CANN/Ascend C原子比较交换API

AtomicCas

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

产品支持情况

功能说明

调用该接口后，可在指定GM地址上进行原子比较，如果和value1相等，则把value2的值赋值到GM上；如果和value1不相等，则GM上的值不变。

函数原型

template <typename T> __aicore__ inline T AtomicCas(__gm__ T *address, T value1, T value2)

参数说明

表 1模板参数说明

表 2参数说明

返回值说明

GM地址上做原子操作前的数据。

约束说明

原子操作涉及标量计算，如果标量计算单元和搬运单元（MTE2/MTE3）在读写GM时存在数据依赖，开发者需要根据实际情况插入同步。

调用示例

extern "C" __global__ __aicore__ void atomic_simple_kernel(__gm__ uint8_t* dst， uint32_t dataSize) { // ... dst_global.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dst_gm), dataSize); LocalTensor<T> dstLocal = inQueueX.AllocTensor<T>(); uint32_t value1 = 1; uint32_t value2 = 2; uint32_t a = AscendC::AtomicCas(reinterpret_cast<__gm__ int32_t *>(dst), value1, value2); // 先执行完原子操作之后才能进行搬运操作，有数据依赖，需要手动插入MTE2等待Scalar的同步 event_t eventIdSToMte2 = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::S_MTE2>()); SetFlag<HardEvent::S_MTE2>(eventIdSToMte2); WaitFlag<HardEvent::S_MTE2>(eventIdSToMte2); DataCopy(dstLocal, dst_global, dataSize); // ... }

假设上述函数在3个核上执行，核1、核2、核3依次调度，结果示例如下：

原GM数据dst: [1,1,1,1,1,...,1] 核1： 原子计算后GM数据dst: [2,1,1,1,1,...,1] 返回值 a: 1 核2： 原子计算后GM数据dst: [2,1,1,1,1,...,1] 返回值 a: 2 核3： 原子计算后GM数据dst: [2,1,1,1,1,...,1] 返回值 a: 2

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit