【BBuf的CUDA笔记】五，解读 PyTorch index_add 操作涉及的优化技术

news2024/11/19 7:36:26

本文把pytorch index_add算子的代码抽取出来放在：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/indexing/index_add_cuda_pytorch_impl.cu 。如果不太熟悉PyTorch的话也可以直接看这个.cu文件，有问题请在这个repo提issue。

0x0. 前言

我们可以在 PyTorch 的文档中找到 torch.index_add_ 的定义（https://pytorch.org/docs/stable/generated/torch.Tensor.index_add_.html#torch.Tensor.index_add_）：

在这里插入图片描述
简单来说就是我们需要根据index的索引完成对当前Tensor dim维度的inplace加和，注意被加数是由另外一个Tensor src决定的。在PyTorch的codebase中搜索index_add，我们可以发现这个操作应用得非常广泛，比如说作为as_strided算子的backward的一部分，作为一些sparse操作的一部分等等。我最近研究了一下，发现PyTorch对index_add算子的cuda kernel进行了较为精细的优化，主要有两个亮点，本篇文章就来学习一下。

顺便提一下，在PyTorch中index_add的cuda kernel实现在https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L712 ，如果你想自己详细读这个代码我建议先编译一下PyTorch再进行调试和阅读，编译PyTorch源码可以参考：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/how-to-complie-pytorch-from-source（这个也是参考PyTorch官方的教程，补充了几个报错的坑） 。

0x1. 亮点1: 按照index的元素个数派发不同的实现

PyTorch优化的出发点是，index_add操作中index这个Tensor是尤其重要，它决定了输入Tensor的哪些位置会被重新赋值，然后index的元素可多可少。如果使用同一套naive的计算逻辑可能会因为重复访问index导致全局内存访问过多，而如果index很大那么为了保证性能kernel又需要满足足够的并行度才可以。为了平衡这两种情况，PyTorch按照index的元素个数实现了2套kernel。这2套kernel的实现在：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L576-L675 。然后根据index元素个数进行dispatch：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L801-L829 。

在这里插入图片描述

我在 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/indexing/indexing_pytorch_explain.cu#L381-L505 这里以PyTorch文档展示的例子（https://pytorch.org/docs/stable/generated/torch.Tensor.index_add_.html#torch.Tensor.index_add_）为例记录了各个中间变量的值，并加上了一些方便理解的注释，感兴趣的可以查看。

我们这里展示一下当index的元素很少的时候的indexFuncSmallIndex kernel实现（代码中的设置是index元素个数少于16）：

// 如果索引的数量很少，我们更喜欢使用这个Kernel来避免重新加载 index。
// 这个kernel实际上适用于几乎所有问题大小的选择，但如果选择的索引数量很大，
// 那么indexFuncLargeIndex Kernel是增加并行度的更好选择。
// 下面的innerSize就是输人的self张量忽略dim维度的切片大小，对于每一个indices[i]，我们都要处理innerSize个元素的copy

// selfAddDim(dstAddDim) = 0
// sourceAddDim(srcAddDim) = 0
// sliceSize(innerSize) = 3
// selfAddDimSize(dstAddDimSize) = 5
// selfNumel(dstNumel) = 15
// selfInfo.sizes(dst): 1, 3, 
// selfInfo.strides(dst): 3, 1,
// sourceInfo.sizes(src): 1, 3, 
// sourceInfo.strides(src): 3, 1
// indexInfo.sizes(indices): 3, 
// indexInfo.strides(indices): 1,

template <typename T, typename IndicesType, typename IndexType, int DstDim, int SrcDim, int IdxDim,
          typename func_t>
