WelfordFinalize

Applicability

Function Usage

Welford is a method for calculating the mean and variance online. On one hand, this method can gradually calculate the mean and variance of all samples without needing to store these samples, making it ideal for processing large-scale data. On the other hand, it requires only a single data traversal, reducing memory access and enhancing computational performance. This API is used for post-processing of the Welford algorithm.

If the reduced axis in the LayerNorm algorithm is large, you can split the reduced axis and use this API and WelfordUpdate together to implement equivalent computation of LayerNorm. This API can be used to implement two computing formulas based on whether there is a tail block after the reduced axis is split:

• Scenarios without the tail block or the counts parameter:

  

  For the Atlas 350 Accelerator Card, the formula for computing the variance is as follows:

  • Scenario without the correction variance coefficient:

    

  • Scenario with the correction variance coefficient:

    

  Mean indicates the mean output, and Var indicates the variance output.

  Mean_i indicates the ith mean value of the input, and Var_i indicates the ith variance of the input. Ab indicates the size of a single computation after the reduced axis is split, Rn indicates the number of times the reduced axis is split based on Ab, indicates the corrected variance coefficient rRecWithCorrection, and indicates the uncorrected variance coefficient rRec.

• Scenarios with the tail block or the counts parameter:

  

  For the Atlas 350 Accelerator Card, the formula for computing the variance is as follows:

  • Scenario without the correction variance coefficient:

    

  • Scenario with the correction variance coefficient:

    

  In addition to the preceding parameters, counts_i indicates the coefficient corresponding to Mean_i, R indicates the length of the original reduced axis that is not split, indicates the corrected coefficient rRecWithCorrection, and indicates the uncorrected variance coefficient rRec.

Prototype

• Pass the temporary space through the sharedTmpBuffer input parameter.
  • Scenarios without the counts parameter

    1 2

    template <bool isReuseSource = false, const WelfordFinalizeConfig& config = WFFINALIZE_DEFAULT_CFG> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, const LocalTensor<uint8_t>& sharedTmpBuffer, WelfordFinalizePara& para)

  • Scenarios with the counts parameter

    1 2

    template <bool isReuseSource = false, const WelfordFinalizeConfig& config = WFFINALIZE_DEFAULT_CFG> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, const LocalTensor<int32_t>& counts, const LocalTensor<uint8_t>& sharedTmpBuffer, WelfordFinalizePara& para)

• Allocate the temporary space through the API framework.
  • Scenarios without the counts parameter

    1 2

    template <bool isReuseSource = false, const WelfordFinalizeConfig& config = WFFINALIZE_DEFAULT_CFG> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, WelfordFinalizePara& para)

  • Scenarios with the counts parameter

    1 2

    template <bool isReuseSource = false, const WelfordFinalizeConfig& config = WFFINALIZE_DEFAULT_CFG> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, const LocalTensor<int32_t>& counts, WelfordFinalizePara& para)

Due to the complex computation involved in the internal implementation of this API, additional temporary space is required to store intermediate variables generated 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 is used for temporary space allocation, 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. The method of obtaining the temporary space size (BufferSize) is as follows: Obtain the required maximum and minimum temporary space sizes using the GetWelfordFinalizeMaxMinTmpSize API provided in WelfordFinalize Tiling. The minimum space can ensure correct functionality, while the maximum space is used to improve performance.

Parameters

1
2
3

struct WelfordFinalizeConfig {
     bool isCorrection = false;
}

1

constexpr WelfordFinalizeConfig WFFINALIZE_DEFAULT_CFG = { false };

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

struct WelfordFinalizePara {
    uint32_t rnLength;
    uint32_t abLength;
    uint32_t headCount;
    uint32_t headCountLength;
    uint32_t tailCount;
    uint32_t tailCountLength;
    float abRec;
    float rRec;
    float rRecWithCorrection;
};

Returns

None

Constraints

• The value of para.abLength must be an integer multiple of 32/sizeof(float).
• The sum of the values of para.headCountLength and para.tailCountLength must be equal to the value of para.abLength.
• The API processing logic is based on the value of the para parameter and does not depend on the shape information of the source operand.
• When para.tailCount is 0, para.tailCountLength cannot be set to a non-zero value.
• The source operand address must not overlap the destination operand address.
• The address of sharedTmpBuffer cannot overlap that of the source or destination operand.

Examples

1
2
3
4

pipe.InitBuffer(sharedTmpBuffer, stackBufferSize);        
AscendC::LocalTensor<uint8_t> tmpLocalTensor = sharedTmpBuffer.Get<uint8_t>();         
struct AscendC::WelfordFinalizePara para = {rnLength, abLength, head, headLength, tail, tailLength, abRec, rRec};
AscendC::WelfordFinalize<false>(meanLocal, varianceLocal, inputMeanLocal, inputVarianceLocal, inputCountsLocal, tmpLocalTensor, para); 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

// meanLocal: destination operand of the mean
// varianceLocal: destination operand of the variance
// inmeanLocal: source operand of the mean
// invarLocal: source operand of the variance
// inputXLocal: source operand
// para: parameter information required for computation, of the WelfordFinalizePara type.

// Do not use the correction coefficient to compute the variance.
static constexpr AscendC::WelfordFinalizeConfig CONFIG = {false};
AscendC::WelfordFinalizePara para = { rnLength, abLength, head, headLength, tail, tailLength, abRec, rRec, rRecWithCorrection};
AscendC::WelfordFinalize<false, CONFIG>(meanLocal, varianceLocal, inmeanLocal, invarLocal, inputXLocal, para);

Input (inmeanLocal, shape: [1, 32]):
[  0.0  1.0  2.0  3.0  4.0  5.0  6.0  7.0  8.0  9.0 10.0 11.0 12.0 13.0 14.0 15.0 
  16.0 17.0 18.0 19.0 20.0 21.0 22.0 23.0 24.0 25.0 26.0 27.0 28.0 29.0 30.0 31.0 ]
Input (invarLocal, shape: [1, 32]):
[  0.0  1.0  2.0  3.0  4.0  5.0  6.0  7.0  8.0  9.0 10.0 11.0 12.0 13.0 14.0 15.0 
  16.0 17.0 18.0 19.0 20.0 21.0 22.0 23.0 24.0 25.0 26.0 27.0 28.0 29.0 30.0 31.0 ]
Input (inputX, shape: [1, 32]):
[ 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 ]
Output (meanLocal):
[ 0.484375 0.       0.       0.       0.       0.       0.       0.      ]
Output (varianceLocal):
[ 9.831062 0.       0.       0.       0.       0.       0.       0.      ]