diff --git a/WORKSPACE b/WORKSPACE index 92790b6fa..d68ca1dfe 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -3,11 +3,32 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") +EIGEN_COMMIT = "12e8d57108c50d8a63605c6eb0144c838c128337" +EIGEN_SHA256 = "f689246e342c3955af48d26ce74ac34d21b579a00675c341721a735937919b02" + + +http_archive( + name = "eigen", + build_file_content = """ +cc_library( + name = "eigen3", + textual_hdrs = glob(["Eigen/**", "unsupported/**"]), + visibility = ["//visibility:public"], +) + """, + sha256 = EIGEN_SHA256, + strip_prefix = "eigen-{commit}".format(commit = EIGEN_COMMIT), + urls = [ + "https://storage.googleapis.com/mirror.tensorflow.org/gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT), + "https://gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT), + ], +) + http_archive( name = "qsim", - sha256 = "d39b9c48866ce4d6a095093ae8059444d649e851219497af99e937a74f1e9a45", - strip_prefix = "qsim-0.9.2-dev-20210317", - urls = ["https://github.com/quantumlib/qsim/archive/v0.9.2-dev+20210317.zip"], + sha256 = "91eb09b2697accab9c0f64d5ea6ff482f4772ce000af867670b4efaf06a35224", + strip_prefix = "qsim-0.10.3-dev-20211001", + urls = ["https://github.com/quantumlib/qsim/archive/refs/tags/v0.10.3-dev+20211001.zip"], ) http_archive( diff --git a/release/BUILD b/release/BUILD index c63613e91..deeae5a42 100644 --- a/release/BUILD +++ b/release/BUILD @@ -39,8 +39,9 @@ sh_binary( "//tensorflow_quantum/core/ops:tfq_unitary_op_py", "//tensorflow_quantum/core/ops:tfq_utility_ops_py", "//tensorflow_quantum/core/ops:tfq_simulate_ops_py", - "//tensorflow_quantum/core/ops/math_ops:inner_product_op_py", "//tensorflow_quantum/core/ops/math_ops:fidelity_op_py", + "//tensorflow_quantum/core/ops/math_ops:inner_product_op_py", + "//tensorflow_quantum/core/ops/math_ops:simulate_mps_py", "//tensorflow_quantum/core/ops/noise:noisy_samples_op_py", "//tensorflow_quantum/core/ops/noise:noisy_expectation_op_py", "//tensorflow_quantum/core/ops/noise:noisy_sampled_expectation_op_py", diff --git a/scripts/import_test.py b/scripts/import_test.py index 9ee4ea8d1..48a4d4b8c 100644 --- a/scripts/import_test.py +++ b/scripts/import_test.py @@ -37,6 +37,7 @@ def test_imports(): # Math ops. _ = tfq.math.inner_product _ = tfq.math.fidelity + _ = tfq.math.mps_1d_expectation # Noisy simulation ops. _ = tfq.noise.expectation diff --git a/tensorflow_quantum/core/ops/math_ops/BUILD b/tensorflow_quantum/core/ops/math_ops/BUILD index 6585ec5b8..f62f47311 100644 --- a/tensorflow_quantum/core/ops/math_ops/BUILD +++ b/tensorflow_quantum/core/ops/math_ops/BUILD @@ -17,6 +17,7 @@ cc_binary( srcs = [ "tfq_inner_product.cc", "tfq_inner_product_grad.cc", + "tfq_simulate_1d_expectation.cc", ], copts = select({ ":windows": [ @@ -65,7 +66,10 @@ cc_binary( "//tensorflow_quantum/core/src:adj_util", "//tensorflow_quantum/core/src:circuit_parser_qsim", "//tensorflow_quantum/core/src:util_qsim", + "@qsim//lib:mps_simulator", + "@qsim//lib:mps_statespace", "@qsim//lib:qsim_lib", + "@eigen//:eigen3", # tensorflow core framework # tensorflow core lib # tensorflow core protos @@ -117,3 +121,22 @@ py_test( "//tensorflow_quantum/python:util", ], ) + +py_library( + name = "simulate_mps_py", + srcs = ["simulate_mps.py"], + data = [":_tfq_math_ops.so"], + deps = [ + "//tensorflow_quantum/core/ops:load_module", + ], +) + +py_test( + name = "simulate_mps_test", + srcs = ["simulate_mps_test.py"], + python_version = "PY3", + deps = [ + ":simulate_mps_py", + "//tensorflow_quantum/python:util", + ], +) diff --git a/tensorflow_quantum/core/ops/math_ops/__init__.py b/tensorflow_quantum/core/ops/math_ops/__init__.py index 14caf50eb..caf9e7648 100644 --- a/tensorflow_quantum/core/ops/math_ops/__init__.py +++ b/tensorflow_quantum/core/ops/math_ops/__init__.py @@ -14,5 +14,6 @@ # ============================================================================== """Module for tfq.core.ops.math_ops.*""" -from tensorflow_quantum.core.ops.math_ops.inner_product_op import inner_product from tensorflow_quantum.core.ops.math_ops.fidelity_op import fidelity +from tensorflow_quantum.core.ops.math_ops.inner_product_op import inner_product +from tensorflow_quantum.core.ops.math_ops.simulate_mps import mps_1d_expectation diff --git a/tensorflow_quantum/core/ops/math_ops/simulate_mps.py b/tensorflow_quantum/core/ops/math_ops/simulate_mps.py new file mode 100644 index 000000000..799fe38c0 --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/simulate_mps.py @@ -0,0 +1,57 @@ +# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Module to register MPS simulation ops.""" +import os +import tensorflow as tf +from tensorflow_quantum.core.ops.load_module import load_module + +MATH_OP_MODULE = load_module(os.path.join("math_ops", "_tfq_math_ops.so")) + + +def mps_1d_expectation(programs, + symbol_names, + symbol_values, + pauli_sums, + bond_dim=4): + """Calculate the expectation value of circuits wrt some operator(s) + + Args: + programs: `tf.Tensor` of strings with shape [batch_size] containing + the string representations of the circuits to be executed. + symbol_names: `tf.Tensor` of strings with shape [n_params], which + is used to specify the order in which the values in + `symbol_values` should be placed inside of the circuits in + `programs`. + symbol_values: `tf.Tensor` of real numbers with shape + [batch_size, n_params] specifying parameter values to resolve + into the circuits specificed by programs, following the ordering + dictated by `symbol_names`. + pauli_sums: `tf.Tensor` of strings with shape [batch_size, n_ops] + containing the string representation of the operators that will + be used on all of the circuits in the expectation calculations. + bond_dim: Integer value used for the bond dimension during simulation. + + Returns: + `tf.Tensor` with shape [batch_size, n_ops] that holds the + expectation value for each circuit with each op applied to it + (after resolving the corresponding parameters in). + """ + return MATH_OP_MODULE.tfq_simulate_mps1d_expectation(programs, + symbol_names, + tf.cast( + symbol_values, + tf.float32), + pauli_sums, + bond_dim=bond_dim) diff --git a/tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py b/tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py new file mode 100644 index 000000000..7682d1f9a --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py @@ -0,0 +1,370 @@ +# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests that specifically target simulate_mps.""" +# Remove PYTHONPATH collisions for protobuf. +# pylint: disable=wrong-import-position +import sys +NEW_PATH = [x for x in sys.path if 'com_google_protobuf' not in x] +sys.path = NEW_PATH +# pylint: enable=wrong-import-position + +import numpy as np +import tensorflow as tf +import cirq +import cirq_google +import sympy + +from tensorflow_quantum.core.ops.math_ops import simulate_mps +from tensorflow_quantum.python import util + + +class SimulateMPS1DExpectationTest(tf.test.TestCase): + """Tests mps_1d_expectation.""" + + def test_simulate_mps_1d_expectation_inputs(self): + """Makes sure that the op fails gracefully on bad inputs.""" + n_qubits = 5 + batch_size = 5 + symbol_names = ['alpha'] + qubits = cirq.GridQubit.rect(1, n_qubits) + circuit_batch = [ + cirq.Circuit( + cirq.X(qubits[0])**sympy.Symbol(symbol_names[0]), + cirq.Z(qubits[1]), + cirq.CNOT(qubits[2], qubits[3]), + cirq.Y(qubits[4])**sympy.Symbol(symbol_names[0]), + ) for _ in range(batch_size) + ] + resolver_batch = [{symbol_names[0]: 0.123} for _ in range(batch_size)] + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + + pauli_sums = util.random_pauli_sums(qubits, 3, batch_size) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'programs must be rank 1'): + # Circuit tensor has too many dimensions. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor([circuit_batch]), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbol_names must be rank 1.'): + # symbol_names tensor has too many dimensions. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), np.array([symbol_names]), + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbol_values must be rank 2.'): + # symbol_values_array tensor has too many dimensions. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + np.array([symbol_values_array]), + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbol_values must be rank 2.'): + # symbol_values_array tensor has too few dimensions. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array[0], + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'pauli_sums must be rank 2.'): + # pauli_sums tensor has too few dimensions. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, util.convert_to_tensor(list(pauli_sums))) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'pauli_sums must be rank 2.'): + # pauli_sums tensor has too many dimensions. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[[x]] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'Unparseable proto'): + # circuit tensor has the right type but invalid values. + simulate_mps.mps_1d_expectation( + ['junk'] * batch_size, symbol_names, symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'Could not find symbol in parameter map'): + # symbol_names tensor has the right type but invalid values. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), ['junk'], + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'qubits not found in circuit'): + # pauli_sums tensor has the right type but invalid values. + new_qubits = [cirq.GridQubit(5, 5), cirq.GridQubit(9, 9)] + new_pauli_sums = util.random_pauli_sums(new_qubits, 2, batch_size) + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in new_pauli_sums])) + + with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # circuits tensor has the wrong type. + simulate_mps.mps_1d_expectation( + [1.0] * batch_size, symbol_names, symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # symbol_names tensor has the wrong type. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), [0.1234], + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.UnimplementedError, ''): + # symbol_values tensor has the wrong type. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + [['junk']] * batch_size, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # pauli_sums tensor has the wrong type. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, [[1.0]] * batch_size) + + with self.assertRaisesRegex(TypeError, 'missing'): + # we are missing an argument. + # pylint: disable=no-value-for-parameter + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array) + # pylint: enable=no-value-for-parameter + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'at least minimum 4'): + # pylint: disable=too-many-function-args + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums]), 1) + + with self.assertRaisesRegex(TypeError, 'Expected int'): + # bond_dim should be int. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums]), []) + + with self.assertRaisesRegex(TypeError, 'positional arguments'): + # pylint: disable=too-many-function-args + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums]), 1, []) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='do not match'): + # wrong op size. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums + ][:int(batch_size * 0.5)])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='do not match'): + # wrong symbol_values size. + simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array[:int(batch_size * 0.5)], + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='cirq.Channel'): + # attempting to use noisy circuit. + noisy_circuit = cirq.Circuit(cirq.depolarize(0.3).on_each(*qubits)) + simulate_mps.mps_1d_expectation( + util.convert_to_tensor([noisy_circuit for _ in pauli_sums]), + symbol_names, symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='not in 1D topology'): + # attempting to use a circuit not in 1D topology + # 0--1--2--3 + # \-4 + circuit_not_1d = cirq.Circuit( + cirq.X(qubits[0])**sympy.Symbol(symbol_names[0]), + cirq.Z(qubits[1])**sympy.Symbol(symbol_names[0]), + cirq.CNOT(qubits[2], qubits[3]), + cirq.CNOT(qubits[2], qubits[4]), + ) + simulate_mps.mps_1d_expectation( + util.convert_to_tensor([circuit_not_1d for _ in pauli_sums]), + symbol_names, symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='not in 1D topology'): + # attempting to use a circuit in 1D topology, which looks in 2D. + # 0--1 + # \-2-\ + # 3--4 == 1--0--2--4--3 + circuit_not_1d = cirq.Circuit( + cirq.CNOT(qubits[0], qubits[1]), + cirq.CNOT(qubits[0], qubits[2]), + cirq.CNOT(qubits[2], qubits[4]), + cirq.CNOT(qubits[3], qubits[4]), + ) + simulate_mps.mps_1d_expectation( + util.convert_to_tensor([circuit_not_1d for _ in pauli_sums]), + symbol_names, symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='Found: 3 qubit gate'): + # attempting to use 3 qubit gate + three_qb_circuit = cirq.Circuit( + cirq.ISWAP(qubits[0], qubits[1]).controlled_by(qubits[2]), + cirq.X.on_each(*qubits)) + simulate_mps.mps_1d_expectation( + util.convert_to_tensor([three_qb_circuit for _ in pauli_sums]), + symbol_names, symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + + res = simulate_mps.mps_1d_expectation( + util.convert_to_tensor([cirq.Circuit() for _ in pauli_sums]), + symbol_names, symbol_values_array.astype(np.float64), + util.convert_to_tensor([[x] for x in pauli_sums])) + self.assertDTypeEqual(res, np.float32) + + def test_simulate_mps_1d_expectation_simple(self): + """Makes sure that the op shows the same result with Cirq.""" + n_qubits = 5 + batch_size = 5 + symbol_names = ['alpha'] + qubits = cirq.GridQubit.rect(1, n_qubits) + circuit_batch = [ + cirq.Circuit( + cirq.X(qubits[0])**sympy.Symbol(symbol_names[0]), + cirq.Z(qubits[1]), + cirq.CNOT(qubits[2], qubits[3]), + cirq.Y(qubits[4])**sympy.Symbol(symbol_names[0]), + ) for _ in range(batch_size) + ] + resolver_batch = [{symbol_names[0]: 0.123} for _ in range(batch_size)] + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + + pauli_sums = [ + cirq.Z(qubits[0]) * cirq.X(qubits[4]) for _ in range(batch_size) + ] + + cirq_result = [ + cirq.Simulator().simulate_expectation_values(c, p, r) + for c, p, r in zip(circuit_batch, pauli_sums, resolver_batch) + ] + # Default bond_dim=4 + mps_result = simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in pauli_sums])) + # Expected value of 0.349... + self.assertAllClose(mps_result, cirq_result) + + def _make_1d_circuit(self, qubits, depth): + """Create a 1d ladder circuit.""" + even_pairs = list(zip(qubits[::2], qubits[1::2])) + odd_pairs = list(zip(qubits[1::2], qubits[2::2])) + ret = cirq.Circuit() + + for _ in range(depth): + # return ret + ret += [(cirq.Y(q)**np.random.random()) for q in qubits] + ret += [ + cirq_google.SycamoreGate()(q0, q1)**np.random.random() + for q0, q1 in even_pairs + ] + ret += [(cirq.Y(q)**np.random.random()) for q in qubits] + ret += [ + cirq_google.SycamoreGate()(q1, q0)**np.random.random() + for q0, q1 in odd_pairs + ] + + return ret + + def test_complex_equality(self): + """Check moderate sized 1d random circuits.""" + batch_size = 10 + qubits = cirq.GridQubit.rect(1, 8) + circuit_batch = [ + self._make_1d_circuit(qubits, 3) for _ in range(batch_size) + ] + + pauli_sums = [[ + cirq.Z(qubits[0]), + cirq.Z(qubits[-1]), + cirq.Z(qubits[0]) * cirq.Z(qubits[-1]), + cirq.Z(qubits[0]) + cirq.Z(qubits[-1]) + ] for _ in range(batch_size)] + symbol_names = [] + resolver_batch = [{} for _ in range(batch_size)] + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + + cirq_result = [ + cirq.Simulator().simulate_expectation_values(c, p, r) + for c, p, r in zip(circuit_batch, pauli_sums, resolver_batch) + ] + mps_result = simulate_mps.mps_1d_expectation( + util.convert_to_tensor(circuit_batch), + symbol_names, + symbol_values_array, + util.convert_to_tensor(pauli_sums), + bond_dim=32) + self.assertAllClose(mps_result, cirq_result, atol=1e-5) + + def test_correctness_empty(self): + """Tests the mps op with empty circuits.""" + + empty_circuit = tf.raw_ops.Empty(shape=(0,), dtype=tf.string) + empty_symbols = tf.raw_ops.Empty(shape=(0,), dtype=tf.string) + empty_values = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.float32) + empty_paulis = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.string) + + out = simulate_mps.mps_1d_expectation(empty_circuit, empty_symbols, + empty_values, empty_paulis, 32) + + self.assertShapeEqual(np.zeros((0, 0)), out) + + +if __name__ == "__main__": + tf.test.main() diff --git a/tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc b/tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc new file mode 100644 index 000000000..03aecae7c --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc @@ -0,0 +1,245 @@ +/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include + +#include "../qsim/lib/circuit.h" +#include "../qsim/lib/formux.h" +#include "../qsim/lib/gate_appl.h" +#include "../qsim/lib/gates_cirq.h" +#include "../qsim/lib/mps_simulator.h" +#include "../qsim/lib/mps_statespace.h" +#include "../qsim/lib/seqfor.h" +#include "../qsim/lib/simmux.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/lib/core/error_codes.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/core/threadpool.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow_quantum/core/ops/parse_context.h" +#include "tensorflow_quantum/core/proto/pauli_sum.pb.h" +#include "tensorflow_quantum/core/proto/program.pb.h" +#include "tensorflow_quantum/core/src/program_resolution.h" +#include "tensorflow_quantum/core/src/util_qsim.h" + +namespace tfq { + +using ::tensorflow::Status; +using ::tfq::proto::PauliSum; +using ::tfq::proto::Program; + +typedef qsim::Cirq::GateCirq QsimGate; +typedef qsim::Circuit QsimCircuit; + +class TfqSimulateMPS1DExpectationOp : public tensorflow::OpKernel { + public: + explicit TfqSimulateMPS1DExpectationOp( + tensorflow::OpKernelConstruction* context) + : OpKernel(context) { + // Get the bond dimension of MPS + // Checked that bond_dim is a positive integer >= 2 by QSim definition. + OP_REQUIRES_OK(context, context->GetAttr("bond_dim", &bond_dim_)); + } + + void Compute(tensorflow::OpKernelContext* context) override { + // TODO (mbbrough): add more dimension checks for other inputs here. + const int num_inputs = context->num_inputs(); + OP_REQUIRES(context, num_inputs == 4, + tensorflow::errors::InvalidArgument(absl::StrCat( + "Expected 4 inputs, got ", num_inputs, " inputs."))); + + // Create the output Tensor. + const int output_dim_batch_size = context->input(0).dim_size(0); + const int output_dim_op_size = context->input(3).dim_size(1); + tensorflow::TensorShape output_shape; + output_shape.AddDim(output_dim_batch_size); + output_shape.AddDim(output_dim_op_size); + + tensorflow::Tensor* output = nullptr; + OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); + auto output_tensor = output->matrix(); + + // Parse program protos. + std::vector programs; + std::vector num_qubits; + std::vector> pauli_sums; + + // TODO: remove endianness workaround introduced here: + // https://github.com/tensorflow/quantum/pull/610 + // once https://github.com/quantumlib/qsim/issues/492 + // is resolved. + OP_REQUIRES_OK(context, + GetProgramsAndNumQubits(context, &programs, &num_qubits, + &pauli_sums, true)); + + std::vector maps; + OP_REQUIRES_OK(context, GetSymbolMaps(context, &maps)); + + OP_REQUIRES(context, programs.size() == maps.size(), + tensorflow::errors::InvalidArgument(absl::StrCat( + "Number of circuits and symbol_values do not match. Got ", + programs.size(), " circuits and ", maps.size(), + " symbol values."))); + + // Construct qsim circuits. + std::vector qsim_circuits(programs.size(), QsimCircuit()); + std::vector fused_circuits(programs.size(), + QsimFusedCircuit({})); + Status parse_status = Status::OK(); + auto p_lock = tensorflow::mutex(); + auto construct_f = [&](int start, int end) { + for (int i = start; i < end; i++) { + Status local = + QsimCircuitFromProgram(programs[i], maps[i], num_qubits[i], + &qsim_circuits[i], &fused_circuits[i]); + // If parsing works, check MPS constraints. + if (local.ok()) { + local = CheckMPSSupported(programs[i]); + } + NESTED_FN_STATUS_SYNC(parse_status, local, p_lock); + } + }; + + const int num_cycles = 1000; + context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( + output_dim_batch_size, num_cycles, construct_f); + OP_REQUIRES_OK(context, parse_status); + + int max_num_qubits = 0; + for (const int num : num_qubits) { + max_num_qubits = std::max(max_num_qubits, num); + } + + // Since MPS simulations have much smaller memory footprint, + // we do not need a ComputeLarge like we do for state vector simulation. + ComputeSmall(num_qubits, max_num_qubits, qsim_circuits, pauli_sums, context, + &output_tensor); + } + + private: + int bond_dim_; + + void ComputeSmall(const std::vector& num_qubits, + const int max_num_qubits, + const std::vector& unfused_circuits, + const std::vector>& pauli_sums, + tensorflow::OpKernelContext* context, + tensorflow::TTypes::Matrix* output_tensor) { + using Simulator = qsim::mps::MPSSimulator; + using StateSpace = Simulator::MPSStateSpace_; + + const int output_dim_op_size = output_tensor->dimension(1); + + Status compute_status = Status::OK(); + auto c_lock = tensorflow::mutex(); + auto DoWork = [&](int start, int end) { + int old_batch_index = -2; + int cur_batch_index = -1; + int largest_nq = 1; + int cur_op_index; + + // Note: ForArgs in MPSSimulator and MPSStateState are currently unused. + // So, this 1 is a dummy for qsim::For. + Simulator sim = Simulator(1); + StateSpace ss = StateSpace(1); + auto sv = ss.Create(largest_nq, bond_dim_); + auto scratch = ss.Create(largest_nq, bond_dim_); + for (int i = start; i < end; i++) { + cur_batch_index = i / output_dim_op_size; + cur_op_index = i % output_dim_op_size; + + const int nq = num_qubits[cur_batch_index]; + + // (#679) Just ignore empty program + auto unfused_gates = unfused_circuits[cur_batch_index].gates; + if (unfused_gates.size() == 0) { + (*output_tensor)(cur_batch_index, cur_op_index) = -2.0; + continue; + } + + if (cur_batch_index != old_batch_index) { + // We've run into a new state vector we must compute. + // Only compute a new state vector when we have to. + if (nq > largest_nq) { + largest_nq = nq; + sv = ss.Create(largest_nq, bond_dim_); + scratch = ss.Create(largest_nq, bond_dim_); + } + // no need to update scratch_state since ComputeExpectationMPSQsim + // will take care of things for us. + ss.SetStateZero(sv); + for (auto gate : unfused_gates) { + // Can't fuse, since this might break nearest neighbor constraints. + qsim::ApplyGate(sim, gate, sv); + } + } + + // Compute expectation values without fusing gates. + float exp_v = 0.0; + NESTED_FN_STATUS_SYNC( + compute_status, + ComputeExpectationQsim(pauli_sums[cur_batch_index][cur_op_index], + sim, ss, sv, scratch, &exp_v, false), + c_lock); + (*output_tensor)(cur_batch_index, cur_op_index) = exp_v; + old_batch_index = cur_batch_index; + } + }; + + const int64_t num_cycles = + 200 * (int64_t(1) << static_cast(max_num_qubits)); + context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( + unfused_circuits.size() * output_dim_op_size, num_cycles, DoWork); + OP_REQUIRES_OK(context, compute_status); + } +}; + +REGISTER_KERNEL_BUILDER( + Name("TfqSimulateMPS1DExpectation").Device(tensorflow::DEVICE_CPU), + TfqSimulateMPS1DExpectationOp); + +REGISTER_OP("TfqSimulateMPS1DExpectation") + .Input("programs: string") + .Input("symbol_names: string") + .Input("symbol_values: float") + .Input("pauli_sums: string") + .Output("expectations: float") + .Attr("bond_dim: int >= 4 = 4") + .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) { + tensorflow::shape_inference::ShapeHandle programs_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &programs_shape)); + + tensorflow::shape_inference::ShapeHandle symbol_names_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &symbol_names_shape)); + + tensorflow::shape_inference::ShapeHandle symbol_values_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 2, &symbol_values_shape)); + + tensorflow::shape_inference::ShapeHandle pauli_sums_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 2, &pauli_sums_shape)); + + tensorflow::shape_inference::DimensionHandle output_rows = + c->Dim(programs_shape, 0); + tensorflow::shape_inference::DimensionHandle output_cols = + c->Dim(pauli_sums_shape, 1); + c->set_output(0, c->Matrix(output_rows, output_cols)); + + return tensorflow::Status::OK(); + }); + +} // namespace tfq diff --git a/tensorflow_quantum/core/ops/parse_context.cc b/tensorflow_quantum/core/ops/parse_context.cc index 62404d421..dd8f941fc 100644 --- a/tensorflow_quantum/core/ops/parse_context.cc +++ b/tensorflow_quantum/core/ops/parse_context.cc @@ -150,7 +150,8 @@ Status GetProgramsAndProgramsToAppend( Status GetProgramsAndNumQubits( OpKernelContext* context, std::vector* programs, std::vector* num_qubits, - std::vector>* p_sums /*=nullptr*/) { + std::vector>* p_sums /*=nullptr*/, + bool swap_endianness /*=false*/) { // 1. Parse input programs // 2. (Optional) Parse input PauliSums // 3. Convert GridQubit locations to integers. @@ -180,10 +181,12 @@ Status GetProgramsAndNumQubits( Program& program = (*programs)[i]; unsigned int this_num_qubits; if (p_sums) { - OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits, - &(p_sums->at(i)))); + OP_REQUIRES_OK(context, + ResolveQubitIds(&program, &this_num_qubits, + &(p_sums->at(i)), swap_endianness)); } else { - OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits)); + OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits, + nullptr, swap_endianness)); } (*num_qubits)[i] = this_num_qubits; } diff --git a/tensorflow_quantum/core/ops/parse_context.h b/tensorflow_quantum/core/ops/parse_context.h index 4cc214385..c811b68c5 100644 --- a/tensorflow_quantum/core/ops/parse_context.h +++ b/tensorflow_quantum/core/ops/parse_context.h @@ -64,7 +64,8 @@ typedef absl::flat_hash_map> SymbolMap; tensorflow::Status GetProgramsAndNumQubits( tensorflow::OpKernelContext* context, std::vector* programs, std::vector* num_qubits, - std::vector>* p_sums = nullptr); + std::vector>* p_sums = nullptr, + bool swap_endianness = false); // Parses Cirq Program protos out of the 'circuit_specs' input Tensor. Also // resolves the QubitIds inside of the Program. This override also parses and diff --git a/tensorflow_quantum/core/src/program_resolution.cc b/tensorflow_quantum/core/src/program_resolution.cc index 4a8c39ea0..82af94d35 100644 --- a/tensorflow_quantum/core/src/program_resolution.cc +++ b/tensorflow_quantum/core/src/program_resolution.cc @@ -83,7 +83,8 @@ Status RegisterQubits( } Status ResolveQubitIds(Program* program, unsigned int* num_qubits, - std::vector* p_sums /*=nullptr*/) { + std::vector* p_sums /*=nullptr*/, + bool swap_endianness /*=false*/) { if (program->circuit().moments().empty()) { // (#679) Just ignore empty program. // Number of qubits in empty programs is zero. @@ -116,6 +117,10 @@ Status ResolveQubitIds(Program* program, unsigned int* num_qubits, id_set.end()); std::sort(ids.begin(), ids.end()); + // reverse endian. + if (swap_endianness) { + std::reverse(ids.begin(), ids.end()); + } absl::flat_hash_map id_to_index; for (size_t i = 0; i < ids.size(); i++) { id_to_index[ids[i].second] = absl::StrCat(i); @@ -128,6 +133,11 @@ Status ResolveQubitIds(Program* program, unsigned int* num_qubits, for (Qubit& qubit : *operation.mutable_qubits()) { qubit.set_id(id_to_index.at(qubit.id())); } + // reverse endian. + if (swap_endianness) { + std::reverse(operation.mutable_qubits()->begin(), + operation.mutable_qubits()->end()); + } // Resolve control qubit ids found in the control_qubits arg. absl::string_view control_qubits = operation.args().at("control_qubits").arg_value().string_value(); @@ -318,4 +328,59 @@ Status ResolveSymbols( return Status::OK(); } +Status CheckMPSSupported(const Program& program) { + // Check if (1) there are only 1-qubit or 2-qubit gates. + // (2) each two qubit gate has neighbor qubits only. + // + // Requires: program have qubit ids resolved. + if (program.circuit().moments().empty()) { + return Status::OK(); + } + + for (auto moment : program.circuit().moments()) { + for (auto operation : moment.operations()) { + // Count the number of qubits in this operation. + auto qs = operation.qubits(); + std::vector qubits(qs.begin(), qs.end()); + std::vector control_ids({}); + + if (operation.args().find("control_qubits") != operation.args().end()) { + absl::string_view control_qubits = + operation.args().at("control_qubits").arg_value().string_value(); + if (!control_qubits.empty()) { + control_ids = absl::StrSplit(control_qubits, ','); + } + } + const int total_num_qubits = qubits.size() + control_ids.size(); + if (total_num_qubits > 2) { + return Status( + tensorflow::error::INVALID_ARGUMENT, + absl::StrCat("1D operations only support 1 and 2 qubit gates. " + "Found: ", + total_num_qubits, " qubit gate.")); + } + + if (total_num_qubits == 2) { + int j = 0; + std::vector qids(2, -1234); + for (; j < qubits.size(); j++) { + (void)absl::SimpleAtoi(qubits[j].id(), &qids[j]); + } + for (; j < 2; j++) { + (void)absl::SimpleAtoi(control_ids[j], &qids[j]); + } + + // Are the two qubits not neighbors? + if (std::abs((int)qids[0] - (int)qids[1]) > 1) { + return Status(tensorflow::error::INVALID_ARGUMENT, + "A program is not in 1D topology. It contains an" + " operation with qubits not neighbors each other."); + } + } + } + } + + return Status::OK(); +} + } // namespace tfq diff --git a/tensorflow_quantum/core/src/program_resolution.h b/tensorflow_quantum/core/src/program_resolution.h index 6889401b2..40d5760dd 100644 --- a/tensorflow_quantum/core/src/program_resolution.h +++ b/tensorflow_quantum/core/src/program_resolution.h @@ -37,7 +37,8 @@ namespace tfq { // The number of qubits in the program is recorded in `num_qubits`. tensorflow::Status ResolveQubitIds( tfq::proto::Program* program, unsigned int* num_qubits, - std::vector* p_sums = nullptr); + std::vector* p_sums = nullptr, + bool swap_endianness = false); // Overload which allows for strict resolution of multiple programs. // Will resolve GridQubits in `program` and then double check that @@ -58,6 +59,9 @@ tensorflow::Status ResolveSymbols( const absl::flat_hash_map>& param_map, tfq::proto::Program* program, bool resolve_all = true); +// Checks if the qubits are in 1D topology. +tensorflow::Status CheckMPSSupported(const tfq::proto::Program& program); + } // namespace tfq #endif // TFQ_CORE_SRC_PROGRAM_RESOLUTION diff --git a/tensorflow_quantum/core/src/program_resolution_test.cc b/tensorflow_quantum/core/src/program_resolution_test.cc index 3cf92ae85..74a200b13 100644 --- a/tensorflow_quantum/core/src/program_resolution_test.cc +++ b/tensorflow_quantum/core/src/program_resolution_test.cc @@ -126,6 +126,61 @@ const std::string valid_symbol_program = R"( } )"; +const std::string three_qubit_op_program = R"( + circuit { + moments { + operations { + qubits { + id: "0_0" + } + qubits { + id: "0_1" + } + qubits { + id: "0_2" + } + } + } + } +)"; + +/* Qubit topology: + 1 -- 0 -- 2 + | + | + 3 +*/ +const std::string resolved_qubit_program_not_1d = R"( + circuit { + moments { + operations { + qubits { + id: "0" + } + qubits { + id: "1" + } + } + operations { + qubits { + id: "0" + } + qubits { + id: "2" + } + } + operations { + qubits { + id: "0" + } + qubits { + id: "3" + } + } + } + } +)"; + TEST(ProgramResolutionTest, ResolveQubitIdsValid) { Program program; unsigned int qubit_count; @@ -512,5 +567,41 @@ TEST(ProgramResolutionTest, ResolveSymbolsStrictFull) { 2.0); } +TEST(ProgramResolutionTest, CheckMPSSupportedEmpty) { + Program empty; + EXPECT_EQ(CheckMPSSupported(empty), Status::OK()); +} + +TEST(ProgramResolutionTest, CheckQubitsIn1DFailedByOpWithMoreThan2Qubits) { + Program program_with_3qubit_op; + ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString( + three_qubit_op_program, &program_with_3qubit_op)); + EXPECT_EQ(CheckMPSSupported(program_with_3qubit_op), + Status(tensorflow::error::INVALID_ARGUMENT, + "1D operations only support 1 and 2 qubit gates. " + "Found: 3 qubit gate.")); +} + +TEST(ProgramResolutionTest, + CheckQubitsIn1DFailedByOpWithMoreThan2QubitsOnControlQubits) { + Program program_with_3qubit_op; + ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString( + valid_program, &program_with_3qubit_op)); + EXPECT_EQ(CheckMPSSupported(program_with_3qubit_op), + Status(tensorflow::error::INVALID_ARGUMENT, + "1D operations only support 1 and 2 qubit gates. " + "Found: 3 qubit gate.")); +} + +TEST(ProgramResolutionTest, CheckQubitsIn1DFailedByNot1DTopology) { + Program program_not_1d; + ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString( + resolved_qubit_program_not_1d, &program_not_1d)); + EXPECT_EQ(CheckMPSSupported(program_not_1d), + Status(tensorflow::error::INVALID_ARGUMENT, + "A program is not in 1D topology. It contains an" + " operation with qubits not neighbors each other.")); +} + } // namespace } // namespace tfq diff --git a/tensorflow_quantum/core/src/util_qsim.h b/tensorflow_quantum/core/src/util_qsim.h index 933f1582d..f5f9c7285 100644 --- a/tensorflow_quantum/core/src/util_qsim.h +++ b/tensorflow_quantum/core/src/util_qsim.h @@ -142,7 +142,8 @@ tensorflow::Status ComputeExpectationQsim(const tfq::proto::PauliSum& p_sum, const SimT& sim, const StateSpaceT& ss, StateT& state, StateT& scratch, - float* expectation_value) { + float* expectation_value, + bool fuse_paulis = true) { // apply the gates of the pauliterms to a copy of the state vector // and add up expectation value term by term. tensorflow::Status status = tensorflow::Status::OK(); @@ -165,8 +166,14 @@ tensorflow::Status ComputeExpectationQsim(const tfq::proto::PauliSum& p_sum, } // copy from src to scratch. ss.Copy(state, scratch); - for (const qsim::GateFused& fused_gate : fused_circuit) { - qsim::ApplyFusedGate(sim, fused_gate, scratch); + if (fuse_paulis) { + for (const qsim::GateFused& fused_gate : fused_circuit) { + qsim::ApplyFusedGate(sim, fused_gate, scratch); + } + } else { + for (const auto& unfused_gate : main_circuit.gates) { + qsim::ApplyGate(sim, unfused_gate, scratch); + } } if (!status.ok()) {