bloqade-circuit 0.2.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

bloqade/squin/op/_wrapper.py

@@ -0,0 +1,105 @@
+from kirin.lowering import wraps
+
+from . import stmts, types
+
+
+@wraps(stmts.Kron)
+def kron(lhs: types.Op, rhs: types.Op) -> types.Op: ...
+
+
+@wraps(stmts.Mult)
+def mult(lhs: types.Op, rhs: types.Op) -> types.Op: ...
+
+
+@wraps(stmts.Scale)
+def scale(op: types.Op, factor: complex) -> types.Op: ...
+
+
+@wraps(stmts.Adjoint)
+def adjoint(op: types.Op) -> types.Op: ...
+
+
+@wraps(stmts.Control)
+def control(op: types.Op, *, n_controls: int) -> types.Op:
+    """
+    Create a controlled operator.
+
+    Note, that when considering atom loss, the operator will not be applied if
+    any of the controls has been lost.
+
+    Args:
+        operator: The operator to apply under the control.
+        n_controls: The number qubits to be used as control.
+
+    Returns:
+        Operator
+    """
+    ...
+
+
+@wraps(stmts.Reset)
+def reset() -> types.Op: ...
+
+
+@wraps(stmts.Identity)
+def identity(*, sites: int) -> types.Op: ...
+
+
+@wraps(stmts.Rot)
+def rot(axis: types.Op, angle: float) -> types.Op: ...
+
+
+@wraps(stmts.ShiftOp)
+def shift(theta: float) -> types.Op: ...
+
+
+@wraps(stmts.PhaseOp)
+def phase(theta: float) -> types.Op: ...
+
+
+@wraps(stmts.X)
+def x() -> types.Op: ...
+
+
+@wraps(stmts.Y)
+def y() -> types.Op: ...
+
+
+@wraps(stmts.Z)
+def z() -> types.Op: ...
+
+
+@wraps(stmts.H)
+def h() -> types.Op: ...
+
+
+@wraps(stmts.S)
+def s() -> types.Op: ...
+
+
+@wraps(stmts.T)
+def t() -> types.Op: ...
+
+
+@wraps(stmts.P0)
+def p0() -> types.Op: ...
+
+
+@wraps(stmts.P1)
+def p1() -> types.Op: ...
+
+
+@wraps(stmts.Sn)
+def spin_n() -> types.Op: ...
+
+
+@wraps(stmts.Sp)
+def spin_p() -> types.Op: ...
+
+
+@wraps(stmts.U3)
+def u(theta: float, phi: float, lam: float) -> types.Op: ...
+
+
+@wraps(stmts.PauliString)
+def pauli_string(*, string: str) -> types.Op: ...

bloqade/squin/op/stdlib.py

@@ -0,0 +1,62 @@
+from kirin import ir
+from kirin.prelude import structural_no_opt
+
+from . import types
+from ._dialect import dialect
+from ._wrapper import h, x, y, z, rot, phase, control
+
+
+@ir.dialect_group(structural_no_opt.add(dialect))
+def op(self):
+    def run_pass(method):
+        pass
+
+    return run_pass
+
+
+@op
+def rx(theta: float) -> types.Op:
+    """Rotation X gate."""
+    return rot(x(), theta)
+
+
+@op
+def ry(theta: float) -> types.Op:
+    """Rotation Y gate."""
+    return rot(y(), theta)
+
+
+@op
+def rz(theta: float) -> types.Op:
+    """Rotation Z gate."""
+    return rot(z(), theta)
+
+
+@op
+def cx() -> types.Op:
+    """Controlled X gate."""
+    return control(x(), n_controls=1)
+
+
+@op
+def cy() -> types.Op:
+    """Controlled Y gate."""
+    return control(y(), n_controls=1)
+
+
+@op
+def cz() -> types.Op:
+    """Control Z gate."""
+    return control(z(), n_controls=1)
+
+
+@op
+def ch() -> types.Op:
+    """Control H gate."""
+    return control(h(), n_controls=1)
+
+
+@op
+def cphase(theta: float) -> types.Op:
+    """Control Phase gate."""
+    return control(phase(theta), n_controls=1)

bloqade/squin/op/stmts.py

@@ -142,6 +142,16 @@ class ShiftOp(PrimitiveOp):
     result: ir.ResultValue = info.result(OpType)
 
 
+@statement(dialect=dialect)
+class Reset(PrimitiveOp):
+    """
+    Reset operator for qubits or wires.
+    """
+
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), FixedSites(1)})
+    result: ir.ResultValue = info.result(OpType)
+
+
 @statement
 class PauliOp(ConstantUnitary):
     pass

bloqade/squin/passes/__init__.py

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

bloqade/squin/passes/stim.py

