Support measuring registers in pyqrack

src/bloqade/pyqrack/qasm2/core.py

@@ -51,12 +51,30 @@
     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 | PyQrackReg = 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, PyQrackReg) and isinstance(carg, CRegister):
+            # TODO: clean up iteration after PyQrackReg is refactored
+            for i in range(qarg.size):
+                qubit = qarg[i]
+
+                # TODO: make this consistent with PyQrackReg __getitem__ ?
+                cbit = CBitRef(carg, i)
+
+                if qubit.is_active():
+                    cbit.set_value(Measurement(qarg.sim_reg.m(qubit.addr)))
+                else:
+                    cbit.set_value(interp.loss_m_result)
         else:
-            carg.set_value(interp.loss_m_result)
+            raise RuntimeError(
+                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/qasm2/_wrappers.py

@@ -59,7 +59,7 @@ def reset(qarg: Qubit) -> None:
 
 
 @wraps(core.Measure)
-def measure(qarg: Qubit, cbit: Bit) -> None:
+def measure(qarg: Qubit | QReg, cbit: Bit | CReg) -> None:
     """
     Measure the qubit `qarg` and store the result in the classical bit `cbit`.
 

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

@@ -46,9 +46,9 @@ class Measure(ir.Statement):
 
     name = "measure"
     traits = frozenset({lowering.FromPythonCall()})
-    qarg: ir.SSAValue = info.argument(QubitType)
+    qarg: ir.SSAValue = info.argument(QubitType | QRegType)
     """qarg (Qubit): The qubit to measure."""
-    carg: ir.SSAValue = info.argument(BitType)
+    carg: ir.SSAValue = info.argument(BitType | CRegType)
     """carg (Bit): The bit to store the result in."""
 
 

test/pyqrack/test_target.py

@@ -111,4 +111,29 @@ def multiple_registers():
     assert True
 
 
-test_target_glob()
+def test_measurement():
+
+    @qasm2.main
+    def measure_register():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+        qasm2.x(q[0])
+        qasm2.cx(q[0], q[1])
+        qasm2.measure(q, c)
+        return c
+
+    @qasm2.main
+    def measure_single_qubits():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+        qasm2.x(q[0])
+        qasm2.cx(q[0], q[1])
+        qasm2.measure(q[0], c[0])
+        qasm2.measure(q[1], c[1])
+        return c
+
+    target = PyQrack(2)
+    result_single = target.run(measure_single_qubits)
+    result_reg = target.run(measure_register)
+
+    assert result_single == result_reg == [reg.Measurement.One, reg.Measurement.One]