Fix the UopToParallel pass for qasm2 constants (#477)

This fixes a bug, where the rewrite bailed if the qubit number was a
`qasm2.expr.ConstInt` rather than a `py.Constant`. That meant that
loaded qasm2 kernels weren't rewritten at all.

Author: david-pl

src/bloqade/qasm2/rewrite/register.py

@@ -2,7 +2,7 @@
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.qasm2.dialects import core
+from bloqade.qasm2.dialects import core, expr
 
 
 class RaiseRegisterRule(RewriteRule):
@@ -26,7 +26,7 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         n_qubits_ref = node.n_qubits
 
         n_qubits = n_qubits_ref.owner
-        if isinstance(n_qubits, py.Constant):
+        if isinstance(n_qubits, py.Constant | expr.ConstInt):
             # case where the n_qubits comes from a constant
             new_n_qubits = n_qubits.from_stmt(n_qubits)
             new_n_qubits.insert_before(first_stmt)

src/bloqade/qasm2/rewrite/uop_to_parallel.py

@@ -8,7 +8,7 @@
 from kirin.analysis.const import lattice
 
 from bloqade.analysis import address
-from bloqade.qasm2.dialects import uop, core, parallel
+from bloqade.qasm2.dialects import uop, core, expr, parallel
 from bloqade.squin.analysis.schedule import StmtDag
 
 
@@ -194,6 +194,8 @@ def move_and_collect_qubit_list(
                     new_qubits.append(new_qubit.result)
                 case core.QRegGet(
                     reg=reg, idx=ir.ResultValue(stmt=py.Constant() as idx)
+                ) | core.QRegGet(
+                    reg=reg, idx=ir.ResultValue(stmt=expr.ConstInt() as idx)
                 ):
                     (new_idx := idx.from_stmt(idx)).insert_before(node)
                     (

test/qasm2/passes/test_uop_to_parallel.py

@@ -38,6 +38,14 @@ def test():
     # add this to raise error if there are broken ssa references
     _, _ = address.AddressAnalysis(test.dialects).run_analysis(test, no_raise=False)
 
+    # check that there's parallel statements now
+    assert any(
+        [
+            isinstance(stmt, qasm2.dialects.parallel.UGate)
+            for stmt in test.callable_region.blocks[0].stmts
+        ]
+    )
+
 
 def test_two():