Merge branch 'main' into david/592-3-fidelity-analysis

Author: weinbe58

src/bloqade/analysis/address/impls.py

@@ -366,16 +366,22 @@ def for_loop(
         if iter_type is None:
             return interp_.eval_fallback(frame, stmt)
 
+        body_values = {}
         for value in iterable:
             with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
                 loop_vars = interp_.frame_call_region(
                     body_frame, stmt, stmt.body, value, *loop_vars
                 )
 
+            for ssa, val in body_frame.entries.items():
+                body_values[ssa] = body_values.setdefault(ssa, val).join(val)
+
             if loop_vars is None:
                 loop_vars = ()
 
             elif isinstance(loop_vars, interp.ReturnValue):
+                frame.set_values(body_frame.entries.keys(), body_frame.entries.values())
                 return loop_vars
 
+        frame.set_values(body_values.keys(), body_values.values())
         return loop_vars

src/bloqade/cirq_utils/noise/model.py

@@ -98,7 +98,10 @@ def __post_init__(self):
 
     @staticmethod
     def validate_moments(moments: Iterable[cirq.Moment]):
-        allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset().gates
+        reset_family = cirq.GateFamily(gate=cirq.ResetChannel, ignore_global_phase=True)
+        allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset(
+            additional_gates=[reset_family]
+        ).gates
 
         for moment in moments:
             for operation in moment:
@@ -117,7 +120,7 @@ def validate_moments(moments: Iterable[cirq.Moment]):
                     )
 
     def parallel_cz_errors(
-        self, ctrls: list[int], qargs: list[int], rest: list[int]
+        self, ctrls: Sequence[int], qargs: Sequence[int], rest: Sequence[int]
     ) -> dict[tuple[float, float, float, float], list[int]]:
         raise NotImplementedError(
             "This noise model doesn't support rewrites on bloqade kernels, but should be used with cirq."
@@ -245,8 +248,15 @@ def noisy_moment(self, moment, system_qubits):
         # Moment with original ops
         original_moment = moment
 
+        no_noise_condition = (
+            len(moment.operations) == 0
+            or cirq.is_measurement(moment.operations[0])
+            or isinstance(moment.operations[0].gate, cirq.ResetChannel)
+            or isinstance(moment.operations[0].gate, cirq.BitFlipChannel)
+        )
+
         # Check if the moment is empty
-        if len(moment.operations) == 0:
+        if no_noise_condition:
             move_noise_ops = []
             gate_noise_ops = []
         # Check if the moment contains 1-qubit gates or 2-qubit gates
@@ -319,20 +329,84 @@ def noisy_moments(
 
         # Split into moments with only 1Q and 2Q gates
         moments_1q = [
-            cirq.Moment([op for op in moment.operations if len(op.qubits) == 1])
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (len(op.qubits) == 1)
+                    and (not cirq.is_measurement(op))
+                    and (not isinstance(op.gate, cirq.ResetChannel))
+                ]
+            )
             for moment in moments
         ]
         moments_2q = [
-            cirq.Moment([op for op in moment.operations if len(op.qubits) == 2])
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (len(op.qubits) == 2) and (not cirq.is_measurement(op))
+                ]
+            )
             for moment in moments
         ]
 
-        assert len(moments_1q) == len(moments_2q)
+        moments_measurement = [
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (cirq.is_measurement(op))
+                    or (isinstance(op.gate, cirq.ResetChannel))
+                ]
+            )
+            for moment in moments
+        ]
+
+        assert len(moments_1q) == len(moments_2q) == len(moments_measurement)
 
         interleaved_moments = []
+
+        def count_remaining_cz_moments(moments_2q):
+            remaining_cz_counts = []
+            count = 0
+            for m in moments_2q[::-1]:
+                if any(isinstance(op.gate, cirq.CZPowGate) for op in m.operations):
+                    count += 1
+                remaining_cz_counts = [count] + remaining_cz_counts
+            return remaining_cz_counts
+
+        remaining_cz_moments = count_remaining_cz_moments(moments_2q)
+
+        pm = 2 * self.sitter_pauli_rates[0]
+        ps = 2 * self.cz_unpaired_pauli_rates[0]
+
+        # probability of a bitflip error for a sitting, unpaired qubit during a move/cz/move cycle.
+        heuristic_1step_bitflip_error: float = (
+            2 * pm * (1 - ps) * (1 - pm) + (1 - pm) ** 2 * ps + pm**2 * ps
+        )
+
         for idx, moment in enumerate(moments_1q):
             interleaved_moments.append(moment)
             interleaved_moments.append(moments_2q[idx])
