金属非原子平行还原

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];
    }

}

我首先认为内存复制到缓冲区/从缓冲区出现问题，但我已经验证到/从 CPU/GPU 与结构正常工作。然后我意识到它与跨线程组同步有关。

有很多关于 CUDA 的示例/文档，但很少有其他任何东西，而且 CUDA 并不总是能很好地转化为 Metal。

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

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

Answer 1

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

我首先尝试在我的结构上使用金属语言中包含的_atomic<T>。这应该工作。在为此苦苦挣扎之后，我终于检查了文档并意识到模板目前被苹果限制为 bool's、int's 和 uint's。

然后我尝试使用 atomic int 来“锁定”关键比较部分，但实际上并未成功保护关键部分。我可能在这个实现中做错了什么，并且可以看到它工作。

然后我简化为返回索引而不是点，这允许我再次在结果上使用 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;
                }
            }
        }
    }
}