first round of meeting feedback implemented

Author: johnzl-777

src/bloqade/squin/passes/stim.py

@@ -12,7 +12,12 @@
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
 
-from bloqade.squin.rewrite import SquinWireToStim, SquinQubitToStim, WrapSquinAnalysis
+from bloqade.squin.rewrite import (
+    SquinWireToStim,
+    SquinQubitToStim,
+    WrapSquinAnalysis,
+    SquinWireIdentityElimination,
+)
 from bloqade.analysis.address import AddressAnalysis
 from bloqade.squin.analysis.nsites import (
     NSitesAnalysis,
@@ -43,6 +48,7 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
                     ),
                     SquinQubitToStim(),
                     SquinWireToStim(),
+                    SquinWireIdentityElimination(),
                 )
             )
             .rewrite(mt.code)

src/bloqade/squin/rewrite/__init__.py

@@ -5,3 +5,6 @@
     AddressAttribute as AddressAttribute,
     WrapSquinAnalysis as WrapSquinAnalysis,
 )
+from .wire_identity_elimination import (
+    SquinWireIdentityElimination as SquinWireIdentityElimination,
+)

src/bloqade/squin/rewrite/qubit_to_stim.py

@@ -6,16 +6,19 @@
 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,
-    get_stim_1q_gate,
     are_sites_compatible,
