SetAtomicMax
功能说明
原子操作函数，设置后续从VECOUT传输到GM的数据是否执行原子比较：将待拷贝的内容和GM已有内容进行比较，将最大值写入GM。
可通过设置模板参数来设定不同的数据类型。
函数原型
1 2 | template <typename T> __aicore__ inline void SetAtomicMax() {} |
参数说明
参数名 |
描述 |
|---|---|
T |
设定不同的数据类型。 Atlas A2训练系列产品/Atlas 800I A2推理产品，支持float/half/int16_t/int32_t/int8_t/bfloat16_t |
返回值
无
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
约束说明
- 使用完后，建议通过SetAtomicNone关闭原子最大操作，以免影响后续相关功能。
- 对于Atlas A2训练系列产品/Atlas 800I A2推理产品，目前无法对bfloat16_t类型设置inf/nan模式。
调用示例
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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | // 本演示示例使用DataCopy从VECOUT搬出到外部dstGlobal时进行原子最大。 #include "kernel_operator.h" static const int data_size = 256; template <typename T> class KernelDataCopyAtomicMax { public: __aicore__ inline KernelDataCopyAtomicMax() {} __aicore__ inline void Init(GM_ADDR src0_gm, GM_ADDR src1_gm, GM_ADDR dst_gm, uint32_t size) { this->size = size; src0Global.SetGlobalBuffer((__gm__ T *)src0_gm); src1Global.SetGlobalBuffer((__gm__ T *)src1_gm); dstGlobal.SetGlobalBuffer((__gm__ T *)dst_gm); pipe.InitBuffer(queueSrc0, 1, size * sizeof(T)); pipe.InitBuffer(queueSrc1, 1, size * sizeof(T)); pipe.InitBuffer(queueDst0, 1, size * sizeof(T)); pipe.InitBuffer(queueDst1, 1, size * sizeof(T)); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> src0local = queueSrc0.AllocTensor<T>(); AscendC::LocalTensor<T> src1local = queueSrc1.AllocTensor<T>(); AscendC::DataCopy(src0local, src0Global, size); AscendC::DataCopy(src1local, src1Global, size); queueSrc0.EnQue(src0local); queueSrc1.EnQue(src1local); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> src0local = queueSrc0.DeQue<T>(); AscendC::LocalTensor<T> src1local = queueSrc1.DeQue<T>(); AscendC::LocalTensor<T> dst0Local = queueDst0.AllocTensor<T>(); AscendC::LocalTensor<T> dst1Local = queueDst1.AllocTensor<T>(); AscendC::Abs(dst0Local, src0local, size); AscendC::Abs(dst1Local, src1local, size); queueDst0.EnQue(dst0Local); queueDst1.EnQue(dst1Local); queueSrc0.FreeTensor(src0local); queueSrc1.FreeTensor(src1local); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> dst0Local = queueDst0.DeQue<T>(); AscendC::LocalTensor<T> dst1Local = queueDst1.DeQue<T>(); AscendC::DataCopy(dstGlobal, dst1Local, size); AscendC::PipeBarrier<PIPE_MTE3>(); AscendC::SetAtomicMax<T>(); AscendC::DataCopy(dstGlobal, dst0Local, size); queueDst0.FreeTensor(dst0Local); queueDst1.FreeTensor(dst1Local); AscendC::SetAtomicNone(); } private: AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> queueSrc0; AscendC::TQue<AscendC::QuePosition::VECIN, 1> queueSrc1; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> queueDst0; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> queueDst1; AscendC::GlobalTensor<T> src0Global, src1Global, dstGlobal; uint32_t size; }; extern "C" __global__ __aicore__ void data_copy_atomic_max_kernel(GM_ADDR src0_gm, GM_ADDR src1_gm, GM_ADDR dst_gm) { KernelDataCopyAtomicMax<half> op; op.Init(src0_gm, src1_gm, dst_gm, data_size); op.Process(); } 每个核的输入数据为: Src0: [1,1,1,1,1,...,1] // 256个1 Src1: [2,2,2,2,2,...,2] // 256个2 最终输出数据: [2,2,2,2,2,...,2] // 256个2 |
父主题： 原子操作