+            # Measurements on Gemini will be at the end, so for circuits with mid-circuit measurements we will insert a
+            # bitflip error proportional to the number of moments left in the circuit to account for the decoherence
+            # that will happen before the final terminal measurement.
+            measured_qubits = []
+            for op in moments_measurement[idx].operations:
+                if cirq.is_measurement(op):
+                    measured_qubits += list(op.qubits)
+            # probability of a bitflip error should be Binomial(moments_left,heuristic_1step_bitflip_error)
+            delayed_measurement_error = (
+                1
+                - (1 - 2 * heuristic_1step_bitflip_error) ** (remaining_cz_moments[idx])
+            ) / 2
+            interleaved_moments.append(
+                cirq.Moment(
+                    cirq.bit_flip(delayed_measurement_error).on_each(measured_qubits)
+                )
+            )
+            interleaved_moments.append(moments_measurement[idx])
 
         interleaved_circuit = cirq.Circuit.from_moments(*interleaved_moments)
 
@@ -368,14 +442,21 @@ def noisy_moment(self, moment, system_qubits):
             "all qubits in the circuit must be defined as cirq.GridQubit objects."
         )
         # Check if the moment is empty
-        if len(moment.operations) == 0:
+        if len(moment.operations) == 0 or cirq.is_measurement(moment.operations[0]):
             move_moments = []
             gate_noise_ops = []
         # Check if the moment contains 1-qubit gates or 2-qubit gates
         elif len(moment.operations[0].qubits) == 1:
-            gate_noise_ops, _ = self._single_qubit_moment_noise_ops(
-                moment, system_qubits
-            )
+            if (
+                (isinstance(moment.operations[0].gate, cirq.ResetChannel))
+                or (cirq.is_measurement(moment.operations[0]))
+                or (isinstance(moment.operations[0].gate, cirq.BitFlipChannel))
+            ):
+                gate_noise_ops = []
+            else:
+                gate_noise_ops, _ = self._single_qubit_moment_noise_ops(
+                    moment, system_qubits
+                )
             move_moments = []
         elif len(moment.operations[0].qubits) == 2:
             cg = OneZoneConflictGraph(moment)

src/bloqade/cirq_utils/parallelize.py

