高效、易用、可拓展我全都要：OneFlow CUDA Elementwise模板库的设计优化思路

// Write ReLU Functor. 
template<typename T>
struct ReluFunctor {
  OF_DEVICE_FUNC T operator()(T x) const {
    const T zero_val = static_cast<T>(0); 
    return (x > zero_val) ? x : zero_val; 
  }
};

// Use CUDA Elementwise Template. 
OF_CUDA_CHECK((cuda::elementwise::Unary(ReluFunctor<T>(), elem_cnt, dx->mut_dptr<T>(),
                                        x->dptr<T>(), ctx->stream()->As<ep::CudaStream>()->cuda_stream())));

constexpr int kBlockSize = 256
constexpr int kNumWaves = 32;

inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
  ...
  /*
  n: The number of the elements. 
  sm_count: The number of the SM. 
  tpm: The maximum resident threads in per multiprocessor. 
  */
  *num_blocks = std::max<int>(1, std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
                                                   sm_count * tpm / kBlockSize * kNumWaves));
  return cudaSuccess;
}

template <typename T>
__global__ void ker_layer_norm(T *ln_res, T *vars, T *means, const T *inp,
                               const T *scale, const T *bias, int hidden_size) {
  // step 0. compute local sum
  float l_sum = 0;
  float l_square_sum = 0;
  const float4 *inp_f4 = (const float4 *)inp + blockIdx.x * hidden_size; // use float4
  for (uint idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
    float4 val = inp_f4[idx];
    ...
  }
}

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;

template<typename T, int pack_size>
union Pack {
  static_assert(sizeof(PackType<T, pack_size>) == sizeof(T) * pack_size, "");
  __device__ Pack() {
    // do nothing
  }
  PackType<T, pack_size> storage;
  T elem[pack_size];
};

constexpr int kMaxPackBytes = 128 / 8;
constexpr int kMaxPackSize = 8;

Unary/Binary/Ternary
  -> xxxFactory
     -> GenericLauncher<...>::Launch
       -> ApplyGeneric(CUDA Kernel)

template<int pack_size, bool tail, typename FactoryT, typename R, typename... IN>
__global__ void __launch_bounds__(kBlockSize)
    ApplyGeneric(FactoryT factory, int64_t n_pack, PackType<R, pack_size>* pack_r,
                 const PackType<IN, pack_size>*... pack_in, int64_t n_tail, R* tail_r,
                 const IN*... tail_in) {
  auto functor = factory();
  const int global_tid = blockIdx.x * kBlockSize + threadIdx.x;
  for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) {
    pack_r[i] = ApplyPack<pack_size, decltype(functor), R, IN...>(
        functor, (FetchPack<IN, pack_size>(pack_in + i).elem)...);
  }
  if (tail && global_tid < n_tail) { tail_r[global_tid] = functor((tail_in[global_tid])...); }
}

template<int pack_size, typename FunctorT, typename R, typename... IN>
__device__
    typename std::enable_if<HasApply2<FunctorT>::value == false, PackType<R, pack_size>>::type
    ApplyPack(const FunctorT& functor, const IN... in[pack_size]) {
  Pack<R, pack_size> ret;
#pragma unroll
  for (int j = 0; j < pack_size; ++j) { ret.elem[j] = functor((in[j])...); }
  return ret.storage;
}

template<int pack_size, typename FunctorT, typename R, typename... IN>
__device__ typename std::enable_if<HasApply2<FunctorT>::value == true && pack_size % 2 == 0,
                                   PackType<R, pack_size>>::type
ApplyPack(const FunctorT& functor, const IN... in[pack_size]) {
  Pack<R, pack_size> ret;
#pragma unroll
  for (int j = 0; j < pack_size; j += 2) { functor.Apply2(ret.elem + j, (in + j)...); }
  return ret.storage;
}

template<typename From>
struct CastFunctor<half, From, typename std::enable_if<!std::is_same<From, half>::value>::type> {
  ...

  __device__ void Apply2(half* to, const From* from) const {
    float2 f2;
    f2.x = static_cast<float>(from[0]);
    f2.y = static_cast<float>(from[1]);
    *reinterpret_cast<half2*>(to) = __float22half2_rn(f2);
  }
};

template<typename FunctorT>
struct SimpleFactory {
  explicit SimpleFactory(FunctorT functor) : tpl(functor) {}
  __device__ FunctorT operator()() const { return tpl; }

 private:
  FunctorT tpl;
};

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

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

// BinaryWithFactory TernaryWithFactory ... 
// Binary Ternary ...