Fix IfElse impl and test known branch execution

Author: david-pl

src/bloqade/analysis/fidelity/analysis.py

@@ -19,23 +19,17 @@ class FidelityAnalysis(AddressAnalysis):
     ## Usage examples
 
     ```
-    from bloqade import qasm2
-    from bloqade.noise import native
-    from bloqade.analysis.fidelity import FidelityAnalysis
-    from bloqade.qasm2.passes.noise import NoisePass
+    from bloqade import squin
 
-    noise_main = qasm2.extended.add(native.dialect)
-
-    @noise_main
+    @squin.kernel
     def main():
-        q = qasm2.qreg(2)
-        qasm2.x(q[0])
+        q = squin.qalloc(1)
+        squin.x(q[0])
+        squin.depolarize(q[0])
         return q
 
-    NoisePass(main.dialects)(main)
-
     fid_analysis = FidelityAnalysis(main.dialects)
-    fid_analysis.run_analysis(main, no_raise=False)
+    fid_analysis.run(main)
 
     gate_fidelities = fid_analysis.gate_fidelities
     qubit_survival_probs = fid_analysis.qubit_survival_fidelities

src/bloqade/analysis/fidelity/impls.py

@@ -1,8 +1,8 @@
 from kirin import interp
-from kirin.analysis import ForwardFrame
+from kirin.analysis import ForwardFrame, const
 from kirin.dialects import scf
 
-from bloqade.analysis.address import Address
+from bloqade.analysis.address import Address, ConstResult
 
 from .analysis import FidelityAnalysis
 
@@ -18,17 +18,27 @@ def if_else(
         current_gate_fidelities = interp_.gate_fidelities
         current_survival_fidelities = interp_.qubit_survival_fidelities
 
-        # TODO: check if the condition is constant and fix the branch in that case
-        # run both branches
+        address_cond = frame.get(stmt.cond)
+
+        # NOTE: if the condition is known at compile time, run specific branch
+        if isinstance(address_cond, ConstResult) and isinstance(
+            const_cond := address_cond.result, const.Value
+        ):
+            body = stmt.then_body if const_cond.data else stmt.else_body
+            with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+                ret = interp_.frame_call_region(body_frame, stmt, body, address_cond)
+                return ret
+
+        # NOTE: runtime condition, evaluate both
         with interp_.new_frame(stmt, has_parent_access=True) as then_frame:
             # NOTE: reset fidelities before stepping into the then-body
             interp_.reset_fidelities()
 
-            interp_.frame_call_region(
+            then_results = interp_.frame_call_region(
                 then_frame,
                 stmt,
                 stmt.then_body,
-                *(interp_.lattice.bottom() for _ in range(len(stmt.args))),
+                address_cond,
             )
             then_fids = interp_.gate_fidelities
             then_survival = interp_.qubit_survival_fidelities
@@ -37,11 +47,11 @@ def if_else(
             # NOTE: reset again before stepping into else-body
             interp_.reset_fidelities()
 
-            interp_.frame_call_region(
+            else_results = interp_.frame_call_region(
                 else_frame,
                 stmt,
                 stmt.else_body,
-                *(interp_.lattice.bottom() for _ in range(len(stmt.args))),
+                address_cond,
             )
 
             else_fids = interp_.gate_fidelities
@@ -60,3 +70,17 @@ def if_else(
             then_survival,
             else_survival,
         )
+
+        # TODO: pick the non-return value
+        if isinstance(then_results, interp.ReturnValue) and isinstance(
+            else_results, interp.ReturnValue
+        ):
+            return interp.ReturnValue(then_results.value.join(else_results.value))
+        elif isinstance(then_results, interp.ReturnValue):
+            ret = else_results
+        elif isinstance(else_results, interp.ReturnValue):
+            ret = then_results
+        else:
+            ret = interp_.join_results(then_results, else_results)
+
+        return ret

test/analysis/fidelity/test_fidelity.py

@@ -344,3 +344,41 @@ def main():
         ]
         + [FidelityRange(0.87, 0.87)] * 4
     )  # squin.correlated_qubit_loss
+
+
+def test_squin_know_if():
+    @squin.kernel
+    def main():
+        x = True
+        q = squin.qalloc(4)
+
+        if x:
+            squin.depolarize(0.1, q[0])
+            squin.qubit_loss(0.1, q[0])
+        else:
+            squin.depolarize(0.1, q[1])
+            squin.qubit_loss(0.1, q[1])
+
+        if not x:
+            squin.depolarize(0.2, q[2])
+            squin.qubit_loss(0.2, q[2])
+        else:
+            squin.depolarize(0.2, q[3])
+            squin.qubit_loss(0.2, q[3])
+
+    fidelity_analysis = FidelityAnalysis(main.dialects)
+    fidelity_analysis.run(main)
+
+    assert fidelity_analysis.gate_fidelities == [
+        FidelityRange(0.9, 0.9),
+        FidelityRange(1.0, 1.0),
+        FidelityRange(1.0, 1.0),
+        FidelityRange(0.8, 0.8),
+    ]
+
+    assert fidelity_analysis.qubit_survival_fidelities == [
+        FidelityRange(0.9, 0.9),
+        FidelityRange(1.0, 1.0),
+        FidelityRange(1.0, 1.0),
+        FidelityRange(0.8, 0.8),
+    ]