Merge branch 'main' into 19-rewrite-from-squin-to-stim

Author: johnzl-777

src/bloqade/analysis/address/impls.py

@@ -192,7 +192,10 @@ def unwrap(
 
         origin_qubit = frame.get(stmt.qubit)
 
-        return (AddressWire(origin_qubit=origin_qubit),)
+        if isinstance(origin_qubit, AddressQubit):
+            return (AddressWire(origin_qubit=origin_qubit),)
+        else:
+            return (Address.top(),)
 
     @interp.impl(squin.wire.Apply)
     def apply(
@@ -201,14 +204,7 @@ def apply(
         frame: ForwardFrame[Address],
         stmt: squin.wire.Apply,
     ):
-
-        origin_qubits = tuple(
-            [frame.get(input_elem).origin_qubit for input_elem in stmt.inputs]
-        )
-        new_address_wires = tuple(
-            [AddressWire(origin_qubit=origin_qubit) for origin_qubit in origin_qubits]
-        )
-        return new_address_wires
+        return frame.get_values(stmt.inputs)
 
     @interp.impl(squin.wire.MeasureAndReset)
     def measure_and_reset(

src/bloqade/analysis/address/lattice.py

@@ -81,5 +81,5 @@ class AddressWire(Address):
 
     def is_subseteq(self, other: Address) -> bool:
         if isinstance(other, AddressWire):
-            return self.origin_qubit == self.origin_qubit
+            return self.origin_qubit == other.origin_qubit
         return False

src/bloqade/noise/native/model.py

@@ -102,10 +102,9 @@ class MoveNoiseModelABC(abc.ABC):
     params: MoveNoiseParams = field(default_factory=MoveNoiseParams)
     """Parameters for calculating move noise."""
 
-    @classmethod
     @abc.abstractmethod
     def parallel_cz_errors(
-        cls, ctrls: List[int], qargs: List[int], rest: List[int]
+        self, ctrls: List[int], qargs: List[int], rest: List[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/pyqrack/__init__.py

@@ -1,7 +1,6 @@
 from .reg import (
     CBitRef as CBitRef,
     CRegister as CRegister,
-    PyQrackReg as PyQrackReg,
     QubitState as QubitState,
     Measurement as Measurement,
     PyQrackQubit as PyQrackQubit,

src/bloqade/pyqrack/noise/native.py

@@ -93,8 +93,7 @@ def atom_loss_channel(
 
         for qarg in active_qubits:
             if interp.rng_state.uniform() <= stmt.prob:
-                sim_reg = qarg.ref.sim_reg
-                sim_reg.force_m(qarg.addr, 0)
+                qarg.sim_reg.m(qarg.addr)
                 qarg.drop()
 
         return ()

src/bloqade/pyqrack/qasm2/core.py

@@ -1,9 +1,12 @@
+from typing import Any
+
 from kirin import interp
+from kirin.interp import InterpreterError
+from kirin.dialects import ilist
 
 from bloqade.pyqrack.reg import (
     CBitRef,
     CRegister,
-    PyQrackReg,
     QubitState,
     Measurement,
     PyQrackQubit,
@@ -19,14 +22,13 @@ def qreg_new(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.QRegNew
     ):
         n_qubits: int = frame.get(stmt.n_qubits)
-        return (
-            PyQrackReg(
-                size=n_qubits,
-                sim_reg=interp.memory.sim_reg,
-                addrs=interp.memory.allocate(n_qubits),
-                qubit_state=[QubitState.Active] * n_qubits,
-            ),
+        qreg = ilist.IList(
+            [
+                PyQrackQubit(i, interp.memory.sim_reg, QubitState.Active)
+                for i in interp.memory.allocate(n_qubits=n_qubits)
+            ]
         )
+        return (qreg,)
 
     @interp.impl(core.CRegNew)
     def creg_new(
@@ -39,7 +41,9 @@ def creg_new(
     def qreg_get(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.QRegGet
     ):
-        return (PyQrackQubit(ref=frame.get(stmt.reg), pos=frame.get(stmt.idx)),)
+        reg = frame.get(stmt.reg)
+        i = frame.get(stmt.idx)
+        return (reg[i],)
 
     @interp.impl(core.CRegGet)
     def creg_get(
@@ -51,12 +55,25 @@ def creg_get(
     def measure(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.Measure
     ):
-        qarg: PyQrackQubit = frame.get(stmt.qarg)
-        carg: CBitRef = frame.get(stmt.carg)
-        if qarg.is_active():
-            carg.set_value(Measurement(qarg.sim_reg.m(qarg.addr)))
+        qarg: PyQrackQubit | ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qarg)
+        carg: CBitRef | CRegister = frame.get(stmt.carg)
+
+        if isinstance(qarg, PyQrackQubit) and isinstance(carg, CBitRef):
+            if qarg.is_active():
+                carg.set_value(Measurement(qarg.sim_reg.m(qarg.addr)))
+            else:
+                carg.set_value(interp.loss_m_result)
+        elif isinstance(qarg, ilist.IList) and isinstance(carg, CRegister):
+            for i, qubit in enumerate(qarg):
+                cbit = CBitRef(carg, i)
+                if qubit.is_active():
+                    cbit.set_value(Measurement(qubit.sim_reg.m(qubit.addr)))
+                else:
+                    cbit.set_value(interp.loss_m_result)
         else:
-            carg.set_value(interp.loss_m_result)
+            raise InterpreterError(
+                f"Expected measure call on either a single qubit and classical bit, or two registers, but got the types {type(qarg)} and {type(carg)}"
+            )
 
         return ()
 

src/bloqade/pyqrack/qasm2/glob.py

@@ -3,7 +3,7 @@
 from kirin import interp
 from kirin.dialects import ilist
 
-from bloqade.pyqrack.reg import PyQrackReg
+from bloqade.pyqrack.reg import PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 from bloqade.qasm2.dialects import glob
 
@@ -12,7 +12,9 @@
 class PyQrackMethods(interp.MethodTable):
     @interp.impl(glob.UGate)
     def ugate(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: glob.UGate):
-        registers: ilist.IList[PyQrackReg, Any] = frame.get(stmt.registers)
+        registers: ilist.IList[ilist.IList[PyQrackQubit, Any], Any] = frame.get(
+            stmt.registers
+        )
         theta, phi, lam = (
             frame.get(stmt.theta),
             frame.get(stmt.phi),

src/bloqade/pyqrack/reg.py

@@ -1,8 +1,8 @@
 import enum
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
-from bloqade.qasm2.types import QReg, Qubit
+from bloqade.qasm2.types import Qubit
 
 if TYPE_CHECKING:
     from pyqrack import QrackSimulator
@@ -45,57 +45,18 @@ class QubitState(enum.Enum):
     Lost = enum.auto()
 
 
-@dataclass(frozen=True)
-class PyQrackReg(QReg):  # TODO: clean up implementation with list base class
-    """Simulation runtime value of a quantum register."""
-
-    size: int
-    """The number of qubits in this register."""
-
-    sim_reg: "QrackSimulator"
-    """The register of the simulator."""
-
-    addrs: tuple[int, ...]
-    """The global addresses of the qubits in this register."""
-
-    qubit_state: List[QubitState]
-    """The state of each qubit in this register."""
-
-    def drop(self, pos: int):
-        """Drop the qubit at the given position in-place.
-
-        Args
-            pos (int): The position of the qubit to drop.
-
-        """
-        assert self.qubit_state[pos] is QubitState.Active, "Qubit already lost"
-        self.qubit_state[pos] = QubitState.Lost
-
-    def __getitem__(self, pos: int):
-        if not 0 <= pos < self.size:
-            raise IndexError("Qubit index out of bounds of register.")
-        return PyQrackQubit(self, pos)
-
-
-@dataclass(frozen=True)
+@dataclass
 class PyQrackQubit(Qubit):
     """The runtime representation of a qubit reference."""
 
-    ref: PyQrackReg
-    """The quantum register that is holding this qubit."""
+    addr: int
+    """The address of this qubit in the quantum register."""
 
-    pos: int
-    """The position of this qubit in the quantum register."""
-
-    @property
-    def sim_reg(self):
-        """The register of the simulator."""
-        return self.ref.sim_reg
+    sim_reg: "QrackSimulator"
+    """The register of the simulator."""
 
-    @property
-    def addr(self) -> int:
-        """The global address of the qubit."""
-        return self.ref.addrs[self.pos]
+    state: QubitState
+    """The state of the qubit (active/lost)"""
 
     def is_active(self) -> bool:
         """Check if the qubit is active.
@@ -104,8 +65,8 @@ def is_active(self) -> bool:
             True if the qubit is active, False otherwise.
 
         """
-        return self.ref.qubit_state[self.pos] is QubitState.Active
+        return self.state is QubitState.Active
 
     def drop(self):
         """Drop the qubit in-place."""
-        self.ref.drop(self.pos)
+        self.state = QubitState.Lost

src/bloqade/qasm2/_wrappers.py

@@ -1,3 +1,5 @@
+from typing import overload
+
 from kirin.lowering import wraps
 
 from .types import Bit, CReg, QReg, Qubit
@@ -58,7 +60,7 @@ def reset(qarg: Qubit) -> None:
     ...
 
 
-@wraps(core.Measure)
+@overload
 def measure(qarg: Qubit, cbit: Bit) -> None:
     """
     Measure the qubit `qarg` and store the result in the classical bit `cbit`.
@@ -70,6 +72,22 @@ def measure(qarg: Qubit, cbit: Bit) -> None:
     ...
 
 
+@overload
+def measure(qarg: QReg, carg: CReg) -> None:
+    """
+    Measure each qubit in the quantum register `qarg` and store the result in the classical register `carg`.
+
+    Args:
+        qarg: The quantum register to measure.
+        carg: The classical bit to store the result in.
+    """
+    ...
+
+
+@wraps(core.Measure)
+def measure(qarg, carg) -> None: ...
+
+
 @wraps(uop.CX)
 def cx(ctrl: Qubit, qarg: Qubit) -> None:
     """

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

@@ -46,10 +46,21 @@ class Measure(ir.Statement):
 
     name = "measure"
     traits = frozenset({lowering.FromPythonCall()})
-    qarg: ir.SSAValue = info.argument(QubitType)
-    """qarg (Qubit): The qubit to measure."""
-    carg: ir.SSAValue = info.argument(BitType)
-    """carg (Bit): The bit to store the result in."""
+    qarg: ir.SSAValue = info.argument(QubitType | QRegType)
+    """qarg (Qubit | QReg): The qubit or quantum register to measure."""
+    carg: ir.SSAValue = info.argument(BitType | CRegType)
+    """carg (Bit | CReg): The bit or register to store the result in."""
+
+    def check_type(self) -> None:
+        qarg_is_qubit = self.qarg.type.is_subseteq(QubitType)
+        carg_is_bit = self.carg.type.is_subseteq(BitType)
+        if (qarg_is_qubit and not carg_is_bit) or (not qarg_is_qubit and carg_is_bit):
+            raise ir.TypeCheckError(
+                self,
+                "Can't perform measurement with single (qu)bit and an entire register!",
+                help="Instead of `measure(qreg[i], creg)` or `measure(qreg, creg[i])`"
+                "use `measure(qreg[i], creg[j])` or `measure(qreg, creg)`, respectively.",
+            )
 
 
 @statement(dialect=dialect)