Fixing extraction of global qubit ids (#237)

Author: weinbe58

src/bloqade/qasm2/passes/noise.py

@@ -34,13 +34,32 @@ class NoisePass(Pass):
     def __post_init__(self):
         self.address_analysis = address.AddressAnalysis(self.dialects)
 
+    def get_qubit_values(self, mt: ir.Method):
+        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_values = {}
+        # Traverse statements in block order to fine the first SSA value for each qubit
+        for block in mt.callable_region.blocks:
+            for stmt in block.stmts:
+                if len(stmt.results) != 1:
+                    continue
+
+                addr = frame.entries.get(result := stmt.results[0])
+                if (
+                    isinstance(addr, address.AddressQubit)
+                    and (index := addr.data) not in qubit_ssa_values
+                ):
+                    qubit_ssa_values[index] = result
+
+        return qubit_ssa_values, frame.entries
+
     def unsafe_run(self, mt: ir.Method):
         result = LiftQubits(self.dialects).unsafe_run(mt)
-        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_value, address_analysis = self.get_qubit_values(mt)
         result = (
             Walk(
                 NoiseRewriteRule(
-                    address_analysis=frame.entries,
+                    qubit_ssa_value=qubit_ssa_value,
+                    address_analysis=address_analysis,
                     noise_model=self.noise_model,
                     gate_noise_params=self.gate_noise_params,
                 ),
@@ -49,5 +68,6 @@ def unsafe_run(self, mt: ir.Method):
             .rewrite(mt.code)
             .join(result)
         )
+
         result = Fixpoint(Walk(DeadCodeElimination())).rewrite(mt.code).join(result)
         return result

src/bloqade/qasm2/rewrite/heuristic_noise.py

@@ -18,22 +18,14 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
     """
 
     address_analysis: Dict[ir.SSAValue, address.Address]
+    qubit_ssa_value: Dict[int, ir.SSAValue]
     gate_noise_params: native.GateNoiseParams = field(
         default_factory=native.GateNoiseParams
     )
     noise_model: native.MoveNoiseModelABC = field(
         default_factory=native.TwoRowZoneModel
     )
 
-    def __post_init__(self):
-        self.qubit_ssa_value: Dict[int, ir.SSAValue] = {}
-        for ssa_value, addr in self.address_analysis.items():
-            if (
-                isinstance(addr, address.AddressQubit)
-                and ssa_value not in self.qubit_ssa_value
-            ):
-                self.qubit_ssa_value[addr.data] = ssa_value
-
     def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
         if isinstance(node, uop.SingleQubitGate):
             return self.rewrite_single_qubit_gate(node)

test/qasm2/passes/test_heuristic_noise.py

@@ -29,7 +29,8 @@ def test_single_qubit_noise():
 
     test_qubit = ir.TestValue(type=qasm2.QubitType)
     address_analysis = {test_qubit: address.AddressQubit(0)}
-    rule = NoiseRewriteRule(address_analysis, noise_params, model)
+    qubit_ssa_value = {0: test_qubit}
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
     rule.qubit_ssa_value[0] = test_qubit
     block = ir.Block(
         [
@@ -70,7 +71,9 @@ def test_parallel_qubit_noise():
         test_qubit: address.AddressQubit,
         qubit_list.result: address.AddressTuple([address.AddressQubit(0)]),
     }
-    rule = NoiseRewriteRule(address_analysis, noise_params, model)
+    qubit_ssa_value = {0: test_qubit}
+
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
     rule.qubit_ssa_value[0] = test_qubit
     block = ir.Block(
         [
@@ -128,7 +131,8 @@ def test_cz_gate_noise():
         ctrl_qubit: address.AddressQubit(0),
         qarg_qubit: address.AddressQubit(1),
     }
-    rule = NoiseRewriteRule(address_analysis, noise_params, model)
+    qubit_ssa_value = {0: ctrl_qubit, 1: qarg_qubit}
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
     rule.qubit_ssa_value[0] = ctrl_qubit
     rule.qubit_ssa_value[1] = qarg_qubit
     block = ir.Block(
@@ -206,7 +210,8 @@ def test_parallel_cz_gate_noise():
         ctrl_list.result: address.AddressTuple([address.AddressQubit(0)]),
         qarg_list.result: address.AddressTuple([address.AddressQubit(1)]),
     }
-    rule = NoiseRewriteRule(address_analysis, noise_params, model)
+    qubit_ssa_value = {0: ctrl_qubit, 1: qarg_qubit}
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
     rule.qubit_ssa_value[0] = ctrl_qubit
     rule.qubit_ssa_value[1] = qarg_qubit
     block = ir.Block(