WelfordUpdate
Applicability
Product |
Supported |
|---|---|
Atlas 350 Accelerator Card |
√ |
√ |
|
√ |
|
x |
|
√ |
|
x |
|
x |
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 pre-processing of the Welford algorithm.
If the reduced axis in the LayerNorm algorithm is large, you can split the axis and use this API and WelfordFinalize together to implement equivalent computation of LayerNorm.
As shown in the following figure, the reduced axis of the data is split. Assume that the shape of each piece of data after splitting is [1, k], and each piece of data is numbered 1, 2, 3, ..., n.

Below is the formula of this API. After data is split, this API is called n times. The following formula is used to compute each piece of split data.


In the formula, the shapes of xi, Meanti, and Mi are all [1, k]. xi indicates the ith data block after splitting. Meanti indicates the mean value of the first i data blocks obtained by calling this API for the ith time. Mi indicates the intermediate result of the variance of the first i data blocks obtained by calling this API for the ith time. (The intermediate result is saved for computing the variance in the following sections.) When this API is called for the first time (i = 1), Meant0 and M0 in the formula are defined by the user as data with the shape of [1, k], and all values of 0.
The following figure shows the computation process of Meantn. After this API is called n times, Meantn and Mn with the shape of [1, k] are obtained. Meantn and Mn are used for subsequent computation by the WelfordFinalize API.

