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)

参数说明

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

struct CumSumInfo
{
    uint32_t outter{0};    // 表示输入数据的外轴长度
    uint32_t inner{0};     // 表示输入数据的内轴长度
};

返回值

无

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

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