在 Metal 中同步网格中的所有线程

Question

我正在尝试为 Metal 中的 n 大小向量编写范数或平方长度函数。为此，我计划让每个线程平方每个元素，然后选择一个线程来汇总所有元素。

这是我当前的内核：

#include <metal_stdlib>
#include <metal_compute>
using namespace metal;

kernel void length_squared(const device float *x [[ buffer(0) ]],
                           device float *s [[ buffer(1) ]],
                           device float *out [[ buffer(2) ]],
                           uint gid [[ thread_position_in_grid ]],
                           uint numElements [[ threads_per_grid ]])
{
    s[gid] = x[gid];// * x[gid];
    simdgroup_barrier(mem_flags::mem_none);
    if(gid == 0){
        for(uint i = 0; i < numElements; i++){
            *out += s[i];
        }
    }
}

不幸的是，此代码无法编译，因为“使用未声明的标识符simdgroup_barrier”。该方法记录在Metal Shading Language Specification 中。

有人遇到过吗？或者知道如何跨网格同步所有线程？ threadgroup_barrier 没有实现完全同步。

我是否错误地处理了这个问题？同步此操作的最佳方法是什么？

Answer 1

SIMD 组小于线程组，因此无法进行同步。

相反，您需要使用parallel reduction 来并行汇总这些值。 Here 是我找到的一些金属代码。

不过，如果您不介意单个线程进行所有求和，您可以运行一个单独的内核，只使用一个线程进行求和。当然，这可能会很慢。