Skip to content

Commit cff0517

Browse files
khalatepradnyacursoragent
khalatepradnya
and
cursoragent
committed
Kernel builder parity
Test restructuring Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Pradnya Khalate <pkhalate@nvidia.com>
1 parent cd3aa6e commit cff0517

7 files changed

Lines changed: 454 additions & 282 deletions

File tree

‎python/cudaq/kernel/kernel_builder.py‎

Lines changed: 22 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1128,15 +1128,16 @@ def mz(self, target, regName=None):
11281128
11291129
Args:
11301130
target (:class:`QuakeValue`): The qubit or qubits to measure.
1131-
register_name (Optional[:obj:`str`]): The optional name to provide the
1132-
results of the measurement. Defaults to an empty string.
1131+
register_name (Optional[:obj:`str`]): The optional name to provide the
1132+
results of the measurement. Defaults to ``None``, in which case
1133+
no register name is attached to the measurement op.
11331134
11341135
Returns:
11351136
:class:`QuakeValue`: A handle to this measurement operation in the MLIR.
11361137
11371138
Note:
1138-
Measurements may be applied both mid-circuit and at the end of
1139-
the circuit. Conditional logic on mid-circuit measurements is no longer
1139+
Measurements may be applied both mid-circuit and at the end of
1140+
the circuit. Conditional logic on mid-circuit measurements is no longer
11401141
supported.
11411142
11421143
```python
@@ -1149,22 +1150,16 @@ def mz(self, target, regName=None):
11491150
```
11501151
"""
11511152
with self.ctx, self.insertPoint, self.loc:
1152-
i1Ty = IntegerType.get_signless(1)
1153-
qubitTy = target.mlirValue.type
1154-
retTy = i1Ty
1155-
measTy = quake.MeasureType.get()
1156-
stdvecTy = cc.StdvecType.get(i1Ty)
1153+
measTy = cc.MeasureHandleType.get()
11571154
if quake.VeqType.isinstance(target.mlirValue.type):
1158-
retTy = stdvecTy
11591155
measTy = cc.StdvecType.get(measTy)
11601156
if regName is not None:
11611157
res = quake.MzOp(measTy, [], [target.mlirValue],
11621158
registerName=StringAttr.get(regName,
11631159
context=self.ctx))
11641160
else:
11651161
res = quake.MzOp(measTy, [], [target.mlirValue])
1166-
disc = quake.DiscriminateOp(retTy, res)
1167-
return self.__createQuakeValue(disc.result)
1162+
return self.__createQuakeValue(res.measOut)
11681163

11691164
def mx(self, target, regName=None):
11701165
"""
@@ -1176,15 +1171,16 @@ def mx(self, target, regName=None):
11761171
11771172
Args:
11781173
target (:class:`QuakeValue`): The qubit or qubits to measure.
1179-
register_name (Optional[:obj:`str`]): The optional name to provide the
1180-
results of the measurement. Defaults to an empty string.
1174+
register_name (Optional[:obj:`str`]): The optional name to provide the
1175+
results of the measurement. Defaults to ``None``, in which case
1176+
no register name is attached to the measurement op.
11811177
11821178
Returns:
11831179
:class:`QuakeValue`: A handle to this measurement operation in the MLIR.
11841180
11851181
Note:
1186-
Measurements may be applied both mid-circuit and at the end of
1187-
the circuit. Conditional logic on mid-circuit measurements is no longer
1182+
Measurements may be applied both mid-circuit and at the end of
1183+
the circuit. Conditional logic on mid-circuit measurements is no longer
11881184
supported.
11891185
11901186
```python
@@ -1196,22 +1192,16 @@ def mx(self, target, regName=None):
11961192
```
11971193
"""
11981194
with self.ctx, self.insertPoint, self.loc:
1199-
i1Ty = IntegerType.get_signless(1)
1200-
qubitTy = target.mlirValue.type
1201-
retTy = i1Ty
1202-
measTy = quake.MeasureType.get()
1203-
stdvecTy = cc.StdvecType.get(i1Ty)
1195+
measTy = cc.MeasureHandleType.get()
12041196
if quake.VeqType.isinstance(target.mlirValue.type):
1205-
retTy = stdvecTy
12061197
measTy = cc.StdvecType.get(measTy)
12071198
if regName is not None:
12081199
res = quake.MxOp(measTy, [], [target.mlirValue],
12091200
registerName=StringAttr.get(regName,
12101201
context=self.ctx))
12111202
else:
12121203
res = quake.MxOp(measTy, [], [target.mlirValue])
1213-
disc = quake.DiscriminateOp(retTy, res)
1214-
return self.__createQuakeValue(disc.result)
1204+
return self.__createQuakeValue(res.measOut)
12151205

12161206
def my(self, target, regName=None):
12171207
"""
@@ -1223,16 +1213,17 @@ def my(self, target, regName=None):
12231213
12241214
Args:
12251215
target (:class:`QuakeValue`): The qubit or qubits to measure.
1226-
register_name (Optional[:obj:`str`]): The optional name to provide the
1227-
results of the measurement. Defaults to an empty string.
1216+
register_name (Optional[:obj:`str`]): The optional name to provide the
1217+
results of the measurement. Defaults to ``None``, in which case
1218+
no register name is attached to the measurement op.
12281219
12291220
Returns:
12301221
:class:`QuakeValue`: A handle to this measurement operation in the MLIR.
12311222
12321223
Note:
1233-
Measurements may be applied both mid-circuit and at the end of
1234-
the circuit. Conditional logic on mid-circuit measurements is no longer
1235-
supported.
1224+
Measurements may be applied both mid-circuit and at the end of
1225+
the circuit. Conditional logic on mid-circuit measurements is no longer
1226+
supported.
12361227
12371228
```python
12381229
# Example:
@@ -1244,22 +1235,16 @@ def my(self, target, regName=None):
12441235
```
12451236
"""
12461237
with self.ctx, self.insertPoint, self.loc:
1247-
i1Ty = IntegerType.get_signless(1)
1248-
qubitTy = target.mlirValue.type
1249-
retTy = i1Ty
1250-
measTy = quake.MeasureType.get()
1251-
stdvecTy = cc.StdvecType.get(i1Ty)
1238+
measTy = cc.MeasureHandleType.get()
12521239
if quake.VeqType.isinstance(target.mlirValue.type):
1253-
retTy = stdvecTy
12541240
measTy = cc.StdvecType.get(measTy)
12551241
if regName is not None:
12561242
res = quake.MyOp(measTy, [], [target.mlirValue],
12571243
registerName=StringAttr.get(regName,
12581244
context=self.ctx))
12591245
else:
12601246
res = quake.MyOp(measTy, [], [target.mlirValue])
1261-
disc = quake.DiscriminateOp(retTy, res)
1262-
return self.__createQuakeValue(disc.result)
1247+
return self.__createQuakeValue(res.measOut)
12631248

12641249
def adjoint(self, otherKernel, *target_arguments):
12651250
"""

‎python/tests/kernel/test_measure_handle.py‎

Lines changed: 2 additions & 230 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,6 @@ def reset_run_clear():
2323
cudaq.reset_target()
2424

2525

26-
def _ir(kernel):
27-
"""Materialize MLIR for ``kernel`` and return it as a string. Triggers
28-
AST bridge codegen (so any frontend diagnostic fires here) without
29-
running the kernel on a simulator -- enough for IR-shape and negative
30-
tests."""
31-
return str(kernel)
32-
33-
3426
def test_host_construction_raises_runtime_error():
3527
with pytest.raises(
3628
RuntimeError,
@@ -45,232 +37,12 @@ def test_host_to_bools_raises_runtime_error():
4537
cudaq.to_bools([])
4638

4739

48-
# Scalar `mz` / `mx` / `my` return a `!cc.measure_handle` and do *not*
49-
# inline a `quake.discriminate`.
50-
51-
52-
def test_scalar_mz_emits_handle_no_discriminate():
53-
54-
@cudaq.kernel
55-
def k():
56-
q = cudaq.qubit()
57-
h = mz(q)
58-
59-
ir = _ir(k)
60-
assert "quake.mz" in ir
61-
assert "!cc.measure_handle" in ir
62-
assert "quake.discriminate" not in ir, (
63-
"scalar mz with no coercion site must not emit a discriminate; "
64-
"got:\n" + ir)
65-
66-
67-
def test_scalar_mx_my_emit_handles():
68-
69-
@cudaq.kernel
70-
def kx():
71-
q = cudaq.qubit()
72-
h = mx(q)
73-
74-
@cudaq.kernel
75-
def ky():
76-
q = cudaq.qubit()
77-
h = my(q)
78-
79-
irx, iry = _ir(kx), _ir(ky)
80-
assert re.search(r"quake\.mx.*-> !cc\.measure_handle", irx), irx
81-
assert re.search(r"quake\.my.*-> !cc\.measure_handle", iry), iry
82-
assert "quake.discriminate" not in irx
83-
assert "quake.discriminate" not in iry
84-
85-
86-
def test_vector_mz_emits_stdvec_of_handles():
87-
88-
@cudaq.kernel
89-
def k():
90-
qv = cudaq.qvector(3)
91-
hs = mz(qv)
92-
93-
ir = _ir(k)
94-
assert "!cc.stdvec<!cc.measure_handle>" in ir
95-
assert "quake.discriminate" not in ir
96-
97-
9840
# ---------------------------------------------------------------------------
99-
# Bool-coercion sites lower to `quake.discriminate`.
41+
# `cudaq.to_integer` rejects a raw `!cc.stdvec<!cc.measure_handle>`
42+
# without an intervening `cudaq.to_bools`.
10043
# ---------------------------------------------------------------------------
10144

10245

103-
def test_bool_coercion_in_if():
104-
105-
@cudaq.kernel
106-
def k():
107-
q = cudaq.qubit()
108-
h = mz(q)
109-
if h:
110-
x(q)
111-
112-
ir = _ir(k)
113-
assert re.search(r"quake\.discriminate.*-> i1", ir), ir
114-
115-
116-
def test_bool_coercion_in_while():
117-
118-
@cudaq.kernel
119-
def k():
120-
q = cudaq.qubit()
121-
h = mz(q)
122-
while h:
123-
x(q)
124-
h = mz(q)
125-
126-
ir = _ir(k)
127-
assert re.search(r"quake\.discriminate.*-> i1", ir), ir
128-
129-
130-
def test_bool_coercion_in_not():
131-
132-
@cudaq.kernel
133-
def k():
134-
q = cudaq.qubit()
135-
h = mz(q)
136-
if not h:
137-
x(q)
138-
139-
ir = _ir(k)
140-
assert re.search(r"quake\.discriminate.*-> i1", ir), ir
141-
142-
143-
def test_bool_coercion_in_bool_call():
144-
145-
@cudaq.kernel
146-
def k() -> bool:
147-
q = cudaq.qubit()
148-
return bool(mz(q))
149-
150-
ir = _ir(k)
151-
assert re.search(r"quake\.discriminate.*-> i1", ir), ir
152-
153-
154-
def test_bool_coercion_in_return():
155-
156-
@cudaq.kernel
157-
def k() -> bool:
158-
q = cudaq.qubit()
159-
h = mz(q)
160-
return h
161-
162-
ir = _ir(k)
163-
assert re.search(r"quake\.discriminate.*-> i1", ir), ir
164-
165-
166-
# ---------------------------------------------------------------------------
167-
# Equality / inequality of two handles
168-
# ---------------------------------------------------------------------------
169-
170-
171-
def test_handle_equality_emits_two_discriminates_and_cmpi():
172-
173-
@cudaq.kernel
174-
def k() -> bool:
175-
qv = cudaq.qvector(2)
176-
return mz(qv[0]) == mz(qv[1])
177-
178-
ir = _ir(k)
179-
n_disc = ir.count("quake.discriminate")
180-
n_mz = ir.count("quake.mz")
181-
assert n_disc >= 2 and n_disc == n_mz, (
182-
f"each mz must contribute exactly one discriminate; "
183-
f"got {n_disc} disc / {n_mz} mz:\n{ir}")
184-
assert "arith.cmpi" in ir, ir
185-
186-
187-
def test_handle_inequality_emits_two_discriminates_and_cmpi():
188-
189-
@cudaq.kernel
190-
def k() -> bool:
191-
qv = cudaq.qvector(2)
192-
return mz(qv[0]) != mz(qv[1])
193-
194-
ir = _ir(k)
195-
assert ir.count("quake.discriminate") >= 2
196-
assert "arith.cmpi" in ir
197-
198-
199-
# ---------------------------------------------------------------------------
200-
# Short-circuit semantics for `and` / `or`
201-
# ---------------------------------------------------------------------------
202-
203-
204-
def test_and_short_circuits_second_discriminate():
205-
206-
@cudaq.kernel
207-
def k():
208-
qv = cudaq.qvector(2)
209-
h0 = mz(qv[0])
210-
h1 = mz(qv[1])
211-
if h0 and h1:
212-
x(qv[0])
213-
214-
ir = _ir(k)
215-
discr = [m.start() for m in re.finditer(r"quake\.discriminate", ir)]
216-
ccif = ir.find("cc.if")
217-
assert len(discr) >= 2, ir
218-
assert ccif != -1 and discr[1] > ccif, (
219-
"second discriminate (RHS of `and`) must be inside the short-circuit "
220-
"branch, not at function entry:\n" + ir)
221-
222-
223-
def test_or_short_circuits_second_discriminate():
224-
225-
@cudaq.kernel
226-
def k():
227-
qv = cudaq.qvector(2)
228-
h0 = mz(qv[0])
229-
h1 = mz(qv[1])
230-
if h0 or h1:
231-
x(qv[0])
232-
233-
ir = _ir(k)
234-
discr = [m.start() for m in re.finditer(r"quake\.discriminate", ir)]
235-
ccif = ir.find("cc.if")
236-
assert len(discr) >= 2, ir
237-
assert ccif != -1 and discr[1] > ccif, (
238-
"second discriminate (RHS of `or`) must be inside the short-circuit "
239-
"branch, not at function entry:\n" + ir)
240-
241-
242-
# ---------------------------------------------------------------------------
243-
# `cudaq.to_bools` (vector discrimination) and composition with
244-
# `cudaq.to_integer`.
245-
# ---------------------------------------------------------------------------
246-
247-
248-
def test_to_bools_lowers_to_vectorized_discriminate():
249-
250-
@cudaq.kernel
251-
def k() -> list[bool]:
252-
qv = cudaq.qvector(3)
253-
return cudaq.to_bools(mz(qv))
254-
255-
ir = _ir(k)
256-
assert "!cc.stdvec<!cc.measure_handle>" in ir
257-
assert re.search(
258-
r"quake\.discriminate.*!cc\.stdvec<!cc\.measure_handle>.*"
259-
r"!cc\.stdvec<i1>", ir), ir
260-
261-
262-
def test_to_integer_composes_with_to_bools():
263-
264-
@cudaq.kernel
265-
def k() -> int:
266-
qv = cudaq.qvector(3)
267-
return cudaq.to_integer(cudaq.to_bools(mz(qv)))
268-
269-
ir = _ir(k)
270-
assert "!cc.stdvec<!cc.measure_handle>" in ir, ir
271-
assert "!cc.stdvec<i1>" in ir, ir
272-
273-
27446
def test_to_integer_rejects_raw_handle_vector():
27547

27648
@cudaq.kernel

0 commit comments

Comments
 (0)