CumSum

Applicability

Function Usage

Performs cumulative sum of an input tensor along a row or column. Each element in the output is the cumulative sum of the element at the corresponding position and all preceding elements along a given row or column.

The formula is as follows:

• Row-by-row accumulation algorithm
  • First axis processing (or row-wise accumulation): The first row remains unchanged, and the subsequent rows are accumulated in sequence. The formula for calculating the element in the ith row and jth column of the output is as follows.

    

    Example: If the input tensor is ([[0, 1, 2], [3, 4, 5]]), the output tensor is ([[0, 1, 2], [3, 5, 7]]).

  • Last axis processing (or column-wise accumulation): The first column remains unchanged, and the subsequent columns are accumulated in sequence. The formula for calculating the element in the ith row and jth column of the output is as follows.

    

    Example: if the input tensor is ([[0, 1, 2], [3, 4, 5]]), the final output tensor is ([[0, 1, 3], [3, 7, 12]]).

• Sklansky binary prefix-sum algorithm

  It is supported only by the Atlas 350 Accelerator Card.

  The Sklansky binary prefix-sum algorithm is based on the parallel prefix-sum logic of Sklansky Adder. Figure 1 shows the one-dimensional binary parallel prefix-sum algorithm. This algorithm can be extended to a cumulative sum algorithm for two-dimensional tensors. Take row-wise accumulation as an example. Figure 2 shows how the extended algorithm works. Row-wise accumulation is implemented by performing parallel accumulation across multiple rows based on the Sklansky binary prefix-sum algorithm.

  Figure 1 Sklansky Adder algorithm
  
  Figure 2 Sklansky binary prefix-sum algorithm
  

Prototype

• Pass the temporary space through the sharedTmpBuffer input parameter.

  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)

• The API framework allocates temporary space.

  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)

Precision conversion is involved in the internal implementation of this API. Therefore, extra temporary space is required to store intermediate variables during computation. The temporary space can be allocated through the API framework or passed by developers through the sharedTmpBuffer input parameter.

• When the API framework allocates temporary space, developers do not need to allocate the space but must reserve the required size for the temporary space.

• When the sharedTmpBuffer input parameter is used for passing the temporary space, the tensor serves as the temporary space. In this case, the API framework is not required for temporary space allocation. This enables developers to manage the sharedTmpBuffer space and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization.

If the API framework is used, developers must reserve the temporary space. If sharedTmpBuffer is used, developers must allocate space for the tensor. To obtain the size of the temporary space (BufferSize) to be reserved, use the API provided in GetCumSumMaxMinTmpSize.

Parameters

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

struct CumSumConfig {
    bool isLastAxis{true};
    bool isReuseSource{false};
    bool outputLastRow{false};
    CumSumAlgorithm algorithm{CumSumAlgorithm::CUMSUM_ALGORITHM_LINEBYLINE};
};
enum class CumSumAlgorithm {    
    CUMSUM_ALGORITHM_LINEBYLINE = 0,    
    CUMSUM_ALGORITHM_SKLANSKY = 1
};

1
2
3
4
5

struct CumSumInfo
{
    uint32_t outter{0};
    uint32_t inner{0};
};

Returns

None

Constraints

• For details about the operand address alignment requirements, see General Address Alignment Restrictions.
• The input supports only the two-dimensional structure.
• The product of cumSumInfo.inner × sizeof(T) must be an integer multiple of 32 bytes.

Examples

For a complete call example, see CumSum operator sample.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

// dstLocal: tensor for storing the computation result
// lastRowLocal: tensor for storing the last row of the computation result
// srcLocal: input tensor involved in computation

// Process data based on the last axis (data accumulated by column) and output the data in the last column. The row-by-row accumulation algorithm is used.
constexpr AscendC::CumSumConfig cumSumConfig{true, false, true, AscendC::CumSumAlgorithm::CUMSUM_ALGORITHM_LINEBYLINE};
// outer: outer axis length
// inner: inner axis length
const AscendC::CumSumInfo cumSumInfo{outer, inner};
AscendC::CumSum<T, cumSumConfig>(dstLocal, lastRowLocal, srcLocal, cumSumInfo);

// Process data based on the first axis (data accumulated by row) and output the data in the last row. The row-by-row accumulation algorithm is used.
constexpr AscendC::CumSumConfig cumSumConfig{false, false, true, AscendC::CumSumAlgorithm::CUMSUM_ALGORITHM_LINEBYLINE};
AscendC::CumSum<T, cumSumConfig>(dstLocal, lastRowLocal, srcLocal, cumSumInfo);

Input (srcLocal): [1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]
Output (dstLocal): [1 2 3 4 5 6 7 8 1 2 3 4 5 6 7 8 1 2 3 4 5 6 7 8 1 2 3 4 5 6 7 8]
Output (lastRowLocal): [8 8 8 8]

Input (srcLocal): [1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]
Output (dstLocal): [1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 3 3 3 3 3 3 3 3 4 4 4 4 4 4 4 4]
Output (lastRowLocal): [4 4 4 4 4 4 4 4]