Brcb
功能说明
给定一个输入张量,每一次取输入张量中的8个数填充到结果张量的8个datablock(32Bytes)中去,每个数对应一个datablock。
函数原型
1 2 |
template <typename T> __aicore__ inline void Brcb(const LocalTensor<T>& dstLocal, const LocalTensor<T>& src0Local, const uint8_t repeatTimes, const BrcbRepeatParams& repeatParams) |
参数说明
参数名 |
描述 |
---|---|
T |
操作数数据类型。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:int16_t/uint16_t/int32_t/uint32_t/half/bfloat16_t/float Atlas推理系列产品AI Core,支持的数据类型为:int16_t/uint16_t/int32_t/uint32_t/half/float |
参数名称 |
输入/输出 |
含义 |
---|---|---|
dstLocal |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 LocalTensor的起始地址需要32字节对齐。 地址需要32bytes对齐。 |
srcLocal |
输入 |
源操作数,连续存储score的elements。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 LocalTensor的起始地址需要32字节对齐。 数据类型和dstLocal保持一致。 |
repeatTimes |
输入 |
指令迭代次数,每次迭代完成8个datablock的数据收集,数据范围:repeatTimes∈[0,255]。 |
repeatParams |
输入 |
指令迭代参数,类型为BrcbRepeatParams,参数说明参考表2。 |
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品AI Core
注意事项
不支持srcLocal与dstLocal为同一块内存地址
针对Atlas推理系列产品AI Core,使用时需要预留8K的Unified Buffer空间,作为接口的临时数据存放区。
调用示例
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 |
#include "kernel_operator.h" class VbrcbCase { public: __aicore__ inline VbrcbCase() {} __aicore__ inline void Init(__gm__ uint8_t *x, __gm__ uint8_t *y) { x_gm.SetGlobalBuffer(reinterpret_cast<__gm__ uint16_t *>(x)); y_gm.SetGlobalBuffer(reinterpret_cast<__gm__ uint16_t *>(y)); tpipe.InitBuffer(vecIn, 1, 16 * sizeof(uint16_t)); tpipe.InitBuffer(vecOut, 1, 256 * sizeof(uint16_t)); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } __aicore__ inline void CopyIn() { auto x_buf = vecIn.AllocTensor<uint16_t>(); AscendC::DataCopy(x_buf, x_gm, 16); vecIn.EnQue(x_buf); } __aicore__ inline void Compute() { auto x_buf = vecIn.DeQue<uint16_t>(); auto y_buf = vecOut.AllocTensor<uint16_t>(); AscendC::Brcb(y_buf, x_buf, 2, {1,8}); vecOut.EnQue(y_buf); vecIn.FreeTensor(x_buf); } __aicore__ inline void CopyOut() { auto y_buf = vecOut.DeQue<uint16_t>(); AscendC::DataCopy(y_gm, y_buf, 256); vecOut.FreeTensor(y_buf); } private: AscendC::GlobalTensor<uint16_t> x_gm; AscendC::GlobalTensor<uint16_t> y_gm; AscendC::TPipe tpipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> vecIn; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> vecOut; }; extern "C" __global__ __aicore__ void vbrcb_uint16_t_16(__gm__ uint8_t *x, __gm__ uint8_t *y) { VbrcbCase op; op.Init(x, y); op.Process(); } |
输入数据(srcGlobal): [1 2 3 ... 16] 输出数据(dstGlobal):[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 ... 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16]