深度学习C++模板推导再炫技：统一CUDA elementwise kernel的调用接口

本文原发自公众号"AI不止算法"，欢迎关注，不定时分享技术感悟和就业指导，原文链接在此

背景

elementwise的概念似乎并没有一个官方说法，为此我特意去wiki上查了一下确实没有，如下图。所以好像属于行内用语，一些公开的博客上称elementwise指两个相同shape的张量（tensor），在对应元素上进行逐位运算。但这是不准确的，后半段是正确的，前半段不对，elementwise不仅指两个相同shape的tensor，而且还指一个或三个甚至四个，只是在AI里面三个和四个用得相对少，但是一个是很多的，并且输出也可以是一个tensor或多个tensor。

举个例子：

一个tensor的elementwise操作包括(又叫unary operation)：++,–,!

两个tensor的elementwise操作包括(又叫binary operation): +, -, *, /

问题

由上述可知，所有elementwise kernel的计算特点都是一样的（在对应元素上逐个apply上述operation），不一样的地方仅仅存在两处：

1. 这个operation是binary中的某某还是unary中的某某还是tenary（三个tensor）中的某某

2. 这个operation是一个input/output tensor还是两个还是三个

那我们能否将所有的这些elementwise kernel用一个wrapper包装起来，在这个wrapper里面处理不同个数的input tensor，再把它们分发到不同的operation进行计算并写到不同个数的output tensor？

wrapper的实现

以下就是我们wrapper的签名，接收inputs和outputs vector，以及具体的operation function（后文简称op），并且把Operation作为模板参数，交给编译器做模板推导，推导出具体的函数类型，以此泛化接收不同格式的function（函数指针，operator()等等)

template <typename Operation, int NumOuts = 1>
void ElementwiseKernelWrapper(const Device &dev,
                       const std::vector<const Tensor*> &ins,
                       std::vector<Tensor*> *outs,
                       Operation func)

因为每个op的输入个数都可能不一样，所以我们还需要萃取出它们的个数，由以下L9可以拿到输入个数（涉及到萃取的东西，以后文章再细谈）

template <typename>
struct FunctionTraits;

template <typename T>
struct FunctionTraits : public FunctionTraits<decltype(&T::operator())> {};

......

template <typename ReturnType, typename... Args>
struct FunctionTraits<ReturnType(Args...)> {
  static const size_t arity = sizeof...(Args);
  using ArgsTuple = std::tuple<Args...>;
};

有了这样一个萃取的工具后，我们再回到Wrapper，如下，L7拿到了输入tensor的个数，在L9，我们就可以直接launch elementwise kernel。

template <typename Operation, int NumOuts = 1>
void ElementwiseKernelWrapper(const Device &dev,
                       const std::vector<const Tensor*> &ins,
                       std::vector<Tensor*> *outs,
                       Operation func){
  using Traits = FunctionTraits<Operation>;
  const int Arity = Traits::arity;
  // ...some macro checks...
  LaunchElementwiseCudaKernel(...);
 }

那我们再来看看LaunchElementwiseCudaKernel的签名和实现，这里的签名与之前多数一样，重点在于实现，实现的难度在于如何handle不同个数的输入和输出，为了接收不同个数的输入和输出，我把输入个数获取提前到了编译期来获取，也就是上面的函数萃取来获取，至于输出个数，可以看出，作为模板参数直接传进来，而后，使用std::array来接收这些输入buffer和输出buffer，见L6和L7，至此，我们获得了可以handle任意个数输入和输出的能力

template <typename Operation, int Arity, int NumOuts, int VecSize>
void LaunchElementwiseCudaKernel(const Device &dev,
                                 const std::vector<const Tensor*> &ins,
                                 std::vector<Tensor*> *outs,
                                 Operation func)
  std::array<Tin*, Arity> ins_data;
  std::array<Tout*, NumOuts> outs_data；

接下来，我们还需要把ins buffer和outs buffer分别拷贝到ins_data和outs_data里面的每个buffer，类似于是以下这么个操作，ps：其实这里也可以用模板推导来完成，不过这个留在下篇文章来讲吧

  for (int i = 0; i < NumOuts; ++i) {
    outs_data[i] = (TOut *)((*outs)[i]));
  }

接下来，求出grid size和block size之后，即可launch真正的cuda kernel

  ElementwiseKernel<Operation, Arity, NumOuts, VecSize>
      <<<grid_size, block_size, 0, stream>>>(
          ins_data, outs_data, func);

那么这个真正的CUDA elementwise kernel又是怎么写的呢，以及模板推导来如何完成ins buffer和outs buffer分别拷贝到ins_data和outs_data里面的每个buffer的，那么假期后的下篇文章再来揭晓，感谢阅读和三连！

本文原发自公众号"AI不止算法"，欢迎关注，不定时分享技术感悟和就业指导，原文链接在此

彩蛋

最后，彩蛋时间

1.最近随着春招和实习的招聘，越来越多来联系我的同学已经拿到了不错的offer，恭喜~

2.各位在校同学可以转到自己的实验室，我准备来应聘教职~为中国高等教育做出贡献！