@@ -0,0 +1,68 @@
+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.squin.rewrite import (
+    SquinWireToStim,
+    SquinQubitToStim,
+    WrapSquinAnalysis,
+    SquinMeasureToStim,
+    SquinWireIdentityElimination,
+)
+from bloqade.analysis.address import AddressAnalysis
+from bloqade.squin.analysis.nsites import (
+    NSitesAnalysis,
+)
+
+
+@dataclass
+class SquinToStim(Pass):
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        fold_pass = Fold(mt.dialects)
+        # propagate constants
+        rewrite_result = fold_pass(mt)
+
+        # Get necessary analysis results to plug into hints
+        address_analysis = AddressAnalysis(mt.dialects)
+        address_frame, _ = address_analysis.run_analysis(mt)
+        site_analysis = NSitesAnalysis(mt.dialects)
+        sites_frame, _ = site_analysis.run_analysis(mt)
+
+        # Wrap Rewrite + SquinToStim can happen w/ standard walk
+        rewrite_result = (
+            Walk(
+                Chain(
+                    WrapSquinAnalysis(
+                        address_analysis=address_frame.entries,
+                        op_site_analysis=sites_frame.entries,
+                    ),
+                    SquinQubitToStim(),
+                    SquinWireToStim(),
+                    SquinMeasureToStim(),  # reduce duplicated logic, can split out even more rules later
+                    SquinWireIdentityElimination(),
+                )
+            )
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
+        rewrite_result = (
+            Fixpoint(
+                Walk(Chain(DeadCodeElimination(), CommonSubexpressionElimination()))
+            )
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
+        return rewrite_result

bloqade/squin/qubit.py

@@ -17,6 +17,8 @@ from kirin.lowering import wraps
 from bloqade.types import Qubit, QubitType
 from bloqade.squin.op.types import Op, OpType
 
+from .lowering import ApplyAnyCallLowering
+
 dialect = ir.Dialect("squin.qubit")
 
 
@@ -34,6 +36,14 @@ class Apply(ir.Statement):
     qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
 
 
+@statement(dialect=dialect)
+class ApplyAny(ir.Statement):
+    # NOTE: custom lowering to deal with vararg calls
+    traits = frozenset({ApplyAnyCallLowering()})
+    operator: ir.SSAValue = info.argument(OpType)
+    qubits: tuple[ir.SSAValue, ...] = info.argument()
+
+
 @statement(dialect=dialect)
 class Broadcast(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
@@ -68,19 +78,6 @@ class MeasureQubitList(ir.Statement):
     result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
 
 
-@statement(dialect=dialect)
-class MeasureAndReset(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-    result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
-
-
-@statement(dialect=dialect)
-class Reset(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-
-
 # NOTE: no dependent types in Python, so we have to mark it Any...
 @wraps(New)
 def new(n_qubits: int) -> ilist.IList[Qubit, Any]:
@@ -95,7 +92,7 @@ def new(n_qubits: int) -> ilist.IList[Qubit, Any]:
     ...
 
 
-@wraps(Apply)
+@overload
 def apply(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
     """Apply an operator to a list of qubits.
 
@@ -112,6 +109,27 @@ def apply(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
     ...
 
 
+@overload
+def apply(operator: Op, *qubits: Qubit) -> None:
+    """Apply and operator to any number of qubits.
+
+    Note, that when considering atom loss, lost qubits will be skipped.
+
+    Args:
+        operator: The operator to apply.
+        *qubits: The qubits to apply the operator to. The number of qubits must
+            match the size of the operator.
+
+    Returns:
+        None
+    """
+    ...
+
+
+@wraps(ApplyAny)
+def apply(operator: Op, *qubits) -> None: ...
+
+
 @overload
 def measure(input: Qubit) -> bool: ...
 @overload
@@ -161,26 +179,3 @@ def broadcast(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> No
         None
     """
     ...
-
-
-@wraps(MeasureAndReset)
-def measure_and_reset(qubits: ilist.IList[Qubit, Any]) -> ilist.IList[bool, Any]:
-    """Measure the qubits in the list and reset them."
-
-    Args:
-        qubits: The list of qubits to measure and reset.
-
-    Returns:
-        list[bool]: The result of the measurement.
-    """
-    ...
-
-
-@wraps(Reset)
-def reset(qubits: ilist.IList[Qubit, Any]) -> None:
-    """Reset the qubits in the list."
-
-    Args:
-        qubits: The list of qubits to reset.
-    """
-    ...

bloqade/squin/rewrite/__init__.py

@@ -0,0 +1,11 @@
+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 .wire_identity_elimination import (
+    SquinWireIdentityElimination as SquinWireIdentityElimination,
+)

bloqade/squin/rewrite/desugar.py

@@ -0,0 +1,65 @@
+from kirin import ir, types
+from kirin.dialects import ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin.qubit import (
+    Apply,
+    ApplyAny,
+    QubitType,
+    MeasureAny,
+    MeasureQubit,
+    MeasureQubitList,
+)
+
+
+class MeasureDesugarRule(RewriteRule):
+    """
+    Desugar measure operations in the circuit.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, MeasureAny):
+            return RewriteResult()
+
+        if node.input.type.is_subseteq(QubitType):
+            node.replace_by(
+                MeasureQubit(
+                    qubit=node.input,
+                )
+            )
+            return RewriteResult(has_done_something=True)
+        elif node.input.type.is_subseteq(ilist.IListType[QubitType, types.Any]):
+            node.replace_by(
+                MeasureQubitList(
+                    qubits=node.input,
+                )
+            )
+            return RewriteResult(has_done_something=True)
+
+        return RewriteResult()
+
+
+class ApplyDesugarRule(RewriteRule):
+    """
+    Desugar apply operators in the kernel.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, ApplyAny):
+            return RewriteResult()
+
+        op = node.operator
+        qubits = node.qubits
+
+        if len(qubits) == 1 and qubits[0].type.is_subseteq(ilist.IListType):
+            # NOTE: already calling with just a single argument that is already an ilist
+            qubits_ilist = qubits[0]
+        else:
+            (qubits_ilist_stmt := ilist.New(values=qubits)).insert_before(node)
+            qubits_ilist = qubits_ilist_stmt.result
+
+        stmt = Apply(operator=op, qubits=qubits_ilist)
+        node.replace_by(stmt)
+        return RewriteResult(has_done_something=True)

bloqade/squin/rewrite/qubit_to_stim.py

@@ -0,0 +1,61 @@
+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 (
+    SQUIN_STIM_GATE_MAPPING,
+    rewrite_Control,
+    insert_qubit_idx_from_address,
+)
+
+
+class SquinQubitToStim(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        match node:
+            case qubit.Apply() | qubit.Broadcast():
+                return self.rewrite_Apply_and_Broadcast(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_Apply_and_Broadcast(
+        self, stmt: qubit.Apply | qubit.Broadcast
+    ) -> RewriteResult:
+        """
+        Rewrite Apply and Broadcast nodes to their stim equivalent statements.
+        """
+
+        # this is an SSAValue, need it to be the actual operator
+        applied_op = stmt.operator.owner
+        assert isinstance(applied_op, op.stmts.Operator)
+
+        if isinstance(applied_op, op.stmts.Control):
+            return rewrite_Control(stmt)
+
+        # 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))
+        if stim_1q_op is None:
+            return RewriteResult()
+
+        address_attr = stmt.qubits.hints.get("address")
+        if address_attr is None:
+            return RewriteResult()
+
+        assert isinstance(address_attr, AddressAttribute)
+        qubit_idx_ssas = insert_qubit_idx_from_address(
+            address=address_attr, stmt_to_insert_before=stmt
+        )
+
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
+        stmt.replace_by(stim_1q_stmt)
+
+        return RewriteResult(has_done_something=True)
+
+
+# put rewrites for measure statements in separate rule, then just have to dispatch

bloqade/squin/rewrite/squin_measure.py

@@ -0,0 +1,73 @@
+# create rewrite rule name SquinMeasureToStim using kirin
+from kirin import ir
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import wire, qubit
+from bloqade.stim.dialects import collapse
+from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
+from bloqade.squin.rewrite.stim_rewrite_util import (
+    is_measure_result_used,
+    insert_qubit_idx_from_address,
+)
+
+
+class SquinMeasureToStim(RewriteRule):
+    """
+    Rewrite squin measure-related statements to stim statements.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        match node:
+            case qubit.MeasureQubit() | qubit.MeasureQubitList() | wire.Measure():
+                return self.rewrite_Measure(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_Measure(
+        self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure
+    ) -> RewriteResult:
+        if is_measure_result_used(measure_stmt):
+            return RewriteResult()
+
+        qubit_idx_ssas = self.get_qubit_idx_ssas(measure_stmt)
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        prob_noise_stmt = py.constant.Constant(0.0)
+        stim_measure_stmt = collapse.MZ(
+            p=prob_noise_stmt.result,
+            targets=qubit_idx_ssas,
+        )
+        prob_noise_stmt.insert_before(measure_stmt)
+        measure_stmt.replace_by(stim_measure_stmt)
+
+        return RewriteResult(has_done_something=True)
+
+    def get_qubit_idx_ssas(
+        self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure
+    ) -> tuple[ir.SSAValue, ...] | None:
+        """
+        Extract the address attribute and insert qubit indices for the given measure statement.
+        """
+        match measure_stmt:
+            case qubit.MeasureQubit():
+                address_attr = measure_stmt.qubit.hints.get("address")
+            case qubit.MeasureQubitList():
+                address_attr = measure_stmt.qubits.hints.get("address")
+            case wire.Measure():
+                address_attr = measure_stmt.wire.hints.get("address")
+            case _:
+                return None
+
+        if address_attr is None:
+            return None
+
+        assert isinstance(address_attr, AddressAttribute)
+
+        qubit_idx_ssas = insert_qubit_idx_from_address(
+            address=address_attr, stmt_to_insert_before=measure_stmt
+        )
+
+        return qubit_idx_ssas