squin to stim rewrite

src/bloqade/analysis/address/impls.py

@@ -206,6 +206,20 @@
     ):
         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

@@ -2,7 +2,7 @@
 
 from kirin import ir, interp
 
-from bloqade.squin import op
+from bloqade.squin import op, wire
 
 from .lattice import (
     NoSites,
@@ -11,6 +11,24 @@
 from .analysis import NSitesAnalysis
 
 
+@wire.dialect.register(key="op.nsites")
+class SquinWire(interp.MethodTable):
+
+    @interp.impl(wire.Apply)
+    def apply(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: wire.Apply):
+
+        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,59 @@
+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
+
+import bloqade.squin.rewrite as squin_rewrite
+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(
+                    squin_rewrite.WrapSquinAnalysis(
+                        address_analysis=address_frame.entries,
+                        op_site_analysis=sites_frame.entries,
+                    ),
+                    squin_rewrite._SquinToStim(),
+                )
+            )
+            .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/qubit.py

@@ -48,6 +48,11 @@ class MeasureAndReset(ir.Statement):
     result: ir.ResultValue = info.result(types.Int)
 
 
+# MZ -> RZ
+# MRZ (never really used, but could have performance benefit)
+#
+
+
 @statement(dialect=dialect)
 class Reset(ir.Statement):
     qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])

src/bloqade/squin/rewrite/__init__.py

@@ -0,0 +1,6 @@
+from .stim import (
+    SitesAttribute as SitesAttribute,
+    AddressAttribute as AddressAttribute,
+    WrapSquinAnalysis as WrapSquinAnalysis,
+    _SquinToStim as _SquinToStim,
+)