如何让Eigen代码跑在GPU/线程池上1

前提

虽说Eigen是一个优秀的并行计算库，但是只跑在单线程环境下在牛鼻和你我写出来的代码也没啥区别。本文需要读者了解GPU并行计算的基础。(这里针对的是Eigen/Tensor，非unsupported中的类可以通过显示的host to device启动cuda核函数)。

先看会源码

大伙先来看源码：
有句话说的好，源码之前了无秘密
/usr/local/include/unsupported/Eigen/CXX11
这个文件夹下的头文件将满足你的所有需求(嘿嘿，起飞咯)
先看一下Tensor头文件：

#ifdef EIGEN_USE_THREADS
#include "ThreadPool"
#endif

#ifdef EIGEN_USE_GPU
#include <iostream>
#include <cuda_runtime.h>
#if __cplusplus >= 201103L
#include <atomic>
#include <unistd.h>
#endif
#endif

这里面涉及了两个重要的宏定义EIGEN_USE_THREADS、EIGEN_USE_GPU
顾名思义我就不多解释了。想要对应的功能就在#include<…Tensor>前先定义一下对应的宏（这里面涉及的头文件错综复杂好在Eigen使用宏定义帮我们理好了，我之前没注意就走了好多弯路）。
有了这些前提就可以正式使用Eigen封装好的ThreadPool/GpuStreamDevice来编程了。

线程池和设备

1. 老规矩，先看源码学会怎么使用再说
   细心的同学应该在Tensor头文件内看到#include "ThreadPool"这个头文件的声明了，所有的秘密都在里面,不细心也没关系，我来带你看。

这里举NonBlockingThreadPool
src/ThreadPool/NonBlockingThreadPool.h:
这里简单说明一下这个玩意是啥，很容易看出来这是Eigen设计的一个非阻塞线程池。它接受num_threads和env作为构造函数参数,前面的num_threads很容易理解，就是线程池的容量，后面那个是初始化的环境。看我列出来的代码最后Eigen已经给我们封装好了一个标准环境，编程的时候直接用那个type就行了。

namespace Eigen {

template <typename Environment>
class NonBlockingThreadPoolTempl : public Eigen::ThreadPoolInterface {
 public:
  typedef typename Environment::Task Task;
  typedef RunQueue<Task, 1024> Queue;

  NonBlockingThreadPoolTempl(int num_threads, Environment env = Environment())
      : env_(env),
        threads_(num_threads),
        queues_(num_threads),
        coprimes_(num_threads),
        waiters_(num_threads),
        blocked_(0),
        spinning_(0),
        done_(false),
        ec_(waiters_) {
    waiters_.resize(num_threads);
    ......
    .....
    typedef NonBlockingThreadPoolTempl<StlThreadEnvironment> NonBlockingThreadPool;

有了线程池这个利器后离正式编程还差一点点。
2. 看一下src/Tensor/TensorDeviceThreadPool.h

struct ThreadPoolDevice {
  // The ownership of the thread pool remains with the caller.
  ThreadPoolDevice(ThreadPoolInterface* pool, int num_cores) : pool_(pool), num_threads_(num_cores) { }

这里的ThreadPoolDevice就用到上面的线程池了，这个类主要管理Tensor对象的内存de/allocate等设备信息。它的构造函数就接受一个ThreadPoolInterface基类指针，使用NonBlocking/sample的话就传对应的对象地址进取就行了。num_cores就是同时运行的线程数了（小于等于线程池保管的线程数）

例子

有了上面基础就可以写一些简单的代码了：
Tensorflow1.15中AdaMax，CPU的实现（这里用NonBlockingThreadPool实现）
上来就王炸，surprise。
下一篇写一下同样的代码怎么跑在GPU上，不过是很类似的，毕竟用现成的库没啥技术含量。以后有机会也写写纯CUDA的实现。

#include<iostream>
#define EIGEN_USE_THREADS
#define EIGEN_TEST_NO_LONGDOUBLE
#define EIGEN_TEST_NO_COMPLEX
#define EIGEN_TEST_FUNC cxx11_tensor_cuda
#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
#define EIGEN_USE_GPU
#include<unsupported/Eigen/CXX11/Tensor>
#define N 40000
//#define CPU
using namespace Eigen;
typedef Eigen::CudaStreamDevice GPUStream;
typedef Eigen::GpuDevice GPUDevice;
typedef Eigen::DenseIndex IndexType;
typedef Eigen::ThreadPoolDevice CPUDevice;
#define NDIMS 1
template<typename T>
using TTensor =  Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using Flat = Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using  ConstScalar = Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>,Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using ConstFlat =  Eigen::TensorMap<
Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned>;

template<typename T>
struct test_cpu{
  Flat<T> call(T* _m,T* _v,T* _var,T* _grad) {
    NonBlockingThreadPool p(4);
    CPUDevice d(&p,4);
    Flat<T> m(_m,N),v(_v,N),var(_var,N);
    const T a = 0.5;
    ConstScalar<T> beta1_power(&a), lr(&a),beta1(&a),
      beta2(&a),epsilon(&a);
    ConstFlat<T> grad(_grad,N);
    m.device(d) += (grad - m)*(T(1) - beta1());
    v.device(d) = (beta2()*v).cwiseMax(grad.abs());
    var.device(d) -= lr() / (T(1) - beta1_power()) * (m/(v + epsilon()));
    return var;
  }
};


int main() {
  float _m[N],_v[N],_var[N],_grad[N];
  for (int i = 0;i < N;i++) {
    _m[i] = (float)rand();
    _v[i] = (float)rand();
    _var[i] = (float)rand();
    _grad[i] = (float)rand();
  }
  test_cpu<float> cpu_gen;
  Flat<float>out_cpu =  cpu_gen.call(_m,_v,_var,_grad);
  std::cout<<out_cpu<<std::endl;
}

可能遇到的问题

我当初用nvcc10编译.cu代码遇到了两个问题
1.TensorForwardDeclarations.h依赖include src/Tensor/TensorDeviceThreadPool.h的问题，修改Tensor头文件代码：
#include “src/Tensor/TensorMacros.h”
//#include “src/Tensor/TensorForwardDeclarations.h”
#include “src/Tensor/TensorMeta.h”
#include “src/Tensor/TensorFunctors.h”
#include “src/Tensor/TensorCostModel.h”
#include “src/Tensor/TensorDeviceDefault.h”
#include “src/Tensor/TensorDeviceThreadPool.h”
#include “src/Tensor/TensorForwardDeclarations.h”
#include “src/Tensor/TensorDeviceCuda.h”
#include “src/Tensor/TensorDeviceSycl.h”
#include “src/Tensor/TensorIndexList.h”
#include “src/Tensor/TensorDimensionList.h”
2.EIGEN_HAS_CUDA_FP16引起的问题
src/Tensor/TensorMeta.h:

#if defined(EIGEN_USE_GPU) && defined(__CUDACC__) && defined(EIGEN_HAS_CUDA_FP16)
template <>
struct PacketType<half, **GpuDevice**> {
  typedef half2 type;
  static const int size = 2;

这里的这个GpuDevice在
==#include “src/Tensor/TensorDeviceCuda.h”==被定义，一但nvcc版本符合要求就会被加入编译（有兴趣可以自己递归着看）导致出现编译错误（src/Tensor/TensorMeta.h在include "src/Tensor/TensorDeviceCuda.h前被include）