Plumb gate tabulation through optimized_for_sycamore (#2697)

 - Convert list of optimizers to functions that return configured
   optimizers
 - plumb through `tolerance` argument
 - plumb through `tabulation` for SYC gate decomposition
   (this should also work for sqrt_iswap but I haven't tested/enabled it)
 - use lru_cache to avoid re-tabulation

cirq/google/optimizers/__init__.py

@@ -14,6 +14,10 @@
 """
 Package for optimizers and gate compilers related to Google-specific devices.
 """
+from cirq.google.optimizers.two_qubit_gates import (
+    gate_product_tabulation,
+    GateTabulation,
+)
 
 from cirq.google.optimizers.convert_to_sycamore_gates import (
     ConvertToSycamoreGates,)

cirq/google/optimizers/optimize_for_sycamore.py

@@ -12,73 +12,123 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """A combination of several optimizations targeting XmonDevice."""
+from functools import lru_cache
 from typing import Callable, cast, List, Optional, TYPE_CHECKING
 
-from cirq import circuits, devices, optimizers
-from cirq.google.optimizers import (convert_to_xmon_gates,
-                                    ConvertToSycamoreGates,
-                                    ConvertToSqrtIswapGates)
+import numpy as np
+
+from cirq import circuits, devices, optimizers, protocols
+from cirq.google import ops as cg_ops
+from cirq.google.optimizers import (
+    convert_to_xmon_gates,
+    ConvertToSycamoreGates,
+    ConvertToSqrtIswapGates,
+    gate_product_tabulation,
+    GateTabulation,
+)
 
 if TYPE_CHECKING:
     import cirq
 
-_TOLERANCE = 1e-5
-
-
-def _merge_rots(c: 'cirq.Circuit'):
-    return optimizers.merge_single_qubit_gates_into_phased_x_z(c, _TOLERANCE)
-
-
-_XMON_OPTIMIZERS: List[Callable[['cirq.Circuit'], None]] = [
-    convert_to_xmon_gates.ConvertToXmonGates().optimize_circuit,
-    optimizers.MergeInteractions(tolerance=_TOLERANCE,
-                                 allow_partial_czs=False).optimize_circuit,
-    _merge_rots,
-    optimizers.EjectPhasedPaulis(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.EjectZ(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.DropNegligible(tolerance=_TOLERANCE).optimize_circuit,
-]
-
-_XMON_OPTIMIZERS_PART_CZ: List[Callable[['cirq.Circuit'], None]] = [
-    convert_to_xmon_gates.ConvertToXmonGates().optimize_circuit,
-    optimizers.MergeInteractions(tolerance=_TOLERANCE,
-                                 allow_partial_czs=True).optimize_circuit,
-    _merge_rots,
-    optimizers.EjectPhasedPaulis(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.EjectZ(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.DropNegligible(tolerance=_TOLERANCE).optimize_circuit,
-]
-
-_SYCAMORE_OPTIMIZERS = [
-    ConvertToSycamoreGates().optimize_circuit,
-    _merge_rots,
-    optimizers.EjectPhasedPaulis(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.EjectZ(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.DropNegligible(tolerance=_TOLERANCE).optimize_circuit,
-]  # type: List[Callable[[circuits.Circuit], None]]
-
-_SQRT_ISWAP_OPTIMIZERS = [
-    ConvertToSqrtIswapGates().optimize_circuit,
-    _merge_rots,
-    optimizers.EjectPhasedPaulis(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.EjectZ(tolerance=_TOLERANCE).optimize_circuit,
-    optimizers.DropNegligible(tolerance=_TOLERANCE).optimize_circuit,
-]  # type: List[Callable[[circuits.Circuit], None]]
+
+def _get_common_cleanup_optimizers(tolerance: float
+                                  ) -> List[Callable[['cirq.Circuit'], None]]:
+    return [
+        optimizers.EjectPhasedPaulis(tolerance=tolerance).optimize_circuit,
+        optimizers.EjectZ(tolerance=tolerance).optimize_circuit,
+        optimizers.DropNegligible(tolerance=tolerance).optimize_circuit,
+    ]
+
+
+def _get_xmon_optimizers(tolerance: float, tabulation: Optional[GateTabulation]
+                        ) -> List[Callable[['cirq.Circuit'], None]]:
+    if tabulation is not None:
+        # coverage: ignore
+        raise ValueError("Gate tabulation not supported for xmon")
+
+    return [
+        convert_to_xmon_gates.ConvertToXmonGates().optimize_circuit,
+        optimizers.MergeInteractions(tolerance=tolerance,
+                                     allow_partial_czs=False).optimize_circuit,
+        lambda c: optimizers.merge_single_qubit_gates_into_phased_x_z(
+            c, tolerance),
+        *_get_common_cleanup_optimizers(tolerance=tolerance),
+    ]
+
+
+def _get_xmon_optimizers_part_cz(tolerance: float,
+                                 tabulation: Optional[GateTabulation]
+                                ) -> List[Callable[['cirq.Circuit'], None]]:
+    if tabulation is not None:
+        # coverage: ignore
+        raise ValueError("Gate tabulation not supported for xmon")
+    return [
+        convert_to_xmon_gates.ConvertToXmonGates().optimize_circuit,
+        optimizers.MergeInteractions(tolerance=tolerance,
+                                     allow_partial_czs=True).optimize_circuit,
+        lambda c: optimizers.merge_single_qubit_gates_into_phased_x_z(
+            c, tolerance),
+        *_get_common_cleanup_optimizers(tolerance=tolerance),
+    ]
+
+
+def _get_sycamore_optimizers(tolerance: float,
+                             tabulation: Optional[GateTabulation]
+                            ) -> List[Callable[['cirq.Circuit'], None]]:
+    return [
+        ConvertToSycamoreGates(tabulation=tabulation).optimize_circuit,
+        lambda c: optimizers.merge_single_qubit_gates_into_phased_x_z(
+            c, tolerance),
+        *_get_common_cleanup_optimizers(tolerance=tolerance),
+    ]
+
+
+def _get_sqrt_iswap_optimizers(tolerance: float,
+                               tabulation: Optional[GateTabulation]
+                              ) -> List[Callable[['cirq.Circuit'], None]]:
+    if tabulation is not None:
+        # coverage: ignore
+        raise ValueError("Gate tabulation not supported for sqrt_iswap")
+    return [
+        ConvertToSqrtIswapGates().optimize_circuit,
+        lambda c: optimizers.merge_single_qubit_gates_into_phased_x_z(
+            c, tolerance),
+        *_get_common_cleanup_optimizers(tolerance=tolerance),
+    ]
+
 
 _OPTIMIZER_TYPES = {
-    'xmon': _XMON_OPTIMIZERS,
-    'xmon_partial_cz': _XMON_OPTIMIZERS_PART_CZ,
-    'sqrt_iswap': _SQRT_ISWAP_OPTIMIZERS,
-    'sycamore': _SYCAMORE_OPTIMIZERS,
+    'xmon': _get_xmon_optimizers,
+    'xmon_partial_cz': _get_xmon_optimizers_part_cz,
+    'sqrt_iswap': _get_sqrt_iswap_optimizers,
+    'sycamore': _get_sycamore_optimizers,
 }
 
 
+@lru_cache()
+def _gate_product_tabulation_cached(optimizer_type: str,
+                                    tabulation_resolution: float
+                                   ) -> GateTabulation:
+    random_state = np.random.RandomState(51)
+    if optimizer_type == 'sycamore':
+        return gate_product_tabulation(protocols.unitary(cg_ops.SYC),
+                                       tabulation_resolution,
+                                       random_state=random_state)
+    else:
+        raise NotImplementedError(
+            f"Gate tabulation not supported for {optimizer_type}")
+
+
 def optimized_for_sycamore(
         circuit: 'cirq.Circuit',
+        *,
         new_device: Optional['cirq.google.XmonDevice'] = None,
         qubit_map: Callable[['cirq.Qid'], devices.GridQubit] = lambda e: cast(
             devices.GridQubit, e),
-        optimizer_type: str = 'sqrt_iswap') -> 'cirq.Circuit':
+        optimizer_type: str = 'sqrt_iswap',
+        tolerance: float = 1e-5,
+        tabulation_resolution: Optional[float] = None,
+) -> 'cirq.Circuit':
     """Optimizes a circuit for Google devices.
 
     Uses a set of optimizers that will compile to the proper gateset for the
@@ -94,14 +144,27 @@ def optimized_for_sycamore(
         optimizer_type: A string defining the optimizations to apply.
             Possible values are  'xmon', 'xmon_partial_cz', 'sqrt_iswap',
             'sycamore'
+        tolerance: The tolerance passed to the various circuit optimization
+            passes.
+        tabulation_resolution: If provided, compute a gateset tabulation
+            with the specified resolution and use it to approximately
+            compile arbitrary two-qubit gates for which an analytic compilation
+            is not known.
     Returns:
         The optimized circuit.
     """
     copy = circuit.copy()
     if optimizer_type not in _OPTIMIZER_TYPES:
         raise ValueError(f'{optimizer_type} is not an allowed type.  Allowed '
                          f'types are: {_OPTIMIZER_TYPES.keys()}')
-    opts = _OPTIMIZER_TYPES[optimizer_type]
+
+    tabulation: Optional[GateTabulation] = None
+    if tabulation_resolution is not None:
+        tabulation = _gate_product_tabulation_cached(optimizer_type,
+                                                     tabulation_resolution)
+
+    opts = _OPTIMIZER_TYPES[optimizer_type](tolerance=tolerance,
+                                            tabulation=tabulation)
     for optimizer in opts:
         optimizer(copy)
 

cirq/google/optimizers/optimize_for_sycamore_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import pytest
 
 import cirq
@@ -21,6 +22,8 @@
 @pytest.mark.parametrize('optimizer_type, gateset', [
     ('sqrt_iswap', cg.SQRT_ISWAP_GATESET),
     ('sycamore', cg.SYC_GATESET),
+    ('xmon', cg.XMON),
+    ('xmon_partial_cz', cg.XMON),
 ])
 def test_optimizer_output_gates_are_supported(optimizer_type, gateset):
     q0, q1 = cirq.LineQubit.range(2)
@@ -41,3 +44,45 @@ def test_invalid_input():
                                cirq.X(q0)**0.2,
                                cirq.Z(q1)**0.2, cirq.measure(q0, q1, key='m'))
         _ = cg.optimized_for_sycamore(circuit, optimizer_type='for_tis_100')
+
+
+def test_tabulation():
+    q0, q1 = cirq.LineQubit.range(2)
+    u = cirq.testing.random_special_unitary(
+        4, random_state=np.random.RandomState(52))
+    circuit = cirq.Circuit(cirq.MatrixGate(u).on(q0, q1))
+    np.testing.assert_allclose(u, cirq.unitary(circuit))
+
+    circuit2 = cg.optimized_for_sycamore(circuit, optimizer_type='sycamore')
+    cirq.testing.assert_allclose_up_to_global_phase(u,
+                                                    cirq.unitary(circuit2),
+                                                    atol=1e-5)
+    assert len(circuit2) == 14
+
+    # sorry for the weak tolerances...
+    circuit3 = cg.optimized_for_sycamore(circuit,
+                                         optimizer_type='sycamore',
+                                         tabulation_resolution=0.01)
+    cirq.testing.assert_allclose_up_to_global_phase(u,
+                                                    cirq.unitary(circuit3),
+                                                    rtol=1e-1,
+                                                    atol=1e-1)
+    assert len(circuit3) == 8
+
+
+def test_no_tabulation():
+    circuit = cirq.Circuit(cirq.X(cirq.LineQubit(0)))
+    with pytest.raises(NotImplementedError):
+        cg.optimized_for_sycamore(circuit,
+                                  optimizer_type='sqrt_iswap',
+                                  tabulation_resolution=0.01)
+
+    with pytest.raises(NotImplementedError):
+        cg.optimized_for_sycamore(circuit,
+                                  optimizer_type='xmon',
+                                  tabulation_resolution=0.01)
+
+    with pytest.raises(NotImplementedError):
+        cg.optimized_for_sycamore(circuit,
+                                  optimizer_type='xmon_partial_cz',
+                                  tabulation_resolution=0.01)

cirq/google/optimizers/optimize_for_xmon.py

@@ -29,7 +29,12 @@ def optimized_for_xmon(
         allow_partial_czs: bool = False,
 ) -> 'cirq.Circuit':
     if allow_partial_czs:
-        return optimized_for_sycamore(circuit, new_device, qubit_map,
-                                      'xmon_partial_cz')
+        return optimized_for_sycamore(circuit,
+                                      new_device=new_device,
+                                      qubit_map=qubit_map,
+                                      optimizer_type='xmon_partial_cz')
     else:
-        return optimized_for_sycamore(circuit, new_device, qubit_map, 'xmon')
+        return optimized_for_sycamore(circuit,
+                                      new_device=new_device,
+                                      qubit_map=qubit_map,
+                                      optimizer_type='xmon')

cirq/google/optimizers/two_qubit_gates/__init__.py

@@ -0,0 +1,4 @@
+from cirq.google.optimizers.two_qubit_gates.gate_compilation import (
+    gate_product_tabulation,
+    GateTabulation,
+)