[python] Bind cudaq.measure_handle, rewire AST bridge, add tests

Python-side counterpart to PR 3b's C++ frontend rewire, per spec
cudaq-spec/proposals/measure_handle.bs (Python API; Operational
Semantics; IR Representation).

Bindings and host stubs
- python/runtime/mlir/py_register_dialects.cpp: bind
  `!cc.measure_handle` so the AST bridge can construct, check, and
  emit the type from Python.
- python/cudaq/kernel_types.py: `cudaq.measure_handle` kernel-type
  stub (mirrors `qubit` / `qvector` / `qview`); host-scope
  construction raises a device-only RuntimeError.
- python/cudaq/__init__.py: re-export `cudaq.measure_handle` and add
  a host stub for `cudaq.to_bools(handles)` raising the same
  device-only error -- kernel-side calls are intercepted by the AST
  bridge.
- python/cudaq/kernel/utils.py: `containsMeasureHandle(ty)`,
  mirroring `cudaq::cc::containsMeasureHandle` from PR 3b
  (`lib/Optimizer/Dialect/CC/CCTypes.cpp`); used by the bridge's
  boundary check.

AST bridge (`PyASTBridge` in `python/cudaq/kernel/ast_bridge.py`)
- mz / mx / my emit `quake.{mz,mx,my}` producing
  `!cc.measure_handle` (scalar) or `!cc.stdvec<!cc.measure_handle>`
  (vector). Discriminate insertion is deferred to coercion sites.
- `__discriminateIfMeasureHandle` inserts `quake.discriminate` at
  every spec-listed coercion site: arithmetic-to-bool (if / while /
  not / and / or), `changeOperandToType` (i1 + i8 + stdvec<i1>),
  Compare (==/!=), explicit `bool(...)`, IfExp test, Assert test;
  surfaces the `discriminating an unbound measure_handle`
  diagnostic for the default-constructed pattern.
- `cudaq.measure_handle()` -> `cc.UndefOp(!cc.measure_handle)` so
  the unbound-handle diagnostic has a recognizable source.
- `cudaq.to_bools(handles)` -> vectorized `quake.discriminate` on
  `!cc.stdvec<!cc.measure_handle>`; legacy
  `cudaq.to_integer(mz(qvec))` is rejected with a targeted
  diagnostic matching PR 3b's C++ rejection.
- Boundary check at entry-point creation: walk the kernel signature
  with `containsMeasureHandle` and reject any parameter / return
  position transitively containing a `measure_handle`. Diagnostic
  matches PR 3b's spec-canonical wording.

Tests
- python/tests/kernel/test_measure_handle.py: new spec-coverage
  suite (host-scope rejections, scalar/vector emission shape, every
  bool-coercion site, `to_bools` lowering, `to_integer(to_bools(...))`
  composition, unbound-handle diagnostic, boundary diagnostic).
- python/tests/mlir/{bug_1775,bug_1777,bug_1875,call_qpu}.py:
  type-rename and structural CHECK refresh for the new bridge IR
  shape.
- python/tests/kernel/test_{assignments,kernel_features,run_kernel,
  to_integer,kernel_shift_operators}.py: pytest migrations off
  the implicit `mz`-returns-bool assumption.

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Pradnya Khalate <pkhalate@nvidia.com>

Author: khalatepradnya

python/cudaq/__init__.py

@@ -204,7 +204,7 @@ def _isinstance(other, _cls=cls, _isinst=py_isinstance):
 parallel = cudaq_runtime.parallel
 
 # Primitive Types (stubs; used only in kernels, parsed to MLIR)
-from .kernel_types import qubit, qvector, qview
+from .kernel_types import measure_handle, qubit, qvector, qview
 
 Pauli = cudaq_runtime.Pauli
 Kernel = PyKernel
@@ -319,6 +319,17 @@ def amplitudes(array_data):
     return numpy.array(array_data, dtype=complex())
 
 
+def to_bools(handles):
+    """Bulk-discriminate a ``list[cudaq.measure_handle]`` into a
+    ``list[bool]``. Device-only: this Python symbol exists so kernel
+    code can call ``cudaq.to_bools(...)``; the AST bridge intercepts
+    the call and lowers it to a vector form ``quake.discriminate`` on
+    ``!cc.stdvec<!cc.measure_handle>``. Host-side invocation raises a
+    ``RuntimeError``.
+    """
+    raise RuntimeError("device-only; usable only inside @cudaq.kernel")
+
+
 def __clearKernelRegistries():
     global globalRegisteredOperations
     globalRegisteredOperations.clear()

