Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2026 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #ifndef __KFC_DUMP_L2NORM_H__
      12              : #define __KFC_DUMP_L2NORM_H__
      13              : 
      14              : #include "kfc_dump_base.h"
      15              : 
      16              : namespace KfcDumpStat {
      17              : 
      18              : template <typename T, typename InputT>
      19          200 : __aicore__ inline void ComputeL2NormSingleType(
      20              :     TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, LocalTensor<InputT>& inputLocal, LocalTensor<uint8_t>& calMidAddr,
      21              :     int64_t curLoop, int64_t curProcessCount, int64_t appendNum, int64_t totalCount)
      22              : {
      23          200 :     LocalTensor<uint8_t> curNormAddr = calMidAddr[sizeof(float)];
      24          200 :     LocalTensor<InputT> workLocal = workQueue.AllocTensor<InputT>();
      25              : 
      26          200 :     Mul(inputLocal, inputLocal, inputLocal, curProcessCount);
      27          200 :     pipe_barrier(PIPE_ALL);
      28              : 
      29          200 :     ReduceSum<InputT>(workLocal, inputLocal, workLocal, curProcessCount);
      30          200 :     pipe_barrier(PIPE_ALL);
      31              : 
      32              :     // 更新最终结果
      33          200 :     LocalTensor<InputT> curNorm = curNormAddr.ReinterpretCast<InputT>();
      34          200 :     if (curLoop == 0) {
      35          189 :         curNorm.SetValue(0, workLocal.GetValue(0));
      36              :     } else {
      37           11 :         curNorm.SetValue(0, curNorm.GetValue(0) + workLocal.GetValue(0));
      38              :     }
      39              : 
      40          200 :     workQueue.FreeTensor(workLocal);
      41          200 : }
      42              : 
      43              : template <typename T>
      44          200 : __aicore__ inline void ComputeL2NormMultiDataType(
      45              :     TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& castXBuf,
      46              :     TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, LocalTensor<uint8_t>& calMidAddr, int64_t curLoop,
      47              :     int64_t curProcessCount, int64_t appendNum, int64_t totalCount)
      48              : {
      49          200 :     LocalTensor<T> x = xQue.DeQue<T>();
      50              : 
      51          200 :     AppendTailPadding<T>(x, curProcessCount, appendNum);
      52              : 
      53              :     if constexpr (std::is_same_v<T, float>) {
      54          157 :         ComputeL2NormSingleType<float>(workQueue, x, calMidAddr, curLoop, curProcessCount, appendNum, totalCount);
      55              :     } else {
      56              :         // 整型/half/bfloat16 及支持的 fp8 类型统一 cast 到 float32 后统计
      57           43 :         LocalTensor<float> float32X = CastXToFloat32<T>(x, castXBuf, curProcessCount);
      58           43 :         AppendTailPaddingFp32(float32X, curProcessCount, appendNum);
      59           43 :         pipe_barrier(PIPE_ALL);
      60           43 :         ComputeL2NormSingleType<float>(
      61              :             workQueue, float32X, calMidAddr, curLoop, curProcessCount, appendNum, totalCount);
      62              :     }
      63              : 
      64          200 :     xQue.FreeTensor(x);
      65          200 : }
      66              : 
      67              : template <typename T>
      68          189 : __aicore__ inline void ProcessL2Norm(
      69              :     TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& calMiddleBuf,
      70              :     TBuf<TPosition::VECCALC>& castXBuf, TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, GlobalTensor<T>& xGm,
      71              :     int64_t innerLoopTime, int64_t tileLengthMean, int64_t tileLengthEnd, int64_t perBlockCount, int64_t blockOffset,
      72              :     int64_t tileNumEnd, int64_t totalCount)
      73              : {
      74          189 :     LocalTensor<uint8_t> calMidAddr =
      75          189 :         calMiddleBuf.Get<uint8_t>()[static_cast<int64_t>(StatClass::STAT_L2NORM) * BLOCK_SIZE];
      76          189 :     LocalTensor<uint8_t> curNormAddr = calMidAddr[sizeof(float)];
      77              : 
      78          201 :     for (int64_t curLoopTime = 0; curLoopTime < innerLoopTime; ++curLoopTime) {
      79           12 :         CopyInX(xQue, xGm, blockOffset + curLoopTime * tileLengthMean, tileLengthMean, perBlockCount);
      80           12 :         int64_t appendNum = 0;
      81           12 :         ComputeL2NormMultiDataType<T>(
      82              :             xQue, castXBuf, workQueue, calMidAddr, curLoopTime, tileLengthMean, appendNum, totalCount);
      83              :     }
      84              : 
      85          189 :     if (tileNumEnd) {
      86          188 :         CopyInX(xQue, xGm, blockOffset + innerLoopTime * tileLengthMean, tileLengthEnd, perBlockCount);
      87          188 :         auto appendNum = CeilAlign(tileLengthEnd, perBlockCount) - tileLengthEnd;
      88          188 :         ComputeL2NormMultiDataType<T>(
      89              :             xQue, castXBuf, workQueue, calMidAddr, innerLoopTime, tileLengthEnd + appendNum, appendNum, totalCount);
      90              :     }
      91              : 
      92          189 :     WriteStatOutput<float>(calMidAddr, curNormAddr);
      93          189 : }
      94              : 
      95              : } // namespace KfcDumpStat
      96              : 
      97              : #endif // __KFC_DUMP_L2NORM_H__