WelfordFinalize

产品支持情况

功能说明

Welford计算是一种在线计算均值和方差的方法。一方面，它可以在不存储所有样本的情况下，逐步计算所有样本的均值和方差，更适合处理海量数据；另一方面，它只需要对数据进行一次遍历，能减少访存次数，提高计算性能。本接口为Welford算法的后处理。

LayerNorm算法中Reduce轴较大的场景，可以通过切分Reduce轴，联合使用本接口与WelfordUpdate，能够实现等效计算LayerNorm。根据Reduce轴切分后是否有尾块，本接口分为如下两种计算公式：

• 不带尾块/不带counts参数场景：

  

  对于Atlas 350 加速卡，方差计算公式如下。

  • 方差系数未修正场景：

    

  • 方差系数修正场景：

    

  其中，Mean为均值输出，Var为方差输出。

  Mean_i代表输入的第i个均值，Var_i代表输入的第i个方差。Ab代表Reduce轴切分后一次计算的大小，Rn代表Reduce轴按Ab拆分的次数，代表未修正的方差系数rRec，代表修正方差系数rRecWithCorrection。

• 带尾块/带counts参数场景：

  

  对于Atlas 350 加速卡，方差计算公式如下。

  • 方差系数未修正场景：

    

  • 方差系数修正场景：

    

  除上述参数含义外，counts_i代表Mean_i对应的系数，R代表未切分的原始Reduce轴长度，代表未修正的方差系数rRec，代表修正系数rRecWithCorrection。

函数原型

• 通过sharedTmpBuffer入参传入临时空间
  • 不带counts参数场景

    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)

  • 带counts参数场景

    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)

• 接口框架申请临时空间
  • 不带counts参数场景

    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)

  • 带counts参数场景

    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)

由于该接口的内部实现中涉及复杂的计算，需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。

• 接口框架申请临时空间，开发者无需申请，但是需要预留临时空间的大小。

• 通过sharedTmpBuffer入参传入，使用该tensor作为临时空间进行处理，接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间，并在接口调用完成后，复用该部分内存，内存不会反复申请释放，灵活性较高，内存利用率也较高。

接口框架申请的方式，开发者需要预留临时空间；通过sharedTmpBuffer传入的情况，开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下：通过WelfordFinalize Tiling中提供的GetWelfordFinalizeMaxMinTmpSize接口获取所需最大和最小临时空间大小，最小空间可以保证功能正确，最大空间用于提升性能。

参数说明

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;
};

返回值说明

无

约束说明

• 接口参数para.abLength的取值必须为32/sizeof(float)的整数倍。
• 接口参数para.headCountLength与para.tailCountLength的和必须等于参数para.abLength。
• 接口处理逻辑以参数para中设置的具体参数值为准，不依赖源操作数的shape信息。
• 接口参数para.tailCount为0时，禁止配置para.tailCountLength为非0值。
• 不支持源操作数与目的操作数地址重叠。
• 不支持sharedTmpBuffer与源操作数和目的操作数地址重叠。

调用示例

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：均值目的操作数
// varianceLocal：方差目的操作数
// inmeanLocal：均值源操作数
// invarLocal：方差源操作数
// inputXLocal：源操作数
// para：计算所需参数信息，WelfordFinalizePara类型

// 计算方差时不使用修正系数
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);

输入数据(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 ]
输入数据(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 ]
输入数据(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 ]
输出数据(meanLocal): 
[ 0.484375 0.       0.       0.       0.       0.       0.       0.      ]
输出数据(varianceLocal): 
[ 9.831062 0.       0.       0.       0.       0.       0.       0.      ]