Fixing tests

Author: weinbe58

src/bloqade/analysis/address/__init__.py

@@ -1,10 +1,15 @@
 from . import impls as impls
 from .lattice import (
+    Bottom as Bottom,
     Address as Address,
+    Unknown as Unknown,
+    AddressReg as AddressReg,
+    UnknownReg as UnknownReg,
     ConstResult as ConstResult,
     AddressQubit as AddressQubit,
+    PartialIList as PartialIList,
+    PartialTuple as PartialTuple,
     UnknownQubit as UnknownQubit,
     PartialLambda as PartialLambda,
-    StaticContainer as StaticContainer,
 )
 from .analysis import AddressAnalysis as AddressAnalysis

src/bloqade/analysis/fidelity/analysis.py

@@ -7,7 +7,7 @@
 from kirin.interp.value import Successor
 from kirin.analysis.forward import ForwardFrame
 
-from ..address import AddressAnalysis
+from ..address import Address, AddressAnalysis
 
 
 class FidelityAnalysis(Forward):
@@ -57,7 +57,7 @@ def main():
 
     _current_atom_survival_probability: list[float] = field(init=False)
 
-    addr_frame: ForwardFrame = field(init=False)
+    addr_frame: ForwardFrame[Address] = field(init=False)
 
     def initialize(self):
         super().initialize()

src/bloqade/qasm2/dialects/core/address.py

@@ -0,0 +1,50 @@
+from kirin import interp
+from kirin.analysis import const
+
+from bloqade.analysis.address import (
+    Address,
+    AddressReg,
+    ConstResult,
+    AddressQubit,
+    UnknownQubit,
+    AddressAnalysis,
+)
+from bloqade.analysis.address.lattice import Bottom, UnknownReg
+
+from .stmts import QRegGet, QRegNew
+from ._dialect import dialect
+
+
+@dialect.register(key="qubit.address")
+class AddressMethodTable(interp.MethodTable):
+
+    @interp.impl(QRegNew)
+    def new(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame[Address],
+        stmt: QRegNew,
+    ):
+        n_qubits = frame.get(stmt.n_qubits)
+        match n_qubits:
+            case ConstResult(const.Value(int() as n)):
+                addr = AddressReg(range(interp.next_address, interp.next_address + n))
+                interp.next_address += n
+                return (addr,)
+            case _:
+                return (UnknownReg(),)
+
+    @interp.impl(QRegGet)
+    def get(self, interp: AddressAnalysis, frame: interp.Frame[Address], stmt: QRegGet):
+        addr = frame.get(stmt.reg)
+        idx = frame.get(stmt.idx)
+
+        match (addr, idx):
+            case (AddressReg(data), ConstResult(const.Value(int() as i))) if (
+                0 <= i < len(data)
+            ):
+                return (AddressQubit(data[i]),)
+            case (UnknownReg(), ConstResult()):
+                return (UnknownQubit(),)
+            case _:
+                return (Bottom(),)

src/bloqade/qasm2/dialects/noise/fidelity.py

@@ -1,7 +1,6 @@
 from kirin import interp
 from kirin.lattice import EmptyLattice
 
-from bloqade.analysis.address import AddressQubit, AddressTuple
 from bloqade.analysis.fidelity import FidelityAnalysis
 
 from .stmts import PauliChannel, CZPauliChannel, AtomLossChannel
@@ -42,10 +41,8 @@ def atom_loss(
         stmt: AtomLossChannel,
     ):
         # NOTE: since AtomLossChannel acts on IList[Qubit], we know the assigned address is a tuple
-        addresses: AddressTuple = interp.addr_frame.get(stmt.qargs)
-
+        addresses = interp.addr_frame.get(stmt.qargs)
+        print(addresses)
         # NOTE: get the corresponding index and reduce survival probability accordingly
-        for qbit_address in addresses.data:
-            assert isinstance(qbit_address, AddressQubit)
-            index = qbit_address.data
+        for index in addresses.data:
             interp._current_atom_survival_probability[index] *= 1 - stmt.prob

src/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 @@ def sitter_errors(
 
     @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

src/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 @@ def rewrite_single_qubit_gate(self, node: uop.SingleQubitGate):
 
     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 @@ def rewrite_global_single_qubit_gate(self, node: glob.UGate):
 
     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 @@ def rewrite_parallel_cz_gate(self, node: parallel.CZ):
         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)
             )

