设置接下来从VECOUT到GM的数据传输是否进行原子累加,可根据参数不同设定不同的累加数据类型。
template <typename T> __aicore__ inline void SetAtomicAdd() {}
__aicore__ inline void SetAtomicNone();
参数名 |
输入/输出 |
描述 |
---|---|---|
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