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的整数
调用示例
#ifdef ASCENDC_CPU_DEBUG #include "tikicpulib.h" #endif #include "kernel_operator.h" #include "../../instrs/common_utils/register_utils.h" using namespace AscendC; template <typename T> class KernelCumSum { public: __aicore__ inline KernelCumSum() {} __aicore__ inline void Init( GM_ADDR src_gm, GM_ADDR dst_gm, GM_ADDR last_row_gm, const CumSumInfo &cumSumParams) { outer = cumSumParams.outter; inner = cumSumParams.inner; src_global.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(src_gm), outer * inner); dst_global.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dst_gm), outer * inner); last_row_global.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(last_row_gm), 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() { LocalTensor<T> srcLocal = inQueueX.AllocTensor<T>(); DataCopy(srcLocal, src_global, outer * inner); inQueueX.EnQue(srcLocal); } __aicore__ inline void Compute() { LocalTensor<T> dstLocal = outQueue.AllocTensor<T>(); LocalTensor<T> lastRowLocal = lastRowQueue.AllocTensor<T>(); LocalTensor<T> srcLocal = inQueueX.DeQue<T>(); static constexpr CumSumConfig cumSumConfig{True, false, true}; const CumSumInfo cumSumInfo{outer, inner}; CumSum<T, cumSumConfig>(dstLocal, lastRowLocal, srcLocal, cumSumInfo); outQueue.EnQue<T>(dstLocal); lastRowQueue.EnQue<T>(lastRowLocal); inQueueX.FreeTensor(srcLocal); } __aicore__ inline void CopyOut() { LocalTensor<T> dstLocal = outQueue.DeQue<T>(); DataCopy(dst_global, dstLocal, outer * inner); outQueue.FreeTensor(dstLocal); LocalTensor<T> lastRowLocal = lastRowQueue.DeQue<T>(); DataCopy(last_row_global, lastRowLocal, inner); lastRowQueue.FreeTensor(lastRowLocal); } private: GlobalTensor<T> src_global; GlobalTensor<T> dst_global; GlobalTensor<T> last_row_global; TPipe pipe; TQue<QuePosition::VECIN, 1> inQueueX; TQue<QuePosition::VECOUT, 1> outQueue; TQue<QuePosition::VECOUT, 1> lastRowQueue; uint32_t outer{1}; uint32_t inner{1}; }; template <typename T> __aicore__ void kernel_cumsum_operator( GM_ADDR src_gm, GM_ADDR dst_gm, GM_ADDR last_row_gm, const CumSumInfo &cumSumParams) { KernelCumSum<T> op; op.Init(src_gm, dst_gm, last_row_gm, cumSumParams); op.Process(); }