diff --git a/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate.h b/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate.h index 57de4d0ff1..d4a7d2c209 100644 --- a/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate.h +++ b/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate.h @@ -10,11 +10,12 @@ namespace functor { template struct ApplyGateFunctor { void operator()(const OpKernelContext* context, const Device& d, T* state, - const T* gate, int nqubits, int target); + const T* gate, int nqubits, int target, + const int32* controls, int ncontrols); }; } // namespace functor } // namespace tensorflow -#endif // KERNEL_APPLY_GATE_H_ \ No newline at end of file +#endif // KERNEL_APPLY_GATE_H_ diff --git a/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate_kernels.cc b/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate_kernels.cc index cf1e655427..814e0d3e3a 100644 --- a/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate_kernels.cc +++ b/src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate_kernels.cc @@ -14,22 +14,42 @@ using thread::ThreadPool; template struct ApplyGateFunctor { void operator()(const OpKernelContext* context, const CPUDevice& d, T* state, - const T* gate, int nqubits, int target) { + const T* gate, int nqubits, int target, + const int32* controls, int ncontrols) { const int64 nstates = std::pow(2, nqubits); - const int64 k = std::pow(2, nqubits - target - 1); + const int64 tk = std::pow(2, nqubits - target - 1); + + int64 cktot = 0; + std::vector cks(ncontrols); + for (int i = 0; i < ncontrols; i++) { + cks[i] = std::pow(2, nqubits - controls[i] - 1); + cktot += cks[i]; + } auto DoWork = [&](int64 t, int64 w) { - for (auto g = t; g < w; g += 2 * k) { - for (auto i = g; i < g + k; i++) { - const auto buffer = state[i]; - state[i] = gate[0] * state[i] + gate[1] * state[i + k]; - state[i + k] = gate[2] * buffer + gate[3] * state[i + k]; + for (auto g = t; g < w; g += 2 * tk) { + for (auto i = g; i < g + tk; i++) { + bool apply = true; + for (const auto &q: cks) { + if (((int64) i / q) % 2) { + apply = false; + break; + } + } + + if (apply) { + const int64 i1 = i + cktot; + const int64 i2 = i1 + tk; + const auto buffer = state[i1]; + state[i1] = gate[0] * state[i1] + gate[1] * state[i2]; + state[i2] = gate[2] * buffer + gate[3] * state[i2]; + } } } }; const ThreadPool::SchedulingParams p( - ThreadPool::SchedulingStrategy::kFixedBlockSize, absl::nullopt, 2 * k); + ThreadPool::SchedulingStrategy::kFixedBlockSize, absl::nullopt, 2 * tk); auto thread_pool = context->device()->tensorflow_cpu_worker_threads()->workers; thread_pool->ParallelFor(nstates, p, DoWork); @@ -48,6 +68,8 @@ class ApplyGateOp : public OpKernel { // grabe the input tensor Tensor state = context->input(0); const Tensor& gate = context->input(1); + const Tensor& controls = context->input(2); + const int ncontrols = controls.flat().size(); // prevent running on GPU OP_REQUIRES( @@ -56,8 +78,11 @@ class ApplyGateOp : public OpKernel { // call the implementation ApplyGateFunctor()(context, context->eigen_device(), - state.flat().data(), gate.flat().data(), - nqubits_, target_); + state.flat().data(), + gate.flat().data(), + nqubits_, target_, + controls.flat().data(), + ncontrols); context->set_output(0, state); } diff --git a/src/qibo/tensorflow/custom_operators/cc/ops/custom_ops.cc b/src/qibo/tensorflow/custom_operators/cc/ops/custom_ops.cc index afea59c7a5..9e4284ce3b 100644 --- a/src/qibo/tensorflow/custom_operators/cc/ops/custom_ops.cc +++ b/src/qibo/tensorflow/custom_operators/cc/ops/custom_ops.cc @@ -16,10 +16,11 @@ REGISTER_OP("ApplyGate") .Attr("T: {complex64, complex128}") .Input("state: T") .Input("gate: T") + .Input("controls: int32") .Attr("nqubits: int") .Attr("target: int") .Output("out: T") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(0)); return Status::OK(); - }); \ No newline at end of file + }); diff --git a/src/qibo/tensorflow/custom_operators/python/ops/qibo_tf_custom_operators.py b/src/qibo/tensorflow/custom_operators/python/ops/qibo_tf_custom_operators.py index 0c9f96b67d..27c36b6e2a 100644 --- a/src/qibo/tensorflow/custom_operators/python/ops/qibo_tf_custom_operators.py +++ b/src/qibo/tensorflow/custom_operators/python/ops/qibo_tf_custom_operators.py @@ -9,4 +9,27 @@ initial_state = custom_module.initial_state # apply_gate operator -apply_gate = custom_module.apply_gate +def apply_gate(state, gate, nqubits, target, controls=[]): + """Applies one-qubit gates to a state vector. + + Modifies ``state`` in-place. + Gates can be controlled to multiple qubits. + + Args: + state (tf.Tensor): State vector of shape ``(2 ** nqubits,)``. + gate (tf.Tensor): Gate matrix of shape ``(2, 2)``. + nqubits (int): Total number of qubits in the state vector. + target (int): Qubit ID that the gate will act on. + Must be smaller than ``nqubits``. + controls (list): List with qubit IDs that the gate will be controlled on. + All qubit IDs must be smaller than ``nqubits``. + + Return: + state (tf.Tensor): State vector of shape ``(2 ** nqubits,)`` after + ``gate`` is applied. + """ + if not (isinstance(controls, list) or isinstance(controls, tuple)): + raise TypeError("Control qubits must be a list or tuple but {} " + "was given.".format(type(controls))) + + return custom_module.apply_gate(state, gate, controls, nqubits, target) diff --git a/src/qibo/tests/test_custom_operators.py b/src/qibo/tests/test_custom_operators.py index dd8f254ffa..86f93ef954 100644 --- a/src/qibo/tests/test_custom_operators.py +++ b/src/qibo/tests/test_custom_operators.py @@ -26,31 +26,86 @@ def apply_operator(dtype): np.testing.assert_allclose(final_state, exact_state) -@pytest.mark.parametrize(("nqubits", "target", "dtype", "compile"), - [(5, 4, np.float32, False), - (4, 2, np.float32, True), - (4, 2, np.float64, False), - (3, 0, np.float64, True), - (8, 5, np.float64, False)]) -def test_apply_gate(nqubits, target, dtype, compile): - """Check that `op.apply_gate` agrees with `tf.einsum`.""" +def tensorflow_random_complex(shape, dtype): + _re = tf.random.uniform(shape, dtype=dtype) + _im = tf.random.uniform(shape, dtype=dtype) + return tf.complex(_re, _im) + + +@pytest.mark.parametrize(("nqubits", "target", "dtype", "compile", "einsum_str"), + [(5, 4, np.float32, False, "abcde,Ee->abcdE"), + (4, 2, np.float32, True, "abcd,Cc->abCd"), + (4, 2, np.float64, False, "abcd,Cc->abCd"), + (3, 0, np.float64, True, "abc,Aa->Abc"), + (8, 5, np.float64, False, "abcdefgh,Ff->abcdeFgh")]) +def test_apply_gate(nqubits, target, dtype, compile, einsum_str): + """Check that ``op.apply_gate`` agrees with ``tf.einsum``.""" def apply_operator(state, gate): return op.apply_gate(state, gate, nqubits, target) - state = tf.complex(tf.random.uniform((2 ** nqubits,), dtype=dtype), - tf.random.uniform((2 ** nqubits,), dtype=dtype)) - gate = tf.complex(tf.random.uniform((2, 2), dtype=dtype), - tf.random.uniform((2, 2), dtype=dtype)) + state = tensorflow_random_complex((2 ** nqubits,), dtype) + gate = tensorflow_random_complex((2, 2), dtype) - einsum_str = {3: "abc,Aa->Abc", - 4: "abcd,Cc->abCd", - 5: "abcde,Ee->abcdE", - 8: "abcdefgh,Ff->abcdeFgh"} target_state = tf.reshape(state, nqubits * (2,)) - target_state = tf.einsum(einsum_str[nqubits], target_state, gate) + target_state = tf.einsum(einsum_str, target_state, gate) target_state = target_state.numpy().ravel() if compile: apply_operator = tf.function(apply_operator) state = apply_operator(state, gate) np.testing.assert_allclose(target_state, state.numpy(), atol=_atol) + + +@pytest.mark.parametrize(("nqubits", "compile"), + [(2, True), (3, False), (4, True), (5, False)]) +def test_apply_gate_cx(nqubits, compile): + """Check ``op.apply_gate`` for multiply-controlled X gates.""" + state = tensorflow_random_complex((2 ** nqubits,), dtype=tf.float64) + + target_state = state.numpy() + gate = np.eye(2 ** nqubits, dtype=target_state.dtype) + gate[-2, -2], gate[-2, -1] = 0, 1 + gate[-1, -2], gate[-1, -1] = 1, 0 + target_state = gate.dot(target_state) + + xgate = tf.cast([[0, 1], [1, 0]], dtype=state.dtype) + controls = list(range(nqubits - 1)) + def apply_operator(state): + return op.apply_gate(state, xgate, nqubits, nqubits - 1, controls) + if compile: + apply_operator = tf.function(apply_operator) + state = apply_operator(state) + + np.testing.assert_allclose(target_state, state.numpy()) + + +@pytest.mark.parametrize(("nqubits", "target", "controls", "einsum_str"), + [(3, 0, [1, 2], "a,Aa->A"), + (4, 3, [0, 1, 2], "a,Aa->A"), + (5, 3, [1], "abcd,Cc->abCd"), + (5, 2, [1, 4], "abc,Bb->aBc"), + (6, 3, [0, 2, 5], "abc,Bb->aBc"), + (6, 3, [0, 2, 4, 5], "ab,Bb->aB")]) +def test_apply_gate_controlled(nqubits, target, controls, einsum_str): + """Check ``op.apply_gate`` for random controlled gates.""" + state = tensorflow_random_complex((2 ** nqubits,), dtype=tf.float64) + gate = tensorflow_random_complex((2, 2), dtype=tf.float64) + + target_state = state.numpy().reshape(nqubits * (2,)) + slicer = nqubits * [slice(None)] + for c in controls: + slicer[c] = 1 + slicer = tuple(slicer) + target_state[slicer] = np.einsum(einsum_str, target_state[slicer], gate) + target_state = target_state.ravel() + + state = op.apply_gate(state, gate, nqubits, target, controls) + np.testing.assert_allclose(target_state, state.numpy()) + + +def test_apply_gate_error(): + """Check that ``TypeError`` is raised for invalid ``controls``.""" + state = tensorflow_random_complex((2 ** 2,), dtype=tf.float64) + gate = tensorflow_random_complex((2, 2), dtype=tf.float64) + with pytest.raises(TypeError): + state = op.apply_gate(state, gate, 2, 0, "a")