From f6a53bdae51dbae65919b1354cb62271eb2fc12a Mon Sep 17 00:00:00 2001
From: Brad Chase <brad@unitary.foundation>
Date: Mon, 10 Feb 2025 16:07:03 -0500
Subject: [PATCH] Add target gate set for cirq benchmarking

---
 benchmarks/scripts/common.py | 96 +++++++++++++++++++++++++++++++++---
 1 file changed, 88 insertions(+), 8 deletions(-)

diff --git a/benchmarks/scripts/common.py b/benchmarks/scripts/common.py
index 1c4ccb0..506ff63 100644
--- a/benchmarks/scripts/common.py
+++ b/benchmarks/scripts/common.py
@@ -4,7 +4,8 @@
 import pandas as pd
 import matplotlib
 from datetime import datetime
-from cirq import optimize_for_target_gateset, ops
+from typing import Dict, Any, List
+from cirq import optimize_for_target_gateset
 from cirq.transformers.target_gatesets import TwoQubitCompilationTargetGateset
 import cirq
 from pytket.circuit import OpType
@@ -107,32 +108,111 @@ def qiskit_compile(qiskit_circuit):
         basis_gates=["rz", "rx", "ry", "h", "cx"],
     )
 
+
 class BenchmarkTargetGateset(TwoQubitCompilationTargetGateset):
-    def __init__(self, atol: float = 1e-8):
-        super().__init__(ops.Rz,ops.Rx, ops.Ry, ops.H, ops.CNOT,name="BenchmarkTargetGateset")
+    """
+    Target gateset for compiling cirq circuits for benchmarking.
+
+    This is based on the source code for the IonQTargetGateset in the
+    cirq.ionq package, but restricts the gateset to only H, CNOT, Rx, Ry, Rz.
+    """
+
+    def __init__(self, *, atol: float = 1e-8):
+        """Initializes BenchmarkTargetGateset
+
+        Args:
+            atol: A limit on the amount of absolute error introduced by the decomposition.
+        """
+        super().__init__(
+            cirq.H,
+            cirq.CNOT,
+            cirq.Rx,
+            cirq.Ry,
+            cirq.Rz,
+            cirq.MeasurementGate,
+            unroll_circuit_op=False,
+        )
         self.atol = atol
 
-    def _decompose_two_qubit_operation(self, op: cirq.Operation, _) -> cirq.OP_TREE:
+    def _decompose_single_qubit_operation(
+        self, op: cirq.Operation, _
+    ) -> cirq.OP_TREE:
+        qubit = op.qubits[0]
+        mat = cirq.unitary(op)
+        for gate in cirq.single_qubit_matrix_to_gates(mat, self.atol):
+            yield gate(qubit)
+
+    def _decompose_two_qubit_operation(
+        self, op: cirq.Operation, _
+    ) -> cirq.OP_TREE:
         if not cirq.has_unitary(op):
             return NotImplemented
         mat = cirq.unitary(op)
         q0, q1 = op.qubits
-        naive = cirq.two_qubit_matrix_to_cz_operations(q0, q1, mat, allow_partial_czs=False)
+        naive = cirq.two_qubit_matrix_to_cz_operations(
+            q0, q1, mat, allow_partial_czs=False
+        )
         temp = cirq.map_operations_and_unroll(
             cirq.Circuit(naive),
             lambda op, _: (
-                [cirq.H(op.qubits[1]), cirq.CNOT(*op.qubits), cirq.H(op.qubits[1])]
+                [
+                    cirq.H(op.qubits[1]),
+                    cirq.CNOT(*op.qubits),
+                    cirq.H(op.qubits[1]),
+                ]
                 if op.gate == cirq.CZ
                 else op
             ),
         )
         return cirq.merge_k_qubit_unitaries(
-            temp, k=1, rewriter=lambda op: self._decompose_single_qubit_operation(op, -1)
+            temp,
+            k=1,
+            rewriter=lambda op: self._decompose_single_qubit_operation(op, -1),
         ).all_operations()
 
+    def _decompose_multi_qubit_operation(
+        self, op: cirq.Operation, _
+    ) -> cirq.OP_TREE:
+        return NotImplemented
+
+    @property
+    def preprocess_transformers(self) -> List["cirq.TRANSFORMER"]:
+        """List of transformers which should be run before decomposing individual operations.
+
+        Decompose to three qubit gates because three qubit gates have different decomposition
+        for all-to-all connectivity between qubits.
+        """
+        return [
+            cirq.create_transformer_with_kwargs(
+                cirq.expand_composite,
+                no_decomp=lambda op: cirq.num_qubits(op) <= 3,
+            )
+        ]
+
+    @property
+    def postprocess_transformers(self) -> List["cirq.TRANSFORMER"]:
+        """List of transformers which should be run after decomposing individual operations."""
+        return [cirq.drop_negligible_operations, cirq.drop_empty_moments]
+
+    def __repr__(self) -> str:
+        return f"BenchmarkTargetGateset(atol={self.atol})"
+
+    def _value_equality_values_(self) -> Any:
+        return self.atol
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        return cirq.obj_to_dict_helper(self, ["atol"])
+
+    @classmethod
+    def _from_json_dict_(cls, atol, **kwargs):
+        return cls(atol=atol)
+
+
 # Cirq compilation
 def cirq_compile(cirq_circuit):
-    return optimize_for_target_gateset(cirq_circuit, gateset=BenchmarkTargetGateset())
+    return optimize_for_target_gateset(
+        cirq_circuit, gateset=BenchmarkTargetGateset()
+    )
 
 
 # Multi-qubit gate count for PyTkets