squin to stim rewrite (#148)

Co-authored-by: Roger-luo <[email protected]>

Author: johnzl-777

src/bloqade/analysis/address/impls.py

@@ -206,6 +206,20 @@ def apply(
     ):
         return frame.get_values(stmt.inputs)
 
+    @interp.impl(squin.wire.MeasureAndReset)
+    def measure_and_reset(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: squin.wire.MeasureAndReset,
+    ):
+
+        # take the address data from the incoming wire
+        # and propagate that forward to the new wire generated.
+        # The first entry can safely be NotQubit because
+        # it's an integer
+        return (NotQubit(), frame.get(stmt.wire))
+
 
 @squin.qubit.dialect.register(key="qubit.address")
 class SquinQubitMethodTable(interp.MethodTable):

src/bloqade/squin/analysis/nsites/__init__.py

@@ -1,6 +1,7 @@
 # Need this for impl registration to work properly!
 from . import impls as impls
 from .lattice import (
+    Sites as Sites,
     NoSites as NoSites,
     AnySites as AnySites,
     NumberSites as NumberSites,

src/bloqade/squin/analysis/nsites/impls.py

@@ -1,6 +1,6 @@
 from kirin import interp
 
-from bloqade.squin import op
+from bloqade.squin import op, wire
 
 from .lattice import (
     NoSites,
@@ -9,6 +9,30 @@
 from .analysis import NSitesAnalysis
 
 
+@wire.dialect.register(key="op.nsites")
+class SquinWire(interp.MethodTable):
+
+    @interp.impl(wire.Apply)
+    @interp.impl(wire.Broadcast)
+    def apply(
+        self,
+        interp: NSitesAnalysis,
+        frame: interp.Frame,
+        stmt: wire.Apply | wire.Broadcast,
+    ):
+
+        return tuple(frame.get(input) for input in stmt.inputs)
+
+    @interp.impl(wire.MeasureAndReset)
+    def measure_and_reset(
+        self, interp: NSitesAnalysis, frame: interp.Frame, stmt: wire.MeasureAndReset
+    ):
+
+        # MeasureAndReset produces both a new wire
+        # and an integer which don't have any sites at all
+        return (NoSites(), NoSites())
+
+
 @op.dialect.register(key="op.nsites")
 class SquinOp(interp.MethodTable):
 

src/bloqade/squin/passes/__init__.py

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

src/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

src/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,
+)

src/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

src/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