finally put everything into a pass

Author: johnzl-777

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/rewrite/__init__.py

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

src/bloqade/squin/rewrite/stim.py

@@ -8,7 +8,7 @@
 
 from bloqade import stim
 from bloqade.squin import op, wire, qubit
-from bloqade.analysis.address import Address
+from bloqade.analysis.address import Address, AddressWire, AddressTuple
 from bloqade.squin.analysis.nsites import Sites
 
 # Probably best to move these attributes to a
@@ -80,13 +80,19 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
 
 @dataclass
-class SquinToStim(RewriteRule):
+class _SquinToStim(RewriteRule):
 
     def get_address(self, value: ir.SSAValue):
-        return value.hints.get("address")
+        try:
+            return value.hints["address"]
+        except KeyError:
+            raise KeyError(f"The address analysis hint for {value} does not exist")
 
     def get_sites(self, value: ir.SSAValue):
-        return value.hints.get("sites")
+        try:
+            return value.hints["sites"]
+        except KeyError:
+            raise KeyError(f"The sites analysis hint for {value} does not exist")
 
     # Go from (most) squin 1Q Ops to stim Ops
     ## X, Y, Z, H, S, (no T!)
@@ -105,7 +111,9 @@ def get_stim_1q_gate(self, squin_op: op.stmts.Operator):
             case op.stmts.Identity():  # enforce sites defined = num wires in
                 return stim.gate.Identity
             case _:
-                return None
+                raise NotImplementedError(
+                    f"The squin operator {squin_op} is not supported in the stim dialect"
+                )
 
     # get the qubit indices from the Apply statement argument
     # wires/qubits
@@ -158,6 +166,8 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
                 return self.rewrite_Apply(node)
             case wire.Wrap():
                 return self.rewrite_Wrap(node)
+            case wire.Measure() | qubit.Measure():
+                return self.rewrite_Measure(node)
             case _:
                 return RewriteResult()
 
@@ -236,3 +246,41 @@ def rewrite_Control(
         stim_stmt.insert_before(apply_stmt_ctrl)
 
         return RewriteResult(has_done_something=True)
+
+    def rewrite_Measure(
+        self, measure_stmt: qubit.Measure | wire.Measure
+    ) -> RewriteResult:
+
+        if isinstance(measure_stmt, qubit.Measure):
+            qubit_ilist_ssa = measure_stmt.qubits
+            # qubits are in an ilist which makes up an AddressTuple
+            address_tuple: AddressTuple = self.get_address(qubit_ilist_ssa).address
+            qubit_idx_ssas = []
+            for qubit_address in address_tuple:
+                qubit_idx = qubit_address.data
+                qubit_idx_stmt = py.constant.Constant(qubit_idx)
+                qubit_idx_stmt.insert_before(measure_stmt)
+                qubit_idx_ssas.append(qubit_idx_stmt.result)
+            qubit_idx_ssas = tuple(qubit_idx_ssas)
+
+        elif isinstance(measure_stmt, wire.Measure):
+            wire_ssa = measure_stmt.wire
+            wire_address: AddressWire = self.get_address(wire_ssa).address
+
+            qubit_idx = wire_address.origin_qubit.data
+            qubit_idx_stmt = py.constant.Constant(qubit_idx)
+            qubit_idx_stmt.insert_before(measure_stmt)
+            qubit_idx_ssas = (qubit_idx_stmt.result,)
+
+        else:
+            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)
+        stim_measure_stmt.insert_before(measure_stmt)
+
+        return RewriteResult(has_done_something=True)

test/squin/stim/stim.py

@@ -1,12 +1,8 @@
 from kirin import ir, types
-from kirin.passes import Fold
-from kirin.rewrite import Walk, Fixpoint, DeadCodeElimination
 from kirin.dialects import py, func, ilist
 
+import bloqade.squin.passes as squin_passes
 from bloqade import qasm2, squin
-from bloqade.analysis import address
-from bloqade.squin.rewrite import SquinToStim, WrapSquinAnalysis
-from bloqade.squin.analysis import nsites
 
 
 def as_int(value: int):
@@ -38,9 +34,6 @@ def gen_func_from_stmts(stmts):
         arg_names=[],
     )
 
-    fold_pass = Fold(extended_dialect)
-    fold_pass(constructed_method)
-
     return constructed_method
 
 
@@ -77,34 +70,8 @@ def test_1q():
 
     constructed_method.print()
 
-    address_frame, _ = address.AddressAnalysis(
-        constructed_method.dialects
-    ).run_analysis(constructed_method, no_raise=False)
-
-    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
-        constructed_method, no_raise=False
-    )
-
-    constructed_method.print(analysis=address_frame.entries)
-    constructed_method.print(analysis=nsites_frame.entries)
-
-    # attempt to wrap analysis results
-
-    wrap_squin_analysis = WrapSquinAnalysis(
-        address_analysis=address_frame.entries, op_site_analysis=nsites_frame.entries
-    )
-    fix_walk_squin_analysis = Fixpoint(Walk(wrap_squin_analysis))
-    rewrite_res = fix_walk_squin_analysis.rewrite(constructed_method.code)
-
-    # attempt rewrite to Stim
-    # Be careful with Fixpoint, can go to infinity until reaches defined threshold
-    squin_to_stim = Walk(SquinToStim())
-    rewrite_res = squin_to_stim.rewrite(constructed_method.code)
-
-    # Get rid of the unused statements
-    dce = Fixpoint(Walk(DeadCodeElimination()))
-    rewrite_res = dce.rewrite(constructed_method.code)
-    print(rewrite_res)
+    squin_to_stim = squin_passes.SquinToStim(constructed_method.dialects)
+    squin_to_stim(constructed_method)
 
     constructed_method.print()
 
@@ -137,34 +104,8 @@ def test_control():
     constructed_method = gen_func_from_stmts(stmts)
     constructed_method.print()
 
-    address_frame, _ = address.AddressAnalysis(
-        constructed_method.dialects
-    ).run_analysis(constructed_method, no_raise=False)
-
-    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
-        constructed_method, no_raise=False
-    )
-
-    constructed_method.print(analysis=address_frame.entries)
-    constructed_method.print(analysis=nsites_frame.entries)
-
-    wrap_squin_analysis = WrapSquinAnalysis(
-        address_analysis=address_frame.entries, op_site_analysis=nsites_frame.entries
-    )
-    fix_walk_squin_analysis = Fixpoint(Walk(wrap_squin_analysis))
-    rewrite_res = fix_walk_squin_analysis.rewrite(constructed_method.code)
-
-    # attempt rewrite to Stim
-    # Be careful with Fixpoint, can go to infinity until reaches defined threshold
-    squin_to_stim = Walk(SquinToStim())
-    rewrite_res = squin_to_stim.rewrite(constructed_method.code)
-
-    constructed_method.print()
-
-    # Get rid of the unused statements
-    dce = Fixpoint(Walk(DeadCodeElimination()))
-    rewrite_res = dce.rewrite(constructed_method.code)
-    print(rewrite_res)
+    squin_to_stim = squin_passes.SquinToStim(constructed_method.dialects)
+    squin_to_stim(constructed_method)
 
     constructed_method.print()