__global__ void indexFuncSmallIndex(cuda::detail::TensorInfo<T, IndexType> dst,
                                    cuda::detail::TensorInfo<T, IndexType> src,
                                    cuda::detail::TensorInfo<IndicesType, IndexType> indices,
                                    int dstAddDim,
                                    int srcAddDim,
                                    IndexType innerSize,
                                    int64_t dstAddDimSize,
                                    int64_t dstNumel,
                                    const func_t& op,
                                    T alpha) {
  // In order to avoid reloading the index that we are copying, load
  // it once to handle all of the points that are being selected, so
  // it can be reused as much as possible. This kernel is chosen when
  // this is a good choice (small number of chosen indices), since
  // re-accessing indices in addition to src elements can be slow.
  // 为了避免重新加载我们正在复制的索引，加载一次以处理所有正在选择的点，以便尽可能地重复使用它。 
  // 当这是一个不错的选择（选择的索引数量很少）时，就会选择这个Kernel，
  // 因为除了 src 元素之外，重新访问索引可能很慢。
  for (IndexType srcIndex = 0; srcIndex < indices.sizes[0]; ++srcIndex) {
    // Lua indices begin at 1
    IndexType dstIndex =
        indices.data[cuda::detail::IndexToOffset<IndicesType, IndexType, IdxDim>::get(srcIndex, indices)];
    CUDA_KERNEL_ASSERT(dstIndex < dstAddDimSize);

    // We stride over the output ignoring the indexed dimension
    // (innerSize), whose offset calculation is handled differently
    for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
         linearIndex < innerSize;
         linearIndex += gridDim.x * blockDim.x) {
      IndexType dstOffset =
          cuda::detail::IndexToOffset<T, IndexType, DstDim>::get(linearIndex, dst);
      dstOffset += dstIndex * dst.strides[dstAddDim];

      IndexType srcOffset =
          cuda::detail::IndexToOffset<T, IndexType, SrcDim>::get(linearIndex, src);
      srcOffset += srcIndex * src.strides[srcAddDim];

      T val = src.data[srcOffset] * alpha;
      op(dst.data, dstOffset, dstNumel, &val);
    }

  }
}

我们可以看到首先有一个for (IndexType srcIndex = 0; srcIndex < indices.sizes[0]; ++srcIndex) 的循环来避免重复加载 index Tensor（这个时候index Tensor信息由indices管理），后续的实验结果也将证明这个优化在 index 元素个数比较小而 self Tensor 比较大的时候是有一定性能提升的。然后选定一个indices[i] 之后就启动一堆线程计算完这个indices[i]对应的 self Tensor的一个切片（linearIndex < innerSize）。

indexFuncLargeIndex Kernel我就不展示了，感兴趣的小伙伴可以直接阅读代码实现。

实现完这两个Kernel之后，我们可以在 https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L753-L778 这里看到PyTorch分别为这两个Kernel设置了不同的GridSize和BlockSize。

// selfAddDim = 0
// sourceAddDim = 0
// sliceSize = 3
// selfAddDimSize = 5
// selfNumel = 15

#define SMALL_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE, SELF_DIM, SOURCE_DIM, IDX_DIM)     \
  indexFuncSmallIndex<TENSOR_TYPE, INDICES_TYPE, TYPE, SELF_DIM, SOURCE_DIM, IDX_DIM>   \
    <<<smallIndexGrid, smallIndexBlock, 0, stream>>>(                                   \
      selfInfo, sourceInfo, indexInfo,                                                  \
      selfAddDim, sourceAddDim, sliceSize, selfAddDimSize,                              \
      selfNumel, reduce_add, alpha_value);                                              \
  C10_CUDA_KERNEL_LAUNCH_CHECK();

#define LARGE_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE,                        \
                    SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR)            \
  indexFuncLargeIndex<TENSOR_TYPE, INDICES_TYPE, TYPE,                      \
                      SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR>          \
    <<<largeIndexGrid, largeIndexBlock, 0, stream>>>(                       \
      selfInfo, sourceInfo, indexInfo,                                      \
      selfAddDim, sourceAddDim, sourceTotalSize,                            \
      (IDX_IS_MAJOR) ? sliceSize : numIndex,                                \
      selfAddDimSize, selfNumel, reduce_add, alpha_value);                  \
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  // small index以正在索引的每个切片的大小为基准来设定GridSize和BlockSize，同时要考虑到需要满足足够多的wave保证利用率
  const dim3 smallIndexGrid(std::min(ceil_div(sliceSize, (ptrdiff_t)128), (ptrdiff_t)(mpc * 8)));
  const dim3 smallIndexBlock(std::min(sliceSize, (ptrdiff_t)128));

  // large index以source 张量的总大小为基准来设定GridSize和BlockSize，同时要考虑到需要满足足够多的wave保证利用率
  const dim3 largeIndexGrid(std::min(ceil_div(sourceTotalSize, (ptrdiff_t)128), (ptrdiff_t)(mpc * 8)));
  const dim3 largeIndexBlock(std::min(sourceTotalSize, (ptrdiff_t)128));

对于index的元素个数比较小也就是smallIndex的情况，线程块的数量由sliceSize来决定，而对于index元素个数比较大也就是largeIndex的时候线程块的数量则由输入Tensor self的总元素数量来决定。我个人感觉这里设置GridSize和BlockSize还是存在一定问题的，在profile的时候ncu对于index比较小并且输入Tensor也不太大的情况会提示grid太小无法充分发挥并行性的问题。建议阅读https://mp.weixin.qq.com/s/1_ao9xM6Qk3JaavptChXew 这篇文章设置更合理的GridSize和BlocSize，或许可以提升smallIndex Kernel的性能。

