CumSum
功能说明
- First轴处理,第一行不变,后面的行依次累加
以tensor([[0, 1, 2], [3, 4, 5]])为例,输出结果是tensor([[0, 1, 2], [3, 5, 7]])
- Last轴处理,第一列不变,后面的列依次累加
以tensor([[0, 1, 2], [3, 4, 5]])为例,输出结果是tensor([[0, 1, 3], [3, 7, 12]])
函数原型
- 通过sharedTmpBuffer入参传入临时空间
1 2
template <typename T, const CumSumConfig &config = defaultCumSumConfig> __aicore__ inline void CumSum(LocalTensor<T> &dstTensor, LocalTensor<T> &lastRowTensor, const LocalTensor<T> &srcTensor, LocalTensor<uint8_t> &sharedTmpBuffer, const CumSumInfo &cumSumInfo)
- 接口框架申请临时空间
1 2
template <typename T, const CumSumConfig &config = defaultCumSumConfig> __aicore__ inline void CumSum(LocalTensor<T> &dstTensor, LocalTensor<T> &lastRowTensor, const LocalTensor<T> &srcTensor, const CumSumInfo &cumSumInfo)
参数说明
参数名 |
描述 |
---|---|
T |
操作数的数据类型。 |
config |
定义CumSum接口编译时config参数 struct CumSumConfig { bool isLastAxis{true}; bool isReuseSource{false}; bool outputLastRow{false}; };
|
参数名 |
输入/输出 |
描述 |
---|---|---|
dstLocal |
输出 |
目的操作数。 类型为LocalTensor。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float |
lastRowTensor |
输出 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 源操作数的数据类型需要与目的操作数保持一致。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float |
srcLocal |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 源操作数的数据类型需要与目的操作数保持一致。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float |
cumSumInfo |
输入 |
srcTensor的shape信息。CumSumInfo类型,具体定义如下: struct CumSumInfo { uint32_t outter{0}; // 表示输入数据的外轴长度 uint32_t inner{0}; // 表示输入数据的内轴长度 }; CumSumInfo.inner*sizeof(T)必须是32字节的整数倍 |
sharedTmpBuffer |
输入 |
临时缓存。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 用于Cumsum内部复杂计算时存储中间变量,由开发者提供。 临时空间大小BufferSize的获取方式请参考GetCumSumMaxMinTmpSize。 |
返回值
无
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品(Ascend 310P处理器)AI Core
约束说明
- 操作数地址偏移对齐要求请参见通用约束。
- 输入input只支持二维结构
- inner必须是32B的整数
调用示例
#include "kernel_operator.h" template <typename T> class KernelCumSum { public: __aicore__ inline KernelCumSum(){} __aicore__ inline void Init( GM_ADDR srcGm, GM_ADDR dstGm, GM_ADDR lastRowGm, const AscendC::CumSumInfo& cumSumParams) { outer = cumSumParams.outter; inner = cumSumParams.inner; srcGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(srcGm), outer * inner); dstGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dstGm), outer * inner); lastRowGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(lastRowGm), inner); pipe.InitBuffer(inQueueX, 1, outer * inner * sizeof(T)); pipe.InitBuffer(outQueue, 1, outer * inner * sizeof(T)); pipe.InitBuffer(lastRowQueue, 1, inner * sizeof(T)); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> srcLocal = inQueueX.AllocTensor<T>(); AscendC::DataCopy(srcLocal, srcGlobal, outer * inner); inQueueX.EnQue(srcLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> dstLocal = outQueue.AllocTensor<T>(); AscendC::LocalTensor<T> lastRowLocal = lastRowQueue.AllocTensor<T>(); AscendC::LocalTensor<T> srcLocal = inQueueX.DeQue<T>(); static constexpr AscendC::CumSumConfig cumSumConfig{true, false, true}; const AscendC::CumSumInfo cumSumInfo{outer, inner}; AscendC::CumSum<T, cumSumConfig>(dstLocal, lastRowLocal, srcLocal, cumSumInfo); outQueue.EnQue<T>(dstLocal); lastRowQueue.EnQue<T>(lastRowLocal); inQueueX.FreeTensor(srcLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> dstLocal = outQueue.DeQue<T>(); AscendC::DataCopy(dstGlobal, dstLocal, outer * inner); outQueue.FreeTensor(dstLocal); AscendC::LocalTensor<T> lastRowLocal = lastRowQueue.DeQue<T>(); AscendC::DataCopy(lastRowGlobal, lastRowLocal, inner); lastRowQueue.FreeTensor(lastRowLocal); } private: AscendC::GlobalTensor<T> srcGlobal; AscendC::GlobalTensor<T> dstGlobal; AscendC::GlobalTensor<T> lastRowGlobal; AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueX; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueue; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> lastRowQueue; uint32_t outer{1}; uint32_t inner{1}; }; template <typename T> __aicore__ inline void kernel_cumsum_operator( GM_ADDR srcGm, GM_ADDR dstGm, GM_ADDR lastRowGm, const AscendC::CumSumInfo &cumSumParams) { KernelCumSum<T> op; op.Init(srcGm, dstGm, lastRowGm, cumSumParams); op.Process(); }