+    is_measure_result_used,
     insert_qubit_idx_from_address,
 )
 
 
 class SquinQubitToStim(RewriteRule):
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        # don't want to alloc dict, change back to if else/match case
         rewrite_methods = {
             qubit.Apply: self.rewrite_Apply_and_Broadcast,
             qubit.Broadcast: self.rewrite_Apply_and_Broadcast,
@@ -31,13 +34,13 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         return rewrite_method(node)
 
-    # handle Control
     def rewrite_Apply_and_Broadcast(
         self, stmt: qubit.Apply | qubit.Broadcast
     ) -> RewriteResult:
         """
         Rewrite Apply and Broadcast nodes to their stim equivalent statements.
         """
+        # get rid of are_sites_compatible, assume program is properly structured
         if not are_sites_compatible(stmt):
             return RewriteResult()
 
@@ -50,7 +53,7 @@ def rewrite_Apply_and_Broadcast(
 
         # need to handle Control through separate means
         # but we can handle X, Y, Z, H, and S here just fine
-        stim_1q_op = get_stim_1q_gate(applied_op)
+        stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
         if stim_1q_op is None:
             return RewriteResult()
 
@@ -75,6 +78,9 @@ def rewrite_Measure(
         self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList
     ) -> RewriteResult:
 
+        if is_measure_result_used(measure_stmt):
+            return RewriteResult()
+
         # qubit_ssa will always be an ilist of qubits
         # but need to be careful with singular vs plural "qubit" attribute name
         if isinstance(measure_stmt, qubit.MeasureQubit):
@@ -132,6 +138,9 @@ def rewrite_MeasureAndReset(
         self, meas_and_reset_stmt: qubit.MeasureAndReset
     ) -> RewriteResult:
 
+        if is_measure_result_used(meas_and_reset_stmt):
+            return RewriteResult()
+
         address_attr = meas_and_reset_stmt.qubits.hints.get("address")
         if address_attr is None:
             return RewriteResult()
@@ -154,3 +163,6 @@ def rewrite_MeasureAndReset(
         meas_and_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/stim_rewrite_util.py

@@ -10,20 +10,14 @@
 from bloqade.squin.analysis.nsites import NumberSites
 from bloqade.squin.rewrite.wrap_analysis import SitesAttribute, AddressAttribute
 
-
-def get_stim_1q_gate(squin_op: op.stmts.Operator):
-    """
-    Map squin 1Q Ops to stim Ops.
-    """
-    gate_mapping = {
-        op.stmts.X: stim.gate.X,
-        op.stmts.Y: stim.gate.Y,
-        op.stmts.Z: stim.gate.Z,
-        op.stmts.H: stim.gate.H,
-        op.stmts.S: stim.gate.S,
-        op.stmts.Identity: stim.gate.Identity,
-    }
-    return gate_mapping.get(type(squin_op))
+SQUIN_STIM_GATE_MAPPING = {
+    op.stmts.X: stim.gate.X,
+    op.stmts.Y: stim.gate.Y,
+    op.stmts.Z: stim.gate.Z,
+    op.stmts.H: stim.gate.H,
+    op.stmts.S: stim.gate.S,
+    op.stmts.Identity: stim.gate.Identity,
+}
 
 
 def insert_qubit_idx_from_address(
@@ -190,3 +184,18 @@ def rewrite_Control(
     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)

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

src/bloqade/squin/rewrite/wire_to_stim.py

@@ -6,9 +6,10 @@
 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,
-    get_stim_1q_gate,
     are_sites_compatible,
+    is_measure_result_used,
     insert_qubit_idx_from_address,
     insert_qubit_idx_from_wire_ssa,
 )
@@ -17,21 +18,17 @@
 class SquinWireToStim(RewriteRule):
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-
-        rewrite_methods = {
-            wire.Apply: self.rewrite_Apply_and_Broadcast,
-            wire.Broadcast: self.rewrite_Apply_and_Broadcast,
-            wire.Wrap: self.rewrite_Wrap,
-            wire.Measure: self.rewrite_Measure,
-            wire.Reset: self.rewrite_Reset,
-            wire.MeasureAndReset: self.rewrite_MeasureAndReset,
-        }
-
-        rewrite_method = rewrite_methods.get(type(node))
-        if rewrite_method is None:
-            return RewriteResult()
-
-        return rewrite_method(node)
+        match node:
+            case wire.Apply() | wire.Broadcast():
+                return self.rewrite_Apply_and_Broadcast(node)
+            case wire.Measure():
+                return self.rewrite_Measure(node)
+            case wire.Reset():
+                return self.rewrite_Reset(node)
+            case wire.MeasureAndReset():
+                return self.rewrite_MeasureAndReset(node)
+            case _:
+                return RewriteResult()
 
     def rewrite_Apply_and_Broadcast(
         self, stmt: wire.Apply | wire.Broadcast
@@ -47,7 +44,7 @@ def rewrite_Apply_and_Broadcast(
         if isinstance(applied_op, op.stmts.Control):
             return rewrite_Control(stmt)
 
-        stim_1q_op = get_stim_1q_gate(applied_op)
+        stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
         if stim_1q_op is None:
             return RewriteResult()
 
@@ -69,21 +66,11 @@ def rewrite_Apply_and_Broadcast(
 
         return RewriteResult(has_done_something=True)
 
-    def rewrite_Wrap(self, wrap_stmt: wire.Wrap) -> RewriteResult:
-
-        # structure at this point should be:
-        ## w = wire.Unwrap(wire)
-        ## wire.Wrap(qubit, w)
-
-        wire_origin_stmt = wrap_stmt.wire.owner
-        if isinstance(wire_origin_stmt, wire.Unwrap):
-            wrap_stmt.delete()
-            return RewriteResult(has_done_something=True)
-
-        return RewriteResult()
-
     def rewrite_Measure(self, measure_stmt: wire.Measure) -> RewriteResult:
 
+        if is_measure_result_used(measure_stmt):
+            return RewriteResult()
+
         wire_ssa = measure_stmt.wire
         address_attr = wire_ssa.hints.get("address")
         if address_attr is None:
@@ -126,6 +113,9 @@ def rewrite_Reset(self, reset_stmt: wire.Reset) -> RewriteResult:
 
     def rewrite_MeasureAndReset(self, meas_and_reset_stmt: wire.MeasureAndReset):
 
+        if is_measure_result_used(meas_and_reset_stmt):
+            return RewriteResult()
+
         address_attr = meas_and_reset_stmt.wire.hints.get("address")
         if address_attr is None:
             return RewriteResult()
@@ -141,6 +131,7 @@ def rewrite_MeasureAndReset(self, meas_and_reset_stmt: wire.MeasureAndReset):
         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)

test/squin/stim/stim.py

@@ -157,6 +157,16 @@ def test_broadcast_wire_1q_application():
     constructed_method.print()
 
 
+# before ANY rewrite, aggressively inline everything, then do the rewrite
+# for Stim pass, need to call validation , check any invoke
+
+# Put one codegen test to stim
+# finish measurement analysis Friday - if painful, ask help from Kai
+# work on other detector rewrite
+
+# later on lower for loop to repeat
+
+
 def test_broadcast_qubit_1q_application():
 
     stmts: list[ir.Statement] = [