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_MAX_OR_MIN_H__
12 : #define __KFC_DUMP_MAX_OR_MIN_H__
13 :
14 : #include "kfc_dump_base.h"
15 :
16 : namespace KfcDumpStat {
17 :
18 : // max/min 统计共用模板，isStatMax 为编译期常量，由编译器完成分支裁剪
19 : template <typename InputT, bool isStatMax>
20 383 : __aicore__ inline void ComputeMaxOrMinSingleType(
21 : TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, LocalTensor<InputT>& inputLocal, LocalTensor<uint8_t>& calMidAddr,
22 : int64_t curLoop, int64_t curProcessCount)
23 : {
24 383 : LocalTensor<uint8_t> curMaxOrMinAddr = calMidAddr[sizeof(float)];
25 383 : LocalTensor<InputT> workLocal = workQueue.AllocTensor<InputT>();
26 383 : LocalTensor<InputT> calMiddleResult = calMidAddr.ReinterpretCast<InputT>();
27 :
28 : if constexpr (isStatMax) {
29 192 : ReduceMax<InputT>(calMiddleResult, inputLocal, workLocal, curProcessCount, false);
30 : } else {
31 191 : ReduceMin<InputT>(calMiddleResult, inputLocal, workLocal, curProcessCount, false);
32 : }
33 383 : pipe_barrier(PIPE_ALL);
34 :
35 : // 更新最终结果
36 383 : LocalTensor<InputT> curMaxOrMin = curMaxOrMinAddr.ReinterpretCast<InputT>();
37 383 : bool needUpdate = curLoop == 0;
38 383 : if (!needUpdate) {
39 : if constexpr (isStatMax) {
40 10 : needUpdate = static_cast<float>(calMiddleResult.GetValue(0)) > static_cast<float>(curMaxOrMin.GetValue(0));
41 : } else {
42 10 : needUpdate = static_cast<float>(calMiddleResult.GetValue(0)) < static_cast<float>(curMaxOrMin.GetValue(0));
43 : }
44 : }
45 383 : if (needUpdate) {
46 363 : curMaxOrMin.SetValue(0, calMiddleResult.GetValue(0));
47 : }
48 383 : pipe_barrier(PIPE_ALL);
49 :
50 383 : workQueue.FreeTensor(workLocal);
51 383 : }
52 :
53 : // int32 走 Max/Min 向量指令，不复用 ReduceMax/ReduceMin 路径
54 : template <bool isStatMax>
55 16 : __aicore__ inline void ComputeInt32MaxOrMin(
56 : LocalTensor<int32_t>& int32X, LocalTensor<int32_t>& x, int64_t curLoop, int64_t curProcessCount,
57 : int64_t tileLengthMean)
58 : {
59 16 : if (curLoop == 0) {
60 14 : int32_t inputVal = x.GetValue(0);
61 14 : Duplicate<int32_t>(int32X, inputVal, tileLengthMean);
62 14 : pipe_barrier(PIPE_ALL);
63 : }
64 :
65 : if constexpr (isStatMax) {
66 8 : Max(int32X, x, int32X, curProcessCount);
67 : } else {
68 8 : Min(int32X, x, int32X, curProcessCount);
69 : }
70 16 : }
71 :
72 : template <typename T, bool isStatMax>
73 399 : __aicore__ inline bool ComputeMaxOrMinMultiDataType(
74 : TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& castXBuf,
75 : TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, LocalTensor<uint8_t>& calMidAddr, int64_t curLoop,
76 : int64_t curProcessCount, int64_t appendNum, int64_t tileLengthMean)
77 : {
78 399 : LocalTensor<T> x = xQue.DeQue<T>();
79 :
80 : // 32B 对齐，max/min 尾块以首元素填充（不改变最值）；fp8 类型延后到 cast 完成后填充
81 399 : if (appendNum > 0) {
82 : if constexpr (!IsFp8Type<T>::value) {
83 300 : T appendValue = x.GetValue(0);
84 1084 : for (int64_t i = curProcessCount - appendNum; i < curProcessCount; ++i) {
85 784 : x.SetValue(i, appendValue);
86 : }
87 : }
88 306 : pipe_barrier(PIPE_ALL);
89 : }
90 :
91 399 : bool isProcessFloat32 = true;
92 : if constexpr (std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>) {
93 : // b8 转 half 后统计
94 26 : LocalTensor<half> float16X = castXBuf.Get<half>();
95 26 : Cast(float16X, x, RoundMode::CAST_NONE, curProcessCount);
96 26 : pipe_barrier(PIPE_ALL);
97 26 : ComputeMaxOrMinSingleType<half, isStatMax>(workQueue, float16X, calMidAddr, curLoop, curProcessCount);
98 26 : isProcessFloat32 = false;
99 : } else if constexpr (std::is_same_v<T, int32_t>) {
100 16 : LocalTensor<int32_t> int32X = castXBuf.Get<int32_t>();
101 16 : ComputeInt32MaxOrMin<isStatMax>(int32X, x, curLoop, curProcessCount, tileLengthMean);
102 16 : isProcessFloat32 = false;
103 : } else if constexpr (std::is_same_v<T, half>) {
104 8 : ComputeMaxOrMinSingleType<half, isStatMax>(workQueue, x, calMidAddr, curLoop, curProcessCount);
105 8 : isProcessFloat32 = false;
106 : } else {
107 : // int16/bfloat16 及支持的 fp8 类型统一 cast 到 float32 后统计
108 349 : LocalTensor<float> float32X = CastXToFloat32<T>(x, castXBuf, curProcessCount);
109 : // int16/bfloat16 在 cast 前已按首元素填充（首元素已随 cast 携带），仅 fp8 需在 cast 后补填；
110 : // 同样以首元素填充保证不改变最值，否则全负数据的 max / 全正数据的 min 会被错误统计为 0
111 : if constexpr (IsFp8Type<T>::value) {
112 18 : AppendTailPaddingFp32(float32X, curProcessCount, appendNum, float32X.GetValue(0));
113 : }
114 349 : pipe_barrier(PIPE_ALL);
115 349 : ComputeMaxOrMinSingleType<float, isStatMax>(workQueue, float32X, calMidAddr, curLoop, curProcessCount);
116 : }
117 :
118 399 : xQue.FreeTensor(x);
119 399 : return isProcessFloat32;
120 : }
121 :
122 : // 将 curMaxOrMinValueAddr 中的最值搬到 calMidAddr 首地址，不同输入类型的目标读取类型不同
123 : template <typename T, bool isStatMax>
124 377 : __aicore__ inline void WriteMaxOrMinOutput(
125 : LocalTensor<uint8_t>& calMidAddr, LocalTensor<uint8_t>& curMaxOrMinValueAddr, TBuf<TPosition::VECCALC>& castXBuf,
126 : bool isProcessFloat32, int64_t tileLengthMean)
127 : {
128 : if constexpr (std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>) {
129 14 : LocalTensor<half> tmpValue = curMaxOrMinValueAddr.ReinterpretCast<half>();
130 14 : int32_t outputValue = static_cast<int32_t>(tmpValue.GetValue(0));
131 14 : LocalTensor<int32_t> localOutput = calMidAddr.ReinterpretCast<int32_t>();
132 14 : localOutput.SetValue(0, outputValue);
133 : } else if constexpr (std::is_same_v<T, int16_t>) {
134 8 : LocalTensor<float> tmpValue = curMaxOrMinValueAddr.ReinterpretCast<float>();
135 8 : int32_t outputValue = static_cast<int32_t>(tmpValue.GetValue(0));
136 8 : LocalTensor<int32_t> localOutput = calMidAddr.ReinterpretCast<int32_t>();
137 8 : localOutput.SetValue(0, outputValue);
138 : } else if constexpr (std::is_same_v<T, int32_t>) {
139 : // int32 结果保留在 castXBuf 中，按标量逐个收尾取最值
140 14 : LocalTensor<int32_t> int32X = castXBuf.Get<int32_t>();
141 14 : int32_t outputValue = int32X.GetValue(0);
142 67584 : for (int32_t i = 1; i < tileLengthMean; ++i) {
143 67570 : int32_t curValue = int32X.GetValue(i);
144 33785 : if (isStatMax && outputValue < curValue) {
145 4704 : outputValue = curValue;
146 33785 : } else if (!isStatMax && outputValue > curValue) {
147 0 : outputValue = curValue;
148 : }
149 : }
150 14 : LocalTensor<int32_t> localOutput = calMidAddr.ReinterpretCast<int32_t>();
151 14 : localOutput.SetValue(0, outputValue);
152 : } else if constexpr (std::is_same_v<T, half>) {
153 8 : if (isProcessFloat32) {
154 0 : LocalTensor<float> tmpValue = curMaxOrMinValueAddr.ReinterpretCast<float>();
155 0 : float outputValue = tmpValue.GetValue(0);
156 0 : LocalTensor<float> localOutput = calMidAddr.ReinterpretCast<float>();
157 0 : localOutput.SetValue(0, outputValue);
158 : } else {
159 8 : LocalTensor<half> tmpValue = curMaxOrMinValueAddr.ReinterpretCast<half>();
160 8 : float outputValue = static_cast<float>(tmpValue.GetValue(0));
161 8 : LocalTensor<float> localOutput = calMidAddr.ReinterpretCast<float>();
162 8 : localOutput.SetValue(0, outputValue);
163 : }
164 : } else {
165 : // float/bfloat16 及支持的 fp8 类型：结果以 float32 存放
166 333 : LocalTensor<float> tmpValue = curMaxOrMinValueAddr.ReinterpretCast<float>();
167 333 : auto outputValue = tmpValue.GetValue(0);
168 333 : LocalTensor<float> localOutput = calMidAddr.ReinterpretCast<float>();
169 333 : localOutput.SetValue(0, outputValue);
170 : }
171 377 : }
172 :
173 : template <typename T, bool isStatMax>
174 377 : __aicore__ inline void ProcessMaxOrMin(
175 : TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& calMiddleBuf,
176 : TBuf<TPosition::VECCALC>& castXBuf, TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, GlobalTensor<T>& xGm,
177 : int64_t innerLoopTime, int64_t tileLengthMean, int64_t tileLengthEnd, int64_t perBlockCount, int64_t blockOffset,
178 : int64_t tileNumEnd)
179 : {
180 377 : StatClass statClass = isStatMax ? StatClass::STAT_MAX : StatClass::STAT_MIN;
181 377 : LocalTensor<uint8_t> calMidAddr = calMiddleBuf.Get<uint8_t>()[static_cast<int64_t>(statClass) * BLOCK_SIZE];
182 377 : LocalTensor<uint8_t> curMaxOrMinValueAddr = calMidAddr[sizeof(float)];
183 377 : bool isProcessFloat32 = true;
184 :
185 401 : for (int64_t curLoopTime = 0; curLoopTime < innerLoopTime; ++curLoopTime) {
186 24 : CopyInX(
187 24 : xQue, xGm, blockOffset + curLoopTime * tileLengthMean, tileLengthMean,
188 : perBlockCount); // tileLengthMean 满足 32B 对齐
189 24 : int64_t appendNum = 0;
190 24 : isProcessFloat32 = ComputeMaxOrMinMultiDataType<T, isStatMax>(
191 : xQue, castXBuf, workQueue, calMidAddr, curLoopTime, tileLengthMean, appendNum, tileLengthMean);
192 : }
193 :
194 377 : if (tileNumEnd) {
195 375 : CopyInX(xQue, xGm, blockOffset + innerLoopTime * tileLengthMean, tileLengthEnd, perBlockCount); // 处理尾块
196 375 : auto appendNum = CeilAlign(tileLengthEnd, perBlockCount) - tileLengthEnd;
197 375 : isProcessFloat32 = ComputeMaxOrMinMultiDataType<T, isStatMax>(
198 : xQue, castXBuf, workQueue, calMidAddr, innerLoopTime, tileLengthEnd + appendNum, appendNum, tileLengthMean);
199 : }
200 :
201 377 : WriteMaxOrMinOutput<T, isStatMax>(calMidAddr, curMaxOrMinValueAddr, castXBuf, isProcessFloat32, tileLengthMean);
202 377 : pipe_barrier(PIPE_ALL);
203 377 : }
204 :
205 : } // namespace KfcDumpStat
206 :
207 : #endif // __KFC_DUMP_MAX_OR_MIN_H__
|
