为自定义 PyTorch 激活函数启动正确数量的 CUDA 块

Question

我目前正在为 PyTorch 的自定义操作（激活）编写 CUDA 内核，但我对任何形式的 GPU 编程都不熟悉。作为参考，我正在学习 Custom C++ & CUDA extension 教程。

我想做的那种操作的简化示例：

假设我有一个输入张量 X_in，它可以是任何形状和暗淡的（例如 (16, 3, 50, 100) 之类的东西）。假设我还有一个一维张量，T（例如，T 可以是形状为 (100,) 的张量）。

对于 X_in 中的每个值，我想计算一个应该小于 len(T) 的“索引”值。然后输出基本上是 T 中该索引的值，乘以或加上某个常数。这类似于“查找 table”操作。

示例内核：

__global__ void inplace_lookup_kernel(
    const scalar_t* __restrict__ T,
    scalar_t* __restrict__ X_in,
    const int N) {

    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    const int idx = int(X_in[i]) % N;

    X_in[i] = 5 * T[idx] - 3;
}

我也希望就地进行操作，这就是为什么要将输出计算到 X_in。

我的问题是，像这样一个要逐点应用于X_in的每个值的操作，如何确定启动一个好的threads/blocks个数的方式？在 Custom C++ & CUDA extension 教程中，他们这样做是：

const int threads = 1024;
const dim3 blocks((state_size + threads - 1) / threads, batch_size);

对于他们的用例，操作（lstm 变体）具有特定的输入格式，因此具有固定的维数，可以从中计算块。

但是我正在编写的操作应该接受任何尺寸和形状的输入。对于这种情况，计算块号的正确方法是什么？

我目前正在做以下事情：

const int threads = 1024;
const int nelems = int(X_in.numel());
const dim3 blocks((nelems + threads - 1) / threads);

但是我是凭直觉这样做的，并不确定。有没有更好或正确的方法来做到这一点？如果我像教程中那样以 blocks(otherdim_size, batch_size) 格式定义块，是否有任何计算优势？

Answer 1

我在这里推测，但是 - 因为你的操作似乎完全是元素的 (w.r.t。X_in);而且您似乎没有使用有趣的特定于 SM 的资源，例如共享内存，也没有使用每个线程的大量寄存器，我认为网格分区并不那么重要。就把X_in看成一个一维数组，根据它在内存中的布局，用一个一维网格，块大小，哦，256，或者512，或者1024。

当然 - 始终尝试您的选择，以确保您不会被意想不到的行为所困扰。

为自定义 PyTorch 激活函数启动正确数量的 CUDA 块

Launching the right number of CUDA blocks for a custom PyTorch activation function

c++

cuda

pytorch

libtorch