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_MEAN_H__
      12              : #define __KFC_DUMP_MEAN_H__
      13              : 
      14              : #include "kfc_dump_base.h"
      15              : 
      16              : namespace KfcDumpStat {
      17              : 
      18              : template <typename T, typename InputT>
      19          199 : __aicore__ inline void ComputeMeanSingleType(
      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          199 :     LocalTensor<uint8_t> curMeanAddr = calMidAddr[sizeof(float)];
      24          199 :     LocalTensor<InputT> workLocal = workQueue.AllocTensor<InputT>();
      25              : 
      26          199 :     ReduceSum<InputT>(workLocal, inputLocal, workLocal, curProcessCount);
      27          199 :     PipeBarrier<PIPE_ALL>();
      28              : 
      29              :     // 更新最终结果
      30          199 :     LocalTensor<InputT> curMean = curMeanAddr.ReinterpretCast<InputT>();
      31          199 :     if (curLoop == 0) {
      32          188 :         curMean.SetValue(0, workLocal.GetValue(0) / totalCount);
      33              :     } else {
      34           11 :         curMean.SetValue(0, curMean.GetValue(0) + (workLocal.GetValue(0) / totalCount));
      35              :     }
      36          199 :     PipeBarrier<PIPE_ALL>();
      37              : 
      38          199 :     workQueue.FreeTensor(workLocal);
      39          199 : }
      40              : 
      41              : template <typename T>
      42          199 : __aicore__ inline void ComputeMeanMultiDataType(
      43              :     TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& castXBuf,
      44              :     TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, LocalTensor<uint8_t>& calMidAddr, int64_t curLoop,
      45              :     int64_t curProcessCount, int64_t appendNum, int64_t totalCount)
      46              : {
      47          199 :     LocalTensor<T> x = xQue.DeQue<T>();
      48              : 
      49          199 :     AppendTailPadding<T>(x, curProcessCount, appendNum);
      50              : 
      51              :     if constexpr (std::is_same_v<T, float>) {
      52          156 :         ComputeMeanSingleType<float>(workQueue, x, calMidAddr, curLoop, curProcessCount, appendNum, totalCount);
      53              :     } else {
      54              :         // 整型/half/bfloat16 及支持的 fp8 类型统一 cast 到 float32 后统计
      55           43 :         LocalTensor<float> float32X = CastXToFloat32<T>(x, castXBuf, curProcessCount);
      56           43 :         AppendTailPaddingFp32(float32X, curProcessCount, appendNum);
      57           43 :         PipeBarrier<PIPE_ALL>();
      58           43 :         ComputeMeanSingleType<float>(workQueue, float32X, calMidAddr, curLoop, curProcessCount, appendNum, totalCount);
      59              :     }
      60              : 
      61          199 :     xQue.FreeTensor(x);
      62          199 : }
      63              : 
      64              : template <typename T>
      65          188 : __aicore__ inline void ProcessMean(
      66              :     TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& calMiddleBuf,
      67              :     TBuf<TPosition::VECCALC>& castXBuf, TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, GlobalTensor<T>& xGm,
      68              :     int64_t innerLoopTime, int64_t tileLengthMean, int64_t tileLengthEnd, int64_t perBlockCount, int64_t blockOffset,
      69              :     int64_t tileNumEnd, int64_t totalCount)
      70              : {
      71          188 :     LocalTensor<uint8_t> calMidAddr =
      72          188 :         calMiddleBuf.Get<uint8_t>()[static_cast<int64_t>(StatClass::STAT_MEAN) * BLOCK_SIZE];
      73          188 :     LocalTensor<uint8_t> curMeanAddr = calMidAddr[sizeof(float)];
      74              : 
      75          200 :     for (int64_t curLoopTime = 0; curLoopTime < innerLoopTime; ++curLoopTime) {
      76           12 :         CopyInX(xQue, xGm, blockOffset + curLoopTime * tileLengthMean, tileLengthMean, perBlockCount);
      77           12 :         int64_t appendNum = 0;
      78           12 :         ComputeMeanMultiDataType<T>(
      79              :             xQue, castXBuf, workQueue, calMidAddr, curLoopTime, tileLengthMean, appendNum, totalCount);
      80              :     }
      81              : 
      82          188 :     if (tileNumEnd) {
      83          187 :         CopyInX(xQue, xGm, blockOffset + innerLoopTime * tileLengthMean, tileLengthEnd, perBlockCount);
      84          187 :         auto appendNum = CeilAlign(tileLengthEnd, perBlockCount) - tileLengthEnd;
      85          187 :         ComputeMeanMultiDataType<T>(
      86              :             xQue, castXBuf, workQueue, calMidAddr, innerLoopTime, tileLengthEnd + appendNum, appendNum, totalCount);
      87              :     }
      88              : 
      89          188 :     WriteStatOutput<float>(calMidAddr, curMeanAddr);
      90          188 : }
      91              : 
      92              : } // namespace KfcDumpStat
      93              : 
      94              : #endif // __KFC_DUMP_MEAN_H__