bloqade-circuit 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

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/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/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/qubit_to_stim.py

@@ -0,0 +1,84 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade import stim
+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 qubit.Reset():
+                return self.rewrite_Reset(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)
+
+    def rewrite_Reset(self, reset_stmt: qubit.Reset) -> RewriteResult:
+        qubit_ilist_ssa = reset_stmt.qubits
+        # qubits are in an ilist which makes up an AddressTuple
+        address_attr = qubit_ilist_ssa.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=reset_stmt
+        )
+
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        stim_rz_stmt = stim.collapse.stmts.RZ(targets=qubit_idx_ssas)
+        reset_stmt.replace_by(stim_rz_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,98 @@
+# 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 import stim
+from bloqade.squin import wire, qubit
+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 qubit.MeasureAndReset() | wire.MeasureAndReset():
+                return self.rewrite_MeasureAndReset(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 = stim.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 rewrite_MeasureAndReset(
+        self, meas_and_reset_stmt: qubit.MeasureAndReset | wire.MeasureAndReset
+    ) -> RewriteResult:
+        if not is_measure_result_used(meas_and_reset_stmt):
+            return RewriteResult()
+
+        qubit_idx_ssas = self.get_qubit_idx_ssas(meas_and_reset_stmt)
+
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        error_p_stmt = py.Constant(0.0)
+        stim_mz_stmt = stim.collapse.MZ(targets=qubit_idx_ssas, p=error_p_stmt.result)
+        stim_rz_stmt = stim.collapse.RZ(
+            targets=qubit_idx_ssas,
+        )
+
+        error_p_stmt.insert_before(meas_and_reset_stmt)
+        stim_mz_stmt.insert_before(meas_and_reset_stmt)
+        meas_and_reset_stmt.replace_by(stim_rz_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

bloqade/squin/rewrite/stim_rewrite_util.py

@@ -0,0 +1,158 @@
+from kirin import ir
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteResult
+
+from bloqade.squin import op, wire, qubit
+from bloqade.stim.dialects import gate
+from bloqade.analysis.address import AddressWire, AddressQubit, AddressTuple
+from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
+
+SQUIN_STIM_GATE_MAPPING = {
+    op.stmts.X: gate.X,
+    op.stmts.Y: gate.Y,
+    op.stmts.Z: gate.Z,
+    op.stmts.H: gate.H,
+    op.stmts.S: gate.S,
+    op.stmts.Identity: gate.Identity,
+}
+
+
+def insert_qubit_idx_from_address(
+    address: AddressAttribute, stmt_to_insert_before: ir.Statement
+) -> tuple[ir.SSAValue, ...] | None:
+    """
+    Extract qubit indices from an AddressAttribute and insert them into the SSA form.
+    """
+    address_data = address.address
+    qubit_idx_ssas = []
+
+    if isinstance(address_data, AddressTuple):
+        for address_qubit in address_data.data:
+            if not isinstance(address_qubit, AddressQubit):
+                return
+            qubit_idx = address_qubit.data
+            qubit_idx_stmt = py.Constant(qubit_idx)
+            qubit_idx_stmt.insert_before(stmt_to_insert_before)
+            qubit_idx_ssas.append(qubit_idx_stmt.result)
+    elif isinstance(address_data, AddressWire):
+        address_qubit = address_data.origin_qubit
+        qubit_idx = address_qubit.data
+        qubit_idx_stmt = py.Constant(qubit_idx)
+        qubit_idx_stmt.insert_before(stmt_to_insert_before)
+        qubit_idx_ssas.append(qubit_idx_stmt.result)
+    else:
+        return
+
+    return tuple(qubit_idx_ssas)
+
+
+def insert_qubit_idx_from_wire_ssa(
+    wire_ssas: tuple[ir.SSAValue, ...], stmt_to_insert_before: ir.Statement
+) -> tuple[ir.SSAValue, ...] | None:
+    """
+    Extract qubit indices from wire SSA values and insert them into the SSA form.
+    """
+    qubit_idx_ssas = []
+    for wire_ssa in wire_ssas:
+        address_attribute = wire_ssa.hints.get("address")
+        if address_attribute is None:
+            return
+        assert isinstance(address_attribute, AddressAttribute)
+        wire_address = address_attribute.address
+        assert isinstance(wire_address, AddressWire)
+        qubit_idx = wire_address.origin_qubit.data
+        qubit_idx_stmt = py.Constant(qubit_idx)
+        qubit_idx_ssas.append(qubit_idx_stmt.result)
+        qubit_idx_stmt.insert_before(stmt_to_insert_before)
+
+    return tuple(qubit_idx_ssas)
+
+
+def insert_qubit_idx_after_apply(
+    stmt: wire.Apply | qubit.Apply | wire.Broadcast | qubit.Broadcast,
+) -> tuple[ir.SSAValue, ...] | None:
+    """
+    Extract qubit indices from Apply or Broadcast statements.
+    """
+    if isinstance(stmt, (qubit.Apply, qubit.Broadcast)):
+        qubits = stmt.qubits
+        address_attribute = qubits.hints.get("address")
+        if address_attribute is None:
+            return
+        assert isinstance(address_attribute, AddressAttribute)
+        return insert_qubit_idx_from_address(
+            address=address_attribute, stmt_to_insert_before=stmt
+        )
+    elif isinstance(stmt, (wire.Apply, wire.Broadcast)):
+        wire_ssas = stmt.inputs
+        return insert_qubit_idx_from_wire_ssa(
+            wire_ssas=wire_ssas, stmt_to_insert_before=stmt
+        )
+
+
+def rewrite_Control(
+    stmt_with_ctrl: qubit.Apply | wire.Apply | qubit.Broadcast | wire.Broadcast,
+) -> RewriteResult:
+    """
+    Handle control gates for Apply and Broadcast statements.
+    """
+    ctrl_op = stmt_with_ctrl.operator.owner
+    assert isinstance(ctrl_op, op.stmts.Control)
+
+    ctrl_op_target_gate = ctrl_op.op.owner
+    assert isinstance(ctrl_op_target_gate, op.stmts.Operator)
+
+    qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt_with_ctrl)
+    if qubit_idx_ssas is None:
+        return RewriteResult()
+
+    # Separate control and target qubits
+    target_qubits = []
+    ctrl_qubits = []
+    for i in range(len(qubit_idx_ssas)):
+        if (i % 2) == 0:
+            ctrl_qubits.append(qubit_idx_ssas[i])
+        else:
+            target_qubits.append(qubit_idx_ssas[i])
+
+    target_qubits = tuple(target_qubits)
+    ctrl_qubits = tuple(ctrl_qubits)
+
+    supported_gate_mapping = {
+        op.stmts.X: gate.CX,
+        op.stmts.Y: gate.CY,
+        op.stmts.Z: gate.CZ,
+    }
+
+    stim_gate = supported_gate_mapping.get(type(ctrl_op_target_gate))
+    if stim_gate is None:
+        return RewriteResult()
+
+    stim_stmt = stim_gate(controls=ctrl_qubits, targets=target_qubits)
+
+    if isinstance(stmt_with_ctrl, (wire.Apply, wire.Broadcast)):
+        # have to "reroute" the input of these statements to directly plug in
+        # to subsequent statements, remove dependency on the current statement
+        for input_wire, output_wire in zip(
+            stmt_with_ctrl.inputs, stmt_with_ctrl.results
+        ):
+            output_wire.replace_by(input_wire)
+
+    stmt_with_ctrl.replace_by(stim_stmt)
+
+    return RewriteResult(has_done_something=True)
+
+
+def is_measure_result_used(
+    stmt: (
+        qubit.MeasureAndReset
+        | qubit.MeasureQubit
+        | qubit.MeasureQubitList
+        | wire.MeasureAndReset
+        | wire.Measure
+    ),
+) -> bool:
+    """
+    Check if the result of a measure statement is used in the program.
+    """
+    return bool(stmt.result.uses)

bloqade/squin/rewrite/wire_identity_elimination.py

@@ -0,0 +1,24 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import wire
+
+
+class SquinWireIdentityElimination(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        """
+        Handle the case where an unwrap feeds a wire directly into a wrap,
+        equivalent to nothing happening/identity operation
+
+        w = unwrap(qubit)
+        wrap(qubit, w)
+        """
+        if isinstance(node, wire.Wrap):
+            wire_origin_stmt = node.wire.owner
+            if isinstance(wire_origin_stmt, wire.Unwrap):
+                node.delete()  # get rid of wrap
+                wire_origin_stmt.delete()  # get rid of the unwrap
+                return RewriteResult(has_done_something=True)
+
+        return RewriteResult()

bloqade/squin/rewrite/wire_to_stim.py

@@ -0,0 +1,73 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade import stim
+from bloqade.squin import op, wire
+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,
+    insert_qubit_idx_from_wire_ssa,
+)
+
+
+class SquinWireToStim(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        match node:
+            case wire.Apply() | wire.Broadcast():
+                return self.rewrite_Apply_and_Broadcast(node)
+            case wire.Reset():
+                return self.rewrite_Reset(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_Apply_and_Broadcast(
+        self, stmt: wire.Apply | wire.Broadcast
+    ) -> RewriteResult:
+
+        # 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)
+
+        stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
+        if stim_1q_op is None:
+            return RewriteResult()
+
+        qubit_idx_ssas = insert_qubit_idx_from_wire_ssa(
+            wire_ssas=stmt.inputs, stmt_to_insert_before=stmt
+        )
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
+
+        # Get the wires from the inputs of Apply or Broadcast,
+        # then put those as the result of the current stmt
+        # before replacing it entirely
+        for input_wire, output_wire in zip(stmt.inputs, stmt.results):
+            output_wire.replace_by(input_wire)
+
+        stmt.replace_by(stim_1q_stmt)
+
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_Reset(self, reset_stmt: wire.Reset) -> RewriteResult:
+        address_attr = reset_stmt.wire.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=reset_stmt
+        )
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        stim_rz_stmt = stim.collapse.stmts.RZ(targets=qubit_idx_ssas)
+        reset_stmt.replace_by(stim_rz_stmt)
+
+        return RewriteResult(has_done_something=True)

bloqade/squin/rewrite/wrap_analysis.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.print.printer import Printer
+
+from bloqade.squin import op, wire
+from bloqade.analysis.address import Address
+from bloqade.squin.analysis.nsites import Sites
+
+
+@wire.dialect.register
+@dataclass
+class AddressAttribute(ir.Attribute):
+
+    name = "Address"
+    address: Address
+
+    def __hash__(self) -> int:
+        return hash(self.address)
+
+    def print_impl(self, printer: Printer) -> None:
+        # Can return to implementing this later
+        printer.print(self.address)
+
+
+@op.dialect.register
+@dataclass
+class SitesAttribute(ir.Attribute):
+
+    name = "Sites"
+    sites: Sites
+
+    def __hash__(self) -> int:
+        return hash(self.sites)
+
+    def print_impl(self, printer: Printer) -> None:
+        # Can return to implementing this later
+        printer.print(self.sites)
+
+
+@dataclass
+class WrapSquinAnalysis(RewriteRule):
+
+    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"):
+            return False
+        else:
+            value.hints["address"] = AddressAttribute(address_analysis_result)
+            value.hints["sites"] = SitesAttribute(op_site_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)

bloqade/squin/wire.py

@@ -6,7 +6,7 @@ circuits. Thus we do not define wrapping functions for the statements in this
 dialect.
 """
 
-from kirin import ir, types, interp, lowering
+from kirin import ir, types, lowering
 from kirin.decl import info, statement
 from kirin.lowering import wraps
 
@@ -112,18 +112,6 @@ class Reset(ir.Statement):
     wire: ir.SSAValue = info.argument(WireType)
 
 
-# Issue where constant propagation can't handle
-# multiple return values from Apply properly
-@dialect.register(key="constprop")
-class ConstPropWire(interp.MethodTable):
-
-    @interp.impl(Apply)
-    @interp.impl(Broadcast)
-    def apply(self, interp, frame, stmt: Apply):
-
-        return frame.get_values(stmt.inputs)
-
-
 @wraps(Unwrap)
 def unwrap(qubit: Qubit) -> Wire: ...