Remove cirq tags in parallelize and deal with tags when lowering to squin (#489)

Closes #488 and #375 .

Author: david-pl

src/bloqade/cirq_utils/__init__.py

@@ -2,7 +2,7 @@
 from .parallelize import (
     transpile as transpile,
     parallelize as parallelize,
-    no_similarity as no_similarity,
+    remove_tags as remove_tags,
     auto_similarity as auto_similarity,
     block_similarity as block_similarity,
     moment_similarity as moment_similarity,

src/bloqade/cirq_utils/parallelize.py

@@ -136,7 +136,7 @@ def auto_similarity(
     return cirq.Circuit(flattened_circuit), weights
 
 
-def no_similarity(circuit: cirq.Circuit) -> cirq.Circuit:
+def remove_tags(circuit: cirq.Circuit) -> cirq.Circuit:
     """
     Removes all tags from the circuit
 
@@ -146,10 +146,11 @@ def no_similarity(circuit: cirq.Circuit) -> cirq.Circuit:
     Returns:
     [0] - cirq.Circuit - the circuit with all tags removed.
     """
-    new_moments = []
-    for moment in circuit.moments:
-        new_moments.append([gate.untagged for gate in moment.operations])
-    return cirq.Circuit(new_moments)
+
+    def remove_tag(op: cirq.Operation, _):
+        return op.untagged
+
+    return cirq.map_operations(circuit, remove_tag)
 
 
 def to_dag_circuit(circuit: cirq.Circuit, can_reorder=None) -> nx.DiGraph:
@@ -399,4 +400,6 @@ def parallelize(
     )
     # Convert the epochs to a cirq circuit.
     moments = map(cirq.Moment, epochs)
-    return cirq.Circuit(moments)
+    circuit = cirq.Circuit(moments)
+
+    return remove_tags(circuit)

src/bloqade/squin/cirq/lowering.py

@@ -144,6 +144,11 @@ def visit_GateOperation(
         qbits = self.lower_qubit_getindices(state, node.qubits)
         return state.current_frame.push(qubit.Apply(operator=op_, qubits=qbits))
 
+    def visit_TaggedOperation(
+        self, state: lowering.State[CirqNode], node: cirq.TaggedOperation
+    ):
+        state.lower(node.untagged)
+
     def lower_measurement(
         self, state: lowering.State[CirqNode], node: cirq.GateOperation
     ):

test/cirq_utils/test_parallelize.py

@@ -6,7 +6,7 @@
 
 from bloqade.cirq_utils import (
     parallelize,
-    no_similarity,
+    remove_tags,
     block_similarity,
     moment_similarity,
 )
@@ -28,7 +28,7 @@ def test1():
     circuit_m, _ = moment_similarity(circuit, weight=1.0)
     # print(circuit_m)
     circuit_b, _ = block_similarity(circuit, weight=1.0, block_id=1)
-    circuit_m2 = no_similarity(circuit_m)
+    circuit_m2 = remove_tags(circuit_m)
     print(circuit_m2)
     circuit2 = parallelize(circuit)
     # print(circuit2)

test/squin/cirq/test_cirq_to_squin.py

@@ -34,6 +34,11 @@ def basic_circuit():
     )
 
 
+def tagged_circuit():
+    q = cirq.LineQubit.range(2)
+    return cirq.Circuit(cirq.H(q[0]).with_tags("FOO"), cirq.CX(*q).with_tags("BAR"))
+
+
 def controlled_gates():
     q0 = cirq.NamedQubit("q0")
     q1 = cirq.NamedQubit("q1")
@@ -166,6 +171,7 @@ def nested_circuit():
     "circuit_f",
     [
         basic_circuit,
+        tagged_circuit,
         controlled_gates,
         parity_gate_circuit,
         phased_gates,