【BBuf的CUDA笔记】四，介绍三个高效实用的CUDA算法实现（OneFlow ElementWise模板，FastAtomicAdd模板，OneFlow UpsampleNearest2d模板）

本文链接：https://blog.csdn.net/just_sort/article/details/128596922

本文介绍了OneFlow的ElementWise模板、FastAtomicAdd和UpsampleNearest2d模板的优化方法和原理。ElementWise模板通过向量化、选择合适的GridSize和BlockSize等技术提高性能。FastAtomicAdd针对CUDA中half类型的atomicAdd进行了优化，利用half2类型提升效率。OneFlowUpsampleNearest2d模板针对2倍上采样进行了优化，减少了取模和映射操作，提升了带宽利用率和计算速度。所有代码已整理为独立的CUDA文件供参考。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

0x0. 前言

如题所述，本篇文章推荐和讲解一下OneFlow ElementWise模板，FastAtomicAdd，OneFlow UpsampleNearest2d模板的用法以及原理。但OneFlow ElementWise模板的用法和原理在【BBuf的CUDA笔记】一，解析OneFlow Element-Wise 算子实现已经讲过了，所以这篇文章里不再赘述，主要讲解后面2个。我将上述三个算法的实现都分别抽出来放到了 https://github.com/BBuf/how-to-optim-algorithm-in-cuda 这个工程的 elementwise/FastAtomicAdd/UpsampleNearest2D 三个文件夹中，并且三个算法的实现都分别只用一个.cu文件进行整理，使用nvcc编译可以使用，有需要的同学请自取。

0x1. OneFlow elementwise模板

将 oneflow 的 elementwise 模板抽出来方便大家使用，这个 elementwise 模板实现了高效的性能和带宽利用率，并且用法非常灵活。完整实验代码见 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/elementwise/elementwise.cu，原理讲解请看：【BBuf 的CUDA笔记】一，解析OneFlow Element-Wise 算子实现。这里以逐点乘（z = x * y，其中x，y，z是形状完全一样的Tensor）为例，性能和带宽的测试情况如下 (A100 PCIE 40G)：

优化手段	数据类型	耗时(us)	带宽利用率
naive elementwise	float	298.46us	85.88%
oneflow elementwise	float	284us	89.42%
naive elementwise	half	237.28us	52.55%
oneflow elementwise	half	140.74us	87.31%

可以看到无论是性能还是带宽，使用 oneflow 的 elementwise 模板相比于原始实现都有较大提升。

涉及到的主要优化技术有向量化数据访问，选取合适的GridSize和BlockSize，循环展开和Grid-Stride Loops等技巧。

模板代码和用法详见： https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/elementwise/elementwise.cu

0x2. FastAtomicAdd

众所周知，atomicAdd是CUDA中非常昂贵的操作，特别是对于half类型来说 atomicAdd 巨慢无比，慢到如果一个算法需要用到 atomicAdd，那么相比于用 half ，转成 float ，再 atomicAdd，再转回去还要慢很多。但是我们有时候不得不去执行half类型的原子加，这个时候怎么能提升性能呢？

PyTorch给出了一个快速原子加的实现（我这里魔改了一下，去掉了一些不需要的参数，完整测试代码见 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/fast_atomic_add_half.cu ）：

// FastAdd is referenced from
// https://github.com/pytorch/pytorch/blob/396c3b1d88d7624938a2bb0b287f2a19f1e89bb4/aten/src/ATen/native/cuda/KernelUtils.cuh#L29
template<typename T, typename std::enable_if<std::is_same<half, T>::value>::type* = nullptr>
__device__ __forceinline__ void FastSpecializedAtomicAdd(T* base, size_t offset,
                                                         const size_t length, T value) {
   
#if ((defined(CUDA_VERSION) && (CUDA_VERSION < 10000)) \
     || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700)))
  atomicAdd(reinterpret_cast<half*>(base) + offset, static_cast<half>(value));
#else
  // Accounts for the chance base falls on an odd 16 bit alignment (ie, not 32 bit aligned)
  __half* target_addr = reinterpret_cast<__half*>(base + offset);
  bool low_byte = (reinterpret_cast<std::uintptr_t>(target_addr) %