Fidelity analysis for for loops

Author: david-pl

src/bloqade/analysis/fidelity/impls.py

@@ -1,12 +1,13 @@
 from kirin import interp
 from kirin.lattice import EmptyLattice
-from kirin.dialects.scf import dialect as scf
-from kirin.dialects.scf.stmts import IfElse
+from kirin.analysis import const
+from kirin.dialects import scf
+from kirin.dialects.scf.stmts import For, Yield, IfElse
 
 from .analysis import FidelityAnalysis
 
 
-@scf.register(key="circuit.fidelity")
+@scf.dialect.register(key="circuit.fidelity")
 class ScfFidelityMethodTable(interp.MethodTable):
 
     @interp.impl(IfElse)
@@ -44,3 +45,40 @@ def if_else(
             interp.current_fidelity = then_fidelity
         else:
             interp.current_fidelity = else_fidelity
+
+    @interp.impl(Yield)
+    def yield_(
+        self, interp: FidelityAnalysis, frame: interp.Frame[EmptyLattice], stmt: Yield
+    ):
+        # NOTE: yield can by definition only contain values, never any stmts, so fidelity cannot decrease
+        return
+
+    @interp.impl(For)
+    def for_loop(
+        self, interp: FidelityAnalysis, frame: interp.Frame[EmptyLattice], stmt: For
+    ):
+
+        if not isinstance(hint := stmt.iterable.hints.get("const"), const.Value):
+            # NOTE: not clear how long this loop is
+            # TODO: should we at least count the body once?
+            return
+
+        current_fidelity = interp.current_fidelity
+
+        # NOTE: we reset the interpreter fidelity and evaluate the fidelity for the body only once
+        interp.current_fidelity = 1.0
+        for s in stmt.body.stmts():
+            stmt_impl = interp.lookup_registry(frame=frame, stmt=s)
+            if stmt_impl is None:
+                continue
+
+            stmt_impl(interp=interp, frame=frame, stmt=s)
+
+        # NOTE: reset current fidelity now in case of 0 iterations
+        loop_body_fidelity = interp.current_fidelity
+        interp.current_fidelity = current_fidelity
+
+        # NOTE: now we simply decrease the fidelity according to the number of iterations
+        iterable = hint.data
+        for _ in iterable:
+            interp.current_fidelity *= loop_body_fidelity

test/analysis/fidelity/test_fidelity.py

@@ -106,3 +106,65 @@ def main_if():
     fid_if_analysis.run_analysis(main_if, no_raise=False)
 
     assert 0 < fid_if_analysis.global_fidelity == fid_analysis.global_fidelity < 1
+
+
+def test_for():
+
+    @noise_main
+    def main():
+        q = qasm2.qreg(1)
+        c = qasm2.creg(1)
+        qasm2.h(q[0])
+        qasm2.measure(q, c)
+
+        # unrolled for loop
+        qasm2.x(q[0])
+        qasm2.x(q[0])
+        qasm2.x(q[0])
+
+        qasm2.measure(q, c)
+
+        return c
+
+    @noise_main
+    def main_for():
+        q = qasm2.qreg(1)
+        c = qasm2.creg(1)
+        qasm2.h(q[0])
+        qasm2.measure(q, c)
+
+        for _ in range(3):
+            qasm2.x(q[0])
+
+        qasm2.measure(q, c)
+        return c
+
+    px = 0.01
+    py = 0.01
+    pz = 0.01
+    p_loss = 0.01
+
+    noise_params = native.GateNoiseParams(
+        global_loss_prob=p_loss,
+        global_px=px,
+        global_py=py,
+        global_pz=pz,
+        local_px=0.002,
+    )
+
+    model = NoiseTestModel()
+    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+    fid_analysis = FidelityAnalysis(main.dialects)
+    fid_analysis.run_analysis(main, no_raise=False)
+
+    model = NoiseTestModel()
+    NoisePass(main_for.dialects, noise_model=model, gate_noise_params=noise_params)(
+        main_for
+    )
+
+    main_for.print()
+
+    fid_for_analysis = FidelityAnalysis(main_for.dialects)
+    fid_for_analysis.run_analysis(main_for, no_raise=False)
+
+    assert 0 < fid_for_analysis.global_fidelity == fid_analysis.global_fidelity < 1