CUDA 高效的 nd-array(张量）切片

CUDA Efficient nd-array(tensor) slicing

本文关键字：张量 切片 nd-array 高效 CUDA      更新时间：2023-10-16

假设我有 3d 数组作为扁平化 1d，大小为[N, M, K].我想像[0:N, 1:M, 0:K]一样处理其中的切片. 我创建了一个帮助程序函数，该函数通过切片数组中的索引对底层数组进行寻址(为简单起见，我只按二维切片(。

#define N somevalue
#define M somevalue
#define K somevalue
// i is an index in sliced array so we need to translate it into original one
template<class T, int FROM>
__device__   __forceinline__ T slice(const T * const __restrict__ x, const size_t i) {
auto batch_size = (M - FROM) * K;
auto batch_index = i / batch_size;
auto offset_0 = i % batch_size;
auto offset_1 = offset_0 / STATES;
auto offset_2 = offset_0 % STATES;
return x[batch_index * M * K + (offset_1 + FROM) * K + offset_2];
}

从NVidia剖析器中，我看到除法和模除法需要大量的计算能力。此外，大小不是 2 的幂，所以我不能直接使用移位技巧。

您有什么建议？ 众所周知，切片是TF中非常常见的操作，那么他们是如何解决的呢？

Cuda 是关于合并内存访问和 simd。而任意切片则恰恰相反。所以答案像往常一样：这取决于。

如果偏移量为并保持 1，请将内存布局更改为 M N K。如果被忽略的条目真的非常稀疏，请使用传统方式并闲置几个线程(是的，这很痛苦，但是一些没有模的threadIdx计算可能会更快(。否则，您将需要计算此双射映射线程/块 id 到元素 id，就像您在问题中写的那样。

有一些方法可以通过其他一些操作来表示模。但通常最好花时间改进内核的其他部分。

相关文章：