how-to-optim-algorithm-in-cuda

how-to-optim-algorithm-in-cuda

how to optimize some algorithm in cuda.

Stars: 1506

Visit
 screenshot

This repository documents how to optimize common algorithms based on CUDA. It includes subdirectories with code implementations for specific optimizations. The optimizations cover topics such as compiling PyTorch from source, NVIDIA's reduce optimization, OneFlow's elementwise template, fast atomic add for half data types, upsample nearest2d optimization in OneFlow, optimized indexing in PyTorch, OneFlow's softmax kernel, linear attention optimization, and more. The repository also includes learning resources related to deep learning frameworks, compilers, and optimization techniques.

README:

how-to-optim-algorithm-in-cuda

我也维护了一个学习深度学习框架(PyTorch和OneFlow)的仓库 https://github.com/BBuf/how-to-learn-deep-learning-framework 以及一个如何学习深度学习编译器(TVM/MLIR/LLVM)的学习仓库 https://github.com/BBuf/tvm_mlir_learn , 有需要的小伙伴可以点一点star

本工程记录如何基于 cuda 优化一些常见的算法。请注意,下面的介绍都分别对应了子目录的代码实现,所以想复现性能的话请查看对应子目录下面的 README 。

0. cuda-mode

一直想系统看一下某个课程系统和科学的学习下 CUDA ,感觉 CUDA-MODE 这个课程能满足我的需求。这个课程是几个 PyTorch 的 Core Dev 搞的,比较系统和专业。不过由于这个课程是 Youtube 上的英语课程,所以要学习和理解这个课程还是需要花不少时间的,我这里记录一下学习这个课程的每一课的笔记,希望可以通过这个笔记帮助对这个课程以及 CUDA 感兴趣的读者更快吸收这个课程的知识。这个课程相比于以前的纯教程更加关注的是我们可以利用 CUDA 做什么事情,而不是让读者陷入到 CUDA 专业术语的细节中,那会非常痛苦。伟大无需多言,感兴趣请阅读本文件夹下的各个课程的学习笔记。

1. how-to-compile-pytorch-from-source

记录如何手动编译 PyTorch 源码,学习 PyTorch 的一些 cuda 实现。

2. reduce

这里记录学习 NIVDIA 的reduce优化官方博客 做的笔记。完整实验代码见这里 , 原理讲解请看:【BBuf的CUDA笔记】三,reduce优化入门学习笔记 。后续又添加了 PyTorch BlockReduce 模板以及在这个模板的基础上额外加了一个数据 Pack ,又获得了一些带宽的提升。详细数据如下:

性能和带宽的测试情况如下 (A100 PCIE 40G):

图片

3. elementwise

将 oneflow 的 elementwise 模板抽出来方便大家使用,这个 elementwise 模板实现了高效的性能和带宽利用率,并且用法非常灵活。完整实验代码见这里 ,原理讲解请看:【BBuf 的CUDA笔记】一,解析OneFlow Element-Wise 算子实现 。这里以逐点乘为例,性能和带宽的测试情况如下 (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 模板相比于原始实现都有较大提升。

4. FastAtomicAdd

实现的脚本是针对half数据类型做向量的内积,用到了atomicAdd,保证数据的长度以及gridsize和blocksize都是完全一致的。一共实现了3个脚本:

  1. https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/atomic_add_half.cu 纯half类型的atomicAdd。
  2. https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/atomic_add_half_pack2.cu half+pack,最终使用的是half2类型的atomicAdd。
  3. https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/fast_atomic_add_half.cu 快速原子加,虽然没有显示的pack,但本质上也是通过对单个half补0使用上了half2的原子加。

性能和带宽的测试情况如下 (A100 PCIE 40G):

原子加方式 性能(us)
纯half类型 422.36ms
pack half2类型 137.02ms
fastAtomicAdd 137.01ms

可以看到使用pack half的方式和直接使用half的fastAtomicAdd方式得到的性能结果一致,均比原始的half的原子加快3-4倍。

5. UpsampleNearest2D

upsample_nearest_2d.cu 展示了 oneflow 对 upsample_nearest2d 的前后向的优化 kernel 的用法,性能和带宽的测试情况如下 (A100 PCIE 40G):

框架 数据类型 Op类型 带宽利用率 耗时
PyTorch Float32 UpsampleNearest2D forward 28.30% 111.42us
PyTorch Float32 UpsampleNearest2D backward 60.16% 65.12us
OneFlow Float32 UpsampleNearest2D forward 52.18% 61.44us
OneFlow Float32 UpsampleNearest2D backward 77.66% 50.56us
PyTorch Float16 UpsampleNearest2D forward 16.99% 100.38us
PyTorch Float16 UpsampleNearest2D backward 31.56% 57.38us
OneFlow Float16 UpsampleNearest2D forward 43.26% 35.36us
OneFlow Float16 UpsampleNearest2D backward 44.82% 40.26us

可以看到基于 oneflow upsample_nearest2d 的前后向的优化 kernel 可以获得更好的带宽利用率和性能。注意这里的 profile 使用的是 oneflow 脚本,而不是 upsample_nearest_2d.cu ,详情请看 UpsampleNearest2D/README.md

6. indexing

在 PyTorch 中对 index_add 做了极致的优化,我这里将 PyTorch 的 index_add 实现 进行了剥离,方便大家应用于其它框架。具体请看 indexing 文件夹的 README 。其中还有和 oneflow 的 index_add 实现的各个 case 的性能比较结果。整体来说 PyTorch 在 index Tensor元素很小,但Tensor很大的情况下有较大的性能提升,其它情况和 OneFlow 基本持平。详情请看 indexing/README.md

7. oneflow-cuda-optimize-skills

OneFlow 深度学习框架中基于 cuda 做的优化工作,动态更新中。

8. FastTransformer

总结 FastTransformer 相关的 cuda 优化技巧。README_BERT.md 总结了 BERT 相关的优化技巧。

9. softmax

学习了oneflow的softmax kernel实现以及Faster Transformer softmax kernel的实现,并以个人的角度分别解析了原理和代码实现,最后对性能做一个对比方便大家直观的感受到oneflow softmax kernel相比于FasterTransformer的优越性。

10. linear-attention

学习一些 linear attention 的 cuda 优化技巧。

图片

11. large-language-model-note

收集了和大语言模型原理,训练,推理,数据标注的相关文章。

12. mlsys-paper

前研的大模型训练相关 AI-Infra 论文收集以及阅读笔记。

13. triton

Triton 学习过程中的代码记录和学习笔记。

14. meagtron-lm

Meagtron-LM 学习笔记。

15. triton-meetup

Triton 中国举办的 Meetup 的slides汇总。点卡这个文件夹也可以找到对应的Meetup的视频回放。

16. ptx-isa

对 CUDA PTX ISA 文档的一个翻译和学习。

17. pytorch-blog-codes

对 PyTorch 团队发布的 cuda 技术的一些学习笔记。

18. cutlass

cutlass 相关的学习笔记。

19. cuda-paper

cuda 相关的 paper 的阅读。

20. 原创学习笔记

21. CUDA 学习资料收集

专栏

文章

Star History

Star History Chart

For Tasks:

Click tags to check more tools for each tasks

For Jobs:

Alternative AI tools for how-to-optim-algorithm-in-cuda

Similar Open Source Tools

For similar tasks

For similar jobs