跳到主要内容

原子函数

原子函数对驻留在全局内存或共享内存中的数据执行读 - 修改 - 写原子操作。原子函数只能在设备函数中使用。

概述​

原子函数是并行编程中的重要工具，用于避免多个线程同时访问共享数据时的竞争条件。MUSA 提供以下类型的原子函数：

  • 算术原子函数: atomicAdd, atomicSub, atomicInc, atomicDec, atomicExch
  • 比较并交换: atomicCAS
  • 位原子操作: atomicAnd, atomicOr, atomicXor
备注

任何原子操作都可以基于 atomicCAS()（比较并交换）实现。

算术原子函数​

atomicAdd()​

int atomicAdd(int* address, int val);
unsigned int atomicAdd(unsigned int* address, unsigned int val);
unsigned long long int atomicAdd(unsigned long long int* address, unsigned long long int val);
float atomicAdd(float* address, float val);
double atomicAdd(double* address, double val);

读取位于全局或共享内存中地址 address 存储的值 old，计算 (old + val)，并将结果存储回同一地址。返回 old。

示例:

__global__ void countKernel(int* counter, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        // 原子增加计数器
        atomicAdd(counter, 1);
    }
}

// 主机代码
int* d_counter;
int h_counter = 0;
musaMalloc(&d_counter, sizeof(int));
musaMemcpy(d_counter, &h_counter, sizeof(int), musaMemcpyHostToDevice);

countKernel<<<gridSize, blockSize>>>(d_counter, 10000);

musaMemcpy(&h_counter, d_counter, sizeof(int), musaMemcpyDeviceToHost);
printf("Counter: %d\n", h_counter);  // 输出：10000

atomicSub()​

int atomicSub(int* address, int val);
unsigned int atomicSub(unsigned int* address, unsigned int val);

读取位于全局或共享内存中地址 address 存储的 32 位值 old，计算 (old - val)，并将结果存储回同一地址。返回 old。

atomicExch()​

int atomicExch(int* address, int val);
unsigned int atomicExch(unsigned int* address, unsigned int val);
unsigned long long int atomicExch(unsigned long long int* address, unsigned long long int val);
float atomicExch(float* address, float val);

读取位于全局或共享内存中地址 address 存储的值 old，将 val 存储到同一地址。返回 old。

atomicInc()​

unsigned int atomicInc(unsigned int* address, unsigned int val);

读取位于全局或共享内存中地址 address 存储的 32 位值 old，计算 ((old >= val) ? 0 : (old+1))，并将结果存储回同一地址。返回 old。

示例 - 循环计数器：

// address 处的值从 0 递增到 val-1，然后回绕到 0
unsigned int old = atomicInc(counter, 100);  // 0 -> 1 -> 2 -> ... -> 99 -> 0 -> ...

atomicDec()​

unsigned int atomicDec(unsigned int* address, unsigned int val);

读取位于全局或共享内存中地址 address 存储的 32 位值 old，计算 (((old == 0) || (old > val)) ? val : (old-1))，并将结果存储回同一地址。返回 old。

atomicCAS()​

int atomicCAS(int* address, int compare, int val);
unsigned int atomicCAS(unsigned int* address, unsigned int compare, unsigned int val);
unsigned long long int atomicCAS(unsigned long long int* address, 
                                  unsigned long long int compare,
                                  unsigned long long int val);
unsigned short int atomicCAS(unsigned short int* address,
                              unsigned short int compare,
                              unsigned short int val);

比较并交换 (Compare And Swap): 读取位于全局或共享内存中地址 address 存储的值 old，计算 (old == compare ? val : old)，并将结果存储到同一地址。返回 old。

示例 - 实现自旋锁：

__device__ void lock(int* mutex) {
    while (atomicCAS(mutex, 0, 1) != 0) {
        // 自旋等待
    }
}

__device__ void unlock(int* mutex) {
    atomicExch(mutex, 0);
}

__global__ void criticalSection(int* mutex, int* shared_data) {
    lock(mutex);
    // 临界区代码
    *shared_data += 1;
    unlock(mutex);
}

