Fixing NoisePass with LiftQubits Pass. (#230)

in issue #226 The problem is because how the qubit -> global index. In
this PR I add an explicit pass that lifts all qubit references to the
top of the entry block.

---------

Co-authored-by: Xiu-zhe (Roger) Luo <[email protected]>

Author: weinbe58

src/bloqade/qasm2/passes/lift_qubits.py

@@ -0,0 +1,26 @@
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    CommonSubexpressionElimination,
+)
+from kirin.passes.hint_const import HintConst
+
+from bloqade.qasm2.rewrite.insert_qubits import InsertGetQubit
+
+
+class LiftQubits(Pass):
+    """This pass lifts the creation of qubits to the block where the register is defined."""
+
+    def unsafe_run(self, mt: ir.Method):
+        result = Walk(InsertGetQubit()).rewrite(mt.code)
+        result = HintConst(self.dialects).unsafe_run(mt).join(result)
+        result = (
+            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        return result

src/bloqade/qasm2/passes/noise.py

@@ -1,19 +1,17 @@
 from dataclasses import field, dataclass
 
 from kirin import ir
-from kirin.passes import Pass, HintConst
+from kirin.passes import Pass
 from kirin.rewrite import (
     Walk,
-    Chain,
     Fixpoint,
-    ConstantFold,
     DeadCodeElimination,
-    CommonSubexpressionElimination,
 )
 
 from bloqade.noise import native
 from bloqade.analysis import address
-from bloqade.qasm2.rewrite.heuristic_noise import InsertGetQubit, NoiseRewriteRule
+from bloqade.qasm2.passes.lift_qubits import LiftQubits
+from bloqade.qasm2.rewrite.heuristic_noise import NoiseRewriteRule
 
 
 @dataclass
@@ -37,8 +35,7 @@ def __post_init__(self):
         self.address_analysis = address.AddressAnalysis(self.dialects)
 
     def unsafe_run(self, mt: ir.Method):
-        result = Walk(InsertGetQubit()).rewrite(mt.code)
-        HintConst(self.dialects).unsafe_run(mt)
+        result = LiftQubits(self.dialects).unsafe_run(mt)
         frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
         result = (
             Walk(
@@ -52,10 +49,5 @@ def unsafe_run(self, mt: ir.Method):
             .rewrite(mt.code)
             .join(result)
         )
-        rule = Chain(
-            ConstantFold(),
-            DeadCodeElimination(),
-            CommonSubexpressionElimination(),
-        )
-        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+        result = Fixpoint(Walk(DeadCodeElimination())).rewrite(mt.code).join(result)
         return result

src/bloqade/qasm2/rewrite/heuristic_noise.py

@@ -3,39 +3,11 @@
 
 from kirin import ir
 from kirin.rewrite import abc as rewrite_abc
-from kirin.dialects import py, ilist
+from kirin.dialects import ilist
 
 from bloqade.noise import native
 from bloqade.analysis import address
-from bloqade.qasm2.dialects import uop, core, glob, parallel
-
-
-class InsertGetQubit(rewrite_abc.RewriteRule):
-
-    def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
-        if (
-            not isinstance(node, core.QRegNew)
-            or not isinstance(n_qubits_stmt := node.n_qubits.owner, py.Constant)
-            or not isinstance(n_qubits := n_qubits_stmt.value.unwrap(), int)
-            or (block := node.parent_block) is None
-        ):
-            return rewrite_abc.RewriteResult()
-
-        n_qubits_stmt.detach()
-        node.detach()
-        if block.first_stmt is None:
-            block.stmts.append(node)
-        else:
-            node.insert_before(block.first_stmt)
-            n_qubits_stmt.insert_before(block.first_stmt)
-
-        for idx_val in range(n_qubits):
-            idx = py.constant.Constant(value=idx_val)
-            qubit = core.QRegGet(node.result, idx=idx.result)
-            qubit.insert_after(node)
-            idx.insert_after(node)
-
-        return rewrite_abc.RewriteResult(has_done_something=True)
+from bloqade.qasm2.dialects import uop, glob, parallel
 
 
 @dataclass

src/bloqade/qasm2/rewrite/insert_qubits.py

@@ -0,0 +1,34 @@
+from kirin import ir
+from kirin.rewrite import abc as rewrite_abc
+from kirin.dialects import py
+
+
+class InsertGetQubit(rewrite_abc.RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+        from bloqade.qasm2 import core
+
+        if (
+            not isinstance(node, core.QRegNew)
+            or not isinstance(n_qubits_stmt := node.n_qubits.owner, py.Constant)
+            or not isinstance(n_qubits := n_qubits_stmt.value.unwrap(), int)
+            or (block := node.parent_block) is None
+        ):
+            return rewrite_abc.RewriteResult()
+
+        n_qubits_stmt.detach()
+        node.detach()
+        if block.first_stmt is None:
+            block.stmts.append(n_qubits_stmt)
+            block.stmts.append(node)
+        else:
+            node.insert_before(block.first_stmt)
+            n_qubits_stmt.insert_before(block.first_stmt)
+
+        for idx_val in range(n_qubits):
+            idx = py.constant.Constant(value=idx_val)
+            qubit = core.QRegGet(node.result, idx=idx.result)
+            qubit.insert_after(node)
+            idx.insert_after(node)
+
+        return rewrite_abc.RewriteResult(has_done_something=True)