Also make QubitId and MeasurementId adhere to broadcasting semantics

Author: david-pl

src/bloqade/pyqrack/squin/qubit.py

@@ -42,15 +42,17 @@ def measure_qubit_list(
     def qubit_id(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.QubitId
     ):
-        qbit: PyQrackQubit = frame.get(stmt.qubit)
-        return (qbit.addr,)
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        ids = ilist.IList([qbit.addr for qbit in qubits])
+        return (ids,)
 
     @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,)
+        measurements: ilist.IList[Measurement, Any] = frame.get(stmt.measurements)
+        ids = ilist.IList([measurement.measurement_id for measurement in measurements])
+        return (ids,)
 
     @interp.impl(qubit.Reset)
     def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Reset):

src/bloqade/qubit/_interface.py

@@ -36,11 +36,13 @@ def measure(qubits: ilist.IList[Qubit, N]) -> ilist.IList[MeasurementResult, N]:
 
 
 @wraps(QubitId)
-def get_qubit_id(qubit: Qubit) -> int: ...
+def get_qubit_id(qubits: ilist.IList[Qubit, N]) -> ilist.IList[int, N]: ...
 
 
 @wraps(MeasurementId)
-def get_measurement_id(measurement: MeasurementResult) -> int: ...
+def get_measurement_id(
+    measurements: ilist.IList[MeasurementResult, N],
+) -> ilist.IList[int, N]: ...
 
 
 @wraps(Reset)

src/bloqade/qubit/stdlib/broadcast.py

@@ -33,3 +33,30 @@ def measure(qubits: ilist.IList[Qubit, N]) -> ilist.IList[MeasurementResult, N]:
             A MeasurementResult can represent both 0 and 1, but also atoms that are lost.
     """
     return _qubit.measure(qubits)
+
+
+@kernel
+def get_qubit_id(qubits: ilist.IList[Qubit, N]) -> ilist.IList[int, N]:
+    """Get the global, unique ID of each qubit in the list.
+
+    Args:
+        qubits (IList[Qubit, N]): The list of qubits of which you want the ID.
+
+    Returns:
+        qubit_ids (IList[int, N]): The list of global, unique IDs of the qubits.
+    """
+    return _qubit.get_qubit_id(qubits)
+
+
+@kernel
+def get_measurement_id(
+    measurements: ilist.IList[MeasurementResult, N],
+) -> ilist.IList[int, N]:
+    """Get the global, unique ID of each of the measurement results in the list.
+
+    Args:
+        measurements (IList[MeasurementResult, N]): The previously taken measurement of which you want to know the ID.
+    Returns:
+        measurement_ids (IList[int, N]): The list of global, unique IDs of the measurements.
+    """
+    return _qubit.get_measurement_id(measurements)

src/bloqade/qubit/stdlib/simple.py

@@ -3,7 +3,6 @@
 from bloqade.types import Qubit, MeasurementResult
 
 from . import broadcast
-from .. import _interface as _qubit
 from .._prelude import kernel
 
 
@@ -43,7 +42,8 @@ def get_qubit_id(qubit: Qubit) -> int:
     Returns:
         qubit_id (int): The global, unique ID of the qubit.
     """
-    return _qubit.get_qubit_id(qubit)
+    ids = broadcast.get_qubit_id(ilist.IList([qubit]))
+    return ids[0]
 
 
 @kernel
@@ -55,4 +55,5 @@ def get_measurement_id(measurement: MeasurementResult) -> int:
     Returns:
         measurement_id (int): The global, unique ID of the measurement.
     """
-    return _qubit.get_measurement_id(measurement)
+    ids = broadcast.get_measurement_id(ilist.IList([measurement]))
+    return ids[0]

src/bloqade/qubit/stmts.py

@@ -26,15 +26,17 @@ class Measure(ir.Statement):
 @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)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, Len])
+    result: ir.ResultValue = info.result(ilist.IListType[types.Int, Len])
 
 
 @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)
+    measurements: ir.SSAValue = info.argument(
+        ilist.IListType[MeasurementResultType, Len]
+    )
+    result: ir.ResultValue = info.result(ilist.IListType[types.Int, Len])
 
 
 @statement(dialect=dialect)