bloqade-circuit 0.6.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/qasm2/dialects/noise/model.py

@@ -1,4 +1,5 @@
 import abc
+from typing import Sequence
 from dataclasses import field, dataclass
 
 
@@ -161,7 +162,7 @@ class MoveNoiseModelABC(abc.ABC):
 
     @abc.abstractmethod
     def parallel_cz_errors(
-        self, ctrls: list[int], qargs: list[int], rest: list[int]
+        self, ctrls: Sequence[int], qargs: Sequence[int], rest: Sequence[int]
     ) -> dict[tuple[float, float, float, float], list[int]]:
         """Takes a set of ctrls and qargs and returns a noise model for all qubits."""
         pass

bloqade/qasm2/emit/base.py

@@ -2,8 +2,9 @@ from abc import ABC
 from typing import Generic, TypeVar, overload
 from dataclasses import field, dataclass
 
-from kirin import ir, idtable
-from kirin.emit import EmitABC, EmitError, EmitFrame
+from kirin import ir, interp, idtable
+from kirin.emit import EmitABC, EmitFrame
+from kirin.worklist import WorkList
 from typing_extensions import Self
 
 from bloqade.qasm2.parse import ast
@@ -15,6 +16,9 @@ EmitNode = TypeVar("EmitNode", bound=ast.Node)
 @dataclass
 class EmitQASM2Frame(EmitFrame[ast.Node | None], Generic[StmtType]):
     body: list[StmtType] = field(default_factory=list)
+    worklist: WorkList[interp.Successor] = field(default_factory=WorkList)
+    block_ref: dict[ir.Block, ast.Node | None] = field(default_factory=dict)
+    _indent: int = 0
 
 
 @dataclass