@@ -119,7 +119,11 @@ def auto_similarity(
     flattened_circuit: list[GateOperation] = list(cirq.flatten_op_tree(circuit))
     weights = {}
     for i in range(len(flattened_circuit)):
+        if not cirq.has_unitary(flattened_circuit[i]):
+            continue
         for j in range(i + 1, len(flattened_circuit)):
+            if not cirq.has_unitary(flattened_circuit[j]):
+                continue
             op1 = flattened_circuit[i]
             op2 = flattened_circuit[j]
             if can_be_parallel(op1, op2):
@@ -297,14 +301,20 @@ def colorize(
     for epoch in epochs:
         oneq_gates = []
         twoq_gates = []
+        nonunitary_gates = []
         for gate in epoch:
-            if len(gate.val.qubits) == 1:
+            if not cirq.has_unitary(gate.val):
+                nonunitary_gates.append(gate.val)
+            elif len(gate.val.qubits) == 1:
                 oneq_gates.append(gate.val)
             elif len(gate.val.qubits) == 2:
                 twoq_gates.append(gate.val)
             else:
                 raise RuntimeError("Unsupported gate type")
 
+        if len(nonunitary_gates) > 0:
+            yield nonunitary_gates
+
         if len(oneq_gates) > 0:
             yield oneq_gates
 

test/analysis/address/test_qubit_analysis.py

@@ -1,7 +1,7 @@
 import pytest
 from util import collect_address_types
 from kirin.analysis import const
-from kirin.dialects import ilist
+from kirin.dialects import scf, ilist
 
 from bloqade import qubit, squin
 from bloqade.analysis import address
@@ -265,3 +265,30 @@ def main():
 
     assert ret == address.AddressReg(data=tuple(range(20)))
     assert analysis.qubit_count == 20
+
+
+def test_for_loop_body_values():
+    @squin.kernel
+    def main():
+        q = squin.qalloc(4)
+        for i in range(1, len(q)):
+            squin.cx(q[0], q[i])
+
+    address_analysis = address.AddressAnalysis(main.dialects)
+    frame, result = address_analysis.run(main)
+    main.print(analysis=frame.entries)
+
+    (for_stmt,) = tuple(
+        stmt for stmt in main.callable_region.walk() if isinstance(stmt, scf.For)
+    )
+
+    for_analysis = [
+        value
+        for stmt in for_stmt.body.walk()
+        for value in frame.get_values(stmt.results)
+    ]
+
+    assert address.AddressQubit(data=0) in for_analysis
+    assert address.ConstResult(const.Value(0)) in for_analysis
+    assert address.ConstResult(const.Value(None)) in for_analysis
+    assert address.Unknown() in for_analysis

test/cirq_utils/noise/test_noise_models.py

@@ -14,7 +14,7 @@
 )
 
 
-def create_ghz_circuit(qubits):
+def create_ghz_circuit(qubits, measurements: bool = False):
     n = len(qubits)
     circuit = cirq.Circuit()
 
@@ -24,26 +24,41 @@ def create_ghz_circuit(qubits):
     # Step 2: CNOT chain from qubit i to i+1
     for i in range(n - 1):
         circuit.append(cirq.CNOT(qubits[i], qubits[i + 1]))
+        if measurements:
+            circuit.append(cirq.measure(qubits[i]))
+            circuit.append(cirq.reset(qubits[i]))
+
+    if measurements:
+        circuit.append(cirq.measure(qubits[-1]))
+        circuit.append(cirq.reset(qubits[-1]))
 
     return circuit
 
 
 @pytest.mark.parametrize(
-    "model,qubits",
+    "model,qubits,measurements",
     [
-        (GeminiOneZoneNoiseModel(), None),
+        (GeminiOneZoneNoiseModel(), None, False),
+        (
+            GeminiOneZoneNoiseModelConflictGraphMoves(),
+            cirq.GridQubit.rect(rows=1, cols=2),
+            False,
+        ),
+        (GeminiTwoZoneNoiseModel(), None, False),
+        (GeminiOneZoneNoiseModel(), None, True),
         (
             GeminiOneZoneNoiseModelConflictGraphMoves(),
             cirq.GridQubit.rect(rows=1, cols=2),
+            True,
         ),
-        (GeminiTwoZoneNoiseModel(), None),
+        (GeminiTwoZoneNoiseModel(), None, True),
     ],
 )
-def test_simple_model(model: cirq.NoiseModel, qubits):
+def test_simple_model(model: cirq.NoiseModel, qubits, measurements: bool):
     if qubits is None:
         qubits = cirq.LineQubit.range(2)
 
-    circuit = create_ghz_circuit(qubits)
+    circuit = create_ghz_circuit(qubits, measurements=measurements)
 
     with pytest.raises(ValueError):
         # make sure only native gate set is supported
@@ -74,13 +89,25 @@ def test_simple_model(model: cirq.NoiseModel, qubits):
         for i in range(4):
             pops_bloqade[i] += abs(ket[i]) ** 2 / nshots
 
-    for pops in (pops_bloqade, pops_cirq):
-        assert math.isclose(pops[0], 0.5, abs_tol=1e-1)
-        assert math.isclose(pops[3], 0.5, abs_tol=1e-1)
-        assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
-        assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
-
-        assert pops[0] < 0.5001
-        assert pops[3] < 0.5001
-        assert pops[1] >= 0.0
-        assert pops[2] >= 0.0
+    if measurements is True:
+        for pops in (pops_bloqade, pops_cirq):
+            assert math.isclose(pops[0], 1.0, abs_tol=1e-1)
+            assert math.isclose(pops[3], 0.0, abs_tol=1e-1)
+            assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
+            assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
+
+            assert pops[0] > 0.99
+            assert pops[3] >= 0.0
+            assert pops[1] >= 0.0
+            assert pops[2] >= 0.0
+    else:
+        for pops in (pops_bloqade, pops_cirq):
+            assert math.isclose(pops[0], 0.5, abs_tol=1e-1)
+            assert math.isclose(pops[3], 0.5, abs_tol=1e-1)
+            assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
+            assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
+
+            assert pops[0] < 0.5001
+            assert pops[3] < 0.5001
+            assert pops[1] >= 0.0
+            assert pops[2] >= 0.0