比如index很小但是输入Tensor只有一个维度的情况下，这个时候PyTorch只会启动一个Block以及一个Thread，这显然是个bad case：

在这里插入图片描述

0x2. 亮点2: 维度压缩减少坐标映射的计算量

index_add里面的第二个优化亮点是对Tensor的维度压缩，对应代码的https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L793， https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L787 ，这个维度压缩是什么意思呢？

假设index_add操作的输入Tensor是三个维度假设形状为(32, 1024, 1024)，而dim设置为0。那么在cuda Kernel中索引位置的时候是可以提前把dim后面的维度给合并起来的（这里使用TensorInfo数据结构来完成，其实本质上就是操作这个TensorInfo对象维护的Tensor的stride和size，具体可见这里的实现：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/CollapseDims.h#L22），这样子原始的输入Tensor的形状就会变成（32, 1024）。这样在indexFuncSmallIndex和indexFuncLargeIndex Kernel里面做坐标映射的时候就可以降低计算量以及降低对全局内存的访问提升带宽。注意，这里的维度压缩也可以压缩dim之前的所有维度为一个维度，这样子最终kernel需要处理的self输入张量维度只有1，2，3三种情况。

虽然这个优化是算法层面的优化，但是也间接让cuda kernel的带宽进行了提升和计算量进行了下降。实际上这个思路也启发了我在oneflow中实现index_add的kernel，我也是间接做了维度压缩。以这个例子来说：

