Merge pull request tensorflow#387 from verult/cirq-repetition-overhead

QSimSimulator: memoizing circuit translations

qsimcirq/qsim_simulator.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import deque
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from cirq import (
@@ -80,7 +81,10 @@ class QSimSimulator(
     SimulatesExpectationValues,
 ):
     def __init__(
-        self, qsim_options: dict = {}, seed: value.RANDOM_STATE_OR_SEED_LIKE = None
+        self,
+        qsim_options: dict = {},
+        seed: value.RANDOM_STATE_OR_SEED_LIKE = None,
+        circuit_memoization_size: int = 0,
     ):
         """Creates a new QSimSimulator using the given options and seed.
 
@@ -97,6 +101,14 @@ def __init__(
                 over when calculating expectation values for a noisy circuit.
                 Note that this does not apply to other simulation types.
             seed: A random state or seed object, as defined in cirq.value.
+            circuit_memoization_size: The number of last translated circuits
+                to be memoized from simulation executions, to eliminate
+                translation overhead. Every simulation will perform a linear
+                search through the list of memoized circuits using circuit
+                equality checks, so a large circuit_memoization_size with large
+                circuits will incur a significant runtime overhead.
+                Note that every resolved parameterization results in a separate
+                circuit to be memoized.
 
         Raises:
             ValueError if internal keys 'c', 'i' or 's' are included in 'qsim_options'.
@@ -109,6 +121,8 @@ def __init__(
         self._prng = value.parse_random_state(seed)
         self.qsim_options = {"t": 1, "f": 2, "v": 0, "r": 1}
         self.qsim_options.update(qsim_options)
+        # Deque of (<original cirq circuit>, <translated qsim circuit>) tuples.
+        self._translated_circuits = deque(maxlen=circuit_memoization_size)
 
     def get_seed(self):
         # Limit seed size to 32-bit integer for C++ conversion.
@@ -225,7 +239,12 @@ def _sample_measure_results(
                     else [ops.IdentityGate(1).on(q) for q in op.qubits]
                     for op in program.moments[i]
                 )
-            options["c"] = program.translate_cirq_to_qsim(ops.QubitOrder.DEFAULT)
+            translator_fn_name = "translate_cirq_to_qsim"
+            options["c"] = self._translate_circuit(
+                program,
+                translator_fn_name,
+                ops.QubitOrder.DEFAULT,
+            )
             options["s"] = self.get_seed()
             final_state = qsim.qsim_simulate_fullstate(options, 0)
             full_results = sim.sample_state_vector(
@@ -241,8 +260,11 @@ def _sample_measure_results(
                     for j, q in enumerate(meas_indices):
                         results[key][i][j] = full_results[i][q]
         else:
-            translator_fn = getattr(program, translator_fn_name)
-            options["c"] = translator_fn(ops.QubitOrder.DEFAULT)
+            options["c"] = self._translate_circuit(
+                program,
+                translator_fn_name,
+                ops.QubitOrder.DEFAULT,
+            )
             for i in range(repetitions):
                 options["s"] = self.get_seed()
                 measurements = sampler_fn(options)
@@ -304,8 +326,11 @@ def compute_amplitudes_sweep(
 
         for prs in param_resolvers:
             solved_circuit = protocols.resolve_parameters(program, prs)
-            translator_fn = getattr(solved_circuit, translator_fn_name)
-            options["c"] = translator_fn(cirq_order)
+            options["c"] = self._translate_circuit(
+                solved_circuit,
+                translator_fn_name,
+                cirq_order,
+            )
             options["s"] = self.get_seed()
             amplitudes = simulator_fn(options)
             trials_results.append(amplitudes)
@@ -380,8 +405,12 @@ def simulate_sweep(
 
         for prs in param_resolvers:
             solved_circuit = protocols.resolve_parameters(program, prs)
-            translator_fn = getattr(solved_circuit, translator_fn_name)
-            options["c"] = translator_fn(cirq_order)
+
+            options["c"] = self._translate_circuit(
+                solved_circuit,
+                translator_fn_name,
+                cirq_order,
+            )
             options["s"] = self.get_seed()
             qubit_map = {qubit: index for index, qubit in enumerate(qsim_order)}
 
@@ -503,8 +532,11 @@ def simulate_expectation_values_sweep(
 
         for prs in param_resolvers:
             solved_circuit = protocols.resolve_parameters(program, prs)
-            translator_fn = getattr(solved_circuit, translator_fn_name)
-            options["c"] = translator_fn(cirq_order)
+            options["c"] = self._translate_circuit(
+                solved_circuit,
+                translator_fn_name,
+                cirq_order,
+            )
             options["s"] = self.get_seed()
 
             if isinstance(initial_state, int):
@@ -514,3 +546,24 @@ def simulate_expectation_values_sweep(
             results.append(evs)
 
         return results
+
+    def _translate_circuit(
+        self,
+        circuit: Any,
+        translator_fn_name: str,
+        qubit_order: ops.QubitOrderOrList,
+    ):
+        # If the circuit is memoized, reuse the corresponding translated
+        # circuit.
+        translated_circuit = None
+        for original, translated in self._translated_circuits:
+            if original == circuit:
+                translated_circuit = translated
+                break
+
+        if translated_circuit is None:
+            translator_fn = getattr(circuit, translator_fn_name)
+            translated_circuit = translator_fn(qubit_order)
+            self._translated_circuits.append((circuit, translated_circuit))
+
+        return translated_circuit

qsimcirq_tests/qsimcirq_test.py

@@ -1376,3 +1376,152 @@ def test_cirq_qsim_global_shift():
     assert cirq.linalg.allclose_up_to_global_phase(
         qsim_result.state_vector(), cirq_result.state_vector()
     )
+
+
+@pytest.mark.parametrize("mode", ["noiseless", "noisy"])
+def test_cirq_qsim_circuit_memoization_compute_amplitudes(mode: str):
+    """Verifies the correctness of simulator functions when
+    circuit_memoization_size is set."""
+    execution_repetitions = 3
+    qsim_sim = qsimcirq.QSimSimulator(circuit_memoization_size=4)
+
+    # Pick qubits.
+    a, b, c, d = [
+        cirq.GridQubit(0, 0),
+        cirq.GridQubit(0, 1),
+        cirq.GridQubit(1, 1),
+        cirq.GridQubit(1, 0),
+    ]
+
+    # Create a circuit
+    cirq_circuit = cirq.Circuit(
+        cirq.X(a) ** 0.5,
+        cirq.Y(b) ** 0.5,
+        cirq.Z(c),
+        cirq.CZ(a, d),
+    )
+
+    if mode == "noisy":
+        cirq_circuit.append(NoiseTrigger().on(a))
+
+    for _ in range(execution_repetitions):
+        result = qsim_sim.compute_amplitudes(cirq_circuit, bitstrings=[0b0100, 0b1011])
+        assert np.allclose(result, [0.5j, 0j])
+
+
+@pytest.mark.parametrize("mode", ["noiseless", "noisy"])
+def test_cirq_qsim_circuit_memoization_simulate(mode: str):
+    execution_repetitions = 3
+    qsim_sim = qsimcirq.QSimSimulator(circuit_memoization_size=4)
+    cirq_sim = cirq.Simulator()
+
+    # Pick qubits.
+    a, b, c, d = [
+        cirq.GridQubit(0, 0),
+        cirq.GridQubit(0, 1),
+        cirq.GridQubit(1, 1),
+        cirq.GridQubit(1, 0),
+    ]
+
+    # Create a circuit.
+    cirq_circuit = cirq.Circuit(
+        cirq.Moment(
+            cirq.X(a) ** 0.5,
+            cirq.H(b),
+            cirq.X(c),
+            cirq.H(d),
+        ),
+        cirq.Moment(
+            cirq.X(a) ** 0.5,
+            cirq.CX(b, c),
+            cirq.S(d),
+        ),
+        cirq.Moment(
+            cirq.I(a),
+            cirq.ISWAP(b, c),
+        ),
+    )
+
+    if mode == "noisy":
+        cirq_circuit.append(NoiseTrigger().on(a))
+
+    cirq_result = cirq_sim.simulate(cirq_circuit, qubit_order=[a, b, c, d])
+    for _ in range(execution_repetitions):
+        result = qsim_sim.simulate(cirq_circuit, qubit_order=[a, b, c, d])
+        assert result.state_vector().shape == (16,)
+        assert cirq.linalg.allclose_up_to_global_phase(
+            result.state_vector(), cirq_result.state_vector()
+        )
+
+
+@pytest.mark.parametrize("mode", ["noiseless", "noisy"])
+def test_cirq_qsim_circuit_memoization_run(mode: str):
+    execution_repetitions = 3
+    qsim_sim = qsimcirq.QSimSimulator(circuit_memoization_size=4)
+
+    # Pick qubits.
+    a, b, c, d = [
+        cirq.GridQubit(0, 0),
+        cirq.GridQubit(0, 1),
+        cirq.GridQubit(1, 1),
+        cirq.GridQubit(1, 0),
+    ]
+
+    # Create a circuit
+    cirq_circuit = cirq.Circuit(
+        cirq.X(a) ** 0.5,
+        cirq.Y(b) ** 0.5,
+        cirq.Z(c),
+        cirq.CZ(a, d),
+        # measure qubits
+        cirq.measure(a, key="ma"),
+        cirq.measure(b, key="mb"),
+        cirq.measure(c, key="mc"),
+        cirq.measure(d, key="md"),
+    )
+    if mode == "noisy":
+        cirq_circuit.append(NoiseTrigger().on(a))
+
+    for _ in range(execution_repetitions):
+        result = qsim_sim.run(cirq_circuit, repetitions=5)
+        for key, value in result.measurements.items():
+            assert value.shape == (5, 1)
+
+
+@pytest.mark.parametrize("mode", ["noiseless", "noisy"])
+def test_cirq_qsim_circuit_memoization_simulate_expectation_values_sweep(mode: str):
+    execution_repetitions = 3
+    qsim_sim = qsimcirq.QSimSimulator(circuit_memoization_size=4)
+    cirq_sim = cirq.Simulator()
+
+    a, b = [cirq.GridQubit(0, 0), cirq.GridQubit(0, 1)]
+
+    x_exp = sympy.Symbol("x_exp")
+    h_exp = sympy.Symbol("h_exp")
+    circuit = cirq.Circuit(
+        cirq.X(a) ** x_exp,
+        cirq.H(b),
+        cirq.H(a) ** h_exp,
+        cirq.H(b) ** h_exp,
+    )
+    params = [
+        {x_exp: 0, h_exp: 0},  # |0+)
+        {x_exp: 1, h_exp: 0},  # |1+)
+        {x_exp: 0, h_exp: 1},  # |+0)
+        {x_exp: 1, h_exp: 1},  # |-0)
+    ]
+    psum1 = cirq.Z(a) + 3 * cirq.X(b)
+    psum2 = cirq.X(a) - 3 * cirq.Z(b)
+
+    if mode == "noisy":
+        circuit.append(NoiseTrigger().on(a))
+
+    cirq_result = cirq_sim.simulate_expectation_values_sweep(
+        circuit, [psum1, psum2], params
+    )
+
+    for _ in range(execution_repetitions):
+        qsim_result = qsim_sim.simulate_expectation_values_sweep(
+            circuit, [psum1, psum2], params
+        )
+        assert cirq.approx_eq(qsim_result, cirq_result, atol=1e-6)