Prototype
- Pass the temporary space through the sharedTmpBuffer input parameter.
1 2
template <typename T, typename U,bool isReuseSource = false, const WelfordUpdateConfig& config = WFUPDATE_DEFAULT_CFG> __aicore__ inline void WelfordUpdate(const LocalTensor<U>& outputMean, const LocalTensor<U>& outputVariance, const LocalTensor<U>& inputMean, const LocalTensor<U>& inputVariance, const LocalTensor<T>& inputX, const LocalTensor<uint8_t>& sharedTmpBuffer, const WelfordUpdateParam& para)
- Allocate the temporary space through the API framework.
1 2
template <typename T, typename U,bool isReuseSource = false, const WelfordUpdateConfig& config = WFUPDATE_DEFAULT_CFG> __aicore__ inline void WelfordUpdate(const LocalTensor<U>& outputMean, const LocalTensor<U>& outputVariance, const LocalTensor<U>& inputMean, const LocalTensor<U>& inputVariance, const LocalTensor<T>& inputX, const WelfordUpdateParam& 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 GetWelfordUpdateMaxMinTmpSize API described in WelfordUpdate Tiling. The minimum space can ensure correct functionality, while the maximum space is used to improve performance.
Parameters
Parameter |
Description |
||||
|---|---|---|---|---|---|
T |
Data type of the inputX operand. For the Atlas 350 Accelerator Card, the supported data types are half, bfloat16_t, and float. For the For the For the |
||||
U |
Data type of the outputMean, outputVariance, inputMean, and inputVariance operands. For the Atlas 350 Accelerator Card, the supported data type is float. For the For the For the |
||||
isReuseSource |
Whether the source operand can be modified. The default value is false. If developers allow the source operand to be modified, enable this parameter, to reduce memory space usage. If this parameter is set to true, the inputX memory space is reused during internal computation of this API to reduce memory space usage. If this parameter is set to false, the inputX memory space is not reused during internal computation of this API. In For details about how to use isReuseSource, see Example 4. |
||||
config |
Reuse relationship between the destination operand and source operand that are not in the specified computation scope range. The WelfordUpdateConfig type is defined as follows:
A configuration example is as follows:
This parameter is used together with the tiling computation API on the kernel. |
Parameter |
Input/Output |
Description |
||
|---|---|---|---|---|
outputMean |
Output |
Destination operand of the mean, corresponding to Meanti in the API formula. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The shape must be the same as that of the source operand inputMean. |
||
outputVariance |
Output |
Destination operand of the intermediate variance result, corresponding to Mi in the API formula. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The shape must be the same as that of the source operand inputVariance. |
||
inputMean |
Input |
Source operand of the mean, corresponding to Meanti-1 in the API formula. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. |
||
inputVariance |
Input |
Source operand of the intermediate variance result, corresponding to Mi-1 in the API formula. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. |
||
inputX |
Input |
Source operand, corresponding to xi in the API formula. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. |
||
sharedTmpBuffer |
Input |
Temporary space. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. This parameter is used to store intermediate variables during complex internal API computation and is provided by developers. For details about how to obtain the temporary space size (BufferSize), see WelfordUpdate Tiling. |
||
para |
Input |
Parameter information required for calculation. The WelfordUpdateParam type is defined as follows:
The shape of each destination operand and source operand is [rnLength, abLength]. |
Returns
None
Constraints
- Currently, the value of para.rnLength can only be 1.
- The value of para.abLength must be an integer multiple of 32/sizeof (T).
- The value of para.abComputeLength must be greater than 0.
- 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
For a complete call example, see WelfordUpdate operator sample.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | // outputMean: updated mean value Meant, with shape [1, abLength] // outputVariance: updated intermediate variance result Mi, with shape [1, abLength] // inputMean: previous mean value Meant-1, used as the input // inputVariance: previous intermediate variance result Mi-1, used as the input // inputX: input data xi at the current time step, with shape [1, abLength] // sharedTmpBuffer: temporary space managed by developers, used for complex internal computation // para: parameter structure that contains the block information and normalization coefficient of the reduced axis // Use the WelfordUpdate API to update the Welford algorithm online. struct AscendC::WelfordUpdateParam para = { nLength, rLength, abComputeLength, 0.3 }; AscendC::WelfordUpdate<T, U, false, WELFORD_UPDATE_ENABLE_INPLACE_CFG>( outputMean, // Output: updated mean value outputVariance, // Output: updated intermediate variance result inputMean, // Input: mean value at the previous moment inputVariance, // Input: intermediate variance result at the previous moment inputX, // Input: current input xi sharedTmpBuffer, // Input: temporary buffer (provided by developers) para // Input: Welford update parameters ); |
Input (inputX, shape: [1, 64]): [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. 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. ] Input (gammaLocal, shape: [64]): [ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25. 26. 27. 28. 29. 30. 31. 32. 33. 34. 35. 36. 37. 38. 39. 40. 41. 42. 43. 44. 45. 46. 47. 48. 49. 50. 51. 52. 53. 54. 55. 56. 57. 58. 59. 60. 61. 62. 63. ] Input (betaLocal, shape: [64]): [ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25. 26. 27. 28. 29. 30. 31. 32. 33. 34. 35. 36. 37. 38. 39. 40. 41. 42. 43. 44. 45. 46. 47. 48. 49. 50. 51. 52. 53. 54. 55. 56. 57. 58. 59. 60. 61. 62. 63. ] Output (meanLocal): [ 0.125 1. 1.875 2.75 3.625 4.5 5.375 6.25 7.125 8. 8.875 9.75 10.625 11.5 12.375 13.25 14.125 15. 15.875 16.75 17.625 18.5 19.375 20.25 21.125 22. 22.875 23.75 24.625 25.5 26.375 27.25 28.125 29. 29.875 35. 36. 37. 38. 39. 40. 41. 42. 43. 44. 45. 46. 47. 48. 49. 50. 51. 52. 53. 54. 55. 56. 57. 58. 59. 60. 61. 62. 63. ] Output (varianceLocal): [8.75000e-01 1.00000e+00 2.87500e+00 6.50000e+00 1.18750e+01 1.90000e+01 2.78750e+01 3.85000e+01 5.08750e+01 6.50000e+01 8.08750e+01 9.85000e+01 1.17875e+02 1.39000e+02 1.61875e+02 1.86500e+02 2.12875e+02 2.41000e+02 2.70875e+02 3.02500e+02 3.35875e+02 3.71000e+02 4.07875e+02 4.46500e+02 4.86875e+02 5.29000e+02 5.72875e+02 6.18500e+02 6.65875e+02 7.15000e+02 7.65875e+02 8.18500e+02 8.72875e+02 9.29000e+02 9.86875e+02 3.50000e+01 3.60000e+01 3.70000e+01 3.80000e+01 3.90000e+01 4.00000e+01 4.10000e+01 4.20000e+01 4.30000e+01 4.40000e+01 4.50000e+01 4.60000e+01 4.70000e+01 4.80000e+01 4.90000e+01 5.00000e+01 5.10000e+01 5.20000e+01 5.30000e+01 5.40000e+01 5.50000e+01 5.60000e+01 5.70000e+01 5.80000e+01 5.90000e+01 6.00000e+01 6.10000e+01 6.20000e+01 6.30000e+01 ]