位原子操作​

atomicAnd()​

int atomicAnd(int* address, int val);
unsigned int atomicAnd(unsigned int* address, unsigned int val);
unsigned long long int atomicAnd(unsigned long long int* address, unsigned long long int val);

读取位于全局或共享内存中地址 address 存储的值 old，计算 (old & val)，并将结果存储回同一地址。返回 old。

atomicOr()​

int atomicOr(int* address, int val);
unsigned int atomicOr(unsigned int* address, unsigned int val);
unsigned long long int atomicOr(unsigned long long int* address, unsigned long long int val);

读取位于全局或共享内存中地址 address 存储的值 old，计算 (old | val)，并将结果存储回同一地址。返回 old。

atomicXor()​

int atomicXor(int* address, int val);
unsigned int atomicXor(unsigned int* address, unsigned int val);
unsigned long long int atomicXor(unsigned long long int* address, unsigned long long int val);

读取位于全局或共享内存中地址 address 存储的值 old，计算 (old ^ val)，并将结果存储回同一地址。返回 old。

高级用法​

基于 atomicCAS 实现自定义原子操作​

// 双精度浮点原子加
__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)));
    // 注意：使用整数比较以避免 NaN 问题 (NaN != NaN)
    } while (assumed != old);

    return __longlong_as_double(old);
}

原子操作实现直方图​

__global__ void histogramKernel(float* data, int* bins, int n, int binCount) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        // 计算 bin 索引
        int bin = min((int)(data[idx] * binCount), binCount - 1);
        // 原子增加 bin 计数
        atomicAdd(&bins[bin], 1);
    }
}

原子操作实现归约​

__global__ void reduceKernel(float* input, float* output, int n) {
    __shared__ float shared[256];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 加载数据
    shared[tid] = (idx < n) ? input[idx] : 0.0f;
    __syncthreads();
    
    // 块内归约
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        __syncthreads();
    }
    
    // 每个块的结果原子加到全局输出
    if (tid == 0) {
        atomicAdd(output, shared[0]);
    }
}

性能注意事项​

注意事项说明
串行化同一地址的原子操作会串行化，影响性能
内存位置共享内存的原子操作比全局内存快
竞争程度线程竞争越激烈，性能影响越大
替代方案考虑先块内归约，再原子操作

优化示例：减少原子操作竞争​

// ❌ 性能较差：所有线程直接原子操作全局内存
__global__ void badHistogram(float* data, int* bins, int n, int binCount) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        int bin = min((int)(data[idx] * binCount), binCount - 1);
        atomicAdd(&bins[bin], 1);  // 高竞争
    }
}

// ✅ 性能较好：先块内归约，再原子操作
__global__ void goodHistogram(float* data, int* bins, int n, int binCount) {
    __shared__ int sharedBins[256];
    
    // 初始化共享内存 bins
    for (int i = threadIdx.x; i < binCount; i += blockDim.x) {
        sharedBins[i] = 0;
    }
    __syncthreads();
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        int bin = min((int)(data[idx] * binCount), binCount - 1);
        atomicAdd(&sharedBins[bin], 1);  // 共享内存原子操作，低延迟
    }
    __syncthreads();
    
    // 合并到全局 bins
    for (int i = threadIdx.x; i < binCount; i += blockDim.x) {
        atomicAdd(&bins[i], sharedBins[i]);  // 每个块只调用一次
    }
}

支持的数据类型和内存区域​

函数intuintullongfloatdoubleshort全局内存共享内存
atomicAdd✅✅✅✅✅❌✅✅
atomicSub✅✅❌❌❌❌✅✅
atomicExch✅✅✅✅❌❌✅✅
atomicInc❌✅❌❌❌❌✅✅
atomicDec❌✅❌❌❌❌✅✅
atomicCAS✅✅✅❌❌✅✅✅
atomicAnd✅✅✅❌❌❌✅✅
atomicOr✅✅✅❌❌❌✅✅
atomicXor✅✅✅❌❌❌✅✅

相关文档​

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。