Add qubit ID and measurement ID statements (#499)

Implements #496 .

I had to adapt the returns from the measurement PyQrack implementation
to actually return a `Measurement` rather than just a boolean, in order
to be able to store an ID on the resulting measurement. Not sure if the
implementation is that clean here.

Author: david-pl

src/bloqade/pyqrack/base.py

@@ -146,7 +146,13 @@ class PyQrackInterpreter(Interpreter, typing.Generic[MemoryType]):
     loss_m_result: Measurement = field(default=Measurement.One, kw_only=True)
     """The value of a measurement result when a qubit is lost."""
 
+    global_measurement_id: int = field(init=False, default=0)
+
     def initialize(self) -> Self:
         super().initialize()
         self.memory.reset()  # reset allocated qubits
         return self
+
+    def set_global_measurement_id(self, m: Measurement):
+        m.measurement_id = self.global_measurement_id
+        self.global_measurement_id += 1

src/bloqade/pyqrack/reg.py

@@ -2,15 +2,20 @@
 from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
+from bloqade.types import MeasurementResult
 from bloqade.qasm2.types import Qubit
 
 if TYPE_CHECKING:
     from pyqrack import QrackSimulator
 
 
-class Measurement(enum.IntEnum):
+class Measurement(MeasurementResult, enum.IntEnum):
     """Enumeration of measurement results."""
 
+    def __init__(self, measurement_id: int = 0) -> None:
+        super().__init__()
+        self.measurement_id = measurement_id
+
     Zero = 0
     One = 1
     Lost = enum.auto()

src/bloqade/pyqrack/squin/noise/native.py

@@ -46,7 +46,7 @@ def n_sites(self) -> int:
         return 1
 
     def apply(self, qubit: PyQrackQubit, adjoint: bool = False) -> None:
-        if random.uniform(0.0, 1.0) < self.p:
+        if random.uniform(0.0, 1.0) <= self.p:
             qubit.state = QubitState.Lost
 
 

src/bloqade/pyqrack/squin/qubit.py

@@ -4,7 +4,7 @@
 from kirin.dialects import ilist
 
 from bloqade.squin import qubit
-from bloqade.pyqrack.reg import QubitState, PyQrackQubit
+from bloqade.pyqrack.reg import QubitState, Measurement, PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 
 from .runtime import OperatorRuntimeABC
@@ -41,9 +41,12 @@ def broadcast(
 
     def _measure_qubit(self, qbit: PyQrackQubit, interp: PyQrackInterpreter):
         if qbit.is_active():
-            return bool(qbit.sim_reg.m(qbit.addr))
+            m = Measurement(bool(qbit.sim_reg.m(qbit.addr)))
         else:
-            return interp.loss_m_result
+            m = Measurement(interp.loss_m_result)
+
+        interp.set_global_measurement_id(m)
+        return m
 
     @interp.impl(qubit.MeasureQubitList)
     def measure_qubit_list(
@@ -63,3 +66,17 @@ def measure_qubit(
         qbit: PyQrackQubit = frame.get(stmt.qubit)
         result = self._measure_qubit(qbit, interp)
         return (result,)
+
+    @interp.impl(qubit.QubitId)
+    def qubit_id(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.QubitId
+    ):
+        qbit: PyQrackQubit = frame.get(stmt.qubit)
+        return (qbit.addr,)
+
+    @interp.impl(qubit.MeasurementId)
+    def measurement_id(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasurementId
+    ):
+        measurement: Measurement = frame.get(stmt.measurement)
+        return (measurement.measurement_id,)

src/bloqade/squin/qubit.py

@@ -79,6 +79,20 @@ class MeasureQubitList(ir.Statement):
     result: ir.ResultValue = info.result(ilist.IListType[MeasurementResultType])
 
 
+@statement(dialect=dialect)
+class QubitId(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    qubit: ir.SSAValue = info.argument(QubitType)
+    result: ir.ResultValue = info.result(types.Int)
+
+
+@statement(dialect=dialect)
+class MeasurementId(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    measurement: ir.SSAValue = info.argument(MeasurementResultType)
+    result: ir.ResultValue = info.result(types.Int)
+
+
 # NOTE: no dependent types in Python, so we have to mark it Any...
 @wraps(New)
 def new(n_qubits: int) -> ilist.IList[Qubit, Any]:
@@ -127,9 +141,10 @@ def measure(input: Any) -> Any:
         input: A qubit or a list of qubits to measure.
 
     Returns:
-        bool | list[bool]: The result of the measurement. If a single qubit is measured,
-            a single boolean is returned. If a list of qubits is measured, a list of booleans
+        MeasurementResult | list[MeasurementResult]: The result of the measurement. If a single qubit is measured,
+            a single result is returned. If a list of qubits is measured, a list of results
             is returned.
+            A MeasurementResult can represent both 0 and 1, but also atoms that are lost.
     """
     ...
 
@@ -169,3 +184,11 @@ def ghz():
         None
     """
     ...
+
+
+@wraps(QubitId)
+def get_qubit_id(qubit: Qubit) -> int: ...
+
+
+@wraps(MeasurementId)
+def get_measurement_id(measurement: MeasurementResult) -> int: ...

test/squin/test_qubit.py

@@ -0,0 +1,48 @@
+from kirin import types
+
+from bloqade import squin
+from bloqade.pyqrack import StackMemorySimulator
+from bloqade.pyqrack.reg import Measurement
+
+
+def test_get_ids():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(3)
+
+        m = squin.qubit.measure(q)
+
+        qid = squin.qubit.get_qubit_id(q[0])
+        mid = squin.qubit.get_measurement_id(m[1])
+        return mid + qid
+
+    main.print()
+    assert main.return_type.is_subseteq(types.Int)
+
+    @squin.kernel
+    def main2():
+        q = squin.qubit.new(2)
+
+        qid = squin.qubit.get_qubit_id(q[0])
+        m1 = squin.qubit.measure(q[qid])
+
+        squin.gate.x(q[qid])
+        m2 = squin.qubit.measure(q[qid])
+
+        m1_id = squin.qubit.get_measurement_id(m1)
+        m2_id = squin.qubit.get_measurement_id(m2)
+
+        if m1_id != 0:
+            # do something that errors
+            squin.gate.x(q[4])
+
+        if m2_id != 1:
+            squin.gate.x(q[4])
+
+        return squin.qubit.measure(q)
+
+    sim = StackMemorySimulator(min_qubits=2)
+    result = sim.run(main2)
+    for i, res in enumerate(result):
+        assert isinstance(res, Measurement)
+        assert res.measurement_id == i + 2