Element-wise 逐点运算

template<size_t pack_size>
bool IsAlignedForPack() {
  return true;
}

template<size_t pack_size, typename T, typename... Args>
bool IsAlignedForPack(const T* ptr, const Args*... others) {
  return reinterpret_cast<uintptr_t>(ptr) % sizeof(Pack<T, pack_size>) == 0
         && IsAlignedForPack<pack_size, Args...>(others...);
}

template<size_t pack_size, typename FactoryT, typename R, typename... IN>
cudaError_t LaunchKernel(FactoryT factory, int64_t n, R* r, const IN*... in) {
  const int64_t n_pack = n / pack_size;
  const int64_t tail_offset = n_pack * pack_size;
  const int64_t n_tail = n - tail_offset;
  int num_blocks;
  {
    cudaError_t err = GetNumBlocks(n_pack, &num_blocks);
    if (err != cudaSuccess) { return err; }
  }
  ApplyGeneric<pack_size, FactoryT, R, IN...><<<num_blocks, kBlockSize, 0>>>(
      factory, n_pack, reinterpret_cast<Packed<R, pack_size>*>(r),
      (reinterpret_cast<const Packed<IN, pack_size>*>(in))..., n_tail, r + tail_offset,
      (in + tail_offset)...);
  return cudaPeekAtLastError();
}

template<typename FactoryT, typename R, typename... IN>
struct GenericLauncher {
  static cudaError_t Launch(FactoryT factory, int64_t n, R* r, const IN*... in) {
    constexpr int max_pack_size = PackSize<R, IN...>();
    if (IsAlignedForPack<max_pack_size, R, IN...>(r, in...)) {
      return LaunchKernel<max_pack_size, FactoryT, R, IN...>(factory, n, r, in...);
    } else {
      return LaunchKernel<1, FactoryT, R, IN...>(factory, n, r, in...);
    }
  }
};

template<typename FactoryT, typename R, typename A>
inline cudaError_t UnaryWithFactory(FactoryT factory, int64_t n, R* r, const A* a) {
  return GenericLauncher<FactoryT, R, A>::Launch(factory, n, r, a);
}

template<typename FunctorT, typename R, typename A>
inline cudaError_t Unary(FunctorT functor, int64_t n, R* r, const A* a) {
  return UnaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a);
}

template<typename FactoryT, typename R, typename A, typename B>
inline cudaError_t BinaryWithFactory(FactoryT factory, int64_t n, R* r, const A* a, const B* b) {
  return GenericLauncher<FactoryT, R, A, B>::Launch(factory, n, r, a, b);
}

template<typename FunctorT, typename R, typename A, typename B>
inline cudaError_t Binary(FunctorT functor, int64_t n, R* r, const A* a, const B* b) {
  return BinaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, b);
}

#include <cuda.h>          // driver API，本文件没用到（启动走 runtime）
#include <cuda_runtime.h>  // 真正需要的：__global__、cudaMalloc/cudaMemcpy、cudaDeviceGetAttribute、dim3
#include <time.h>          // 【未使用】main 里没有计时；本机测时用的是 cudaEvent
#include <algorithm>       // 用到的只有 std::max / std::min（GetNumBlocks :30）
#include <iostream>        // 【未使用】打印用的是 printf
#include <cuda_fp16.h>     // half / half2、__hmul2；本文件 main 实例化的就是 half
using namespace std;       // 下面用了 std::max/std::min（带 std:: 前缀），其实可以不写这行

#define N 32 * 1024 * 1024   // 元素个数 33,554,432。宏没加括号，本文件里的用法（乘、加、除）恰好都安全
// elementwise implementation copyed from https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh
constexpr int kBlockSize = 256;   // 每 block 线程数；kernel 上的 __launch_bounds__ 也是它（:115）
constexpr int kNumWaves = 32;     // GetNumBlocks 的"波数"：允许开的 block 数 = 能同时驻留的份数 × 32


