bloqade-circuit 0.4.3__tar.gz → 0.4.4__tar.gz

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     - id: isort
       name: isort (python)
 - repo: https://github.com/psf/black
-  rev: 24.10.0
+  rev: 25.1.0
   hooks:
     - id: black
 - repo: https://github.com/charliermarsh/ruff-pre-commit

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.4.3
+Version: 0.4.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.4.3 → bloqade_circuit-0.4.4}/pyproject.toml

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

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/squin/analysis/nsites/impls.py

@@ -1,4 +1,6 @@
 from kirin import interp
+from kirin.dialects import scf
+from kirin.dialects.scf.typeinfer import TypeInfer as ScfTypeInfer
 
 from bloqade.squin import op, wire
 
@@ -78,3 +80,8 @@ class SquinOp(interp.MethodTable):
     def scale(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Scale):
         op_sites = frame.get(stmt.op)
         return (op_sites,)
+
+
+@scf.dialect.register(key="op.nsites")
+class ScfSquinOp(ScfTypeInfer):
+    pass

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/squin/op/__init__.py

@@ -30,6 +30,9 @@ from ._wrapper import (
     shift as shift,
     spin_n as spin_n,
     spin_p as spin_p,
+    sqrt_x as sqrt_x,
+    sqrt_y as sqrt_y,
+    sqrt_z as sqrt_z,
     adjoint as adjoint,
     control as control,
     identity as identity,

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/squin/op/_wrapper.py

@@ -69,6 +69,18 @@ def y() -> types.Op: ...
 def z() -> types.Op: ...
 
 
+@wraps(stmts.SqrtX)
+def sqrt_x() -> types.Op: ...
+
+
+@wraps(stmts.SqrtY)
+def sqrt_y() -> types.Op: ...
+
+
+@wraps(stmts.S)
+def sqrt_z() -> types.Op: ...
+
+
 @wraps(stmts.H)
 def h() -> types.Op: ...
 

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/squin/op/stmts.py

@@ -142,7 +142,12 @@ class Reset(PrimitiveOp):
 
 
 @statement
-class PauliOp(ConstantUnitary):
+class CliffordOp(ConstantUnitary):
+    pass
+
+
+@statement
+class PauliOp(CliffordOp):
     pass
 
 
@@ -173,6 +178,19 @@ class Z(PauliOp):
     pass
 
 
+@statement(dialect=dialect)
+class SqrtX(ConstantUnitary):
+    pass
+
+
+@statement(dialect=dialect)
+class SqrtY(ConstantUnitary):
+    pass
+
+
+# NOTE no SqrtZ since its equal to S
+
+
 @statement(dialect=dialect)
 class H(ConstantUnitary):
     pass

bloqade_circuit-0.4.4/src/bloqade/squin/rewrite/U3_to_clifford.py

@@ -0,0 +1,149 @@
+# create rewrite rule name SquinMeasureToStim using kirin
+import math
+from typing import List, Tuple, Callable
+
+import numpy as np
+from kirin import ir
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import op, qubit
+
+
+def sdag() -> list[ir.Statement]:
+    return [_op := op.stmts.S(), op.stmts.Adjoint(op=_op.result, is_unitary=True)]
+
+
+# (theta, phi, lam)
+U3_HALF_PI_ANGLE_TO_GATES: dict[
+    tuple[int, int, int], Callable[[], Tuple[List[ir.Statement], ...]]
+] = {
+    (0, 0, 0): lambda: ([op.stmts.Identity(sites=1)],),
+    (0, 0, 1): lambda: ([op.stmts.S()],),
+    (0, 0, 2): lambda: ([op.stmts.Z()],),
+    (0, 0, 3): lambda: (sdag(),),
+    (1, 0, 0): lambda: ([op.stmts.SqrtY()],),
+    (1, 0, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()]),
+    (1, 0, 2): lambda: ([op.stmts.H()],),
+    (1, 0, 3): lambda: (sdag(), [op.stmts.SqrtY()]),
+    (1, 1, 0): lambda: ([op.stmts.SqrtY()], [op.stmts.S()]),
+    (1, 1, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()], [op.stmts.S()]),
+    (1, 1, 2): lambda: ([op.stmts.Z()], [op.stmts.SqrtY()], [op.stmts.S()]),
+    (1, 1, 3): lambda: (sdag(), [op.stmts.SqrtY()], [op.stmts.S()]),
+    (1, 2, 0): lambda: ([op.stmts.SqrtY()], [op.stmts.Z()]),
+    (1, 2, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()], [op.stmts.Z()]),
+    (1, 2, 2): lambda: ([op.stmts.Z()], [op.stmts.SqrtY()], [op.stmts.Z()]),
+    (1, 2, 3): lambda: (sdag(), [op.stmts.SqrtY()], [op.stmts.Z()]),
+    (1, 3, 0): lambda: ([op.stmts.SqrtY()], sdag()),
+    (1, 3, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()], sdag()),
+    (1, 3, 2): lambda: ([op.stmts.Z()], [op.stmts.SqrtY()], sdag()),
+    (1, 3, 3): lambda: (sdag(), [op.stmts.SqrtY()], sdag()),
+    (2, 0, 0): lambda: ([op.stmts.Y()],),
+    (2, 0, 1): lambda: ([op.stmts.S()], [op.stmts.Y()]),
+    (2, 0, 2): lambda: ([op.stmts.Z()], [op.stmts.Y()]),
+    (2, 0, 3): lambda: (sdag(), [op.stmts.Y()]),
+}
+
+
+def equivalent_u3_para(
+    theta_half_pi: int, phi_half_pi: int, lam_half_pi: int
+) -> tuple[int, int, int]:
+    """
+    1. Assume all three angles are in the range [0, 4].
+    2. U3(theta, phi, lam) = -U3(2pi-theta, phi+pi, lam+pi).
+    """
+    return ((4 - theta_half_pi) % 4, (phi_half_pi + 2) % 4, (lam_half_pi + 2) % 4)
+
+
+class SquinU3ToClifford(RewriteRule):
+    """
+    Rewrite squin U3 statements to clifford when possible.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if isinstance(node, (qubit.Apply, qubit.Broadcast)):
+            return self.rewrite_ApplyOrBroadcast_onU3(node)
+        else:
+            return RewriteResult()
+
+    def get_constant(self, node: ir.SSAValue) -> float | None:
+        if isinstance(node.owner, py.Constant):
+            # node.value is a PyAttr, need to get the wrapped value out
+            return node.owner.value.unwrap()
+        else:
+            return None
+
+    def resolve_angle(self, angle: float) -> int | None:
+        """
+        Normalize the angle to be in the range [0, 2π).
+        """
+        # convert to 0.0~1.0, in unit of pi/2
+        angle_half_pi = angle / math.pi * 2.0
+
+        mod = angle_half_pi % 1.0
+        if not (np.isclose(mod, 0.0) or np.isclose(mod, 1.0)):
+            return None
+
+        else:
+            return round((angle / math.tau) % 1 * 4) % 4
+
+    def rewrite_ApplyOrBroadcast_onU3(
+        self, node: qubit.Apply | qubit.Broadcast
+    ) -> RewriteResult:
+        """
+        Rewrite Apply and Broadcast nodes to their clifford equivalent statements.
+        """
+        if not isinstance(node.operator.owner, op.stmts.U3):
+            return RewriteResult()
+
+        gates = self.decompose_U3_gates(node.operator.owner)
+
+        if len(gates) == 0:
+            return RewriteResult()
+
+        for stmt_list in gates:
+            for gate_stmt in stmt_list[:-1]:
+                gate_stmt.insert_before(node)
+
+            oper = stmt_list[-1]
+            oper.insert_before(node)
+            new_node = node.__class__(operator=oper.result, qubits=node.qubits)
+            new_node.insert_before(node)
+
+        node.delete()
+
+        # rewrite U3 to clifford gates
+        return RewriteResult(has_done_something=True)
+
+    def decompose_U3_gates(self, node: op.stmts.U3) -> Tuple[List[ir.Statement], ...]:
+        """
+        Rewrite U3 statements to clifford gates if possible.
+        """
+        theta = self.get_constant(node.theta)
+        phi = self.get_constant(node.phi)
+        lam = self.get_constant(node.lam)
+
+        if theta is None or phi is None or lam is None:
+            return ()
+
+        theta_half_pi: int | None = self.resolve_angle(theta)
+        phi_half_pi: int | None = self.resolve_angle(phi)
+        lam_half_pi: int | None = self.resolve_angle(lam)
+
+        if theta_half_pi is None or phi_half_pi is None or lam_half_pi is None:
+            return ()
+
+        angles_key = (theta_half_pi, phi_half_pi, lam_half_pi)
+        if angles_key not in U3_HALF_PI_ANGLE_TO_GATES:
+            angles_key = equivalent_u3_para(*angles_key)
+            if angles_key not in U3_HALF_PI_ANGLE_TO_GATES:
+                return ()
+
+        gates_stmts = U3_HALF_PI_ANGLE_TO_GATES.get(angles_key)
+
+        # no consistent gates, then:
+        assert (
+            gates_stmts is not None
+        ), "internal error, U3 gates not found for angles: {}".format(angles_key)
+
+        return gates_stmts()

bloqade_circuit-0.4.4/src/bloqade/squin/rewrite/__init__.py

@@ -0,0 +1,7 @@
+from .wrap_analysis import (
+    SitesAttribute as SitesAttribute,
+    AddressAttribute as AddressAttribute,
+    WrapOpSiteAnalysis as WrapOpSiteAnalysis,
+    WrapAddressAnalysis as WrapAddressAnalysis,
+)
+from .remove_dangling_qubits import RemoveDeadRegister as RemoveDeadRegister

bloqade_circuit-0.4.4/src/bloqade/squin/rewrite/remove_dangling_qubits.py

@@ -0,0 +1,19 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import qubit
+
+
+class RemoveDeadRegister(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, qubit.New):
+            return RewriteResult()
+
+        if bool(node.result.uses):
+            return RewriteResult()
+        else:
+            node.delete()
+
+        return RewriteResult(has_done_something=True)

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/squin/rewrite/wrap_analysis.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from dataclasses import dataclass
 
 from kirin import ir
@@ -40,33 +41,47 @@ class SitesAttribute(ir.Attribute):
 
 
 @dataclass
-class WrapSquinAnalysis(RewriteRule):
+class WrapAnalysis(RewriteRule):
 
+    @abstractmethod
+    def wrap(self, value: ir.SSAValue) -> bool:
+        pass
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        has_done_something = any(self.wrap(arg) for arg in node.args)
+        return RewriteResult(has_done_something=has_done_something)
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        has_done_something = any(self.wrap(result) for result in node.results)
+        return RewriteResult(has_done_something=has_done_something)
+
+
+@dataclass
+class WrapAddressAnalysis(WrapAnalysis):
     address_analysis: dict[ir.SSAValue, Address]
-    op_site_analysis: dict[ir.SSAValue, Sites]
 
     def wrap(self, value: ir.SSAValue) -> bool:
         address_analysis_result = self.address_analysis[value]
-        op_site_analysis_result = self.op_site_analysis[value]
 
-        if value.hints.get("address") and value.hints.get("sites"):
+        if value.hints.get("address") is not None:
             return False
-        else:
-            value.hints["address"] = AddressAttribute(address_analysis_result)
-            value.hints["sites"] = SitesAttribute(op_site_analysis_result)
+
+        value.hints["address"] = AddressAttribute(address_analysis_result)
 
         return True
 
-    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
-        has_done_something = False
-        for arg in node.args:
-            if self.wrap(arg):
-                has_done_something = True
-        return RewriteResult(has_done_something=has_done_something)
 
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        has_done_something = False
-        for result in node.results:
-            if self.wrap(result):
-                has_done_something = True
-        return RewriteResult(has_done_something=has_done_something)
+@dataclass
+class WrapOpSiteAnalysis(WrapAnalysis):
+
+    op_site_analysis: dict[ir.SSAValue, Sites]
+
+    def wrap(self, value: ir.SSAValue) -> bool:
+        op_site_analysis_result = self.op_site_analysis[value]
+
+        if value.hints.get("sites") is not None:
+            return False
+
+        value.hints["sites"] = SitesAttribute(op_site_analysis_result)
+
+        return True

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/stim/__init__.py

@@ -1,4 +1,4 @@
-from . import emit as emit, parse as parse, dialects as dialects
+from . import emit as emit, parse as parse, passes as passes, dialects as dialects
 from .groups import main as main
 from ._wrappers import (
     h as h,

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/stim/dialects/auxiliary/stmts/const.py

@@ -47,7 +47,7 @@ class ConstFloat(ir.Statement):
 
 @statement(dialect=dialect)
 class ConstBool(ir.Statement):
-    """IR Statement representing a constant float value."""
+    """IR Statement representing a constant boolean value."""
 
     name = "constant.bool"
     traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})

{bloqade_circuit-0.4.3 → bloqade_circuit-0.4.4}/src/bloqade/stim/dialects/gate/stmts/__init__.py

@@ -1,4 +1,10 @@
 from .pp import SPP as SPP
+from .base import (
+    Gate as Gate,
+    TwoQubitGate as TwoQubitGate,
+    SingleQubitGate as SingleQubitGate,
+    ControlledTwoQubitGate as ControlledTwoQubitGate,
+)
 from .control_2q import CX as CX, CY as CY, CZ as CZ
 from .clifford_1q import (
     H as H,

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

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

bloqade_circuit-0.4.4/src/bloqade/stim/passes/squin_to_stim.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+
+from kirin.passes import Fold
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    DeadCodeElimination,
+    CommonSubexpressionElimination,
+)
+from kirin.ir.method import Method
+from kirin.passes.abc import Pass
+from kirin.rewrite.abc import RewriteResult
+
+from bloqade.stim.groups import main as stim_main_group
+from bloqade.stim.rewrite import (
+    SquinWireToStim,
+    PyConstantToStim,
+    SquinQubitToStim,
+    SquinMeasureToStim,
+    SquinWireIdentityElimination,
+)
+from bloqade.squin.rewrite import RemoveDeadRegister
+
+
+@dataclass
+class SquinToStim(Pass):
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        fold_pass = Fold(mt.dialects)
+        # propagate constants
+        rewrite_result = fold_pass(mt)
+
+        # Assume that address analysis and
+        # wrapping has been done before this pass!
+
+        # Wrap Rewrite + SquinToStim can happen w/ standard walk
+        rewrite_result = (
+            Walk(
+                Chain(
+                    SquinQubitToStim(),
+                    SquinWireToStim(),
+                    SquinMeasureToStim(),  # 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)
+        )
+
+        # 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
+        rewrite_result = (
+            Fixpoint(
+                Walk(
+                    Chain(
+                        DeadCodeElimination(),
+                        CommonSubexpressionElimination(),
+                        RemoveDeadRegister(),
+                    )
+                )
+            )
+            .rewrite(mt.code)
+            .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_circuit-0.4.3/src/bloqade/squin → bloqade_circuit-0.4.4/src/bloqade/stim}/rewrite/__init__.py

@@ -1,11 +1,7 @@
 from .wire_to_stim import SquinWireToStim as SquinWireToStim
 from .qubit_to_stim import SquinQubitToStim as SquinQubitToStim
 from .squin_measure import SquinMeasureToStim as SquinMeasureToStim
-from .wrap_analysis import (
-    SitesAttribute as SitesAttribute,
-    AddressAttribute as AddressAttribute,
-    WrapSquinAnalysis as WrapSquinAnalysis,
-)
+from .py_constant_to_stim import PyConstantToStim as PyConstantToStim
 from .wire_identity_elimination import (
     SquinWireIdentityElimination as SquinWireIdentityElimination,
 )

bloqade_circuit-0.4.4/src/bloqade/stim/rewrite/py_constant_to_stim.py

@@ -0,0 +1,42 @@
+from kirin import ir
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.stim.dialects import auxiliary
+
+# py.constant.int -> stim.const.ConstInt
+# py.constant.float -> stimt.const.ConstFloat
+# py.constant -> stim.const.ConstBool
+#
+
+
+class PyConstantToStim(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        match node:
+            case py.constant.Constant():
+                return self.rewrite_PyConstant(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_PyConstant(self, node: py.constant.Constant) -> RewriteResult:
+
+        # node.value is a PyAttr, need to get the
+        # wrapped value out
+        wrapped_value = node.value.unwrap()
+
+        if isinstance(wrapped_value, int):
+            stim_const = auxiliary.ConstInt(value=wrapped_value)
+        elif isinstance(wrapped_value, float):
+            stim_const = auxiliary.ConstFloat(value=wrapped_value)
+        elif isinstance(wrapped_value, bool):
+            stim_const = auxiliary.ConstBool(value=wrapped_value)
+        elif isinstance(wrapped_value, str):
+            stim_const = auxiliary.ConstStr(value=wrapped_value)
+        else:
+            return RewriteResult()
+
+        node.replace_by(stim_const)
+
+        return RewriteResult(has_done_something=True)

{bloqade_circuit-0.4.3/src/bloqade/squin → bloqade_circuit-0.4.4/src/bloqade/stim}/rewrite/qubit_to_stim.py

@@ -2,8 +2,9 @@ from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.squin import op, qubit
-from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
-from bloqade.squin.rewrite.stim_rewrite_util import (
+from bloqade.squin.rewrite import AddressAttribute
+from bloqade.stim.dialects import gate
+from bloqade.stim.rewrite.util import (
     SQUIN_STIM_GATE_MAPPING,
     rewrite_Control,
     insert_qubit_idx_from_address,
@@ -34,6 +35,15 @@ class SquinQubitToStim(RewriteRule):
         if isinstance(applied_op, op.stmts.Control):
             return rewrite_Control(stmt)
 
+        # check if its adjoint, assume its canonicalized so no nested adjoints.
+        is_conj = False
+        if isinstance(applied_op, op.stmts.Adjoint):
+            if not applied_op.is_unitary:
+                return RewriteResult()
+
+            is_conj = True
+            applied_op = applied_op.op.owner
+
         # need to handle Control through separate means
         # but we can handle X, Y, Z, H, and S here just fine
         stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
@@ -41,9 +51,11 @@ class SquinQubitToStim(RewriteRule):
             return RewriteResult()
 
         address_attr = stmt.qubits.hints.get("address")
+
         if address_attr is None:
             return RewriteResult()
 
+        # sometimes you can get a whole AddressReg...
         assert isinstance(address_attr, AddressAttribute)
         qubit_idx_ssas = insert_qubit_idx_from_address(
             address=address_attr, stmt_to_insert_before=stmt
@@ -52,7 +64,10 @@ class SquinQubitToStim(RewriteRule):
         if qubit_idx_ssas is None:
             return RewriteResult()
 
-        stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
+        if isinstance(stim_1q_op, gate.stmts.Gate):
+            stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas), dagger=is_conj)
+        else:
+            stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
         stmt.replace_by(stim_1q_stmt)
 
         return RewriteResult(has_done_something=True)

{bloqade_circuit-0.4.3/src/bloqade/squin → bloqade_circuit-0.4.4/src/bloqade/stim}/rewrite/squin_measure.py

@@ -4,9 +4,9 @@ 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
-from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
-from bloqade.squin.rewrite.stim_rewrite_util import (
+from bloqade.stim.rewrite.util import (
     is_measure_result_used,
     insert_qubit_idx_from_address,
 )