如何使用两个`u32`缓冲区在`u64`上正确模拟`atomicAdd`？

Question

我正在尝试对

u64

进行原子操作。但由于不支持，该数字存储在两个

u32

缓冲区中

问题是我不确定如何正确执行

atomicAdd

来模拟它对

u64

产生的效果。同时避免其他线程在加载和存储值之间修改内存。

我目前的想法是这样的：

fn tou64(value: u32) -> vec2u {
        return vec2u(u32(value / BASE), value % BASE);
}

fn add(a: vec2u, b: vec2u) -> vec2u {
    let x = a.x + b.x + u32((a.y + b.y) / BASE);
    let y = (a.y + b.y) % BASE;
    return vec2u(x, y);
}

fn main() {
// .....

// convert the value from u32 to 2-buffer representation of u64
let b: vec2u = tou64(value);
// fetch the old value from the 2 buffers
var a = vec2u(0); 
a.x = atomicLoad(&buffer[index]);
a.y = atomicLoad(&buffer[index+1]);
// add the value to the buffer value
let result = add(a, b);
// store back the buffer results 
atomicStore(&buffer[index], result.x);
atomicStore(&buffer[index+1], result.y);
}

仅当没有其他线程同时修改缓冲区时，这才有效

index

。但除此之外，这是一个非常薄弱的实现。线程 1 可以更改

buffer[index+1]

的值，而线程 2 仅读取旧的

buffer[index]

值和新的

buffer[index+1]

值

编辑：在CUDA指南中，注意到：

注意任何原子操作都可以基于atomicCAS()（Compare And Swap）来实现

并且提供了 AtomicAdd on double 的示例

#if __CUDA_ARCH__ < 600
__device__ double atomicAdd(double* address, double val)
{
    unsigned long long int* address_as_ull =
                              (unsigned long long int*)address;
    unsigned long long int old = *address_as_ull, assumed;

    do {
        assumed = old;
        old = atomicCAS(address_as_ull, assumed,
                        __double_as_longlong(val +
                               __longlong_as_double(assumed)));

    // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
    } while (assumed != old);

    return __longlong_as_double(old);
}
#endif

这也可以使用

atomicExchange

应用于 webgpu 吗？ 这个答案展示了如何对用户定义的类型进行原子操作。除了 webgpu 之外，我怎样才能做类似的事情？

Answer 1

正如其他人所说，简单的答案是否定的。 使用 WGSL 提供的原语没有简单的方法来获取 64 位添加。

但是，如果您的数据表现足够良好，则可以解决此问题。 例如，不要使用单个 64 位加法器，而是使用 4 个 16 位加法器。 让您的程序将 64 位数字拆分为 (x & 0xFFFF)、((x >> 16) & 0xFFFF)、((x >> 32) & 0xFFFF) 和 ((x >> 48) & 0xFFFF)。 以原子方式将每个片段添加到四个位置 a0、a1、a2、a3。

然后您的 JavaScript 程序可以将它们组合在一起。 a0 + (a1 << 16) + (a2 << 32) + (a3 << 48).

如果您执行的加法少于 2^16 次，则此方法有效，这样就不会溢出。 如果您正在进行 2^24 次加法，则将数据分成 8 个 8 位块。

本质上你只是让 GPU 做加法，然后让 JS 处理进位。