矩阵乘法中的扭曲级同步

Question

我正在尝试在 CUDA 中实现方阵乘法，并使用扭曲级基元优化点积的求和部分。我以前使用了一种幼稚的方法，但现在我尝试使用

__shfl_down_sync

来加快求和速度。这是初始代码：

for (int i = 0; i < tileSize; i++)
{
    sum += a[threadIdx.y * tileSize + i] * b[i * tileSize + threadIdx.x];
}
__syncthreads();

当tileSize设置为32时，我认为我可以通过使用扭曲级基元计算总和来进行优化。这是我尝试过的修改后的代码：

int mult = a[threadIdx.y * tileSize + threadIdx.x] * b[threadIdx.x * tileSize + threadIdx.y];
for (int i = tileSize >> 1; i > 0; i >>= 1)
{
    sum += __shfl_down_sync(-1, mult, i);
}

但是，我用这种方法没有得到正确的结果。我真的很感激任何更正或见解。谢谢！

Answer 1

如果要优化求和，需要确保编译器可以展开循环。

即：转换

loop:
    do stuff

进入：

    dostuff
    dostuff
    ...
    dostuff

这样循环开销就消失了。

这可以通过将循环的所有元素设置为编译时间常数来轻松完成。

template <typename T>
T warp_reduce_sum(const T a) {
    auto sum = a;
    for (auto i = 1; i < 32; i *= 2) { 
        sum += __shfl_down_sync(-1u, a, i);
    }
    return sum;
}

编译器会将其展开为 5 个连续的

shfl

指令，从而消除 for 循环。如果循环的元素不是编译时常量，则它无法执行此操作，这可能就是您看不到任何改进的原因。

但是，如果您有 GTX3000 (Ampere) 或更新的 GPU，那么您可以使用以下代码比

shfl

减少效果更好：

template <typename T>
T warp_reduce_sum(const T a) {
    auto sum = a;
#if __CUDA_ARCH__ < 800
    for (auto i = 1; i < 32; i *= 2) { 
        sum += __shfl_down_sync(-1u, a, i);
    }
#else
    sum = __reduce_add_sync(-1u, a);
#endif
    return sum;
}

reduce

的运行速度比 shfl 范围快大约 3 倍，但只有当 all 线程参与时，如果只有某些线程执行加法（如果

threadmask != -1

），速度会更慢。

在 Ampere 之前的硬件上，通过在

shfl

内存上使用atomicAdd，您可以做得比

__shared__

更好。 （请参阅：https://www.youtube.com/watch?v=47go5WpbULM&t=33s）。

//reduce 1024 values, adjust as needed to reduce fewer items.
__global__ void reduce_atom() {
    auto volatile x = threadIdx.x;
    const auto warpid = threadIdx.x / 32;
    const auto laneid = threadIdx.x % 32;
    __shared__ int smax[32];
    for (auto run = 0; run < 10; run++) {
        const auto volatile start = clock64();
        smax[warpid] = 0;
        __syncthreads();
        atomicAdd(&smax[laneid], x);
        __syncthreads();
        if (threadIdx.x < 32) {
            x = smax[threadIdx.x];
            __syncwarp();
            atomicAdd(&smax[0], x);
        }
        __syncwarp();
        const auto volatile end = clock64();
        if (threadIdx.x == 0) {
            printf("atomic: run: %i, max: %i, time: %i\n", run, smax[0], int(end-start));
        }
    }
}

如果您想查看

shfl

、

atomic

和

reduce

减少之间的时间差，可以复制粘贴 Off Godbolt 的完整代码。
https://cuda.godbolt.org/z/rfK395o9v

如果您在代码中调用上述方法，编译器将内联它，因为 nvcc 对于内联非常积极，这也将消除调用开销。

这仍然把大象🐘留在房间里。

a[threadIdx.y * tileSize + i] * b[i * tileSize + threadIdx.x];
  ^^^^^^^^^^^                                    ^^^^^^^^^^^

数组

a

的访问模式是通过

threadIdx.y

访问，而

b

是通过

threadIdx.x

访问。这意味着一个数组具有一种非常低效的跨步访问模式，因为它会导致数据从内存中读取多次，而不是只读取一次。
由于从 L2 缓存读取大约需要 400 个周期，而从主内存读取大约需要两倍的时间，因此这很容易淹没通过更好的减少而带来的任何加速。

您可以通过以下方式解决此问题：首先将

a

和

b

的相关部分读取到

__shared__

内存缓冲区中，然后从

a

内存中处理

b

和

__shared__

。

参见：https://stackoverflow.com/a/18856054/650492 以及：https://www.cstechera.com/2016/03/tiled-matrix-multiplication-using-shared-memory-in-cuda.html