gather算子的CUDA编程和算子测试

谨慎付费（看不懂试读博客不要订阅）

于 2025-01-19 22:38:19 发布

阅读量191

点赞数

CC 4.0 BY-SA版权

分类专栏：高性能计算文章标签： CUDA

本文链接：https://blog.csdn.net/forrestguang/article/details/145249199

高性能计算专栏收录该内容

54 篇文章 ¥239.90 ¥399.90

订阅专栏

知乎介绍参考添加链接描述
完整测试框架参考本人仓库
添加链接描述

gather算子的onnx定义参考添加链接描述，该算子的主要变换参考下图：
在这里插入图片描述
这里我们不妨以input = [A, dimsize, D]， indices = [B,C], axis = 1举例子，此时对应的output形状是[A,B,C,D]，并且根据gather算子定义，我们知道output[i,j,k,s] = input[i,indices[j,k], s]，下面我们引入一些定义：
outputSize = ABCD, inputSize = AD×dimsize， indSize = BC。
那么此时我们可以看出，首先input和output共享一个othersize = AD，我们可以这么思考，对于output来说，可以把output分成两部分，看成一个[othersize, indSize]的矩阵，同样的，input也可以分成两部分，看成一个[othersize, dimsize]的矩阵，因此output和input共享了othersize部分的数据，后面的主要数学变换集中在后面部分。
除此之外，不管axis在哪，对于input和output来说，对应的stride是一样的，这里的stride表示axis对应的步长，比如说此时output形状为[A,B,C,D]，对应的全局索引是i(BCD) + j(CD) + k(D) + s = i（BCD） + （j ×C+ k）×D + s，而input对应的全局索引是i(dimsize ×D) + indices[j,k]×D+s，此时的stride对应的就是D，也就是说，不管axis是多