UnPad
函数功能
对height * width的二维Tensor在width方向上进行unpad,如果Tensor的width非32B对齐,则不支持调用本接口unpad。本接口具体功能场景如下:Tensor的width已32B对齐,以half为例,如16*16,进行UnPad,变成16*15。
函数原型
由于该接口的内部实现中涉及复杂的计算,需要额外的临时空间来存储计算过程中的中间变量。临时空间大小BufferSize的获取方法:通过UnPad Tiling中提供的GetUnPadMaxMinTmpSize接口获取所需最大和最小临时空间大小,最小空间可以保证功能正确,最大空间用于提升性能。
临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式,因此UnPad接口的函数原型有两种:
- 通过sharedTmpBuffer入参传入临时空间
1 2
template <typename T> __aicore__ inline void UnPad(const LocalTensor<T> &dstTensor, const LocalTensor<T> &srcTensor, UnPadParams &unPadParams, LocalTensor<uint8_t> &sharedTmpBuffer, UnPadTiling &tiling)
该方式下开发者需自行申请并管理临时内存空间并管理,并在接口调用完成后,复用该部分内存,内存不会反复申请释放,灵活性较高,内存利用率也较高。
- 接口框架申请临时空间
1 2
template <typename T> __aicore__ inline void UnPad(const LocalTensor<T> &dstTensor, const LocalTensor<T> &srcTensor, UnPadParams &unPadParams, UnPadTiling &tiling)
该方式下开发者无需申请,但是需要预留临时空间的大小。
参数说明
接口 |
功能 |
---|---|
T |
操作数的数据类型。 |
参数名称 |
输入/输出 |
含义 |
---|---|---|
dstTensor |
输出 |
目的操作数,类型为LocalTensor,shape为二维,LocalTensor数据结构的定义请参考LocalTensor。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:int16_t/uint16_t/half/int32_t/uint32_t/float Atlas推理系列产品AI Core,支持的数据类型为:int16_t/uint16_t/half/int32_t/uint32_t/float |
srcTensor |
输入 |
源操作数,类型为LocalTensor,shape为二维,LocalTensor数据结构的定义请参考LocalTensor。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:int16_t/uint16_t/half/int32_t/uint32_t/float Atlas推理系列产品AI Core,支持的数据类型为:int16_t/uint16_t/half/int32_t/uint32_t/float |
UnPadParams |
输入 |
UnPad详细参数,UnPadParams数据类型,具体结构体参数说明如下:
UnPadParams结构体的定义如下: struct UnPadParams { uint16_t leftPad = 0; uint16_t rightPad = 0; }; |
sharedTmpBuffer |
输入 |
共享缓冲区,用于存放API内部计算产生的临时数据。该方式开发者可以自行管理sharedTmpBuffer内存空间,并在接口调用完成后,复用该部分内存,内存不会反复申请释放,灵活性较高,内存利用率也较高。共享缓冲区大小的获取方式请参考UnPad Tiling。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:int16_t/uint16_t/half/int32_t/uint32_t/float Atlas推理系列产品AI Core,支持的数据类型为:int16_t/uint16_t/half/int32_t/uint32_t/float |
tiling |
输入 |
计算所需tiling信息,Tiling信息的获取请参考UnPad Tiling。 |
返回值
无
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品AI Core
注意事项
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
本样例:Tensor的width已32B对齐,以half为例,如16*16,进行Pad,变成16*15。输入数据类型均为half。
#include "kernel_operator.h" template <typename T> class KernelUnPad { public: __aicore__ inline KernelUnPad() {} __aicore__ inline void Init(GM_ADDR dstGm, GM_ADDR srcGm, uint16_t heightIn, uint16_t widthIn, uint16_t oriWidthIn, AscendC::UnPadParams &unPadParamsIn, const UnPadTiling &tilingData) { height = heightIn; width = widthIn; oriWidth = oriWidthIn; unPadParams = unPadParamsIn; srcGlobal.SetGlobalBuffer((__gm__ T *)srcGm); dstGlobal.SetGlobalBuffer((__gm__ T *)dstGm); pipe.InitBuffer(inQueueSrcVecIn, 1, height * width * sizeof(T)); pipe.InitBuffer(inQueueSrcVecOut, 1, height * (width - unPadParams.leftPad - unPadParams.rightPad) * sizeof(T)); tiling = tilingData; } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> srcLocal = inQueueSrcVecIn.AllocTensor<T>(); AscendC::DataCopy(srcLocal, srcGlobal, height * width); inQueueSrcVecIn.EnQue(srcLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> dstLocal = inQueueSrcVecIn.DeQue<T>(); AscendC::LocalTensor<T> srcOutLocal = inQueueSrcVecOut.AllocTensor<T>(); AscendC::UnPad(srcOutLocal, dstLocal, unPadParams, tiling); inQueueSrcVecOut.EnQue(srcOutLocal); inQueueSrcVecIn.FreeTensor(dstLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> srcOutLocalDe = inQueueSrcVecOut.DeQue<T>(); AscendC::DataCopy(dstGlobal, srcOutLocalDe, height * (width - unPadParams.leftPad - unPadParams.rightPad)); inQueueSrcVecOut.FreeTensor(srcOutLocalDe); } private: AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrcVecIn; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> inQueueSrcVecOut; AscendC::GlobalTensor<T> srcGlobal; AscendC::GlobalTensor<T> dstGlobal; uint16_t height; uint16_t width; uint16_t oriWidth; AscendC::UnPadParams unPadParams; UnPadTiling tiling; }; extern "C" __global__ __aicore__ void kernel_unpad_half_16_16_16(GM_ADDR src_gm, GM_ADDR dst_gm, __gm__ uint8_t *tiling) { GET_TILING_DATA(tilingData, tiling); KernelUnPad<half> op; AscendC::UnPadParams unPadParams{0, 1}; op.Init(dst_gm, src_gm, 16, 16, 16, unPadParams, tilingData.unpadTilingData); op.Process(); }