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入参传入临时空间

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)

接口框架申请临时空间

template <typename T, const CumSumConfig &config = defaultCumSumConfig>
__aicore__ inline void CumSum(LocalTensor<T> &dstTensor, LocalTensor<T> &lastRowTensor, const LocalTensor<T> &srcTensor, const CumSumInfo &cumSumInfo)

参数说明

表1 模板参数说明

参数名

描述

操作数的数据类型。

config

定义CumSum接口编译时config参数

struct CumSumConfig {
    bool isLastAxis{true};
    bool isReuseSource{false};
    bool outputLastRow{false};
};

isLastAxis：取值为true表示计算按last轴处理，取值为false表示计算按first轴处理；
isReuseSource：是否可以复用srcLocal的内存空间；
outputLastRow：是否输出最后一行数据。

表2 接口参数说明

参数名

输入/输出

描述

dstLocal

输出

目的操作数。按first轴或last轴处理，输入元素的累加和。

类型为LocalTensor。

Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：half/float

Atlas推理系列产品AI Core，支持的数据类型为：half/float

lastRowTensor

输出

目的操作数。模板参数config中的outputLastRow参数取值为true时，输出的最后一行数据。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

源操作数的数据类型需要与目的操作数保持一致。

Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：half/float

Atlas推理系列产品AI Core，支持的数据类型为：half/float

srcLocal

输入

源操作数。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

源操作数的数据类型需要与目的操作数保持一致。

Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：half/float

Atlas推理系列产品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推理系列产品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();
}

父主题： CumSum