文档
注册
评分
提单
论坛
小AI

SetAtomicAdd

功能说明

调用该接口后,可对后续的从VECOUT/L0C/L1到GM的数据传输开启原子累加,通过模板参数设定不同的累加数据类型。

函数原型

1
2
template <typename T>
__aicore__ inline void SetAtomicAdd() {}

参数说明

表1 模板参数说明

参数名

描述

T

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

Atlas 训练系列产品,支持的数据类型为:float;支持的数据通路为VECOUT->GM。

Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为float/half/int16_t;支持的数据通路为VECOUT->GM。

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为float/half/int16_t/int32_t/int8_t/bfloat16_t;支持的数据通路为VECOUT/L0C/L1->GM。

Atlas 200/500 A2推理产品,支持的数据类型为float/half/int16_t/int32_t;支持的数据通路为VECOUT/L0C/L1->GM

返回值

支持的型号

Atlas 训练系列产品

Atlas推理系列产品(Ascend 310P处理器)AI Core

Atlas A2训练系列产品/Atlas 800I A2推理产品

Atlas 200/500 A2推理产品

注意事项

  • 累加操作完成后,建议通过SetAtomicNone关闭原子累加,以免影响后续相关指令功能。
  • 该指令执行前不会对GM的数据做清零操作,开发者根据实际的算子逻辑判断是否需要清零,如果需要自行进行清零操作。

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
// 本演示示例使用DataCopy从VECOUT搬出到外部dstGlobal时进行原子累加。
// 本例子结果经过3个vector核累加得到,并且输出(dstGM)上的元素在调用核函数之前被初始化为0。
#include "kernel_operator.h"

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()
    {
        AscendC::LocalTensor<float> src0Local = inQueueSrc0.AllocTensor<float>();
        AscendC::LocalTensor<float> src1Local = inQueueSrc1.AllocTensor<float>();
        AscendC::DataCopy(src0Local, src0Global, 256);
        AscendC::DataCopy(src1Local, src1Global, 256);
        inQueueSrc0.EnQue(src0Local);
        inQueueSrc1.EnQue(src1Local);
    }
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<float> src0Local = inQueueSrc0.DeQue<float>();
        AscendC::LocalTensor<float> src1Local = inQueueSrc1.DeQue<float>();
        AscendC::LocalTensor<float> dstLocal = outQueueDst.AllocTensor<float>();
        AscendC::Add(dstLocal, src0Local, src1Local, 256);
        outQueueDst.EnQue(dstLocal);
        inQueueSrc0.FreeTensor(src0Local);
        inQueueSrc1.FreeTensor(src1Local);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<float> dstLocal = outQueueDst.DeQue<float>();
        AscendC::SetAtomicAdd<float>();
        AscendC::DataCopy(dstGlobal, dstLocal, 256);
        AscendC::SetAtomicNone();
        outQueueDst.FreeTensor(dstLocal);
    }
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc0, inQueueSrc1;
    AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueueDst;
    AscendC::GlobalTensor<float> src0Global, src1Global, dstGlobal;
};

extern "C" __global__ __aicore__ void set_atomic_add_ops_kernel(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
{
    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
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词