x = torch.randn(32, 1024, 1024).to("cuda")
t = torch.randn(15, 1024, 1024).to("cuda")
index = torch.randint(0, 32, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

使用ncu在a100 pcie 40g上profile，我发现使用了维度压缩优化之后将这个cuda kernel从接近300+us的运行速度提升到了180+ us。

0x3. 实战性能表现

我这里对比了一下PyTorch的index_add和oneflow中index_add的性能表现。做性能profile的时候，我使用了以下脚本：

import torch

x = torch.randn(32*1024*1024).to("cuda")
t = torch.randn(15).to("cuda")
index = torch.randint(0, 1024, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32*1024, 1024).to("cuda")
t = torch.randn(15, 1024).to("cuda")
index = torch.randint(0, 1024, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32, 1024, 1024).to("cuda")
t = torch.randn(15, 1024, 1024).to("cuda")
index = torch.randint(0, 32, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32*1024*1024).to("cuda")
t = torch.randn(1024).to("cuda")
index = torch.randint(0, 1024, (1024,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32*1024, 1024).to("cuda")
t = torch.randn(1024, 1024).to("cuda")
index = torch.randint(0, 1024, (1024,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

测试环境为 A100 PCIE 40G，测试结果如下：

框架self tensor的shapedimsource shapeindex shape速度
PyTorch(32 * 1024 *1024,)0(15)(15)17.15us
OneFlow(32 * 1024 *1024,)0(15)(15)12us
PyTorch(32 * 1024, 1024)0(15, 1024)(15)27.78us
OneFlow(32 * 1024, 1024,)0(15, 1024)(15)26.98us
PyTorch(32, 1024, 1024)0(15, 1024, 1024)(15)186.88us
OneFlow(32 * 1024 *1024,)0(15, 1024, 1024)(15)247.10us
PyTorch(32 * 1024 *1024,)0(1024)(1024)7.9us
OneFlow(32 * 1024 *1024,)0(1024)(1024)7.79us
PyTorch(32 * 1024, 1024,)0(1024, 1024)(1024)27.87us
OneFlow(32 * 1024, 1024,)0(1024, 1024)(1024)28.67us

整体来说 PyTorch 在 index Tensor元素很小，但Tensor很大的情况下相比于oneflow有一些性能提升，其它情况和 OneFlow 基本持平，也有一些case是慢于oneflow比如index很小但是输入Tensor只有一个维度的情况下，这个时候PyTorch只会启动一个Block以及一个Thread，这显然是个bad case。OneFlow的index_add实现在 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/index_add_kernel.cu ，我们并没有针对 index 的大小来单独派发kernel，所以在某些case上性能暂时比PyTorch低一些，后续有需求的话可以继续优化下。

0x4. 总结

我这里相对粗糙的学习了一下调研PyTorch index_add这个算子的cuda实现的优化技术。但PyTorch的这个index_add实现仍然有一些改进空间，比如IndexToOffset的实现有取模操作，这个可以改成一次乘法和减法，可以节省计算指令。然后index_add 的两个kernel来说，GridSize和BlockSize并不是很合理，有改进空间。

0x5. 相关链接

  • https://github.com/pytorch/pytorch
  • https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/index_add_kernel.cu
  • https://github.com/BBuf/how-to-optim-algorithm-in-cuda

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.coloradmin.cn/o/163481.html

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈，一经查实，立即删除！

相关文章

2022，我感受到了CSDN不平凡

2022，我感受到了CSDN不平凡

最初注册CSDN&#xff0c;只是因为老师的要求&#xff0c;负责教C课程的老师让同学们注册CSDN&#xff0c;并经常更新自己的博客。虽然注册了CSDN的博客&#xff0c;也写了几篇博客文章&#xff0c;可最初我并不理解老师为什么要让我们注册&#xff0c;可是随着在CSDN驻留时间的…
阅读更多...
运行时数据区

运行时数据区

目录 一、概述 1.1、数据区 1.2、JAVA线程数据区 二、线程 2.1、JVM线程概述 2.2、JVM系统线程 三、PC寄存器 3.1、寄存器概述 3.2、作用 3.3、常见问题 一、概述 1.1、数据区 内存是非常重要的系统资源&#xff0c;是硬盘和CPU的中间仓库及桥梁&#xff0c;承载着操…
阅读更多...
vulnhub DC系列 DC-5

vulnhub DC系列 DC-5

总结&#xff1a; 下载地址 DC-5.zip (Size: 521 MB)Download: http://www.five86.com/downloads/DC-5.zipDownload (Mirror): https://download.vulnhub.com/dc/DC-5.zip使用方法:解压后&#xff0c;使用vm直接打开ova文件。 漏洞利用 信息收集 这里还是使用DC-1的方法 1.给靶…
阅读更多...
AR Foundation

AR Foundation

AR Session 在一个AR应用中有且只允许存在一个AR Session 包括两个组件 &#xff1a;AR Session &#xff08;用于管理 Session&#xff09;、AR Input Manager (用于管理输入的一些信息) AR Session&#xff08;用于管理 Session&#xff09; 作用&#xff1a;管理AR应用状…
阅读更多...
01.数据的存储

01.数据的存储

1. 数据类型介绍 1&#xff09;基本的内置类型&#xff1a;char //字符数据类型 short //短整型 int //整形 long //长整型 long long //更长的整形 float //单精度浮点数 double //双精度浮点数 2&#xff09;类型的意义&#xff1a; 使用这个类型开辟内存空间的大小&#xff…
阅读更多...
Docker快速入门自用笔记

Docker快速入门自用笔记

1. Docker - 介绍 不同Linux的内核一致。 2. Docker - 与虚拟机的不同 3. Docker - Docker架构 镜像&#xff08;只读&#xff09;&#xff1a;应用程序及其所需依赖、函数库、环境、配置等文件打包在一起&#xff0c;称为镜像。 容器&#xff1a;镜像中的应用程序运行后…
阅读更多...
一文看懂人机对话

一文看懂人机对话

人机对话概述 人机对话是指&#xff0c;让机器理解和运用自然语言实现人机通信的技术&#xff0c;如图1所示&#xff0c;通过人机对话交互&#xff0c;用户可以查询信息&#xff0c;如示例中的第一轮对话&#xff0c;用户查询天气信息&#xff0c;用户也可以和机器机型聊天&am…
阅读更多...
Pytorch深度学习【十三】

Pytorch深度学习【十三】

LeNet网络 基本结构图 构造思路 先用卷积层来学习图片空间信息池化层降低敏感度全连接层来转换到类别空间 代码实现 import torch from torch import nn from d2l import torch as d2l class Reshape(nn.Module):def forward(self ,x):return x.view(-1, 1, 28, 28) # vie…
阅读更多...
ElasticSearch7 Kibana集群安装

ElasticSearch7 Kibana集群安装

文章目录ElasticSearch安装下载安装包基础环境安装JDK安装修改Linux配置安装ES启动报错bootstrap check failure [1] of [1]: memory locking requested for elasticsearch process but memory is not lockedKibana安装Kibana简介Kibana下载Kibana安装ElasticSearch安装 下载安…
阅读更多...
RS485接口电路设计

RS485接口电路设计

RS485接口是串口的一种&#xff0c;常常用在一些工业控制中&#xff0c;485通信是半双工的通信&#xff0c;一条总线最多可连接256个485设备。RS485信号速率最高可以到10Mbps&#xff0c;传输距离最大能到1000多米485接口包含一对差分信号A和B&#xff0c;如下图所示RS485标准规…
阅读更多...
springcloud3 Nacos的集群搭建

springcloud3 Nacos的集群搭建

一 nacos的集群搭建说明 1.1 版本选择的说明 本案例采用的版本是1.4.2 &#xff0c;当选用其他版本均有不同的问题&#xff1a; 1.4.4 nginx代理后&#xff0c;输入密码登录进去。 2.2.0 &#xff0c;2.1.2等 2.x版本均是 3个节点的启动之后&#xff0c;只能访问一个节点。…
阅读更多...
David Murray 加入 VeraViews 担任需求总监

David Murray 加入 VeraViews 担任需求总监

近日&#xff0c;VeraViews宣布任命 David Murray 为需求总监。 VeraViews 是一个旨在提供透明和可审计的无效流量 (IVT) 预防的广告技术平台&#xff0c;很高兴地宣布任命David Murray为 VeraViews 需求总监。 大卫是伦敦的一位商业专家&#xff0c;他在英国许多最大、商业上…
阅读更多...
三种方式二叉树求从根到叶子结点的所有路径的方法整理

三种方式二叉树求从根到叶子结点的所有路径的方法整理

二叉树求从根到叶子结点的所有路径的方法整理 1. 利用递归和回溯方法求解 思路&#xff1a; 利用递归和回溯的方法求解 首先将当前结点加入到path中&#xff0c;然后判断是否为叶子结点&#xff0c;如果为叶子结点&#xff0c;则保存path路径如果不是叶子结点&#xff0c;则…
阅读更多...
IDEA 开发工具

IDEA 开发工具

文章目录IDEA 开发工具一、IDEA 概述二、IDEA 下载和安装1. IDEA 下载2. IDEA 安装三、IDEA 中的第一个代码1. IDEA 项目结构介绍2. IDEA 中的第一个代码1&#xff09;创建 Project 项目2&#xff09;创建 Module 模块3&#xff09;创建 class 类4&#xff09;在类中编写代码5&…
阅读更多...
C++指针

C++指针

指针的基本概念 作用&#xff1a;通过指针间接访问内存。 内存编号从0开始&#xff0c;一般使用16进制表示。 可以利用指针变量保存地址。 指针变量定义和使用 数据类型*变量名 #include<iostream> using namespace std;int main(){//定义指针int a10;//指针定义语…
阅读更多...
绕线机-排线伺服速度解算FC（比例随动编程应用）

绕线机-排线伺服速度解算FC（比例随动编程应用）

绕线机模型和算法详细讲解，专栏已有几篇介绍文章，这篇属于补充和升华。在介绍这篇文章算法之前大家有必要了解什么是"随动系统"。相关链接如下： 绕线机机械结构相关介绍： S7-200 SMART PLC和V20变频器绕线机控制应用_RXXW_Dor的博客-CSDN博客_plc绕线机绕线机P…
阅读更多...
循迹小车基本原理和代码实现

循迹小车基本原理和代码实现

目录 一、循迹模块使用 二、接线方式 三、循迹小车原理 四、代码实现 一、循迹模块使用 1、TCRT5000传感器的红外发射二极管不断发射红外线&#xff0c;当发出的红外线没有被反射回来或被反射回来但强度不太够时&#xff0c;红外接收管一直处于关断状态&#xff0c;此时模…
阅读更多...
将本地SpringBoot微服务制作成Docker Image镜像包然后运行

将本地SpringBoot微服务制作成Docker Image镜像包然后运行

将本地SpringBoot微服务制作成Docker Image镜像包然后运行 文章目录将本地SpringBoot微服务制作成Docker Image镜像包然后运行1&#xff09;环境准备2&#xff09;验证jar包3&#xff09;制作images镜像包4&#xff09;注意1&#xff09;环境准备 准备一个SpringBoot项目&…
阅读更多...
音视频笔记

音视频笔记

音视频笔记 基础理论知识 雷神的博客 [总结]视音频编解码技术零基础学习方法 视频压缩 H264码流结构 码流结构&#xff1a;原来你是这样的H264 H265码流结构 音视频基础&#xff1a;H265/HEVC&码流结构 FFMPEG 雷神的博客 [总结]FFMPEG视音频编解码零基础学习方法…
阅读更多...
1583_AURIX_TC275_SMU的控制以及FSP

1583_AURIX_TC275_SMU的控制以及FSP

全部学习汇总&#xff1a; GreyZhang/g_TC275: happy hacking for TC275! (github.com) SMU的软件控制接口主要是实现了一些控制命令&#xff0c;用于控制SMU的状态机以及FSP。具体的内容在上面的一部分表格以及接下来的一页中的表格中有介绍。 涉及到ed命令&#xff0c;用几个…
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023