如何扩展CUDA.@atomic操作支持自定义结构体类型

CUDA.jl原生@atomic默认仅支持硬件原生实现原子操作的基础数值类型（整数、浮点数等），无法直接作用于自定义结构体。针对Dual数场景，有三类可落地的实现方式，按性能和易用性排序如下：

方案1：逐字段原生原子操作（性能最优） #

如果自定义结构体运算可以拆分为独立的基础字段运算（比如Dual的加法本质是x、y分量各自相加），不需要对整个结构体做原子操作，直接对每个字段调用原生@atomic即可，性能和基础类型原子操作完全一致，没有额外开销。
修改后的可运行示例：

using CUDA

struct Dual
    x::Float64
    y::Float64
end
import Base:+
+(a::Dual, b::Dual) = Dual(a.x+b.x, a.y+b.y)

cu0 = CuArray([Dual(1, 2), Dual(2,3)])
cu1 = CuArray([Dual(1, 2), Dual(2,3)])
indexes = CuArray([1, 1])

function my_kernel(dst, src, idx)
    index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    @inbounds if index <= length(idx)
        pos = idx[index]
        # 对两个字段分别做原子加，完全调用硬件原生实现
        CUDA.@atomic dst[pos].x += src[index].x
        CUDA.@atomic dst[pos].y += src[index].y
    end
    return nothing
end

@cuda threads=100 my_kernel(cu0, cu1, indexes)

这个方案没有架构兼容问题，只要字段是支持原子操作的基础类型就能运行，是优先推荐的实现方式。

方案2：基于CAS循环手写原子逻辑（通用方案） #

如果必须对整个结构体做原子操作（比如结构体字段存在运算依赖，无法拆分），可以利用CUDA硬件支持的比较并交换（CAS）原语，通过循环实现任意纯值类型的原子操作。
使用这个方案的前提是自定义结构体满足两个要求：

• 是纯值类型（isbitstype(T) == true），没有堆引用字段，内存布局固定
• 内存对齐到自身大小，总位宽不超过硬件支持的CAS位宽：32/64位CAS全架构支持，128位CAS需要sm_70（伏特架构）及以上GPU
  核心实现代码如下，可直接替换原kernel中的@atomic行：

function my_kernel(dst, src, idx)
    index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    @inbounds if index <= length(idx)
        pos = idx[index]
        old = dst[pos]
        while true
            new = old + src[index]
            # 调用硬件CAS：如果目标位置值等于old，就替换为new，返回位置上的旧值
            res = CUDA.atomic_cas!(pointer(dst, pos), old, new)
            res === old && break
            old = res
        end
    end
    return nothing
end

这个实现的性能在冲突率不高的场景下和原生原子操作差距很小，唯一缺点是需要手写循环逻辑，不能直接用@atomic的简洁语法。

方案3：扩展CUDA.jl原子操作接口，兼容@atomic语法 #

如果想保留原代码里CUDA.@atomic的写法，可以直接扩展CUDA.jl的原子操作分派方法，让宏自动识别自定义类型，调用你实现的CAS逻辑。
只需要在运行kernel前添加如下方法定义即可：

import CUDA: atomic_add!
function atomic_add!(ptr::CuDevicePtr{Dual}, val::Dual)
    old = unsafe_load(ptr)
    while true
        new = old + val
        res = CUDA.atomic_cas!(ptr, old, new)
        res === old && return old
        old = res
    end
end

扩展完成后，最初写的kernel代码不需要做任何修改，CUDA.@atomic dst[idx[index]] = dst[idx[index]] + src[index]就可以正常运行。如果需要支持减法、乘法等其他原子操作，按照相同逻辑实现对应的atomic_sub!、atomic_mul!方法即可。

注意事项 #

• 不要为非isbits类型实现上述原子逻辑，会直接触发GPU内存访问错误
• 如果结构体总大小超过128位，硬件不支持原生CAS，需要用软件自旋锁实现，性能会大幅下降，这种场景强烈建议用方案1拆分字段实现
• CAS循环存在理论上的活锁可能，但GPU实际运行中冲突概率极低，不会影响正常使用

内容的提问来源于stack exchange，提问作者Timschneider7339 Schneider