python/cudaq/kernel/ast_bridge.py

@@ -24,7 +24,7 @@
 from cudaq.mlir.ir import (ComplexType, F32Type, F64Type, IntegerType, Context,
                            Module)
 from cudaq.mlir._mlir_libs._quakeDialects import register_all_dialects
-from cudaq.kernel_types import qubit, qvector, qview
+from cudaq.kernel_types import measure_handle, qubit, qvector, qview
 
 State = cudaq_runtime.State
 pauli_word = cudaq_runtime.pauli_word
@@ -41,6 +41,23 @@
 globalRegisteredTypes = cudaq_runtime.DataClassRegistry
 
 
+def containsMeasureHandle(ty):
+    """Return True iff ``ty`` is ``!cc.measure_handle`` or transitively
+    contains one.
+    """
+    if cc.MeasureHandleType.isinstance(ty):
+        return True
+    if cc.PointerType.isinstance(ty):
+        return containsMeasureHandle(cc.PointerType.getElementType(ty))
+    if cc.ArrayType.isinstance(ty):
+        return containsMeasureHandle(cc.ArrayType.getElementType(ty))
+    if cc.StdvecType.isinstance(ty):
+        return containsMeasureHandle(cc.StdvecType.getElementType(ty))
+    if cc.StructType.isinstance(ty):
+        return any(containsMeasureHandle(t) for t in cc.StructType.getTypes(ty))
+    return False
+
+
 def getMLIRContext():
     """
     This code creates an MLIRContext singleton for this python process. We do
@@ -383,6 +400,8 @@ def emitFatalErrorOverride(msg):
                     return quake.RefType.get()
                 if annotation.attr == 'pauli_word':
                     return cc.CharspanType.get()
+                if annotation.attr == 'measure_handle':
+                    return cc.MeasureHandleType.get()
 
             if annotation.value.id in ['numpy', 'np']:
                 if annotation.attr in ['array', 'ndarray']:
@@ -675,6 +694,8 @@ def mlirTypeFromPyType(argType, ctx, **kwargs):
         return quake.RefType.get(ctx)
     if argType == pauli_word:
         return cc.CharspanType.get(ctx)
+    if argType == measure_handle:
+        return cc.MeasureHandleType.get(ctx)
 
     if 'argInstance' in kwargs:
         argInstance = kwargs['argInstance']
@@ -754,6 +775,9 @@ def mlirTypeToPyType(argType):
     if cc.CharspanType.isinstance(argType):
         return pauli_word
 
+    if cc.MeasureHandleType.isinstance(argType):
+        return measure_handle
+
     if cc.StdvecType.isinstance(argType):
         eleTy = cc.StdvecType.getElementType(argType)
         if cc.CharspanType.isinstance(eleTy):

python/cudaq/kernel/utils.py

@@ -35,6 +35,29 @@ def __new__(cls, *args, **kwargs):
         raise KernelTypeError(cls)
 
 
+class measure_handle(KernelType):
+    """
+    A handle to a measurement event recorded inside a CUDA-Q kernel.
+
+    Returned by ``mz`` / ``mx`` / ``my`` inside an ``@cudaq.kernel`` body
+    (scalar form on a single qubit; vector form on a ``qvector`` /
+    ``qview``). The classical outcome is read by coercing the handle to
+    ``bool`` in any Python ``bool`` context, and the AST bridge inserts a
+    ``quake.discriminate`` at the coercion site.
+    ``cudaq.to_bools(handles)`` is the bulk counterpart on a
+    ``list[measure_handle]``.
+
+    Instantiating ``cudaq.measure_handle()`` at host scope raises
+    ``RuntimeError`` (it is device-only).
+    """
+
+    def __new__(cls, *args, **kwargs):
+        raise RuntimeError("device-only; usable only inside @cudaq.kernel")
+
+    def __init__(self) -> None:
+        ...
+
+
 class qubit(KernelType):
     """
     The qubit is the primary unit of information in a quantum computer.

python/cudaq/kernel_types.py

@@ -200,6 +200,18 @@ static void registerCCDialectAndTypes(nanobind::module_ &m) {
           },
           nanobind::arg("cls"), nanobind::arg("context") = nanobind::none());
 
