与金属的非原子平行还原

Question

我刚刚进入并行缩减的世界。我正在尝试用 Metal 实现这个。我已经能够使用原子类型和 atomic_fetch_* 函数成功编写一个简单的版本。

我现在正在尝试对非原子变量（一个简单的结构）做一些类似的事情。

定义如下：

struct Point2
{
    int x;
    int y;
};

使用这样的内核函数：

kernel void compareX(const device Point2 *array [[ buffer(0) ]],
                 device Point2 *result [[ buffer(1) ]],
                 uint id [[ thread_position_in_grid ]],
                 uint tid [[ thread_index_in_threadgroup ]],
                 uint bid [[ threadgroup_position_in_grid ]],
                 uint blockDim [[ threads_per_threadgroup ]]) {

    threadgroup Point2 shared_memory[THREADGROUP_SIZE];

    uint i = bid * blockDim + tid;
    shared_memory[tid] = array[i];

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // reduction in shared memory
    for (uint s = 1; s < blockDim; s *= 2) {
        if (tid % (2 * s) == 0 && shared_memory[tid + s].x < shared_memory[tid].x) {

            shared_memory[tid] = shared_memory[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (0 == tid ) {
///THIS IS NOT CORRECT
        result[0] = shared_memory[0];
    }

}

我一开始以为复制 to/from 缓冲区的内存出了问题，但我已经验证 to/from CPU/GPU 与结构一起正常工作。然后我意识到它与跨线程组同步有关。

CUDA 有很多 examples/doc，但其他的很少，而且 CUDA 并不总是能很好地转换为 Metal。

在没有原子类型的情况下，如何实现跨线程组同步？

内核正在尝试获取输入数组中的最小点。现在，由于写入命令，结果会在执行过程中发生变化。

Answer 1

这可能不是最正确或最佳的解决方案。但这是我在为此苦苦挣扎了一段时间后想出的。如果其他人找到更好的解决方案，请post！这也可能与不同版本的 Metal 过时。

我首先尝试在我的结构上使用 Metal 语言中包含的 _atomic<T>。这应该有效。经过一番努力之后，我终于检查了文档，意识到模板目前被苹果限制为 bool、int 和 uint。

然后我尝试使用 atomic int 来 "lock" 关键比较部分，但没有成功保护关键部分。我很可能在这个实现上做错了什么，并且可以看到它在工作。

然后我简化为 return 索引而不是点，这允许我再次对结果使用 atomic_int。有点作弊，仍然使用原子来减少。但是它有效，所以我可以继续前进。

这是内核现在的样子：


kernel void compareX(const device Point2 *array [[ buffer(0) ]],
                     device atomic_int *result [[ buffer(1) ]],
                     uint id [[ thread_position_in_grid ]],
                     uint tid [[ thread_index_in_threadgroup ]],
                     uint bid [[ threadgroup_position_in_grid ]],
                     uint blockDim [[ threads_per_threadgroup ]]) {

    threadgroup int shared_memory[THREADGROUP_SIZE];
    uint i = bid * blockDim + tid;
    shared_memory[tid] = i;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    for (uint s = 1; s < blockDim; s *= 2) {
        if (tid % (2 * s) == 0) {
            // aggregate the index to our smallest value in shared_memory
            if ( array[shared_memory[tid + s]].x < array[shared_memory[tid]].x) {
                shared_memory[tid] = shared_memory[tid + s];
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    if (0 == tid ) {
        // get the current index so we can test against that
        int current = atomic_load_explicit(result, memory_order_relaxed);

        if( array[shared_memory[0]].x < array[current].x) {
            while(!atomic_compare_exchange_weak_explicit(result, &current, shared_memory[0], memory_order_relaxed, memory_order_relaxed)) {
                // another thread won. Check if we still need to set it.
                if (array[shared_memory[0]].x > array[current].x) {
                    // they won, and have a smaller value, ignore our best result
                    break;
                }
            }
        }
    }
}

与金属的非原子平行还原

Non Atomic Parallel Reduction with Metal

parallel-processing

gpgpu

metal