@@ -37,18 +41,13 @@ class EmitQASM2Base(
         return self
 
     def initialize_frame(
-        self, code: ir.Statement, *, has_parent_access: bool = False
+        self, node: ir.Statement, *, has_parent_access: bool = False
     ) -> EmitQASM2Frame[StmtType]:
-        return EmitQASM2Frame(code, has_parent_access=has_parent_access)
-
-    def run_method(
-        self, method: ir.Method, args: tuple[ast.Node | None, ...]
-    ) -> tuple[EmitQASM2Frame[StmtType], ast.Node | None]:
-        return self.run_callable(method.code, (ast.Name(method.sym_name),) + args)
+        return EmitQASM2Frame(node, has_parent_access=has_parent_access)
 
     def emit_block(self, frame: EmitQASM2Frame, block: ir.Block) -> ast.Node | None:
         for stmt in block.stmts:
-            result = self.eval_stmt(frame, stmt)
+            result = self.frame_eval(frame, stmt)
             if isinstance(result, tuple):
                 frame.set_values(stmt.results, result)
         return None
@@ -70,5 +69,11 @@ class EmitQASM2Base(
         node: ast.Node | None,
     ) -> A | B:
         if not isinstance(node, typ):
-            raise EmitError(f"expected {typ}, got {type(node)}")
+            raise TypeError(f"expected {typ}, got {type(node)}")
         return node
+
+    def reset(self):
+        pass
+
+    def eval_fallback(self, frame: EmitQASM2Frame, node: ir.Statement):
+        return tuple(None for _ in range(len(node.results)))

bloqade/qasm2/emit/gate.py

@@ -3,11 +3,12 @@ from dataclasses import field, dataclass
 from kirin import ir, types, interp
 from kirin.dialects import py, func, ilist
 from kirin.ir.dialect import Dialect as Dialect
+from typing_extensions import Self
 
 from bloqade.types import QubitType
 from bloqade.qasm2.parse import ast
 
-from .base import EmitError, EmitQASM2Base, EmitQASM2Frame
+from .base import EmitQASM2Base, EmitQASM2Frame
 
 
 def _default_dialect_group():
@@ -18,9 +19,13 @@ def _default_dialect_group():
 
 @dataclass
 class EmitQASM2Gate(EmitQASM2Base[ast.UOp | ast.Barrier, ast.Gate]):
-    keys = ["emit.qasm2.gate"]
+    keys = ("emit.qasm2.gate",)
     dialects: ir.DialectGroup = field(default_factory=_default_dialect_group)
 
+    def initialize(self) -> Self:
+        super().initialize()
+        return self
+
 
 @ilist.dialect.register(key="emit.qasm2.gate")
 class Ilist(interp.MethodTable):
@@ -45,7 +50,7 @@ class Func(interp.MethodTable):
 
     @interp.impl(func.Call)
     def emit_call(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Call):
-        raise EmitError("cannot emit dynamic call")
+        raise RuntimeError("cannot emit dynamic call")
 
     @interp.impl(func.Invoke)
     def emit_invoke(
@@ -55,7 +60,7 @@ class Func(interp.MethodTable):
         if len(stmt.results) == 1 and stmt.results[0].type.is_subseteq(types.NoneType):
             ret = (None,)
         elif len(stmt.results) > 0:
-            raise EmitError(
+            raise RuntimeError(
                 "cannot emit invoke with results, this "
                 "is not compatible QASM2 gate routine"
                 " (consider pass qreg/creg by argument)"
@@ -67,10 +72,9 @@ class Func(interp.MethodTable):
                 qparams.append(frame.get(arg))
             else:
                 cparams.append(frame.get(arg))
-
         frame.body.append(
             ast.Instruction(
-                name=ast.Name(stmt.callee.sym_name),
+                name=ast.Name(stmt.callee.__getattribute__("sym_name")),
                 params=cparams,
                 qargs=qparams,
             )
@@ -80,9 +84,8 @@ class Func(interp.MethodTable):
     @interp.impl(func.Lambda)
     @interp.impl(func.GetField)
     def emit_err(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
-        raise EmitError(f"illegal statement {stmt.name} for QASM2 gate routine")
+        raise RuntimeError(f"illegal statement {stmt.name} for QASM2 gate routine")
 
     @interp.impl(func.Return)
-    @interp.impl(func.ConstantNone)
     def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         return ()

bloqade/qasm2/emit/main.py

@@ -1,8 +1,10 @@
+from typing import List, cast
 from dataclasses import dataclass
 
 from kirin import ir, interp
 from kirin.dialects import cf, scf, func
 from kirin.ir.dialect import Dialect as Dialect
+from typing_extensions import Self
 
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.dialects.uop import SingleQubitGate, TwoQubitCtrlGate
@@ -14,26 +16,124 @@ from ..dialects.core.stmts import Reset, Measure
 
 @dataclass
 class EmitQASM2Main(EmitQASM2Base[ast.Statement, ast.MainProgram]):
-    keys = ["emit.qasm2.main", "emit.qasm2.gate"]
+    keys = ("emit.qasm2.main", "emit.qasm2.gate")
     dialects: ir.DialectGroup
 
+    def initialize(self) -> Self:
+        super().initialize()
+        return self
+
+    def eval_fallback(self, frame: EmitQASM2Frame, node: ir.Statement):
+        return tuple(None for _ in range(len(node.results)))
+
 
 @func.dialect.register(key="emit.qasm2.main")
 class Func(interp.MethodTable):
+    @interp.impl(func.Invoke)
+    def invoke(self, emit: EmitQASM2Main, frame: EmitQASM2Frame, node: func.Invoke):
+        name = emit.callables.get(node.callee.code)
+        if name is None:
+            name = emit.callables.add(node.callee.code)
+            emit.callable_to_emit.append(node.callee.code)
+
+        if isinstance(node.callee.code, GateFunction):
+            c_params: list[ast.Expr] = []
+            q_args: list[ast.Bit | ast.Name] = []
+
+            for arg in node.args:
+                val = frame.get(arg)
+                if val is None:
+                    raise interp.InterpreterError(f"missing mapping for arg {arg}")
+                if isinstance(val, (ast.Bit, ast.Name)):
+                    q_args.append(val)
+                elif isinstance(val, ast.Expr):
+                    c_params.append(val)
+
+            instr = ast.Instruction(
+                name=ast.Name(name) if isinstance(name, str) else name,
+                params=c_params,
+                qargs=q_args,
+            )
+            frame.body.append(instr)
+            return ()
+
+        callee_name_node = ast.Name(name) if isinstance(name, str) else name
+        args = tuple(frame.get_values(node.args))
+        _, call_expr = emit.call(node.callee.code, callee_name_node, *args)
+        if call_expr is not None:
+            frame.body.append(call_expr)
+        return ()
 
     @interp.impl(func.Function)
     def emit_func(
         self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: func.Function
     ):
         from bloqade.qasm2.dialects import glob, parallel
+        from bloqade.qasm2.emit.gate import EmitQASM2Gate
+
+        if isinstance(stmt, GateFunction):
+            return ()
+
+        func_name = emit.callables.get(stmt)
+        if func_name is None:
+            func_name = emit.callables.add(stmt)
+
+        for block in stmt.body.blocks:
+            frame.current_block = block
+            for s in block.stmts:
+                frame.current_stmt = s
+                stmt_results = emit.frame_eval(frame, s)
+                if isinstance(stmt_results, tuple):
+                    if len(stmt_results) != 0:
+                        frame.set_values(s._results, stmt_results)
+                    continue
+
+        gate_defs: list[ast.Gate] = []
+
+        gate_emitter = EmitQASM2Gate(dialects=emit.dialects).initialize()
+        gate_emitter.callables = emit.callables
+
+        while emit.callable_to_emit:
+            callable_node = emit.callable_to_emit.pop()
+            if callable_node is None:
+                break
+
+            if isinstance(callable_node, GateFunction):
+                with gate_emitter.eval_context():
+                    with gate_emitter.new_frame(
+                        callable_node, has_parent_access=False
+                    ) as gate_frame:
+                        gate_result = gate_emitter.frame_eval(gate_frame, callable_node)
+                        gate_obj = None
+                        if isinstance(gate_result, tuple) and len(gate_result) > 0:
+                            maybe = gate_result[0]
+                            if isinstance(maybe, ast.Gate):
+                                gate_obj = maybe
+
+                        if gate_obj is None:
+                            name = emit.callables.get(
+                                callable_node
+                            ) or emit.callables.add(callable_node)
+                            prefix = getattr(emit.callables, "prefix", "") or ""
+                            emit_name = (
+                                name[len(prefix) :]
+                                if prefix and name.startswith(prefix)
+                                else name
+                            )
+                            gate_obj = ast.Gate(
+                                name=emit_name, cparams=[], qparams=[], body=[]
+                            )
+
+                        gate_defs.append(gate_obj)
 
-        emit.run_ssacfg_region(frame, stmt.body, ())
         if emit.dialects.data.intersection((parallel.dialect, glob.dialect)):
             header = ast.Kirin([dialect.name for dialect in emit.dialects])
         else:
             header = ast.OPENQASM(ast.Version(2, 0))
 
-        emit.output = ast.MainProgram(header=header, statements=frame.body)
+        full_body = gate_defs + frame.body
+        stmt_list = cast(List[ast.Statement], full_body)
+        emit.output = ast.MainProgram(header=header, statements=stmt_list)
         return ()
 
 

bloqade/qasm2/emit/target.py

@@ -115,8 +115,8 @@ class QASM2:
             ParallelToUOp(dialects=entry.dialects)(entry)
 
         Py2QASM(entry.dialects)(entry)
-        target_main = EmitQASM2Main(self.main_target)
-        target_main.run(entry, ())
+        target_main = EmitQASM2Main(self.main_target).initialize()
+        target_main.run(entry)
 
         main_program = target_main.output
         assert main_program is not None, f"failed to emit {entry.sym_name}"
@@ -127,9 +127,13 @@ class QASM2:
 
         if self.custom_gate:
             cg = CallGraph(entry)
-            target_gate = EmitQASM2Gate(self.gate_target)
+            target_gate = EmitQASM2Gate(self.gate_target).initialize()
 
-            for _, fn in cg.defs.items():
+            for _, fns in cg.defs.items():
+                if len(fns) != 1:
+                    raise ValueError("Incorrect callgraph")
+
+                (fn,) = fns
                 if fn is entry:
                     continue
 
@@ -146,7 +150,7 @@ class QASM2:
 
                 Py2QASM(fn.dialects)(fn)
 
-                target_gate.run(fn, tuple(ast.Name(name) for name in fn.arg_names[1:]))
+                target_gate.run(fn)
                 assert target_gate.output is not None, f"failed to emit {fn.sym_name}"
                 extra.append(target_gate.output)
 

bloqade/qasm2/groups.py

@@ -1,6 +1,6 @@
 from kirin import ir, passes
 from kirin.prelude import structural_no_opt
-from kirin.dialects import scf, func, ilist, lowering
+from kirin.dialects import scf, func, ilist, ssacfg, lowering
 
 from bloqade.qasm2.dialects import (
     uop,
@@ -15,7 +15,7 @@ from bloqade.qasm2.dialects import (
 from bloqade.qasm2.rewrite.desugar import IndexingDesugarPass
 
 
-@ir.dialect_group([uop, func, expr, lowering.func, lowering.call])
+@ir.dialect_group([uop, func, expr, lowering.func, lowering.call, ssacfg])
 def gate(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
@@ -58,6 +58,7 @@ def gate(self):
         func,
         lowering.func,
         lowering.call,
+        ssacfg,
     ]
 )
 def main(self):

bloqade/qasm2/parse/lowering.py

@@ -450,7 +450,6 @@ class QASM2(lowering.LoweringABC[ast.Node]):
                 func.Invoke(
                     callee=value,
                     inputs=tuple(params + qargs),
-                    kwargs=tuple(),
                 )
             )
 

bloqade/qasm2/passes/fold.py

@@ -1,27 +1,12 @@
 from dataclasses import field, dataclass
 
 from kirin import ir
-from kirin.passes import Pass, TypeInfer
-from kirin.rewrite import (
-    Walk,
-    Chain,
-    Inline,
-    Fixpoint,
-    WrapConst,
-    Call2Invoke,
-    ConstantFold,
-    CFGCompactify,
-    InlineGetItem,
-    InlineGetField,
-    DeadCodeElimination,
-    CommonSubexpressionElimination,
-)
-from kirin.analysis import const
-from kirin.dialects import scf, ilist
+from kirin.passes import Pass
 from kirin.ir.method import Method
 from kirin.rewrite.abc import RewriteResult
 
 from bloqade.qasm2.dialects import expr
+from bloqade.rewrite.passes import AggressiveUnroll
 
 from .unroll_if import UnrollIfs
 
@@ -30,71 +15,27 @@ from .unroll_if import UnrollIfs
 class QASM2Fold(Pass):
     """Fold pass for qasm2.extended"""
 
-    constprop: const.Propagate = field(init=False)
     inline_gate_subroutine: bool = True
     unroll_ifs: bool = True
+    aggressive_unroll: AggressiveUnroll = field(init=False)
 
     def __post_init__(self):
-        self.constprop = const.Propagate(self.dialects)
-        self.typeinfer = TypeInfer(self.dialects)
+        def inline_simple(node: ir.Statement):
+            if isinstance(node, expr.GateFunction):
+                return self.inline_gate_subroutine
 
-    def unsafe_run(self, mt: Method) -> RewriteResult:
-        result = RewriteResult()
-        frame, _ = self.constprop.run_analysis(mt)
-        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
-        rule = Chain(
-            ConstantFold(),
-            Call2Invoke(),
-            InlineGetField(),
-            InlineGetItem(),
-            DeadCodeElimination(),
-            CommonSubexpressionElimination(),
-        )
-        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+            return True
 
-        result = (
-            Walk(
-                Chain(
-                    scf.unroll.PickIfElse(),
-                    scf.unroll.ForLoop(),
-                    scf.trim.UnusedYield(),
-                )
-            )
-            .rewrite(mt.code)
-            .join(result)
+        self.aggressive_unroll = AggressiveUnroll(
+            self.dialects, inline_simple, no_raise=self.no_raise
         )
 
-        if self.unroll_ifs:
-            UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
-
-        # run typeinfer again after unroll etc. because we now insert
-        # a lot of new nodes, which might have more precise types
-        self.typeinfer.unsafe_run(mt)
-        result = (
-            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
-            .rewrite(mt.code)
-            .join(result)
-        )
-
-        def inline_simple(node: ir.Statement):
-            if isinstance(node, expr.GateFunction):
-                return self.inline_gate_subroutine
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
 
-            if not isinstance(node.parent_stmt, (scf.For, scf.IfElse)):
-                return True  # always inline calls outside of loops and if-else
+        if self.unroll_ifs:
+            result = UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
 
-            # inside loops and if-else, only inline simple functions, i.e. functions with a single block
-            if (trait := node.get_trait(ir.CallableStmtInterface)) is None:
-                return False  # not a callable, don't inline to be safe
-            region = trait.get_callable_region(node)
-            return len(region.blocks) == 1
+        result = self.aggressive_unroll.unsafe_run(mt).join(result)
 
-        result = (
-            Walk(
-                Inline(inline_simple),
-            )
-            .rewrite(mt.code)
-            .join(result)
-        )
-        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
         return result

bloqade/qasm2/passes/glob.py

@@ -51,7 +51,7 @@ class GlobalToUOP(Pass):
     """
 
     def generate_rule(self, mt: ir.Method) -> GlobalToUOpRule:
-        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        frame, _ = address.AddressAnalysis(mt.dialects).run(mt)
         return GlobalToUOpRule(frame.entries)
 
     def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
@@ -105,7 +105,7 @@ class GlobalToParallel(Pass):
     """
 
     def generate_rule(self, mt: ir.Method) -> GlobalToParallelRule:
-        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        frame, _ = address.AddressAnalysis(mt.dialects).run(mt)
         return GlobalToParallelRule(frame.entries)
 
     def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:

bloqade/qasm2/passes/noise.py

@@ -55,7 +55,7 @@ class NoisePass(Pass):
         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)
+        frame, _ = self.address_analysis.run(mt)
         qubit_ssa_values = {}
         # Traverse statements in block order to fine the first SSA value for each qubit
         for block in mt.callable_region.blocks:

bloqade/qasm2/passes/parallel.py

@@ -28,7 +28,6 @@ from bloqade.qasm2.rewrite import (
     UOpToParallelRule,
     ParallelToGlobalRule,
     SimpleOptimalMergePolicy,
-    RydbergGateSetRewriteRule,
 )
 from bloqade.squin.analysis import schedule
 
@@ -64,7 +63,7 @@ class ParallelToUOp(Pass):
     """
 
     def generate_rule(self, mt: ir.Method) -> ParallelToUOpRule:
-        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        frame, _ = address.AddressAnalysis(mt.dialects).run(mt)
 
         id_map = {}
 
@@ -151,16 +150,19 @@ class UOpToParallel(Pass):
             return result
 
         if self.rewrite_to_native_first:
+            # NOTE: this import also imports cirq, so we do it locally here
+            from bloqade.qasm2.rewrite.native_gates import RydbergGateSetRewriteRule
+
             result = (
                 Fixpoint(Walk(RydbergGateSetRewriteRule(self.dialects)))
                 .rewrite(mt.code)
                 .join(result)
             )
 
-        frame, _ = self.constprop.run_analysis(mt)
+        frame, _ = self.constprop.run(mt)
         result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
 
-        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        frame, _ = address.AddressAnalysis(mt.dialects).run(mt)
         dags = schedule.DagScheduleAnalysis(
             mt.dialects, address_analysis=frame.entries
         ).get_dags(mt)
@@ -191,7 +193,7 @@ class ParallelToGlobal(Pass):
 
     def generate_rule(self, mt: ir.Method) -> ParallelToGlobalRule:
         address_analysis = address.AddressAnalysis(mt.dialects)
-        frame, _ = address_analysis.run_analysis(mt)
+        frame, _ = address_analysis.run(mt)
         return ParallelToGlobalRule(frame.entries)
 
     def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:

bloqade/qasm2/rewrite/__init__.py

@@ -3,7 +3,6 @@ from .glob import (
     GlobalToParallelRule as GlobalToParallelRule,
 )
 from .register import RaiseRegisterRule as RaiseRegisterRule
-from .native_gates import RydbergGateSetRewriteRule as RydbergGateSetRewriteRule
 from .parallel_to_uop import ParallelToUOpRule as ParallelToUOpRule
 from .uop_to_parallel import (
     MergePolicyABC as MergePolicyABC,

bloqade/qasm2/rewrite/noise/heuristic_noise.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple, cast
+from typing import Dict, List, Tuple
 from dataclasses import field, dataclass
 
 from kirin import ir
@@ -55,7 +55,7 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
 
     def rewrite_global_single_qubit_gate(self, node: glob.UGate):
         addrs = self.address_analysis[node.registers]
-        if not isinstance(addrs, address.AddressTuple):
+        if not isinstance(addrs, address.PartialIList):
             return rewrite_abc.RewriteResult()
 
         qargs = []
@@ -74,10 +74,7 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
 
     def rewrite_parallel_single_qubit_gate(self, node: parallel.RZ | parallel.UGate):
         addrs = self.address_analysis[node.qargs]
-        if not isinstance(addrs, address.AddressTuple):
-            return rewrite_abc.RewriteResult()
-
-        if not all(isinstance(addr, address.AddressQubit) for addr in addrs.data):
+        if not isinstance(addrs, address.AddressReg):
             return rewrite_abc.RewriteResult()
 
         assert isinstance(node.qargs, ir.ResultValue)
@@ -178,18 +175,11 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
         qargs = self.address_analysis[node.qargs]
 
         has_done_something = False
-        if (
-            isinstance(ctrls, address.AddressTuple)
-            and all(isinstance(addr, address.AddressQubit) for addr in ctrls.data)
-            and isinstance(qargs, address.AddressTuple)
-            and all(isinstance(addr, address.AddressQubit) for addr in qargs.data)
+        if isinstance(ctrls, address.AddressReg) and isinstance(
+            qargs, address.AddressReg
         ):
-            ctrl_qubits = list(
-                map(lambda addr: cast(address.AddressQubit, addr).data, ctrls.data)
-            )
-            qarg_qubits = list(
-                map(lambda addr: cast(address.AddressQubit, addr).data, qargs.data)
-            )
+            ctrl_qubits = tuple(ctrls.data)
+            qarg_qubits = tuple(qargs.data)
             rest = sorted(
                 set(self.qubit_ssa_value.keys()) - set(ctrl_qubits + qarg_qubits)
             )

bloqade/qasm2/rewrite/parallel_to_glob.py

@@ -3,7 +3,6 @@ from dataclasses import dataclass
 
 from kirin import ir
 from kirin.rewrite import abc
-from kirin.analysis import const
 from kirin.dialects import ilist
 
 from bloqade.analysis import address
@@ -20,28 +19,24 @@ class ParallelToGlobalRule(abc.RewriteRule):
             return abc.RewriteResult()
 
         qargs = node.qargs
-        qarg_addresses = self.address_analysis.get(qargs, None)
+        qargs_address = self.address_analysis.get(qargs, address.Unknown())
 
-        if isinstance(qarg_addresses, address.AddressReg):
-            # NOTE: we only have an AddressReg if it's an entire register, definitely rewrite that
-            return self._rewrite_parallel_to_glob(node)
-
-        if not isinstance(qarg_addresses, address.AddressTuple):
+        if not isinstance(qargs_address, address.AddressReg):
             return abc.RewriteResult()
 
-        idxs, qreg = self._find_qreg(qargs.owner, set())
+        qregs = self._get_all_qreg(qargs.owner)
 
-        if qreg is None:
-            # NOTE: no unique register found
+        if len(qregs) != 1:
             return abc.RewriteResult()
 
-        if not isinstance(hint := qreg.n_qubits.hints.get("const"), const.Value):
-            # NOTE: non-constant number of qubits
+        qreg = next(iter(qregs))
+
+        qreg_address = self.address_analysis.get(qreg, address.Unknown())
+
+        if not isinstance(qreg_address, address.AddressReg):
             return abc.RewriteResult()
 
-        n = hint.data
-        if len(idxs) != n:
-            # NOTE: not all qubits of the register are there
+        if set(qargs_address.data) != set(qreg_address.data):
             return abc.RewriteResult()
 
         return self._rewrite_parallel_to_glob(node)
@@ -53,6 +48,24 @@ class ParallelToGlobalRule(abc.RewriteRule):
         node.replace_by(global_u)
         return abc.RewriteResult(has_done_something=True)
 
+    @staticmethod
+    def _get_all_qreg(owner: ir.Statement | ir.Block):
+        stack = [owner]
+        qregs: set[ir.SSAValue] = set()
+        while stack:
+            current = stack.pop()
+
+            if isinstance(current, core.stmts.QRegGet):
+                stack.append(current.reg.owner)
+            elif isinstance(current, ilist.New):
+                for val in current.values:
+                    stack.append(val.owner)
+
+            elif isinstance(current, core.QRegNew):
+                qregs.add(current.result)
+
+        return qregs
+
     @staticmethod
     def _find_qreg(
         qargs_owner: ir.Statement | ir.Block, idxs: set