SetAtomicAdd
功能说明
设置接下来从VECOUT到GM,L0C到GM,L1到GM的数据传输是否进行原子累加,可根据参数不同设定不同的累加数据类型。
函数原型
参数说明
参数名 |
输入/输出 |
描述 |
|---|---|---|
T |
输入 |
设定不同的累加数据类型。 Atlas 训练系列产品,支持float Atlas 训练系列产品,通路支持VECOUT到GM Atlas推理系列产品AI Core,支持float/half/int16_t Atlas推理系列产品AI Core,通路支持VECOUT到GM Atlas A2训练系列产品/Atlas 800I A2推理产品,支持float/half/int16_t/int32_t/int8_t/bfloat16_t Atlas A2训练系列产品/Atlas 800I A2推理产品,通路支持VECOUT到GM,L0C到GM,L1到GM |
返回值
无
支持的型号
Atlas 训练系列产品
Atlas推理系列产品AI Core
Atlas A2训练系列产品/Atlas 800I A2推理产品
注意事项
- 使用完后,建议关闭原子累加,以免影响后续相关指令功能。
- 该指令执行前不会对GM的数据做清零操作,开发者根据实际的算子逻辑判断是否需要清零,如果需要自行进行清零操作。
调用示例
本演示示例模拟3个核进行数据处理,使用DataCopy从VECOUT搬出到外部dstGlobal时进行原子累加。
namespace AscendC {
class KernelSetAtomicAdd {
public:
__aicore__ inline KernelSetAtomicAdd() {}
__aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
{
src0Global.SetGlobalBuffer((__gm__ float*)src0Gm);
src1Global.SetGlobalBuffer((__gm__ float*)src1Gm);
dstGlobal.SetGlobalBuffer((__gm__ float*)dstGm);
pipe.InitBuffer(inQueueSrc0, 1, 256 * sizeof(float));
pipe.InitBuffer(inQueueSrc1, 1, 256 * sizeof(float));
pipe.InitBuffer(outQueueDst, 1, 256 * sizeof(float));
}
__aicore__ inline void Process()
{
CopyIn();
Compute();
CopyOut();
}
private:
__aicore__ inline void CopyIn()
{
LocalTensor<float> src0Local = inQueueSrc0.AllocTensor<float>();
LocalTensor<float> src1Local = inQueueSrc1.AllocTensor<float>();
DataCopy(src0Local, src0Global, 256);
DataCopy(src1Local, src1Global, 256);
inQueueSrc0.EnQue(src0Local);
inQueueSrc1.EnQue(src1Local);
}
__aicore__ inline void Compute()
{
LocalTensor<float> src0Local = inQueueSrc0.DeQue<float>();
LocalTensor<float> src1Local = inQueueSrc1.DeQue<float>();
LocalTensor<float> dstLocal = outQueueDst.AllocTensor<float>();
Add(dstLocal, src0Local, src1Local, 256);
outQueueDst.EnQue(dstLocal);
inQueueSrc0.FreeTensor(src0Local);
inQueueSrc1.FreeTensor(src1Local);
}
__aicore__ inline void CopyOut()
{
LocalTensor<float> dstLocal = outQueueDst.DeQue<float>();
SetAtomicAdd<float>();
DataCopy(dstGlobal, dstLocal, 256);
SetAtomicNone();
outQueueDst.FreeTensor(dstLocal);
}
private:
TPipe pipe;
TQue<QuePosition::VECIN, 1> inQueueSrc0, inQueueSrc1;
TQue<QuePosition::VECOUT, 1> outQueueDst;
GlobalTensor<float> src0Global, src1Global, dstGlobal;
};
} // namespace AscendC
extern "C" __global__ __aicore__ void set_atomic_add_ops_kernel(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
{
AscendC::KernelSetAtomicAdd op;
op.Init(src0Gm, src1Gm, dstGm);
op.Process();
}
每个核的输入数据为:
Src0: [1,1,1,1,1,...,1] // 256个1
Src1: [1,1,1,1,1,...,1] // 256个1
最终输出数据: [6,6,6,6,6,...,6] // 256个6
父主题: 原子操作