// 【pack 的存储类型】aligned_storage 给出"大小 = pack_size*sizeof(T)、对齐 = 同样大"的裸内存块。
//   注意它只是 POD 字节，没有成员可以访问，所以下面还要 union 一个 T elem[] 视图。
template<typename T, int pack_size>
struct GetPackType {
  using type = typename std::aligned_storage<pack_size * sizeof(T), pack_size * sizeof(T)>::type;
};

template<typename T, int pack_size>
using PackType = typename GetPackType<T, pack_size>::type;  


// 【Pack：只在 host 侧用来算 sizeof】union 让 storage 和 elem[] 共用同一块内存。
//   本文件里 Pack<T,pack_size> 唯一的用处是 IsAlignedForPack(:143) 里的 sizeof —— 拿它当
//   "一个包占几字节"的标尺。Packed（下面那个）才是 kernel 里真正用的类型。
template<typename T, int pack_size>
union Pack {
  static_assert(sizeof(PackType<T, pack_size>) == sizeof(T) * pack_size, "");   // 编译期自证：aligned_storage 没被额外 padding
  __device__ Pack() {
    // do nothing   ← 什么都不做的默认构造：既避免隐式构造产生清零代码，也让这个类型不能在 host 侧默认构造
  }
  PackType<T, pack_size> storage;   // 裸字节视图（决定对齐）
  T elem[pack_size];                // 元素视图
};


// 【Packed：kernel 里真正用的"一包"】alignas 是关键：没有它，Packed<half,8> 的 alignof 只有 2，
//   编译器就发不出 128 位访存指令（本机实测 PTX 是 ld.global.v4.b32，SASS 是 LDG.E.128）。
//   匿名 union 只有一个成员，作用是让 Packed<T,N> 能直接当数组用（.elem[j]），
//   同时保持是 trivially-copyable 的聚合体。
template<typename T, int pack_size>
struct alignas(sizeof(T) * pack_size) Packed {
  __device__ Packed() {
    // do nothing
  }
  union {
    T elem[pack_size];
  };
};

constexpr int kMaxPackBytes = 128 / 8;   // 16 字节 = 一条向量访存指令搬的量（128 bit）
constexpr int kMaxPackSize = 8;          // 上限 8 个元素（half 正好 16 字节；float 会先被 16/4=4 卡住）

constexpr int Min(int a, int b) { return a < b ? a : b; }   // constexpr 版 min（没用 std::min，是为了在模板实参里能用）

template<typename T>
constexpr int PackSize() {
  return Min(kMaxPackBytes / sizeof(T), kMaxPackSize);   // half→8、float→4、double→2（本机打印验证）
}


template<typename T, typename U, typename... Args>
constexpr int PackSize() {
  return Min(PackSize<T>(), PackSize<U, Args...>());   // 递归取"输出和所有输入"里的最小包宽：任一操作数是 double 就整体降到 2
}


// 【编译期问一句："你这个 functor 有没有 Apply2 成员？"】
//   test(decltype(&C::Apply2)) 能成立就匹配第一个重载（返回 1 字节的 one），
//   否则靠 ... 匹配第二个（返回 2 字节的 two）。sizeof 一比即得。
//   为什么要问：有 Apply2 就能一次算 2 个元素（half2 成对乘），没有就只能逐个算。
template<typename T>
class HasApply2 {
  typedef char one;
  struct two {
    char x[2];
  };

  template<typename C>
  static one test(decltype(&C::Apply2));   // 只有 C::Apply2 存在时这个声明才合法（SFINAE）
  template<typename C>
  static two test(...);                    // 否则退到这里

 public:
  enum { value = sizeof(test<T>(0)) == sizeof(char) };   // 1 = 有，0 = 没有
};



// 【ApplyPack 版本 A：functor 有 Apply2 且包宽是偶数】一次给 2 个元素的地址，让 functor 成对算。
//   形参 `const Packed<IN, pack_size>... in` 是变参模板的展开：每个输入一个包，按值传（16 字节，进寄存器）。
//   调用点 (:122) 传的是 (pack_in[i])...，即第 i 包的所有输入。
template<int pack_size, typename FunctorT, typename R, typename... IN>
__device__ typename std::enable_if<HasApply2<FunctorT>::value == true && pack_size % 2 == 0,
                                   Packed<R, pack_size>>::type
ApplyPack(const FunctorT& functor, const Packed<IN, pack_size>... in) {
  Packed<R, pack_size> ret;                    // 未初始化（默认构造什么都不做），下面逐元素写满
#pragma unroll
  for (int j = 0; j < pack_size; j += 2) { functor.Apply2(ret.elem + j, (in.elem + j)...); }   // 每次推进 2；传的是元素数组里的下标 j，不是包号
  return ret;
}

// 【ApplyPack 版本 B：没有 Apply2，或包宽是奇数】逐元素调 functor 的 operator()。
//   两个版本的函数体一样长，选哪个由 enable_if 决定；条件互补（|| 与 && 取反），不会同时可选。
template<int pack_size, typename FunctorT, typename R, typename... IN>
__device__ typename std::enable_if<HasApply2<FunctorT>::value == false || pack_size % 2 != 0,
                                   Packed<R, pack_size>>::type
ApplyPack(const FunctorT& functor, const Packed<IN, pack_size>... in) {
  Packed<R, pack_size> ret;
#pragma unroll
  for (int j = 0; j < pack_size; ++j) { ret.elem[j] = functor((in.elem[j])...); }   // 每个元素调一次 operator()，参数是所有输入的第 j 个
  return ret;
}



// 【kernel 本体】grid-stride 走包 + 标量补尾巴。
//   【形参】factory：functor 工厂（见 SimpleFactory）；n_pack：包数；pack_r/pack_in：包视图的指针；
//   n_tail：尾巴元素数；tail_r/tail_in：普通元素视图的指针（指向 tail_offset 处）。
//   同一个数据缓冲区用两套指针看，就是这段代码里最容易被绕晕的地方。
template<int pack_size, typename FactoryT, typename R, typename... IN>
__global__ void __launch_bounds__(kBlockSize)     // 告诉 ptxas 每 block 最多 256 线程（本机实测 36 寄存器，占满 6 block/SM）
    ApplyGeneric(FactoryT factory, int64_t n_pack, Packed<R, pack_size>* pack_r,
                 const Packed<IN, pack_size>*... pack_in, int64_t n_tail, R* tail_r,
                 const IN*... tail_in) {
  auto functor = factory();                      // 在这里、在 device 侧构造 functor —— 这就是传 factory 而不直接传 functor 的原因
  const int global_tid = blockIdx.x * kBlockSize + threadIdx.x;   // 线程全局编号；注意它用的是 kBlockSize 常量，不是 blockDim.x
  for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) {   // grid-stride：步长 = 本次启动的总线程数（本机 n=2^25 时是 2359296）
    pack_r[i] = ApplyPack<pack_size, decltype(functor), R, IN...>(functor, (pack_in[i])...);   // 读第 i 包所有输入 → 算 → 写第 i 包输出
  }
  if (global_tid < n_tail) { tail_r[global_tid] = functor((tail_in[global_tid])...); }   // 尾巴：前 n_tail 个线程各补 1 个元素，用标量路径
}


// 【functor 工厂】把 functor 装进一个可拷贝的小对象，kernel 里 factory() 再取出来。
//   直接传 functor 需要它的类型在 host 和 device 两侧都能构造；工厂把这个动作推到 device 侧。
//   这里 operator() 是 __device__ 的，所以 SimpleFactory<T> 本身在 host 侧构造（:193 那次），
//   取出模板值只在 device 侧发生。
template<typename FunctorT>
struct SimpleFactory {
  explicit SimpleFactory(FunctorT functor) : tpl(functor) {}   // host 侧调用（避免隐式转换）
  __device__ FunctorT operator()() const { return tpl; }       // device 侧：按值返回一份 functor

 private:
  FunctorT tpl;
};