BroadCast
功能说明
将输入按照输出shape进行广播。
比如A的shape为(2,1)，广播的目标shape为(2,16)，则会将原来的一列扩展为相同的16列。
输入数据: [[ 1] [ 2]] 输出数据: [[ 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]]
实现原理
以float类型，ND格式，[m, 1]广播到[m, k]为例，描述BroadCast高阶API内部算法框图，如下图所示。
计算过程分为如下几步，均在Vector上进行：
- brcb步骤：将每个元素广播为一个datablock；
- Copy步骤：将每个datablock均复制为多个datablock，k对齐场景下即为结果y；
- 对于k非对齐的场景，再使用GatherMask截取[m, k]个元素， 其中k'表示k向上对齐32B的大小。
函数原型
- 通过sharedTmpBuffer入参传入临时空间
1 2
template <typename T, int32_t dim, int32_t axis, bool isReuseSource = false> __aicore__ inline void BroadCast(LocalTensor<T> &dstLocal, const LocalTensor<T> &srcLocal, const uint32_t dstShape[dim], const uint32_t srcShape[dim], LocalTensor<uint8_t> &sharedTmpBuffer)
- 接口框架申请临时空间
1 2
template <typename T, int32_t dim, int32_t axis, bool isReuseSource = false> __aicore__ inline void BroadCast(LocalTensor<T> &dstLocal, const LocalTensor<T> &srcLocal, const uint32_t dstShape[dim], const uint32_t srcShape[dim])
该接口需要额外的临时空间来存储计算过程中的中间变量。临时空间支持开发者通过sharedTmpBuffer入参传入和接口框架申请两种方式。
- 通过sharedTmpBuffer入参传入，使用该tensor作为临时空间进行处理，接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间，并在接口调用完成后，复用该部分内存，内存不会反复申请释放，灵活性较高，内存利用率也较高。
- 接口框架申请临时空间，开发者无需申请，但是需要预留临时空间的大小。
通过sharedTmpBuffer传入的情况，开发者需要为tensor申请空间；接口框架申请的方式，开发者需要预留临时空间。临时空间大小BufferSize的获取方式如下：通过GetBroadCastMaxMinTmpSize中提供的接口获取需要预留空间范围的大小。
参数说明
接口 |
功能 |
|---|---|
T |
数据类型: uint8_t/int8_t/half/float |
dim |
输入/输出tensor的维度，目前仅支持1维和2维 |
axis |
要广播的维度，目前仅支持0和1 |
isReuseSource |
是否允许修改源操作数。该参数预留，传入默认值false即可。 |
参数名 |
输入/输出 |
描述 |
|---|---|---|
dstTensor |
输出 |
目的操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas推理系列产品AI Core，支持的数据类型为：uint8_t/int8_t/half/float Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：uint8_t/int8_t/half/float |
srcTensor |
输入 |
源操作数。 源操作数的数据类型需要与目的操作数保持一致。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas推理系列产品AI Core，支持的数据类型为：uint8_t/int8_t/half/float Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：uint8_t/int8_t/half/float |
dstShape |
输入 |
输出tensor的shape：uint32_t类型的数组，长度为1或者2， 输入/输出的shape维度要一致。 |
srcShape |
输入 |
输入tensor的shape：uint32_t类型的数组，长度为1或者2， 输入/输出的shape维度要一致。 |
sharedTmpBuffer |
输入 |
临时缓存。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 用于SwiGLU内部复杂计算时存储中间变量，由开发者提供。 临时空间大小BufferSize的获取方式请参考BroadCast Tiling。 |
返回值
无
支持的型号
Atlas推理系列产品AI Core
Atlas A2训练系列产品/Atlas 800I A2推理产品
约束说明
- 操作数地址偏移对齐要求请参见通用约束。
- 不支持源操作数与目的操作数地址重叠。
- 当前仅支持ND格式的输入，不支持其他格式。
- dim目前仅支持1或者2， axis目前仅支持0或者1。
- 对于Atlas推理系列产品AI Core，在dim=2，axis=1时，srcShape[0]32B对齐
- 在dim=2，axis=0时，要求srchShape[1]32B对齐
调用示例
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 |
#include "kernel_operator.h" template <typename T, int32_t dim, int32_t axis> class KernelBroadCast { public: __aicore__ inline KernelBroadCast() {} __aicore__ inline void Init( GM_ADDR srcGm, GM_ADDR dstGm, const uint32_t dstShape[dim], const uint32_t srcShape[dim]) { for (uint32_t i = 0; i < dim; i++) { srcSize *= srcShape[i]; dstSize *= dstShape[i]; } srcGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(srcGm), srcSize); dstGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dstGm), dstSize); pipe.InitBuffer(inQueueX, 1, srcSize * sizeof(T)); pipe.InitBuffer(outQueue, 1, dstSize * sizeof(T)); dstShape_ = dstShape; srcShape_ = srcShape; } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> srcLocal = inQueueX.AllocTensor<T>(); AscendC::DataCopy(srcLocal, srcGlobal, srcSize); inQueueX.EnQue(srcLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> dstLocal = outQueue.AllocTensor<T>(); AscendC::LocalTensor<T> srcLocal = inQueueX.DeQue<T>(); AscendC::BroadCast<T, dim, axis>(dstLocal, srcLocal, dstShape_, srcShape_); outQueue.EnQue<T>(dstLocal); inQueueX.FreeTensor(srcLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> dstLocal = outQueue.DeQue<T>(); AscendC::DataCopy(dstGlobal, dstLocal, dstSize); outQueue.FreeTensor(dstLocal); } private: AscendC::GlobalTensor<T> srcGlobal; AscendC::GlobalTensor<T> dstGlobal; AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueX; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueue; const uint32_t *dstShape_{nullptr}; const uint32_t *srcShape_{nullptr}; int32_t srcSize{1}; int32_t dstSize{1}; }; template <typename T, int32_t dim, int32_t axis> __aicore__ void kernel_broadcast_operator( GM_ADDR srcGm, GM_ADDR dstGm, const uint32_t dstShape[dim], const uint32_t srcShape[dim]) { KernelBroadCast<T, dim, axis> op; op.Init(srcGm, dstGm, dstShape, srcShape); op.Process(); } |
输入数据(srcTensor): [[ 1] [ 2] [ 3] [ 4] [ 5] [ 6] [ 7] [ 8] [ 9] [10] [11] [12] [13] [14] [15] [16]] 输出数据(dstLocal): [[ 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] [ 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3] [ 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4] [ 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5] [ 6 6 6 6 6 6 6 6 6 6 6 6 6 6 6 6] [ 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7] [ 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8] [ 9 9 9 9 9 9 9 9 9 9 9 9 9 9 9 9] [10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10] [11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11] [12 12 12 12 12 12 12 12 12 12 12 12 12 12 12 12] [13 13 13 13 13 13 13 13 13 13 13 13 13 13 13 13] [14 14 14 14 14 14 14 14 14 14 14 14 14 14 14 14] [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]]
