Fix qubit.apply lowering with multiple qubit indexing operations (#381)

When more than one qubit indexing operation is present, the lowering for
`squin.qubit.apply` falls through which causes issues in the squin to
stim infrastructure.

To be precise, a statement like:
```python
qubit.apply(op.some_gate, [list_of_qubits[i]])
```
Is fine but @liupengy19 helped identify that something like:
```python
qubit.apply(op.some_gate, [q[i], q[i+1])
```
causes problems.

With help from @david-pl and @kaihsin I was able to figure out that some
of the existing lowering logic just needed to be made a bit more
flexible (:

Author: johnzl-777

src/bloqade/squin/rewrite/desugar.py

@@ -54,7 +54,9 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         qubits = node.qubits
 
         if len(qubits) > 1 and all(q.type.is_subseteq(QubitType) for q in qubits):
-            (qubits_ilist_stmt := ilist.New(qubits)).insert_before(node)
+            (qubits_ilist_stmt := ilist.New(qubits)).insert_before(
+                node
+            )  # qubits is just a tuple of SSAValues
             qubits_ilist = qubits_ilist_stmt.result
 
         elif len(qubits) == 1 and qubits[0].type.is_subseteq(QubitType):
@@ -76,34 +78,44 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
                 return RewriteResult()
 
             is_ilist = isinstance(qbit_stmt := qubits[0].stmt, ilist.New)
+
             if is_ilist:
-                if len(qbit_stmt.values) != 1:
-                    return RewriteResult()
 
-                if not isinstance(
-                    qbit_getindex_result := qbit_stmt.values[0], ir.ResultValue
+                if not all(
+                    isinstance(qbit_getindex_result, ir.ResultValue)
+                    for qbit_getindex_result in qbit_stmt.values
                 ):
                     return RewriteResult()
 
-                qbit_getindex = qbit_getindex_result.stmt
+                # Get the parent statement that the qubit came from
+                # (should be a GetItem instance, see logic below)
+                qbit_getindices = [
+                    qbit_getindex_result.stmt
+                    for qbit_getindex_result in qbit_stmt.values
+                ]
             else:
-                qbit_getindex = qubits[0].stmt
+                qbit_getindices = [qubit.stmt for qubit in qubits]
 
-            if not isinstance(qbit_getindex, py.indexing.GetItem):
+            if any(
+                not isinstance(qbit_getindex, py.indexing.GetItem)
+                for qbit_getindex in qbit_getindices
+            ):
                 return RewriteResult()
 
-            if not qbit_getindex.obj.type.is_subseteq(
-                ilist.IListType[QubitType, types.Any]
+            # The GetItem should have been applied on something that returns an IList of Qubits
+            if any(
+                not qbit_getindex.obj.type.is_subseteq(
+                    ilist.IListType[QubitType, types.Any]
+                )
+                for qbit_getindex in qbit_getindices
             ):
                 return RewriteResult()
 
             if is_ilist:
-                values = qbit_stmt.values
+                qubits_ilist = qbit_stmt.result
             else:
-                values = [qubits[0]]
-
-            (qubits_ilist_stmt := ilist.New(values=values)).insert_before(node)
-            qubits_ilist = qubits_ilist_stmt.result
+                (qubits_ilist_stmt := ilist.New(values=[qubits[0]])).insert_before(node)
+                qubits_ilist = qubits_ilist_stmt.result
         else:
             return RewriteResult()
 

test/squin/test_sugar.py

@@ -46,7 +46,7 @@ def main():
         h = squin.op.h()
         x = squin.op.x()
 
-        # test applying to lest with getindex
+        # test applying to list with getindex
         squin.qubit.apply(x, [q[0]])
 
         # test apply with ast.Name
@@ -89,3 +89,18 @@ def main():
 
     assert math.isclose(abs(ket[0]) ** 2, 1, abs_tol=1e-7)
     assert ket[1] == ket[2] == ket[3] == 0
+
+
+def test_apply_in_for_loop_index_multiple_index():
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(3)
+        squin.qubit.apply(squin.op.h(), q[0])
+        cx = squin.op.cx()
+        for i in range(2):
+            squin.qubit.apply(cx, [q[i], q[i + 1]])
+
+    sim = StackMemorySimulator(min_qubits=3)
+    ket = sim.state_vector(main)
+    assert math.isclose(abs(ket[0]) ** 2, 0.5, abs_tol=1e-5)

test/stim/passes/stim_reference_programs/qubit/for_loop.stim

@@ -0,0 +1,4 @@
+
+H 0
+CX 0 1
+CX 1 2

test/stim/passes/test_squin_noise_to_stim.py

@@ -3,8 +3,9 @@
 from kirin import ir
 
 from bloqade.squin import op, noise, qubit, kernel
+from bloqade.stim.passes import SquinToStimPass
 
-from .test_squin_qubit_to_stim import codegen as _codegen, run_address_and_stim_passes
+from .test_squin_qubit_to_stim import codegen as _codegen
 
 
 def codegen(mt: ir.Method) -> str:
@@ -30,7 +31,7 @@ def test():
         qubit.apply(channel, q[0])
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("apply_pauli_channel_1.stim")
     assert codegen(test) == expected_stim_program
 
@@ -44,7 +45,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("broadcast_pauli_channel_1.stim")
     assert codegen(test) == expected_stim_program
 
@@ -58,7 +59,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program(
         "broadcast_pauli_channel_1_many_qubits.stim"
     )
@@ -76,7 +77,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program(
         "broadcast_pauli_channel_1_reuse.stim"
     )
@@ -110,7 +111,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("broadcast_pauli_channel_2.stim")
     assert codegen(test) == expected_stim_program
 
@@ -143,7 +144,7 @@ def test():
         qubit.broadcast(channel, [q[2], q[3]])
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program(
         "broadcast_pauli_channel_2_reuse_on_4_qubits.stim"
     )
@@ -159,7 +160,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("broadcast_depolarize2.stim")
     assert codegen(test) == expected_stim_program
 
@@ -173,7 +174,7 @@ def test():
         qubit.apply(channel, q[0])
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("apply_depolarize1.stim")
     assert codegen(test) == expected_stim_program
 
@@ -187,7 +188,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("broadcast_depolarize1.stim")
     assert codegen(test) == expected_stim_program
 
@@ -202,7 +203,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program(
         "broadcast_iid_bit_flip_channel.stim"
     )
@@ -219,7 +220,7 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program(
         "broadcast_iid_phase_flip_channel.stim"
     )
@@ -236,6 +237,6 @@ def test():
         qubit.broadcast(channel, q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     expected_stim_program = load_reference_program("broadcast_iid_y_flip_channel.stim")
     assert codegen(test) == expected_stim_program

test/stim/passes/test_squin_qubit_to_stim.py

@@ -2,15 +2,12 @@
 import math
 
 from kirin import ir
-from kirin.rewrite import Walk
 from kirin.dialects import py
 
 from bloqade import squin
 from bloqade.squin import op, noise, qubit, kernel
 from bloqade.stim.emit import EmitStimMain
 from bloqade.stim.passes import SquinToStimPass
-from bloqade.squin.rewrite import WrapAddressAnalysis
-from bloqade.analysis.address import AddressAnalysis
 
 
 # Taken gratuitously from Kai's unit test
@@ -38,12 +35,6 @@ def load_reference_program(filename):
         return f.read()
 
 
-def run_address_and_stim_passes(test: ir.Method):
-    addr_frame, _ = AddressAnalysis(test.dialects).run_analysis(test)
-    Walk(WrapAddressAnalysis(address_analysis=addr_frame.entries)).rewrite(test.code)
-    SquinToStimPass(test.dialects)(test)
-
-
 def test_qubit():
     @kernel
     def test():
@@ -57,7 +48,7 @@ def test():
         squin.qubit.measure(ql)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     base_stim_prog = load_reference_program("qubit.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()
@@ -74,7 +65,7 @@ def test():
         squin.qubit.measure(q[0])
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     base_stim_prog = load_reference_program("qubit_reset.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()
@@ -91,7 +82,7 @@ def test():
         squin.qubit.measure(ql)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     base_stim_prog = load_reference_program("qubit_broadcast.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()
@@ -111,7 +102,7 @@ def test():
         squin.qubit.measure(ql)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
     base_stim_prog = load_reference_program("qubit_loss.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()
@@ -129,7 +120,7 @@ def test():
         squin.qubit.measure(q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
 
     base_stim_prog = load_reference_program("u3_to_clifford.stim")
 
@@ -144,7 +135,7 @@ def test():
         qubit.broadcast(op.sqrt_x(), q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
 
     assert codegen(test).strip() == "SQRT_X 0"
 
@@ -157,6 +148,22 @@ def test():
         qubit.broadcast(op.sqrt_y(), q)
         return
 
-    run_address_and_stim_passes(test)
+    SquinToStimPass(test.dialects)(test)
 
     assert codegen(test).strip() == "SQRT_Y 0"
+
+
+def test_for_loop_rewrite():
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(3)
+        squin.qubit.apply(squin.op.h(), q[0])
+        cx = squin.op.cx()
+        for i in range(2):
+            squin.qubit.apply(cx, [q[i], q[i + 1]])
+
+    SquinToStimPass(main.dialects)(main)
+    base_stim_prog = load_reference_program("for_loop.stim")
+
+    assert codegen(main) == base_stim_prog.rstrip()