oneflow element-wise代码解析

/*
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/



#ifndef ONEFLOW_CORE_CUDA_ELEMENTWISE_H_
#define ONEFLOW_CORE_CUDA_ELEMENTWISE_H_

#include <cuda_runtime.h>
#include <cstdint>
#include <algorithm>
#include <type_traits>

namespace oneflow {

namespace cuda {

namespace elementwise {

constexpr int kBlockSize = 256;
constexpr int kNumWaves = 32;

inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
  int dev;
  {
    cudaError_t err = cudaGetDevice(&dev);
    if (err != cudaSuccess) { return err; }
  }
  int sm_count;
  {
    cudaError_t err = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
    if (err != cudaSuccess) { return err; }
  }
  int tpm;
  {
    cudaError_t err = cudaDeviceGetAttribute(&tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev);
    if (err != cudaSuccess) { return err; }
  }
  *num_blocks = std::max<int>(1, std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
                                                   sm_count * tpm / kBlockSize * kNumWaves));
  return cudaSuccess;
}

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];
};

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;
constexpr int kMaxPackSize = 8;

constexpr int Min(int a, int b) { return a < b ? a : b; }

template<typename T>
constexpr int PackSize() {
  return Min(kMaxPackBytes / sizeof(T), kMaxPackSize);
}

template<typename T, typename U, typename... Args>
constexpr int PackSize() {
  return Min(PackSize<T>(), PackSize<U, Args...>());
}

template<typename T>
class HasApply2 {
  typedef char one;
  struct two {
    char x[2];
  };

  template<typename C>
  static one test(decltype(&C::Apply2));
  template<typename C>
  static two test(...);

 public:
  enum { value = sizeof(test<T>(0)) == sizeof(char) };
};

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)...); }
  return ret;
}

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])...); }
  return ret;
}

template<int pack_size, typename FactoryT, typename R, typename... IN>
__global__ void __launch_bounds__(kBlockSize)
    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();
  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, (pack_in[i])...);
  }
  if (global_tid < n_tail) { tail_r[global_tid] = functor((tail_in[global_tid])...); }
}

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

 private:
  FunctorT tpl;
};

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, cudaStream_t stream) {
  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, stream>>>(
      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,
                            cudaStream_t stream) {
    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..., stream);
    } else {
      return LaunchKernel<1, FactoryT, R, IN...>(factory, n, r, in..., stream);
    }
  }
};

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);
}

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,
                                     cudaStream_t stream) {
  return GenericLauncher<FactoryT, R, A, B>::Launch(factory, n, r, a, b, stream);
}

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,
                          cudaStream_t stream) {
  return BinaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, b, stream);
}

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

template<typename FunctorT, typename R, typename A, typename B, typename C>
inline cudaError_t Ternary(FunctorT functor, int64_t n, R* r, const A* a, const B* b, const C* c,
                           cudaStream_t stream) {
  return TernaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, b, c, stream);
}

}  // namespace elementwise

}  // namespace cuda

}  // namespace oneflow

#endif  // ONEFLOW_CORE_CUDA_ELEMENTWISE_H_

oneflow是国产开源深度学习框架，以极致的性能和分布式易用性著称。核心理念是：静态编译与运行时调度分离。使得算子层面做出非常激进的编译期优化。
Element Wise正是这种设计哲学的集中体现
Oneflow在200行代码实现了向量化访存，自适应Grid分配，编译器多态，模板特化等多项技术

Element-wise的和兴思想

Element-wise操作的计算密度很低。一次内存读取只伴随少量浮点运算。因此GPU算力不是瓶颈，而是带宽/也就是如何用最快方式把数据从全局内存搬运到寄存器里面。
Onflow确立了三大支柱：
-memory Pakcing：强制128bit对齐读写，出发编译器和硬件使用向量化访存指令
-SFINAE 编译期优化：当定义的算符支持Apply2接口时，自动切换到更高效的half2等SIMD指令路径
-Adaptive Grid Sizing：动态查询当前GPU的SM数量和最大线程数，计算出恰好能掩盖延迟，又不浪费资源的线程网络

架构师视角

用户接口层

框架要服务号算法工程师，因此用户些element-wise算子时不应该感知 线程块大小是多少，指针是否对齐等底层细节。只想描述数学逻辑，然后给一个函数执行。Oneflow给出了最简洁的接口： Unary, Binar4y, Ternary例如：

// 以 Binary (二元计算) 为例
// 用户自定义一个二元运算
struct AddFunctor {
    __device__ float operator()(float a, float b) {
        return a + b;
    }
};

// 调用 OneFlow 的 Binary 接口
float *d_r, *d_a, *d_b;
int64_t n = 1 << 20;
cudaStream_t stream = 0;
oneflow::cuda::elementwise::Binary(AddFunctor{}, n, d_r, d_a, d_b, stream);

这样，三行核心代码，一个 operator() 用户就将CUDA核函数的烦恼全部丢给了OneFlow
某些算子状态在运行时决定，例如Dropout需要随机种子，这时可以使用 工厂模式， 将一组参数打包成工厂对象，待到GPU端真正执行时再生成最终的算符，对应接口称为：UnaryWithFactory / BinaryWithFactory / TernaryWithFactor

分发层

编译期计算最大打包宽度

接口函数内部会直接调用 GenericLauncher::launch(…) 这是性能决策的枢纽，核心任务时选择pack_size. 核心手段时向量化访存，GPU提供了 LDG.128 和 STG.128指令，一次可以搬运16个字节，但是向量化访存需要严格对齐内存
PackSize 根据所有输入输出类型，在编译期算出安全的最大打爆元素数。它受限于两个元素：
- kMaxPackBytes = 16 字节(128位)，因为GPU的向量化加载指令一次最多搬运128位
- kMaxPackSize=8，即最高一次打包8个元素
以 float 为例，每个 4 字节，16 / 4 = 4，所以 float 的最大 pack 大小是 4。half（2 字节）则是 min(16/2, 8) = 8。这样，一条 128-bit 指令就能加载/存储 4 个 float 或 8 个 half，极大降低指令发射数。

运行期检查内存对齐

向量化访存要求数据地址与访问狂赌对齐IsAlignedForPack 在运行时检查所有输入输出指针是否满徐alignof(Packed<T, pack_size>)（即 sizeof(T) * pack_size）

template<typename FactoryT, typename R, typename... IN>
struct GenericLauncher {
  static cudaError_t Launch(...) {
    // 1. 计算最大可能的打包大小 (最大 16 字节)
    constexpr int max_pack_size = PackSize<R, IN...>();

    // 2. 动态检查指针地址是否满足对齐要求
    if (IsAlignedForPack<max_pack_size, R, IN...>(r, in...)) {
      // 性能模式：按 max_pack_size (如 4 或 8) 启动 Kernel
      return LaunchKernel<max_pack_size, ...>(...);
    } else {
      // 安全回退模式：地址不对齐，按标量 (pack_size = 1) 启动 Kernel
      return LaunchKernel<1, ...>(...);
    }
  }
};

如果地址对齐，就享用向量化红利；如果不对齐，则安全回退到pack_size=1的标量模式。这种防御性变成值得学习

网格优化层

在去顶了Packsize之后，准备启动kernel，此时需要决定分配多少个线程块
对于新手来说，一般就是 blocks = (n + blockSize - 1) / blockSize。这样在数据较小时会喂不饱 GPU，数据量大又会淹没调度器。OneFlow的方法是自适应，引入了GetNemBlocks 算法

constexpr int kBlockSize = 256;
constexpr int kNumWaves = 32;

inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
    int dev, sm_count, tpm;
    cudaGetDevice(&dev);
    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
    cudaDeviceGetAttribute(&tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev);

    *num_blocks = std::max<int>(1, std::min<int64_t>(
        (n + kBlockSize - 1) / kBlockSize,
        sm_count * tpm / kBlockSize * kNumWaves
    ));
    return cudaSuccess;
}

kNumWaves = 32 是一个经验值：让每个SM上常驻kNumWaves 个线程块，以便在访存延迟时立即切换到另一波线程执行。上限 sm_count * tpm / kBlockSize * kNumWaves 确保不会因过度订阅空耗调度资源。无论用的时V100，H100还是4070，这段逻辑都能自动适配，得到最适合的网格大小

内核实现

Grid-stride Loop

template<int pack_size, ...>
__global__ void ApplyGeneric(..., int64_t n_pack, Packed<R, pack_size>* pack_r, ..., int64_t n_tail, R* tail_r, ...) {

  auto functor = factory();
  const int global_tid = blockIdx.x * kBlockSize + threadIdx.x;

  // 1. 主循环：处理 Pack 数据 (向量化读取 -> 计算 -> 向量化写入)
  for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) {
    pack_r[i] = ApplyPack<pack_size, ...>(functor, (pack_in[i])...);
  }

  // 2. 尾部处理：处理无法被 pack_size 整除的剩余标量
  if (global_tid < n_tail) {
    tail_r[global_tid] = functor((tail_in[global_tid])...);
  }
}

分离了打包区与尾部区，确保所有的数据都会被处理

架构师的黑魔法

打包究竟是什么？

template<typename T, int pack_size>
struct alignas(sizeof(T) * pack_size) Packed {
  union {
    T elem[pack_size];
  };
};

alignas 告诉编译器：这个结构体必须按照 sizeof(T) * pack_size 对齐，当我们把输入输出指针 reinterpret_cast 成 Packed<R, pack_size>* 时，编译器会自动生成 128-bit 的加载/存储指令，而不是四次 32-bit 的标量指令。因此自然会调用最高效的宽总线访存指令。

SFINAE

ApplyPack 有两个重载，通过 std::enable_if 在编译期选择

• 默认路径：一个简单的for 循环，诸葛调用functor(elem…)
• 优化路径：当检测到FunctorT 拥有Apply2方法且pack_size % 2 == 0 时，启用这个版本

  template<int pack_size, typename FunctorT, ...>
  __device__ enable_if<HasApply2<FunctorT>::value && 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)...);
      }
      return ret;
  }

用户视角

在此处给出一些应用的典型场景

基本运算

#include <cuda_runtime.h>
#include "elementwise.h"

struct MulFunctor {
    __device__ float operator()(float a, float b) { return a * b; }
};

void launch_mul(float *d_out, const float *d_a, const float *d_b, int64_t n, cudaStream_t stream) {
    oneflow::cuda::elementwise::Binary(MulFunctor{}, n, d_out, d_a, d_b, stream);
}

struct ReLUFunctor {
    __device__ float operator()(float x) { return fmaxf(x, 0.f); }
};

void launch_relu(float *d_out, const float *d_in, int64_t n, cudaStream_t stream) {
    oneflow::cuda::elementwise::Unary(ReLUFunctor{}, n, d_out, d_in, stream);
}

带运行时参数的运算

struct ScaleAddFunctor {
    float alpha;
    __device__ float operator()(float a, float b) const {
        return alpha * a + b;
    }
};

struct ScaleAddFactory {
    float alpha;
    __device__ ScaleAddFunctor operator()() const {
        return ScaleAddFunctor{alpha};
    }
};

void launch_scale_add(float *d_out, const float *d_a, const float *d_b,
                      int64_t n, float alpha, cudaStream_t stream) {
    oneflow::cuda::elementwise::BinaryWithFactory(
        ScaleAddFactory{alpha}, n, d_out, d_a, d_b, stream);
}

把运行时参数存进工厂，核函数再生产出一个带状态的functor/这也是许多复杂算子(比如带随机种子的Dropout)的标准做法

Apply2加速

struct AddFunctor {
    template<typename T>
    __device__ T operator()(T a, T b) const { return a + b; }

    template<typename T>
    __device__ void Apply2(T* out, const T* a, const T* b) const {
        out[0] = a[0] + b[0];
        out[1] = a[1] + b[1];
    }
};