src/bloqade/qasm2/rewrite/parallel_to_glob.py

@@ -3,7 +3,6 @@
 
 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 @@ def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
             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 @@ def _rewrite_parallel_to_glob(node: parallel.UGate) -> abc.RewriteResult:
         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

src/bloqade/qasm2/rewrite/parallel_to_uop.py

@@ -21,16 +21,10 @@ def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
 
     def get_qubit_ssa(self, ilist_ref: ir.SSAValue) -> Optional[List[ir.SSAValue]]:
         addr = self.address_analysis.get(ilist_ref)
-        if not isinstance(addr, address.AddressTuple):
+        if not isinstance(addr, address.AddressReg):
             return None
 
-        ids = []
-        for ele in addr.data:
-            if not isinstance(ele, address.AddressQubit):
-                return None
-
-            ids.append(ele.data)
-
+        ids = addr.data
         return [self.id_map[ele] for ele in ids]
 
     def rewrite_cz(self, node: ir.Statement):

src/bloqade/stim/rewrite/squin_noise.py

@@ -9,7 +9,7 @@
 from bloqade.squin import noise as squin_noise
 from bloqade.stim.dialects import noise as stim_noise
 from bloqade.stim.rewrite.util import insert_qubit_idx_from_address
-from bloqade.analysis.address.lattice import AddressTuple
+from bloqade.analysis.address.lattice import AddressReg, AddressQubit
 from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
 
 
@@ -42,19 +42,24 @@ def rewrite_NoiseChannel(
             if not isinstance(qubit_address_attr, AddressAttribute):
                 return RewriteResult()
 
-            address_tuple = qubit_address_attr.address
+            address_reg = qubit_address_attr.address
 
-            if not isinstance(address_tuple, AddressTuple):
+            if not isinstance(address_reg, AddressReg):
                 return RewriteResult()
 
             qubit_idx_ssas_list = [
-                insert_qubit_idx_from_address(AddressAttribute(address=address), stmt)
-                for address in address_tuple.data
+                insert_qubit_idx_from_address(
+                    AddressAttribute(address=AddressQubit(address)), stmt
+                )
+                for address in address_reg.data
             ]
             if None in qubit_idx_ssas_list:
                 return RewriteResult()
 
             for qubit_idx_ssas in qubit_idx_ssas_list:
+                assert (
+                    qubit_idx_ssas is not None
+                ), "qubit_idx_ssas should not be None here"
                 stim_stmt = rewrite_method(stmt, tuple(qubit_idx_ssas))
                 stim_stmt.insert_before(stmt)
             stmt.delete()

src/bloqade/stim/rewrite/util.py

@@ -2,7 +2,7 @@
 from kirin.dialects import py
 
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.analysis.address import AddressReg, AddressWire, AddressQubit, AddressTuple
+from bloqade.analysis.address import AddressReg, AddressQubit
 
 
 def create_and_insert_qubit_idx_stmt(
@@ -22,14 +22,7 @@ def insert_qubit_idx_from_address(
     address_data = address.address
     qubit_idx_ssas = []
 
-    if isinstance(address_data, AddressTuple):
-        for address_qubit in address_data.data:
-            if not isinstance(address_qubit, AddressQubit):
-                return
-            create_and_insert_qubit_idx_stmt(
-                address_qubit.data, stmt_to_insert_before, qubit_idx_ssas
-            )
-    elif isinstance(address_data, AddressReg):
+    if isinstance(address_data, AddressReg):
         for qubit_idx in address_data.data:
             create_and_insert_qubit_idx_stmt(
                 qubit_idx, stmt_to_insert_before, qubit_idx_ssas
@@ -38,11 +31,6 @@ def insert_qubit_idx_from_address(
         create_and_insert_qubit_idx_stmt(
             address_data.data, stmt_to_insert_before, qubit_idx_ssas
         )
-    elif isinstance(address_data, AddressWire):
-        address_qubit = address_data.origin_qubit
-        create_and_insert_qubit_idx_stmt(
-            address_qubit.data, stmt_to_insert_before, qubit_idx_ssas
-        )
     else:
         return