+  mlir_type_subclass(ccMod, "MeasureHandleType",
+                     [](MlirType type) {
+                       return mlir::isa<cudaq::cc::MeasureHandleType>(
+                           unwrap(type));
+                     })
+      .def_classmethod(
+          "get",
+          [](nanobind::object cls, MlirContext context) {
+            return wrap(cudaq::cc::MeasureHandleType::get(unwrap(context)));
+          },
+          nanobind::arg("cls"), nanobind::arg("context") = nanobind::none());
+
   mlir_type_subclass(
       ccMod, "StateType",
       [](MlirType type) { return mlir::isa<quake::StateType>(unwrap(type)); })

python/runtime/mlir/py_register_dialects.cpp

@@ -1150,25 +1150,23 @@ def test1() -> list[bool]:
     assert 'variable defined in parent scope cannot be modified' in str(e.value)
     assert '(offending source -> c = qs[1])' in str(e.value)
 
-    # TODO: The reason we cannot currently support this is
-    # because we store measurement results as values in the
-    # symbol table. This should be changed and supported when
-    # we do the change to properly distinguish measurement
-    # types from booleans.
-    with pytest.raises(RuntimeError) as e:
-
-        @cudaq.kernel
-        def test2() -> bool:
-            qs = cudaq.qvector(2)
-            res = mz(qs[0])
-            if True:
-                x(qs[1])
-                res = mz(qs[1])
-            return res
+    # Reassigning a `measure_handle`-typed variable across scopes is
+    # supported now that `mz` returns `cudaq.measure_handle` instead of
+    # `bool`: the symbol-table slot has handle type, the inner-scope
+    # store binds a fresh handle, and the bool-coercion at `return res`
+    # discriminates exactly once. Previously this case was disallowed
+    # because measurement results were stored as raw `i1` values in the
+    # symbol table.
+    @cudaq.kernel
+    def test2() -> bool:
+        qs = cudaq.qvector(2)
+        res = mz(qs[0])
+        if True:
+            x(qs[1])
+            res = mz(qs[1])
+        return res
 
-        test2()
-    assert 'variable defined in parent scope cannot be modified' in str(e.value)
-    assert '(offending source -> res = mz(qs[1]))' in str(e.value)
+    test2()
 
 
 def test_var_scopes():
@@ -1391,7 +1389,8 @@ def fct():
                 x(q)
 
         fct()
-        return i, mz(q)
+        # FIXME: aggregate-element typing does not currently auto-discriminate, so coerce explicitly.
+        return i, bool(mz(q))
 
     out = cudaq.run(test1, True, False, shots_count=10)
     assert all(res == (True, False) for res in out)

python/tests/kernel/test_assignments.py

@@ -2727,7 +2727,7 @@ def test_mid_circuit_measurements():
 
     @cudaq.kernel
     def callee(register: cudaq.qview) -> list[bool]:
-        result = [0, 0, 0, 0, 0, 0, 0, 0]
+        result = [False, False, False, False, False, False, False, False]
         for i in range(4):
             j = i * 2
             if i % 2 == 0:

python/tests/kernel/test_kernel_features.py

@@ -76,7 +76,7 @@ def test_run_with_integer_left_shift_operator():
     @cudaq.kernel
     def kernel(n: int) -> int:
         q = cudaq.qvector(n)
-        m = mz(q)
+        m = cudaq.to_bools(mz(q))
         r = 0
         for i in range(n):
             r = r & (m[i] << i)
@@ -96,7 +96,7 @@ def test_run_with_non_integer_left_shift_operator():
         @cudaq.kernel
         def kernel(n: int) -> int:
             q = cudaq.qvector(n)
-            m = mz(q)
+            m = cudaq.to_bools(mz(q))
             r = 0
             for i in range(n):
                 r = r & (m[i] << 1.0)
@@ -113,7 +113,7 @@ def test_run_with_integer_right_shift_operator():
     @cudaq.kernel
     def kernel(n: int) -> int:
         q = cudaq.qvector(n)
-        m = mz(q)
+        m = cudaq.to_bools(mz(q))
         r = 0
         for i in range(n):
             r = r & (m[i] >> i)
@@ -131,7 +131,7 @@ def test_run_with_integer_bitwise_or_operator():
     @cudaq.kernel
     def kernel(n: int) -> int:
         q = cudaq.qvector(n)
-        m = mz(q)
+        m = cudaq.to_bools(mz(q))
         r = 0
         for i in range(n):
             r = r | (m[i] >> i)
@@ -149,7 +149,7 @@ def test_run_with_integer_bitwise_xor_operator():
     @cudaq.kernel
     def kernel(n: int) -> int:
         q = cudaq.qvector(n)
-        m = mz(q)
+        m = cudaq.to_bools(mz(q))
         r = 0
         for i in range(n):
             r = r ^ (m[i] >> i)