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_NAN_H__
      12              : #define __KFC_DUMP_NAN_H__
      13              : 
      14              : #include "kfc_dump_base.h"
      15              : 
      16              : namespace KfcDumpStat {
      17              : 
      18              : template <typename T, typename InputT>
      19          172 : __aicore__ inline void ComputeNanSingleType(
      20              :     TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, TBuf<TPosition::VECCALC>& maskBuf,
      21              :     TBuf<TPosition::VECCALC>& cacheBuf1, LocalTensor<InputT>& inputLocal, LocalTensor<uint8_t>& calMidAddr,
      22              :     int64_t curLoop, int64_t curProcessCount)
      23              : {
      24          172 :     LocalTensor<uint8_t> curNanAddr = calMidAddr[sizeof(float)]; // 四个字节偏移作为当前最大值存放的地址
      25          172 :     LocalTensor<InputT> workLocal = workQueue.AllocTensor<InputT>();
      26          172 :     LocalTensor<uint8_t> compareResult = maskBuf.Get<uint8_t>();
      27              : 
      28              :     // 单核模板字节数不超过 8192 字节。Nan Inf 统计的数据类型字节数最少 2 字节。
      29              :     // 因此最多处理 4096 个元素，最多需要 512 个 uint8，即 16 个 BLOCK_SIZE
      30          172 :     Compare(compareResult, inputLocal, inputLocal, CMPMODE::EQ, curProcessCount);
      31              : 
      32          172 :     InputT inputVal(0);
      33          172 :     Duplicate<InputT>(workLocal, inputVal, curProcessCount);
      34          172 :     LocalTensor<InputT> selectOnesResult = cacheBuf1.Get<InputT>();
      35          172 :     pipe_barrier(PIPE_ALL);
      36              : 
      37          172 :     Select(
      38              :         selectOnesResult, compareResult, workLocal, static_cast<InputT>(1), SELMODE::VSEL_TENSOR_SCALAR_MODE,
      39              :         curProcessCount);
      40          172 :     pipe_barrier(PIPE_ALL);
      41              : 
      42          172 :     ReduceSum<InputT>(selectOnesResult, selectOnesResult, selectOnesResult, curProcessCount);
      43          172 :     pipe_barrier(PIPE_ALL);
      44              : 
      45              :     // 更新最终结果
      46          172 :     LocalTensor<int32_t> curNan = curNanAddr.ReinterpretCast<int32_t>();
      47          172 :     if (curLoop == 0) {
      48          170 :         curNan.SetValue(0, static_cast<int32_t>(selectOnesResult.GetValue(0)));
      49              :     } else {
      50            2 :         curNan.SetValue(0, curNan.GetValue(0) + static_cast<int32_t>(selectOnesResult.GetValue(0)));
      51              :     }
      52          172 :     pipe_barrier(PIPE_ALL);
      53              : 
      54          172 :     workQueue.FreeTensor(workLocal);
      55          172 : }
      56              : 
      57              : // nan 统计计算回调：供 ComputeFloatStatSkeleton 在 cast 完成后调用
      58              : template <typename T>
      59              : struct NanComputeFunc {
      60          172 :     __aicore__ inline void operator()(
      61              :         TQue<QuePosition::VECOUT, BUFFER_NUM>& workQue, TBuf<TPosition::VECCALC>& mask,
      62              :         TBuf<TPosition::VECCALC>& cache1, LocalTensor<float>& inputLocal, LocalTensor<uint8_t>& midAddr, int64_t loop,
      63              :         int64_t processCount) const
      64              :     {
      65          172 :         ComputeNanSingleType<T, float>(workQue, mask, cache1, inputLocal, midAddr, loop, processCount);
      66          172 :     }
      67              : };
      68              : 
      69              : template <typename T>
      70          199 : __aicore__ inline void ComputeNanMultiDataType(
      71              :     TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& castXBuf,
      72              :     TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, TBuf<TPosition::VECCALC>& maskBuf,
      73              :     TBuf<TPosition::VECCALC>& cacheBuf1, LocalTensor<uint8_t>& calMidAddr, int64_t curLoop, int64_t curProcessCount,
      74              :     int64_t appendNum)
      75              : {
      76          199 :     ComputeFloatStatSkeleton<T>(
      77              :         xQue, castXBuf, workQueue, maskBuf, cacheBuf1, calMidAddr, curLoop, curProcessCount, appendNum,
      78              :         NanComputeFunc<T>());
      79          199 : }
      80              : 
      81              : template <typename T>
      82          188 : __aicore__ inline void ProcessNan(
      83              :     TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& calMiddleBuf,
      84              :     TBuf<TPosition::VECCALC>& castXBuf, TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, GlobalTensor<T>& xGm,
      85              :     TBuf<TPosition::VECCALC>& maskBuf, TBuf<TPosition::VECCALC>& cacheBuf1, int64_t innerLoopTime,
      86              :     int64_t tileLengthMean, int64_t tileLengthEnd, int64_t perBlockCount, int64_t blockOffset, int64_t tileNumEnd,
      87              :     uint64_t xDtypeSize)
      88              : {
      89              :     // 获取计算结果保存的地址
      90          188 :     LocalTensor<uint8_t> calMidAddr =
      91          188 :         calMiddleBuf.Get<uint8_t>()[static_cast<int64_t>(StatClass::STAT_NAN) * BLOCK_SIZE];
      92          188 :     LocalTensor<uint8_t> curNanAddr = calMidAddr[sizeof(float)];
      93              : 
      94          200 :     for (int64_t curLoopTime = 0; curLoopTime < innerLoopTime; ++curLoopTime) {
      95           12 :         CopyInX(xQue, xGm, blockOffset + curLoopTime * tileLengthMean, tileLengthMean, perBlockCount);
      96           12 :         int64_t appendNum = 0;
      97           12 :         ComputeNanMultiDataType<T>(
      98              :             xQue, castXBuf, workQueue, maskBuf, cacheBuf1, calMidAddr, curLoopTime, tileLengthMean, appendNum);
      99              :     }
     100              : 
     101          188 :     if (tileNumEnd) {
     102          187 :         CopyInX(xQue, xGm, blockOffset + innerLoopTime * tileLengthMean, tileLengthEnd, perBlockCount);
     103          187 :         auto appendNum = CeilAlign(tileLengthEnd * xDtypeSize, COMMAND_SIZE) / xDtypeSize - tileLengthEnd;
     104          187 :         ComputeNanMultiDataType<T>(
     105          187 :             xQue, castXBuf, workQueue, maskBuf, cacheBuf1, calMidAddr, innerLoopTime, tileLengthEnd + appendNum,
     106              :             appendNum);
     107              :     }
     108              : 
     109          188 :     WriteStatOutput<int32_t>(calMidAddr, curNanAddr);
     110          188 : }
     111              : 
     112              : } // namespace KfcDumpStat
     113              : 
     114              : #endif // __KFC_DUMP_NAN_H__