Support cirq classical control when lowering a circuit to squin (#315)

Author: david-pl

src/bloqade/squin/cirq/lowering.py

@@ -5,7 +5,7 @@
 import cirq
 from kirin import ir, lowering
 from kirin.rewrite import Walk, CFGCompactify
-from kirin.dialects import py, ilist
+from kirin.dialects import py, scf, ilist
 
 from .. import op, noise, qubit
 
@@ -150,10 +150,79 @@ def lower_measurement(
     ):
         if len(node.qubits) == 1:
             qbit = self.lower_qubit_getindex(state, node.qubits[0])
-            return state.current_frame.push(qubit.MeasureQubit(qbit))
+            stmt = state.current_frame.push(qubit.MeasureQubit(qbit))
+        else:
+            qbits = self.lower_qubit_getindices(state, node.qubits)
+            stmt = state.current_frame.push(qubit.MeasureQubitList(qbits))
 
-        qbits = self.lower_qubit_getindices(state, node.qubits)
-        return state.current_frame.push(qubit.MeasureQubitList(qbits))
+        key = node.gate.key
+        if isinstance(key, cirq.MeasurementKey):
+            key = key.name
+
+        state.current_frame.defs[key] = stmt.result
+        return stmt
+
+    def visit_ClassicallyControlledOperation(
+        self, state: lowering.State[CirqNode], node: cirq.ClassicallyControlledOperation
+    ):
+        conditions: list[ir.SSAValue] = []
+        for outcome in node.classical_controls:
+            key = outcome.key
+            if isinstance(key, cirq.MeasurementKey):
+                key = key.name
+            measurement_outcome = state.current_frame.defs[key]
+
+            if measurement_outcome.type.is_subseteq(ilist.IListType):
+                # NOTE: there is currently no convenient ilist.any method, so we need to use foldl
+                # with a simple function that just does an or
+
+                def bool_op_or(x: bool, y: bool) -> bool:
+                    return x or y
+
+                f_code = state.current_frame.push(
+                    lowering.Python(self.dialects).python_function(bool_op_or)
+                )
+                fn = ir.Method(
+                    mod=None,
+                    py_func=bool_op_or,
+                    sym_name="bool_op_or",
+                    arg_names=[],
+                    dialects=self.dialects,
+                    code=f_code,
+                )
+                f_const = state.current_frame.push(py.constant.Constant(fn))
+                init_val = state.current_frame.push(py.Constant(False)).result
+                condition = state.current_frame.push(
+                    ilist.Foldl(f_const.result, measurement_outcome, init=init_val)
+                ).result
+            else:
+                condition = measurement_outcome
+
+            conditions.append(condition)
+
+        if len(conditions) == 1:
+            condition = conditions[0]
+        else:
+            condition = state.current_frame.push(
+                py.boolop.And(conditions[0], conditions[1])
+            ).result
+            for next_cond in conditions[2:]:
+                condition = state.current_frame.push(
+                    py.boolop.And(condition, next_cond)
+                ).result
+
+        then_stmt = self.visit(state, node.without_classical_controls())
+
+        assert isinstance(
+            then_stmt, ir.Statement
+        ), f"Expected operation of classically controlled node {node} to be lowered to a statement, got type {type(then_stmt)}. \
+        Please report this issue!"
+
+        # NOTE: remove stmt from parent block
+        then_stmt.detach()
+        then_body = ir.Block((then_stmt,))
+
+        return state.current_frame.push(scf.IfElse(condition, then_body=then_body))
 
     def visit_SingleQubitPauliStringGateOperation(
         self,

test/squin/cirq/test_cirq_to_squin.py

@@ -159,3 +159,59 @@ def test_circuit(circuit_f, run_sim: bool = False):
     sim = DynamicMemorySimulator()
     ket = sim.state_vector(kernel=kernel)
     print(ket)
+
+
+def test_classical_control(run_sim: bool = False):
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.H(q[0]),
+        cirq.measure(q[0]),
+        cirq.X(q[1]).with_classical_controls("q(0)"),
+        cirq.measure(q[1]),
+    )
+
+    print(circuit)
+
+    if run_sim:
+        simulator = cirq.Simulator()
+        simulator.run(circuit, repetitions=1)
+
+    kernel = squin.cirq.load_circuit(circuit)
+    kernel.print()
+
+
+def test_classical_control_register():
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.H(q[0]),
+        cirq.measure(q, key="test"),
+        cirq.X(q[1]).with_classical_controls("test"),
+        cirq.measure(q[1]),
+    )
+
+    print(circuit)
+
+    kernel = squin.cirq.load_circuit(circuit)
+    kernel.print()
+
+
+def test_multiple_classical_controls(run_sim: bool = False):
+    q = cirq.LineQubit.range(2)
+    q2 = cirq.GridQubit(0, 1)
+    circuit = cirq.Circuit(
+        cirq.H(q[0]),
+        cirq.H(q2),
+        cirq.measure(q, key="test"),
+        cirq.measure(q2),
+        cirq.X(q[1]).with_classical_controls("test", "q(0, 1)"),
+        cirq.measure(q[1]),
+    )
+
+    print(circuit)
+
+    if run_sim:
+        sim = cirq.Simulator()
+        sim.run(circuit)
+
+    kernel = squin.cirq.load_circuit(circuit)
+    kernel.print()