SetAtomicAdd

功能说明

设置接下来从VECOUT到GM的数据传输是否进行原子累加,可根据参数不同设定不同的累加数据类型。

函数原型

参数说明

表1 接口参数说明

参数名

输入/输出

描述

T

输入

设定不同的累加数据类型。

Atlas 训练系列产品,支持float

Atlas推理系列产品AI Core,支持float/half/int16_t

Atlas A2训练系列产品,支持float/half/int16_t/int32_t/int8_t

返回值

支持的型号

Atlas 训练系列产品

Atlas推理系列产品AI Core

Atlas A2训练系列产品

约束说明

调用示例

本演示示例模拟3个核进行数据处理,使用DataCopy从VECIN搬出到外部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