Infer QASM for classical KeyConditions in circuit (#6871)

* Infer qasm for classical keyconditions

* add qasm3.0 option

* comments

---------

Co-authored-by: Michael Hucka <[email protected]>

Author: daxfohl

cirq-core/cirq/circuits/qasm_output.py

@@ -197,13 +197,14 @@ def __init__(
         meas_key_id_map, meas_comments = self._generate_measurement_ids()
         self.meas_comments = meas_comments
         qubit_id_map = self._generate_qubit_ids()
+        self.cregs = self._generate_cregs(meas_key_id_map)
         self.args = protocols.QasmArgs(
             precision=precision,
             version=version,
             qubit_id_map=qubit_id_map,
             meas_key_id_map=meas_key_id_map,
+            meas_key_bitcount={k: v[0] for k, v in self.cregs.items()},
         )
-        self.cregs = self._generate_cregs()
 
     def _generate_measurement_ids(self) -> Tuple[Dict[str, str], Dict[str, Optional[str]]]:
         # Pick an id for the creg that will store each measurement
@@ -227,7 +228,7 @@ def _generate_measurement_ids(self) -> Tuple[Dict[str, str], Dict[str, Optional[
     def _generate_qubit_ids(self) -> Dict['cirq.Qid', str]:
         return {qubit: f'q[{i}]' for i, qubit in enumerate(self.qubits)}
 
-    def _generate_cregs(self) -> Dict[str, tuple[int, str]]:
+    def _generate_cregs(self, meas_key_id_map: Dict[str, str]) -> Dict[str, tuple[int, str]]:
         """Pick an id for the creg that will store each measurement
 
         This function finds the largest measurement using each key.
@@ -239,7 +240,7 @@ def _generate_cregs(self) -> Dict[str, tuple[int, str]]:
         cregs: Dict[str, tuple[int, str]] = {}
         for meas in self.measurements:
             key = protocols.measurement_key_name(meas)
-            meas_id = self.args.meas_key_id_map[key]
+            meas_id = meas_key_id_map[key]
 
             if self.meas_comments[key] is not None:
                 comment = f'  // Measurement: {self.meas_comments[key]}'

cirq-core/cirq/ops/classically_controlled_operation.py

@@ -219,4 +219,4 @@ def _qasm_(self, args: 'cirq.QasmArgs') -> Optional[str]:
         subop_qasm = protocols.qasm(self._sub_operation, args=args)
         if not self._conditions:
             return subop_qasm
-        return f'if ({self._conditions[0].qasm}) {subop_qasm}'
+        return f'if ({protocols.qasm(self._conditions[0], args=args)}) {subop_qasm}'

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -196,7 +196,7 @@ def test_diagram_subcircuit_layered():
     )
 
 
-def test_qasm():
+def test_qasm_sympy_condition():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(
         cirq.measure(q0, key='a'),
@@ -222,6 +222,29 @@ def test_qasm():
     )
 
 
+def test_qasm_key_condition():
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.measure(q0, key='a'), cirq.X(q1).with_classical_controls('a'))
+    qasm = cirq.qasm(circuit)
+    assert (
+        qasm
+        == f"""// Generated from Cirq v{cirq.__version__}
+
+OPENQASM 2.0;
+include "qelib1.inc";
+
+
+// Qubits: [q(0), q(1)]
+qreg q[2];
+creg m_a[1];
+
+
+measure q[0] -> m_a[0];
+if (m_a==1) x q[1];
+"""
+    )
+
+
 def test_qasm_no_conditions():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(

cirq-core/cirq/protocols/qasm.py

@@ -38,6 +38,7 @@ def __init__(
         version: str = '2.0',
         qubit_id_map: Optional[Dict['cirq.Qid', str]] = None,
         meas_key_id_map: Optional[Dict[str, str]] = None,
+        meas_key_bitcount: Optional[Dict[str, int]] = None,
     ) -> None:
         """Inits QasmArgs.
 
@@ -49,11 +50,14 @@ def __init__(
             qubit_id_map: A dictionary mapping qubits to qreg QASM identifiers.
             meas_key_id_map: A dictionary mapping measurement keys to creg QASM
                 identifiers.
+            meas_key_bitcount: A dictionary with of bits for each measurement
+                key.
         """
         self.precision = precision
         self.version = version
         self.qubit_id_map = {} if qubit_id_map is None else qubit_id_map
         self.meas_key_id_map = {} if meas_key_id_map is None else meas_key_id_map
+        self.meas_key_bitcount = {} if meas_key_bitcount is None else meas_key_bitcount
 
     def _format_number(self, value) -> str:
         """OpenQASM 2.0 does not support '1e-5' and wants '1.0e-5'"""

cirq-core/cirq/value/condition.py

@@ -14,7 +14,7 @@
 
 import abc
 import dataclasses
-from typing import Mapping, Tuple, TYPE_CHECKING, FrozenSet
+from typing import Mapping, Tuple, TYPE_CHECKING, FrozenSet, Optional
 
 import sympy
 
@@ -47,6 +47,9 @@ def resolve(self, classical_data: 'cirq.ClassicalDataStoreReader') -> bool:
     def qasm(self):
         """Returns the qasm of this condition."""
 
+    def _qasm_(self, args: 'cirq.QasmArgs', **kwargs) -> Optional[str]:
+        return self.qasm
+
     def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]) -> 'cirq.Condition':
         condition = self
         for k in self.keys:
@@ -115,6 +118,22 @@ def _from_json_dict_(cls, key, **kwargs):
     def qasm(self):
         raise ValueError('QASM is defined only for SympyConditions of type key == constant.')
 
+    def _qasm_(self, args: 'cirq.QasmArgs', **kwargs) -> Optional[str]:
+        args.validate_version('2.0', '3.0')
+        key_str = str(self.key)
+        if key_str not in args.meas_key_id_map:
+            raise ValueError(f'Key "{key_str}" not in QasmArgs.meas_key_id_map.')
+        key = args.meas_key_id_map[key_str]
+        # QASM 3.0 supports !=, so we return it directly.
+        if args.version == '3.0':
+            return f'{key}!=0'
+        # QASM 2.0 only has == operator, so we must limit to single-bit measurement keys == 1.
+        if key not in args.meas_key_bitcount:
+            raise ValueError(f'Key "{key}" not in QasmArgs.meas_key_bitcount.')
+        if args.meas_key_bitcount[str(key)] != 1:
+            raise ValueError('QASM is defined only for single-bit classical conditions.')
+        return f'{key}==1'
+
 
 @dataclasses.dataclass(frozen=True)
 class SympyCondition(Condition):

cirq-core/cirq/value/condition_test.py

@@ -71,6 +71,35 @@ def test_key_condition_qasm():
         _ = cirq.KeyCondition(cirq.MeasurementKey('a')).qasm
 
 
+def test_key_condition_qasm_protocol():
+    cond = cirq.KeyCondition(cirq.MeasurementKey('a'))
+    args = cirq.QasmArgs(meas_key_id_map={'a': 'm_a'}, meas_key_bitcount={'m_a': 1})
+    qasm = cirq.qasm(cond, args=args)
+    assert qasm == 'm_a==1'
+
+
+def test_key_condition_qasm_protocol_v3():
+    cond = cirq.KeyCondition(cirq.MeasurementKey('a'))
+    args = cirq.QasmArgs(meas_key_id_map={'a': 'm_a'}, version='3.0')
+    qasm = cirq.qasm(cond, args=args)
+    assert qasm == 'm_a!=0'
+
+
+def test_key_condition_qasm_protocol_invalid_args():
+    cond = cirq.KeyCondition(cirq.MeasurementKey('a'))
+    args = cirq.QasmArgs()
+    with pytest.raises(ValueError, match='Key "a" not in QasmArgs.meas_key_id_map.'):
+        _ = cirq.qasm(cond, args=args)
+    args = cirq.QasmArgs(meas_key_id_map={'a': 'm_a'})
+    with pytest.raises(ValueError, match='Key "m_a" not in QasmArgs.meas_key_bitcount.'):
+        _ = cirq.qasm(cond, args=args)
+    args = cirq.QasmArgs(meas_key_id_map={'a': 'm_a'}, meas_key_bitcount={'m_a': 2})
+    with pytest.raises(
+        ValueError, match='QASM is defined only for single-bit classical conditions.'
+    ):
+        _ = cirq.qasm(cond, args=args)
+
+
 def test_sympy_condition_with_keys():
     c = init_sympy_condition.replace_key(key_a, key_b)
     assert c.keys == (key_b,)