Multi-controlled gates in apply_gate

src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate.h

@@ -10,11 +10,12 @@ namespace functor {
 template <typename Device, typename T>
 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_
+#endif  // KERNEL_APPLY_GATE_H_

src/qibo/tensorflow/custom_operators/cc/kernels/apply_gate_kernels.cc

@@ -14,22 +14,42 @@ using thread::ThreadPool;
 template <typename T>
 struct ApplyGateFunctor<CPUDevice, T> {
   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<int64> 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<int32>().size();
 
     // prevent running on GPU
     OP_REQUIRES(
@@ -56,8 +78,11 @@ class ApplyGateOp : public OpKernel {
 
     // call the implementation
     ApplyGateFunctor<Device, T>()(context, context->eigen_device<Device>(),
-                                  state.flat<T>().data(), gate.flat<T>().data(),
-                                  nqubits_, target_);
+                                  state.flat<T>().data(),
+                                  gate.flat<T>().data(),
+                                  nqubits_, target_,
+                                  controls.flat<int32>().data(),
+                                  ncontrols);
 
     context->set_output(0, state);
   }

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();
-    });
+    });

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)

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")