GPU custom kernels

src/qibo/gates.py

@@ -1,12 +1,6 @@
 from qibo.config import BACKEND_NAME
 if BACKEND_NAME == "tensorflow":
     import tensorflow as tf
-    if tf.config.list_physical_devices("GPU"):
-        # If GPU is available use Tensorflow gates
-        from qibo.tensorflow.gates import *
-    else:
-        # For CPU use custom operator gates
-        from qibo.tensorflow.cgates import *
-
+    from qibo.tensorflow.cgates import *
 else:
     raise NotImplementedError("Only Tensorflow backend is implemented.")

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

@@ -48,53 +48,76 @@ template <typename Device, typename T>
 struct BaseOneQubitGateFunctor {
   virtual void apply(T& state1, T& state2, const T* gate = NULL) const;
 
-  void operator()(const OpKernelContext* context, const Device& d,
-                  T* state,       //!< Total state vector.
-                  int nqubits,    //!< Total number of qubits in the state.
-                  int target,     //!< Target qubit id.
-                  int ncontrols,  //!< Number of qubits that the gate is controlled on.
-                  const int32* controls,  //!< List of control qubits ids sorted in decreasing order.
-                  const T* gate = NULL    //!< Gate matrix (used only by)
-                 ) const;
+  virtual void nocontrolwork(const Device& d, int numBlocks, int blockSize,
+                             T* state, const T* gate, long tk, int m) const;
+
+  virtual void singlecontrolwork(const Device& d, int numBlocks, int blockSize,
+                                 T* state, const T* gate, long tk,
+                                 long k1, long k2, int m1, int m2) const;
+
+  virtual void multicontrolwork(const Device& d, int numBlocks, int blockSize,
+                                T* state, const T* gate, long tk,
+                                int m, int ncontrols,
+                                const int* controls, int nqubits, int target) const;
+
+  void operator()(
+      const OpKernelContext* context, const Device& d,
+      T* state,       //!< Total state vector.
+      int nqubits,    //!< Total number of qubits in the state.
+      int target,     //!< Target qubit id.
+      int ncontrols,  //!< Number of qubits that the gate is controlled on.
+      const int32*
+          controls,  //!< List of control qubits ids sorted in decreasing order.
+      const int32* tensor_controls,  //!< List of control qubits ids sorted in
+                                     //!< decreasing order and stored on Device.
+      const T* gate = NULL           //!< Gate matrix (used only by)
+  ) const;
 };
 
 template <typename Device, typename T>
-struct ApplyGateFunctor: BaseOneQubitGateFunctor<Device, T> {};
+struct ApplyGateFunctor : BaseOneQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
-struct ApplyXFunctor: BaseOneQubitGateFunctor<Device, T> {};
+struct ApplyXFunctor : BaseOneQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
-struct ApplyYFunctor: BaseOneQubitGateFunctor<Device, T> {};
+struct ApplyYFunctor : BaseOneQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
-struct ApplyZFunctor: BaseOneQubitGateFunctor<Device, T> {};
+struct ApplyZFunctor : BaseOneQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
-struct ApplyZPowFunctor: BaseOneQubitGateFunctor<Device, T> {};
+struct ApplyZPowFunctor : BaseOneQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
 struct BaseTwoQubitGateFunctor {
   virtual void apply(T* state, int64 i, int64 tk1, int64 tk2,
                      const T* gate = NULL) const;
 
+  virtual void nocontrolwork(const Device& d, int numBlocks, int blockSize,
+                             T* state, const T* gate, long ctk1, long ctk2,
+                             long tk1, long tk2, int m1, int m2) const;
+
+  virtual void multicontrolwork(const Device& d, int numBlocks, int blockSize,
+                                T* state, const T* gate, long ctk1, long ctk2,
+                                long tk1, long tk2, int m1, int m2,
+                                int ncontrols, const int* controls, int nqubits,
+                                int t1, int t2) const;
+
   void operator()(const OpKernelContext* context, const Device& d, T* state,
-                  int nqubits,
-                  int target1,
-                  int target2,
-                  int ncontrols,
-                  const int32* controls,
-                  const T* gate = NULL);
+                  int nqubits, int target1, int target2, int ncontrols,
+                  const int32* controls, const int32* tensor_controls,
+                  const T* gate = NULL) const;
 };
 
 template <typename Device, typename T>
-struct ApplyTwoQubitGateFunctor: BaseTwoQubitGateFunctor<Device, T> {};
+struct ApplyTwoQubitGateFunctor : BaseTwoQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
-struct ApplyFsimFunctor: BaseTwoQubitGateFunctor<Device, T> {};
+struct ApplyFsimFunctor : BaseTwoQubitGateFunctor<Device, T> {};
 
 template <typename Device, typename T>
-struct ApplySwapFunctor: BaseTwoQubitGateFunctor<Device, T> {};
+struct ApplySwapFunctor : BaseTwoQubitGateFunctor<Device, T> {};
 
 }  // namespace functor