Star1431
2162
Code
Code
Issues956
Pull Requests781
Actions
Discussion
Wiki
Analysis
Star1431
2162
Merged
perf(fused_cross_entropy_loss_with_max_sum): 小V场景性能优化（平坦打包快路+预取流水+裸vdiv） #14270
AlfengYuanAdd on 2 days ago
perf(fused_cross_entropy_loss_with_max_sum): 小V场景性能优化（平坦打包快路+预取流水+裸vdiv） #14270
Merged
AlfengYuanAdd on 2 days ago
Merged fromalfengyuan/ops-nnissue225-fce-perfintocann/ops-nnmaster
Total: 4 files changed+475-20
Mloss/fused_cross_entropy_loss_with_max_sum/op_host/fused_cross_entropy_loss_with_max_sum_tiling_arch35.cpp+22-2
@@ -170,13 +170,32 @@ ge::graphStatus FusedCrossEntropyLossWithMaxSumRegbaseTiling::DoTilingFull()
170 tilingData_->vLen = v;170 tilingData_->vLen = v;
171 tilingData_->elementsNumber = 0;171 tilingData_->elementsNumber = 0;
172 tilingData_->vocabDtypeId = vocabDtypeId;172 tilingData_->vocabDtypeId = vocabDtypeId;
173+ // 小V快路：按"多行打包成平坦向量"重排UB预算，行数由host按V与dtype算出塞满UB
174+ tilingData_->aPerLoop = FUSED_CE_MAX_SUM_A_PER_LOOP;
175+ if (v <= FUSED_CE_MAX_SUM_SMALL_V_MAX) {
176+ // 每行开销：vocab双缓冲 + softmax双缓冲 + vocab的fp32展开副本 + 行向量buffer
177+ // （max/sum/pred/loss四队列深2支撑跨tile预取流水 + invSum单buffer，共9个）
178+ constexpr int64_t SCALAR_BYTES_PER_ROW = 9 * FLOAT_BYTES;
179+ constexpr int64_t SMALL_V_A_ALIGN = 128;
180+ constexpr int64_t SMALL_V_A_MAX = 8192;
181+ int64_t bytesPerRow = DOUBLE_BUFFER_NUM * v * dtypeSize + DOUBLE_BUFFER_NUM * v * FLOAT_BYTES +
182+ v * FLOAT_BYTES + SCALAR_BYTES_PER_ROW;
183+ // 向量访存按VL对齐多预留一个向量寄存器的余量，避免尾块越界读
184+ int64_t budget = ubSize_ - RESERVED_UB -
185+ DOUBLE_BUFFER_NUM * FUSED_CE_MAX_SUM_V_ALIGN * (dtypeSize + FLOAT_BYTES);
186+ int64_t aPerLoop = FloorAlign(budget / bytesPerRow, SMALL_V_A_ALIGN);
187+ OP_CHECK_IF(aPerLoop < SMALL_V_A_ALIGN,
188+ OP_LOGE(context_, "ub is not enough for smallV, aPerLoop is %ld", aPerLoop),
189+ return ge::GRAPH_FAILED);
190+ tilingData_->aPerLoop = aPerLoop > SMALL_V_A_MAX ? SMALL_V_A_MAX : aPerLoop;
191+ }
173 192 
174 constexpr uint64_t tilingKey = 0; // TILINGKEY_FULL193 constexpr uint64_t tilingKey = 0; // TILINGKEY_FULL
175 OP_LOGI(context_,194 OP_LOGI(context_,
176 "DoTilingFull: bt is %ld, v is %ld, vPerLoop is %ld, formerCoreNum is %ld, formerRows is %ld, "195 "DoTilingFull: bt is %ld, v is %ld, vPerLoop is %ld, formerCoreNum is %ld, formerRows is %ld, "
177- "latterRows is %ld, vCores is %ld, vChunk is %ld, tilingKey is %lu",196+ "latterRows is %ld, vCores is %ld, vChunk is %ld, aPerLoop is %ld, tilingKey is %lu",
178 bt, v, vPerLoop, tilingData_->formerCoreNum, tilingData_->formerRows, tilingData_->latterRows,197 bt, v, vPerLoop, tilingData_->formerCoreNum, tilingData_->formerRows, tilingData_->latterRows,
179- tilingData_->vCores, tilingData_->vChunk, tilingKey);198+ tilingData_->vCores, tilingData_->vChunk, tilingData_->aPerLoop, tilingKey);
180 context_->SetTilingKey(tilingKey);199 context_->SetTilingKey(tilingKey);
181 return ge::GRAPH_SUCCESS;200 return ge::GRAPH_SUCCESS;
182}201}
@@ -207,6 +226,7 @@ ge::graphStatus FusedCrossEntropyLossWithMaxSumRegbaseTiling::DoTilingForMemory(
207 tilingData_->elementsNumber = elementsNumber;226 tilingData_->elementsNumber = elementsNumber;
208 tilingData_->vCores = 1; // 省显存路径只有bt一维，不切v227 tilingData_->vCores = 1; // 省显存路径只有bt一维，不切v
209 tilingData_->vChunk = 0;228 tilingData_->vChunk = 0;
229+ tilingData_->aPerLoop = 0;
210 230 
211 constexpr uint64_t tilingKey = 1; // TILINGKEY_MEMORY231 constexpr uint64_t tilingKey = 1; // TILINGKEY_MEMORY
212 OP_LOGI(context_,232 OP_LOGI(context_,
Mloss/fused_cross_entropy_loss_with_max_sum/op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_ar.h+333-18
@@ -45,6 +45,26 @@ constexpr static AscendC::Reg::DivSpecificMode castDivFt32False = {
45 false,45 false,
46 AscendC::DivAlgo::PRECISION_1ULP_FTZ_FALSE,46 AscendC::DivAlgo::PRECISION_1ULP_FTZ_FALSE,
47};47};
48+// 小V快路的 inv_sum=1/sum_exp 用裸 vdiv：sum_exp 非正常(次正规)时 inv_sum 本就溢出为 inf，
49+// 补偿与否不影响输出；裸 vdiv 省掉补偿序列，吞吐显著更高
50+constexpr static AscendC::Reg::DivSpecificMode castDivIntrinsic = {
51+ AscendC::Reg::MaskMergeMode::ZEROING,
52+ false,
53+ AscendC::DivAlgo::INTRINSIC,
54+};
55+// 小V快路Gather索引构建用CastTrait（仿softmax_focal_loss_grad_nd_vec.h惯用法）
56+constexpr static AscendC::Reg::CastTrait castTraitS32ToFp32 = {
57+ AscendC::Reg::RegLayout::ZERO,
58+ AscendC::Reg::SatMode::NO_SAT,
59+ AscendC::Reg::MaskMergeMode::ZEROING,
60+ AscendC::RoundMode::CAST_NONE,
61+};
62+constexpr static AscendC::Reg::CastTrait castTraitFp32ToS32Floor = {
63+ AscendC::Reg::RegLayout::ZERO,
64+ AscendC::Reg::SatMode::NO_SAT,
65+ AscendC::Reg::MaskMergeMode::ZEROING,
66+ AscendC::RoundMode::CAST_FLOOR,
67+};
48 68 
49template <typename T, bool fullPath>69template <typename T, bool fullPath>
50class FusedCrossEntropyLossWithMaxSumRegBase {70class FusedCrossEntropyLossWithMaxSumRegBase {
@@ -72,11 +92,32 @@ private:
72 int64_t aRows, int64_t vCur);92 int64_t aRows, int64_t vCur);
73 __aicore__ inline void CopyOutSoftmaxTile(const LocalTensor<float>& ub, int64_t rowBase, int64_t vOff,93 __aicore__ inline void CopyOutSoftmaxTile(const LocalTensor<float>& ub, int64_t rowBase, int64_t vOff,
74 int64_t aRows, int64_t vCur);94 int64_t aRows, int64_t vCur);
95+ // 小V快路（vLen<=FUSED_CE_MAX_SUM_SMALL_V_MAX）：多行打包成平坦向量，
96+ // max/inv按Gather(l/V)展开，vocab与softmax均按整块平坦搬运；
97+ // 循环按"预取下一块输入→再等本块输出"排布，使MTE与VEC跨tile重叠
98+ __aicore__ inline void ProcessFullSmallV();
99+ __aicore__ inline void PrefetchSmallVTile(int64_t tileIdx);
100+ __aicore__ inline int64_t SmallVTileRows(int64_t tileIdx);
101+ __aicore__ inline void ComputeLossInvSumLoop(const LocalTensor<float>& sumExpUb,
102+ const LocalTensor<float>& predictedUb,
103+ const LocalTensor<float>& lossUb, const LocalTensor<float>& invSumUb,
104+ int64_t len);
105+ __aicore__ inline void CastVocabToFp32(const LocalTensor<T>& srcUb, const LocalTensor<float>& dstUb, int64_t len);
106+ // softmax平坦计算：源为fp32（原样/预转副本）或b16（内联cast），由S区分
107+ template <typename S>
108+ __aicore__ inline void SoftmaxFlatLoop(__ubuf__ S* src, const LocalTensor<float>& maxUb,
109+ const LocalTensor<float>& invSumUb, const LocalTensor<float>& outUb,
110+ int64_t aRows, int64_t g);
111+ __aicore__ inline void ComputeSoftmaxFlat(const LocalTensor<T>& vocabUb, const LocalTensor<float>& maxUb,
112+ const LocalTensor<float>& invSumUb, const LocalTensor<float>& outUb,
113+ int64_t aRows);
75 // 省显存路径：仅loss114 // 省显存路径：仅loss
76 __aicore__ inline void ProcessMemory();115 __aicore__ inline void ProcessMemory();
77 __aicore__ inline void ComputeLossChunk(const LocalTensor<float>& sumExpUb, const LocalTensor<float>& predictedUb,116 __aicore__ inline void ComputeLossChunk(const LocalTensor<float>& sumExpUb, const LocalTensor<float>& predictedUb,
78 const LocalTensor<float>& lossUb, int64_t len);117 const LocalTensor<float>& lossUb, int64_t len);
79- __aicore__ inline void LoadVocabToFp32(__ubuf__ T* src, AscendC::Reg::RegTensor<float>& dst,118+ // 加载len个元素并转fp32：S=float直接DIST_NORM，S=b16则UNPACK_B16+Cast（源地址须32B对齐）
119+ template <typename S>
120+ __aicore__ inline void LoadVocabToFp32(__ubuf__ S* src, AscendC::Reg::RegTensor<float>& dst,
80 AscendC::Reg::MaskReg& preg, uint32_t offset);121 AscendC::Reg::MaskReg& preg, uint32_t offset);
81 122 
82private:123private:
@@ -95,14 +136,16 @@ private:
95 GlobalTensor<float> softmaxGm_;136 GlobalTensor<float> softmaxGm_;
96 GlobalTensor<float> lossGm_;137 GlobalTensor<float> lossGm_;
97 138 
98- // 完整路径buffer139+ // 完整路径buffer（小V快路四标量队列按深度2做跨tile预取流水；大V通路单缓冲，num=1在[1,bufNum]内合法）
99 TQue<QuePosition::VECIN, 2> vocabQue_;140 TQue<QuePosition::VECIN, 2> vocabQue_;
100 TQue<QuePosition::VECOUT, 2> softmaxQue_;141 TQue<QuePosition::VECOUT, 2> softmaxQue_;
101- TQue<QuePosition::VECIN, 1> logitsMaxQue_;142+ TQue<QuePosition::VECIN, 2> logitsMaxQue_;
102- TQue<QuePosition::VECIN, 1> sumExpScalarQue_;143+ TQue<QuePosition::VECIN, 2> sumExpScalarQue_;
103- TQue<QuePosition::VECIN, 1> predictedScalarQue_;144+ TQue<QuePosition::VECIN, 2> predictedScalarQue_;
104- TQue<QuePosition::VECOUT, 1> lossScalarQue_;145+ TQue<QuePosition::VECOUT, 2> lossScalarQue_;
105 TBuf<> invSumBuf_;146 TBuf<> invSumBuf_;
147+ // 小V快路buffer：vocab的fp32平坦副本（Gather索引在寄存器内经Arange/Cast/Div构建）
148+ TBuf<> vocabFp32Buf_;
106 // 省显存路径buffer149 // 省显存路径buffer
107 TQue<QuePosition::VECIN, 2> sumExpQue_;150 TQue<QuePosition::VECIN, 2> sumExpQue_;
108 TQue<QuePosition::VECIN, 2> predictedQue_;151 TQue<QuePosition::VECIN, 2> predictedQue_;
@@ -111,6 +154,7 @@ private:
111 int64_t myRows_ = 0;154 int64_t myRows_ = 0;
112 int64_t rowStart_ = 0;155 int64_t rowStart_ = 0;
113 int64_t vPartIdx_ = 0; // v切分时本核在一行内负责的v分片序号（0..vCores-1），不切分恒为0156 int64_t vPartIdx_ = 0; // v切分时本核在一行内负责的v分片序号（0..vCores-1），不切分恒为0
157+ bool smallV_ = false; // vLen<=FUSED_CE_MAX_SUM_SMALL_V_MAX时走小V快路
114};158};
115 159 
116template <typename T, bool fullPath>160template <typename T, bool fullPath>
@@ -139,13 +183,32 @@ __aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::Init
139 vocabGm_.SetGlobalBuffer((__gm__ T*)vocabParallelLogits);183 vocabGm_.SetGlobalBuffer((__gm__ T*)vocabParallelLogits);
140 logitsMaxGm_.SetGlobalBuffer((__gm__ float*)logitsMax);184 logitsMaxGm_.SetGlobalBuffer((__gm__ float*)logitsMax);
141 softmaxGm_.SetGlobalBuffer((__gm__ float*)softmaxLogits);185 softmaxGm_.SetGlobalBuffer((__gm__ float*)softmaxLogits);
142- pipe_->InitBuffer(vocabQue_, 2, A_PER_LOOP * tl_->vPerLoop * sizeof(T));186+ smallV_ = tl_->vLen > 0 && tl_->vLen <= FUSED_CE_MAX_SUM_SMALL_V_MAX;
143- pipe_->InitBuffer(softmaxQue_, 2, A_PER_LOOP * tl_->vPerLoop * sizeof(float));187+ if (smallV_) {
144- pipe_->InitBuffer(logitsMaxQue_, 1, SCALAR_BUF_BYTES);188+ // tile按VL对齐预留，保证尾块向量访存的越界读/零写在已分配区内；
145- pipe_->InitBuffer(sumExpScalarQue_, 1, SCALAR_BUF_BYTES);189+ // vocab源缓冲额外多留一个VL：V>1时G*v<VL（如V=3组步长48<64），末组满载
146- pipe_->InitBuffer(predictedScalarQue_, 1, SCALAR_BUF_BYTES);190+ // 向量读会越过tileElems，padding把这次越界读收敛在自有buffer内（掩码位不参与计算）。
147- pipe_->InitBuffer(lossScalarQue_, 1, SCALAR_BUF_BYTES);191+ // 该开销已含在host预算的DOUBLE_BUFFER_NUM*V_ALIGN*(dtypeSize+4B)余量中，不挤占aPerLoop；
148- pipe_->InitBuffer(invSumBuf_, SCALAR_BUF_BYTES);192+ // max/sum/pred/loss四个行向量队列深2，支撑跨tile预取流水
193+ int64_t tileElems = ops::CeilDiv(tl_->aPerLoop * tl_->vLen, static_cast<int64_t>(VL_FP32)) * VL_FP32;
194+ int64_t rowPad = ops::CeilDiv(tl_->aPerLoop, static_cast<int64_t>(VL_FP32)) * VL_FP32;
195+ pipe_->InitBuffer(vocabQue_, 2, (tileElems + VL_FP32) * sizeof(T));
196+ pipe_->InitBuffer(softmaxQue_, 2, tileElems * sizeof(float));
197+ pipe_->InitBuffer(logitsMaxQue_, 2, rowPad * sizeof(float));
198+ pipe_->InitBuffer(sumExpScalarQue_, 2, rowPad * sizeof(float));
199+ pipe_->InitBuffer(predictedScalarQue_, 2, rowPad * sizeof(float));
200+ pipe_->InitBuffer(lossScalarQue_, 2, rowPad * sizeof(float));
atomgit-bot
atomgit-botatomgit-bot2 days ago

🟡 Medium Priority

变更行：fused_cross_entropy_loss_with_max_sum_ar.h:194-197 在小V快路里对 logitsMaxQue_ / sumExpScalarQue_ / predictedScalarQue_ / lossScalarQue_ 调用 pipe_->InitBuffer(que, 2, rowPad * sizeof(float))，但这 4 个队列在 142-145 行声明为 TQue<QuePosition::VECIN, 1>（lossScalarQue_ 为 TQue<QuePosition::VECOUT, 1>），即模板深度参数 bufNum=1。

契约/影响：Ascend C 的 TQue 第二个模板参数是队列深度（可缓存的 Tensor 个数），TPipe::InitBuffer(TQue&, num, len) 的 num 语义上是 [1, bufNum] 范围内的 buffer 数，TQue 内部状态（BufState 数组/事件索引）按 depth 分配。传 num=2 > depth=1 属于越界使用：轻则在框架断言处直接失败/只分配 1 个 buffer 而代码按 2 个流水理解，重则状态数组越界写破坏相邻成员，导致 kernel 编译期报错或运行期数据错乱。

一致性证据：本仓所有双缓冲用法都把深度与 num 对齐——loss/smooth_l1_loss_grad/op_kernel/arch35/smooth_l1_loss_grad.h:142-145 声明 TQue<QuePosition::VECIN, 2> 后才 214-217 行 InitBuffer(..., 2, ...)；loss/sigmoid_focal_loss/op_kernel/arch35/sigmoid_focal_loss_kernel.h:778 声明深度 2 后 391 行传 2。本文件同一函数的 else 分支（200-208 行）也是深度 1 配 num=1。只有新增的小V分支出现深度 1 配 2 的单边偏离。

likedislike
Incorrect?
201+ pipe_->InitBuffer(invSumBuf_, rowPad * sizeof(float));
202+ pipe_->InitBuffer(vocabFp32Buf_, (tileElems + VL_FP32) * sizeof(float));
203+ } else {
204+ pipe_->InitBuffer(vocabQue_, 2, A_PER_LOOP * tl_->vPerLoop * sizeof(T));
205+ pipe_->InitBuffer(softmaxQue_, 2, A_PER_LOOP * tl_->vPerLoop * sizeof(float));
206+ pipe_->InitBuffer(logitsMaxQue_, 1, SCALAR_BUF_BYTES);
207+ pipe_->InitBuffer(sumExpScalarQue_, 1, SCALAR_BUF_BYTES);
208+ pipe_->InitBuffer(predictedScalarQue_, 1, SCALAR_BUF_BYTES);
209+ pipe_->InitBuffer(lossScalarQue_, 1, SCALAR_BUF_BYTES);
210+ pipe_->InitBuffer(invSumBuf_, SCALAR_BUF_BYTES);
211+ }
149 } else {212 } else {
150 pipe_->InitBuffer(sumExpQue_, 2, tl_->elementsNumber * sizeof(float));213 pipe_->InitBuffer(sumExpQue_, 2, tl_->elementsNumber * sizeof(float));
151 pipe_->InitBuffer(predictedQue_, 2, tl_->elementsNumber * sizeof(float));214 pipe_->InitBuffer(predictedQue_, 2, tl_->elementsNumber * sizeof(float));
@@ -160,6 +223,10 @@ __aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::Proc
160 return;223 return;
161 }224 }
162 if constexpr (fullPath) {225 if constexpr (fullPath) {
226+ if (smallV_) {
227+ ProcessFullSmallV();
228+ return;
229+ }
163 ProcessFull();230 ProcessFull();
164 } else {231 } else {
165 ProcessMemory();232 ProcessMemory();
@@ -177,6 +244,88 @@ __aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::Proc
177 }244 }
178}245}
179 246 
247+// 小V快路的tile行数（尾块不足aPerLoop）
248+template <typename T, bool fullPath>
249+__aicore__ inline int64_t FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::SmallVTileRows(int64_t tileIdx)
250+{
251+ int64_t rows = myRows_ - tileIdx * tl_->aPerLoop;
252+ return rows < tl_->aPerLoop ? rows : tl_->aPerLoop;
253+}
254+ 
255+// 下发tileIdx块的全部搬入（3个标量行向量 + vocab整块平坦拷贝），仅入队不等待
256+template <typename T, bool fullPath>
257+__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::PrefetchSmallVTile(int64_t tileIdx)
258+{
259+ int64_t rowBase = rowStart_ + tileIdx * tl_->aPerLoop;
260+ int64_t aRows = SmallVTileRows(tileIdx);
261+ LocalTensor<float> maxUb = logitsMaxQue_.AllocTensor<float>();
262+ LocalTensor<float> sumUb = sumExpScalarQue_.AllocTensor<float>();
263+ LocalTensor<float> predUb = predictedScalarQue_.AllocTensor<float>();
264+ CopyInScalar(maxUb, logitsMaxGm_, rowBase, aRows);
265+ CopyInScalar(sumUb, sumExpGm_, rowBase, aRows);
266+ CopyInScalar(predUb, predictedGm_, rowBase, aRows);
267+ logitsMaxQue_.EnQue(maxUb);
268+ sumExpScalarQue_.EnQue(sumUb);
269+ predictedScalarQue_.EnQue(predUb);
270+ 
271+ LocalTensor<T> vocabUb = vocabQue_.AllocTensor<T>();
272+ DataCopyExtParams vocabCopyParams{1, static_cast<uint32_t>(aRows * tl_->vLen * sizeof(T)), 0, 0, 0};
273+ DataCopyPadExtParams<T> vocabPadParams{false, 0, 0, 0};
274+ DataCopyPad(vocabUb, vocabGm_[rowBase * tl_->vLen], vocabCopyParams, vocabPadParams);
275+ vocabQue_.EnQue(vocabUb);
276+}
277+ 
278+// 小V快路主循环：先下发本块全部VEC计算、再预取下一块输入（MTE与VEC重叠）、最后等本块输出
279+template <typename T, bool fullPath>
280+__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ProcessFullSmallV()
281+{
282+ int64_t aLoops = ops::CeilDiv(myRows_, tl_->aPerLoop);
283+ PrefetchSmallVTile(0);
284+ for (int64_t i = 0; i < aLoops; i++) {
285+ int64_t rowBase = rowStart_ + i * tl_->aPerLoop;
286+ int64_t aRows = SmallVTileRows(i);
287+ 
288+ // 1. 本块输入就绪
289+ LocalTensor<float> maxUb = logitsMaxQue_.DeQue<float>();
290+ LocalTensor<float> sumUb = sumExpScalarQue_.DeQue<float>();
291+ LocalTensor<float> predUb = predictedScalarQue_.DeQue<float>();
292+ 
293+ // 2. loss = log(sum_exp) - predicted 与 inv_sum = 1 / sum_exp（按VL分块向量化）
294+ LocalTensor<float> invSumUb = invSumBuf_.Get<float>();
295+ LocalTensor<float> lossUb = lossScalarQue_.AllocTensor<float>();
296+ ComputeLossInvSumLoop(sumUb, predUb, lossUb, invSumUb, aRows);
297+ lossScalarQue_.EnQue(lossUb);
298+ 
299+ // 3. softmax = exp(vocab - max) * inv_sum，G行一组打包成G*vLen条lane平坦计算；
300+ // b16且组内偏移满足32B对齐时加载即内联cast，否则先整tile预转fp32
301+ LocalTensor<T> vocabUb = vocabQue_.DeQue<T>();
302+ LocalTensor<float> outUb = softmaxQue_.AllocTensor<float>();
303+ ComputeSoftmaxFlat(vocabUb, maxUb, invSumUb, outUb, aRows);
304+ vocabQue_.FreeTensor(vocabUb);
305+ softmaxQue_.EnQue(outUb);
306+ 
307+ // 5. 本块输入已消费完（VEC侧由事件保证读完才复用），释放后立即预取下一块，
308+ // 使下一块的MTE搬入与本块仍在执行的VEC计算重叠
309+ logitsMaxQue_.FreeTensor(maxUb);
310+ sumExpScalarQue_.FreeTensor(sumUb);
311+ predictedScalarQue_.FreeTensor(predUb);
312+ if (i + 1 < aLoops) {
313+ PrefetchSmallVTile(i + 1);
314+ }
315+ 
316+ // 6. 等本块输出并写回
317+ LocalTensor<float> lossOutUb = lossScalarQue_.DeQue<float>();
318+ DataCopyExtParams lossCopyParams{1, static_cast<uint32_t>(aRows * sizeof(float)), 0, 0, 0};
319+ DataCopyPad(lossGm_[rowBase], lossOutUb, lossCopyParams);
320+ lossScalarQue_.FreeTensor(lossOutUb);
321+ 
322+ LocalTensor<float> outGmUb = softmaxQue_.DeQue<float>();
323+ DataCopyExtParams outCopyParams{1, static_cast<uint32_t>(aRows * tl_->vLen * sizeof(float)), 0, 0, 0};
324+ DataCopyPad(softmaxGm_[rowBase * tl_->vLen], outGmUb, outCopyParams);
325+ softmaxQue_.FreeTensor(outGmUb);
326+ }
327+}
328+ 
180template <typename T, bool fullPath>329template <typename T, bool fullPath>
181__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ProcessRowTile(int64_t rowBase,330__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ProcessRowTile(int64_t rowBase,
182 int64_t aRows)331 int64_t aRows)
@@ -290,15 +439,16 @@ __aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::Copy
290}439}
291 440 
292template <typename T, bool fullPath>441template <typename T, bool fullPath>
442+template <typename S>
293__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::LoadVocabToFp32(443__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::LoadVocabToFp32(
294- __ubuf__ T* src, AscendC::Reg::RegTensor<float>& dst, AscendC::Reg::MaskReg& preg, uint32_t offset)444+ __ubuf__ S* src, AscendC::Reg::RegTensor<float>& dst, AscendC::Reg::MaskReg& preg, uint32_t offset)
295{445{
296- if constexpr (IsSameType<T, float>::value) {446+ if constexpr (IsSameType<S, float>::value) {
297 AscendC::Reg::DataCopy<float, AscendC::Reg::LoadDist::DIST_NORM>(dst, src + offset);447 AscendC::Reg::DataCopy<float, AscendC::Reg::LoadDist::DIST_NORM>(dst, src + offset);
298 } else {448 } else {
299- AscendC::Reg::RegTensor<T> b16Reg;449+ AscendC::Reg::RegTensor<S> b16Reg;
300- AscendC::Reg::DataCopy<T, AscendC::Reg::LoadDist::DIST_UNPACK_B16>(b16Reg, src + offset);450+ AscendC::Reg::DataCopy<S, AscendC::Reg::LoadDist::DIST_UNPACK_B16>(b16Reg, src + offset);
301- AscendC::Reg::Cast<float, T, castTraitB16ToFp32>(dst, b16Reg, preg);451+ AscendC::Reg::Cast<float, S, castTraitB16ToFp32>(dst, b16Reg, preg);
302 }452 }
303}453}
304 454 
@@ -363,6 +513,171 @@ __aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::Copy
363 DataCopyPad(softmaxGm_[rowBase * tl_->vLen + vOff], ub, copyParams);513 DataCopyPad(softmaxGm_[rowBase * tl_->vLen + vOff], ub, copyParams);
364}514}
365 515 
516+template <typename T, bool fullPath>
517+__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ComputeLossInvSumLoop(
518+ const LocalTensor<float>& sumExpUb, const LocalTensor<float>& predictedUb, const LocalTensor<float>& lossUb,
519+ const LocalTensor<float>& invSumUb, int64_t len)
520+{
521+ __ubuf__ float* sumAddr = (__ubuf__ float*)sumExpUb.GetPhyAddr();
522+ __ubuf__ float* predAddr = (__ubuf__ float*)predictedUb.GetPhyAddr();
523+ __ubuf__ float* lossAddr = (__ubuf__ float*)lossUb.GetPhyAddr();
524+ __ubuf__ float* invAddr = (__ubuf__ float*)invSumUb.GetPhyAddr();
525+ // 与ComputeLossChunk同构：满块迭代复用编译期满掩码，尾块单独0/1次迭代
526+ uint16_t fullLoops = static_cast<uint16_t>(len / static_cast<int64_t>(VL_FP32));
527+ uint16_t totalLoops = static_cast<uint16_t>(ops::CeilDiv(len, static_cast<int64_t>(VL_FP32)));
528+ uint32_t tailCount = static_cast<uint32_t>(len) - fullLoops * VL_FP32;
529+ __VEC_SCOPE__
530+ {
531+ AscendC::Reg::RegTensor<float> sumReg, predReg, logReg, oneReg, invReg;
532+ AscendC::Reg::MaskReg fullMask = AscendC::Reg::CreateMask<float, AscendC::Reg::MaskPattern::ALL>();
533+ uint32_t offset = 0;
534+ // oneReg为常量1向量，与循环无关，外提避免每轮重复广播
535+ AscendC::Reg::Duplicate(oneReg, 1.0f, fullMask);
536+ for (uint16_t i = 0; i < fullLoops; i++) {
537+ AscendC::Reg::DataCopy(sumReg, sumAddr + offset);
538+ AscendC::Reg::DataCopy(predReg, predAddr + offset);
539+ AscendC::Reg::Log<float, &castLogFt32False>(logReg, sumReg, fullMask);
540+ AscendC::Reg::Sub(logReg, logReg, predReg, fullMask);
541+ AscendC::Reg::DataCopy(lossAddr + offset, logReg, fullMask);
542+ AscendC::Reg::Div<float, &castDivIntrinsic>(invReg, oneReg, sumReg, fullMask);
543+ AscendC::Reg::DataCopy(invAddr + offset, invReg, fullMask);
544+ offset += VL_FP32;
545+ }
546+ for (uint16_t i = fullLoops; i < totalLoops; i++) {
547+ uint32_t tail = tailCount;
548+ AscendC::Reg::MaskReg tailMask = AscendC::Reg::UpdateMask<float>(tail);
549+ AscendC::Reg::DataCopy(sumReg, sumAddr + offset);
550+ AscendC::Reg::DataCopy(predReg, predAddr + offset);
551+ AscendC::Reg::Log<float, &castLogFt32False>(logReg, sumReg, tailMask);
552+ AscendC::Reg::Sub(logReg, logReg, predReg, tailMask);
553+ AscendC::Reg::DataCopy(lossAddr + offset, logReg, tailMask);
554+ AscendC::Reg::Div<float, &castDivIntrinsic>(invReg, oneReg, sumReg, tailMask);
555+ AscendC::Reg::DataCopy(invAddr + offset, invReg, tailMask);
556+ }
557+ }
558+}
559+ 
560+template <typename T, bool fullPath>
561+__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::CastVocabToFp32(
562+ const LocalTensor<T>& srcUb, const LocalTensor<float>& dstUb, int64_t len)
563+{
564+ __ubuf__ T* srcAddr = (__ubuf__ T*)srcUb.GetPhyAddr();
565+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstUb.GetPhyAddr();
566+ uint16_t fullLoops = static_cast<uint16_t>(len / static_cast<int64_t>(VL_FP32));
567+ uint16_t totalLoops = static_cast<uint16_t>(ops::CeilDiv(len, static_cast<int64_t>(VL_FP32)));
568+ uint32_t tailCount = static_cast<uint32_t>(len) - fullLoops * VL_FP32;
569+ __VEC_SCOPE__
570+ {
571+ AscendC::Reg::RegTensor<float> fp32Reg;
572+ AscendC::Reg::MaskReg fullMask = AscendC::Reg::CreateMask<float, AscendC::Reg::MaskPattern::ALL>();
573+ uint32_t offset = 0;
574+ for (uint16_t i = 0; i < fullLoops; i++) {
575+ LoadVocabToFp32(srcAddr, fp32Reg, fullMask, offset);
576+ AscendC::Reg::DataCopy(dstAddr + offset, fp32Reg, fullMask);
577+ offset += VL_FP32;
578+ }
579+ for (uint16_t i = fullLoops; i < totalLoops; i++) {
580+ uint32_t tail = tailCount;
581+ AscendC::Reg::MaskReg tailMask = AscendC::Reg::UpdateMask<float>(tail);
582+ LoadVocabToFp32(srcAddr, fp32Reg, tailMask, offset);
583+ AscendC::Reg::DataCopy(dstAddr + offset, fp32Reg, tailMask);
584+ }
585+ }
586+}
587+ 
588+template <typename T, bool fullPath>
589+__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ComputeSoftmaxFlat(
590+ const LocalTensor<T>& vocabUb, const LocalTensor<float>& maxUb, const LocalTensor<float>& invSumUb,
591+ const LocalTensor<float>& outUb, int64_t aRows)
592+{
593+ const int64_t v = tl_->vLen;
594+ int64_t g = VL_FP32 / v;
595+ while ((g * v) % (UB_BLOCK_SIZE / sizeof(float)) != 0) {
596+ g--;
597+ }
598+ if constexpr (IsSameType<T, float>::value) {
599+ SoftmaxFlatLoop<float>((__ubuf__ float*)vocabUb.GetPhyAddr(), maxUb, invSumUb, outUb, aRows, g);
600+ } else if ((g * v) % (UB_BLOCK_SIZE / sizeof(T)) == 0) {
601+ // b16组内偏移 r0*v*2B 满足32B对齐（V=5/7不满足）：加载时内联cast，
602+ // 省掉整tile预转fp32的一趟UB读写，数值与预转bit级一致（同一条Cast指令）
603+ SoftmaxFlatLoop<T>((__ubuf__ T*)vocabUb.GetPhyAddr(), maxUb, invSumUb, outUb, aRows, g);
604+ } else {
605+ LocalTensor<float> vocabFp32Ub = vocabFp32Buf_.Get<float>();
606+ CastVocabToFp32(vocabUb, vocabFp32Ub, aRows * v);
607+ SoftmaxFlatLoop<float>((__ubuf__ float*)vocabFp32Ub.GetPhyAddr(), maxUb, invSumUb, outUb, aRows, g);
608+ }
609+}
610+ 
611+template <typename T, bool fullPath>
612+template <typename S>
613+__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::SoftmaxFlatLoop(
614+ __ubuf__ S* src, const LocalTensor<float>& maxUb, const LocalTensor<float>& invSumUb,
615+ const LocalTensor<float>& outUb, int64_t aRows, int64_t g)
616+{
617+ __ubuf__ float* maxAddr = (__ubuf__ float*)maxUb.GetPhyAddr();
618+ __ubuf__ float* invAddr = (__ubuf__ float*)invSumUb.GetPhyAddr();
619+ __ubuf__ float* outAddr = (__ubuf__ float*)outUb.GetPhyAddr();
620+ const int64_t v = tl_->vLen;
621+ // g由调用方算好传入：每组G行打包成G*v条lane，组内偏移r0*v*4B须保持32B对齐（G*v为8的倍数），
622+ // 且G*v<=VL；G>=8保证maxAddr+r0（4B*r0）同样32B对齐
623+ uint16_t gU16 = static_cast<uint16_t>(g);
624+ uint16_t aRowsU16 = static_cast<uint16_t>(aRows);
625+ // VF内不写运行时if分支（见ComputeSoftmaxTile处注释）：V=1与V>1分两个独立的VEC域
626+ if (v == 1) {
627+ __VEC_SCOPE__
628+ {
629+ AscendC::Reg::RegTensor<float> maxReg, invReg, xReg, yReg;
630+ for (uint16_t r0 = 0; r0 < aRowsU16; r0 += gU16) {
631+ uint16_t rows = (aRowsU16 - r0) < gU16 ? (aRowsU16 - r0) : gU16;
632+ uint32_t lanes = static_cast<uint32_t>(rows);
633+ AscendC::Reg::MaskReg mask = AscendC::Reg::UpdateMask<float>(lanes);
634+ // V=1无需展开：max/inv与vocab逐元素对位，直接连续加载
635+ AscendC::Reg::DataCopy(maxReg, maxAddr + r0);
636+ AscendC::Reg::DataCopy(invReg, invAddr + r0);
637+ // 满载VL（UB tile按VL对齐预留，越界读在已分配区内），尾块由compute/store掩码收敛
638+ LoadVocabToFp32(src, xReg, mask, r0);
639+ AscendC::Reg::Sub(yReg, xReg, maxReg, mask);
640+ AscendC::Reg::Exp(yReg, yReg, mask);
641+ AscendC::Reg::Mul(yReg, yReg, invReg, mask);
642+ AscendC::Reg::DataCopy(outAddr + r0, yReg, mask);
643+ }
644+ }
645+ } else {
646+ __VEC_SCOPE__
647+ {
648+ AscendC::Reg::RegTensor<float> maxReg, invReg, xReg, yReg;
649+ AscendC::Reg::RegTensor<int32_t> idxReg;
650+ AscendC::Reg::MaskReg fullMask = AscendC::Reg::CreateMask<float, AscendC::Reg::MaskPattern::ALL>();
651+ // 索引在寄存器内构建：idx[l] = floor(l / v)。整型 Reg::Div 无截断语义保证（全仓无
652+ // 先例），按 softmax_focal_loss_grad 惯用法走 fp32 真除 + CAST_FLOOR：l < VL 且
653+ // v <= 8 时 l/v 与最近整数的距离 >= 1/8，远超 0.5ulp，floor 结果必精确
654+ AscendC::Reg::RegTensor<float> idxFReg;
655+ AscendC::Reg::RegTensor<float> vFReg;
656+ AscendC::Reg::Arange(idxReg, 0);
657+ AscendC::Reg::Cast<float, int32_t, castTraitS32ToFp32>(idxFReg, idxReg, fullMask);
658+ AscendC::Reg::Duplicate(vFReg, static_cast<float>(v), fullMask);
659+ AscendC::Reg::Div(idxFReg, idxFReg, vFReg, fullMask);
660+ AscendC::Reg::Cast<int32_t, float, castTraitFp32ToS32Floor>(idxReg, idxFReg, fullMask);
661+ for (uint16_t r0 = 0; r0 < aRowsU16; r0 += gU16) {
662+ uint16_t rows = (aRowsU16 - r0) < gU16 ? (aRowsU16 - r0) : gU16;
663+ uint32_t lanes = static_cast<uint32_t>(rows) * static_cast<uint32_t>(v);
664+ AscendC::Reg::MaskReg mask = AscendC::Reg::UpdateMask<float>(lanes);
665+ uint32_t off = static_cast<uint32_t>(r0) * static_cast<uint32_t>(v);
666+ // V>1用Gather按idx[l]=l/v把每行的max/inv展开到G*v条lane
667+ AscendC::Reg::Gather<float, float, uint32_t>(maxReg, maxAddr + r0,
668+ (AscendC::Reg::RegTensor<uint32_t>&)idxReg, mask);
669+ AscendC::Reg::Gather<float, float, uint32_t>(invReg, invAddr + r0,
670+ (AscendC::Reg::RegTensor<uint32_t>&)idxReg, mask);
671+ LoadVocabToFp32(src, xReg, mask, off);
J
Jjohan_lee2 days ago

【review】 问题分类：正确性 位置：loss/fused_cross_entropy_loss_with_max_sum/op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_ar.h:673，代码：LoadVocabToFp32(src, xReg, mask, off); 问题：FP32 且 V=3 时 g=16，组步长为 48 个 float。当前 aPerLoop 是 128 的倍数，满 tile 的 vocab 缓冲长度已经整除 64，按 VL 向上对齐不增加尾部空间。最后一组从 tileElems-48 开始，LoadVocabToFp32 仍用不带掩码的 DIST_NORM 读取 64 个 float。 可能导致的结果：最后一次向量搬入会越过当前 vocab 缓冲区 16 个 float，读到相邻 UB 区域。这里的 mask 只约束后续计算和写回，无法约束这次读。只覆盖 V=1 的新增 tiling 用例不会触发该边界。 建议的修改方法：按最后一组的实际起点为每个输入槽预留一个完整向量的可读范围，并把额外空间计入 Host UB 预算，或为不足 VL 的组使用受掩码约束的搬入。增加 V=3、每核至少一个满 tile 的边界用例。

likedislike
672+ AscendC::Reg::Sub(yReg, xReg, maxReg, mask);
673+ AscendC::Reg::Exp(yReg, yReg, mask);
674+ AscendC::Reg::Mul(yReg, yReg, invReg, mask);
675+ AscendC::Reg::DataCopy(outAddr + off, yReg, mask);
676+ }
677+ }
678+ }
679+}
680+ 
366template <typename T, bool fullPath>681template <typename T, bool fullPath>
367__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ProcessMemory()682__aicore__ inline void FusedCrossEntropyLossWithMaxSumRegBase<T, fullPath>::ProcessMemory()
368{683{
Mloss/fused_cross_entropy_loss_with_max_sum/op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_tiling_data.h+5-0
@@ -20,6 +20,10 @@
20constexpr int64_t FUSED_CE_MAX_SUM_A_PER_LOOP = 8;20constexpr int64_t FUSED_CE_MAX_SUM_A_PER_LOOP = 8;
21// v维分块列数的对齐粒度，需不小于一个向量寄存器可容纳的fp32元素数（64）21// v维分块列数的对齐粒度，需不小于一个向量寄存器可容纳的fp32元素数（64）
22constexpr int64_t FUSED_CE_MAX_SUM_V_ALIGN = 64;22constexpr int64_t FUSED_CE_MAX_SUM_V_ALIGN = 64;
23+// 小V快路上限：vLen不超过该值时走"多行打包+Gather展开max/inv"的行主序平坦通路。
24+// 大V通路按行广播（每行一次DIST_BRC_B32），小V时每行仅1~8个有效lane，广播与
25+// 2-D小搬运（blockLen=V*sizeof(T)）的固定开销占满流水，实测带宽退化两个数量级。
26+constexpr int64_t FUSED_CE_MAX_SUM_SMALL_V_MAX = 8;
23 27 
24// vocab_parallel_logits dtype id（kernel运行时分发用）28// vocab_parallel_logits dtype id（kernel运行时分发用）
25constexpr int64_t FUSED_CE_MAX_SUM_DTYPE_FP32 = 0;29constexpr int64_t FUSED_CE_MAX_SUM_DTYPE_FP32 = 0;
@@ -36,6 +40,7 @@ struct FusedCrossEntropyLossWithMaxSumRegBaseTilingData {
36 int64_t vocabDtypeId; // vocab_parallel_logits的dtype：0=fp32, 1=fp16, 2=bf1640 int64_t vocabDtypeId; // vocab_parallel_logits的dtype：0=fp32, 1=fp16, 2=bf16
37 int64_t vCores; // 每行在v维上切分的核数（1=不切分）；bt<coreNum且v>vPerLoop时>141 int64_t vCores; // 每行在v维上切分的核数（1=不切分）；bt<coreNum且v>vPerLoop时>1
38 int64_t vChunk; // v切分时每核负责的v长度（已按FUSED_CE_MAX_SUM_V_ALIGN对齐）；不切分时为vLen42 int64_t vChunk; // v切分时每核负责的v长度（已按FUSED_CE_MAX_SUM_V_ALIGN对齐）；不切分时为vLen
43+ int64_t aPerLoop; // 小V快路下每个UB分块处理的行数（大V通路恒为FUSED_CE_MAX_SUM_A_PER_LOOP）
39};44};
40 45 
41#endif // _FUSED_CROSS_ENTROPY_LOSS_WITH_MAX_SUM_REGBASE_TILING_DATA_H_46#endif // _FUSED_CROSS_ENTROPY_LOSS_WITH_MAX_SUM_REGBASE_TILING_DATA_H_
Mloss/fused_cross_entropy_loss_with_max_sum/tests/ut/op_host/test_fused_cross_entropy_loss_with_max_sum_arch3510_tiling.cpp+115-0
@@ -137,6 +137,121 @@ TEST_F(FusedCrossEntropyLossWithMaxSumArch3510Tiling, test_regbase_tiling_full_f
137 EXPECT_EQ(tilingData.vChunk, 4096);137 EXPECT_EQ(tilingData.vChunk, 4096);
138}138}
139 139 
140+TEST_F(FusedCrossEntropyLossWithMaxSumArch3510Tiling, test_regbase_tiling_full_fp16_small_v)
141+{
142+ fe::PlatFormInfos platFormInfo;
143+ map<string, string> socInfos;
144+ map<string, string> aicoreSpec;
145+ map<string, string> intrinsics;
146+ map<string, string> socVersionInfos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
147+ InitRegbasePlatForm(platFormInfo, socInfos, aicoreSpec, intrinsics, socVersionInfos);
148+ 
149+ FusedCeMaxSumUtCompileInfo compileInfo;
150+ gert::StorageShape input0 = {{1024}, {1024}};
151+ gert::StorageShape input1 = {{1024}, {1024}};
152+ gert::StorageShape input2 = {{1024}, {1024}};
153+ gert::StorageShape input3 = {{1024}, {1024}};
154+ gert::StorageShape input4 = {{1024}, {1024}};
155+ gert::StorageShape input5 = {{1024, 1}, {1024, 1}};
156+ gert::StorageShape output0 = {{1024}, {1024}};
157+ gert::StorageShape output1 = {{1024, 1}, {1024, 1}};
158+ 
159+ auto param = gert::TilingData::CreateCap(4096);
160+ ASSERT_NE(param, nullptr);
161+ auto workspaceHolder = gert::ContinuousVector::Create<size_t>(4096);
162+ auto wsSize = reinterpret_cast<gert::ContinuousVector*>(workspaceHolder.get());
163+ auto holder = gert::TilingContextFaker()
164+ .NodeIoNum(6, 2)
165+ .IrInstanceNum({1, 1, 1, 1, 1, 1})
166+ .InputShapes({&input0, &input1, &input2, &input3, &input4, &input5})
167+ .OutputShapes({&output0, &output1})
168+ .CompileInfo(&compileInfo)
169+ .PlatformInfo(reinterpret_cast<char*>(&platFormInfo))
170+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
171+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
172+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
173+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
174+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
175+ .NodeInputTd(5, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
176+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
177+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
178+ .TilingData(param.get())
179+ .Workspace(wsSize)
180+ .Build();
181+ 
182+ FusedCrossEntropyLossWithMaxSumRegBaseTilingData tilingData;
183+ // 小V快路(v=1<=8)：vPerLoop按常规公式取64；aPerLoop按平坦通路UB预算
184+ // 每行52B(vocab双缓冲4+softmax双缓冲8+fp32副本4+行向量9buffer36)
185+ // = FloorAlign((253952-2048-2*64*(2+4))/52, 128) = FloorAlign(251136/52, 128) = FloorAlign(4829.5, 128) = 4736
186+ // （2*64*(dtypeSize+4B)余量吸收vocab双缓冲+fp32副本各一个VL的防越界padding与尾块对齐）
187+ DoRegbaseTiling(holder, platFormInfo, socInfos, aicoreSpec, intrinsics, socVersionInfos, 0, 64, tilingData);
188+ EXPECT_EQ(tilingData.formerCoreNum, 0);
189+ EXPECT_EQ(tilingData.formerRows, 17);
190+ EXPECT_EQ(tilingData.latterRows, 16);
191+ EXPECT_EQ(tilingData.vPerLoop, 64);
192+ EXPECT_EQ(tilingData.vLen, 1);
193+ EXPECT_EQ(tilingData.vCores, 1);
194+ EXPECT_EQ(tilingData.vChunk, 1);
195+ EXPECT_EQ(tilingData.aPerLoop, 4736);
196+}
197+ 
198+TEST_F(FusedCrossEntropyLossWithMaxSumArch3510Tiling, test_regbase_tiling_full_fp32_small_v3)
199+{
200+ fe::PlatFormInfos platFormInfo;
201+ map<string, string> socInfos;
202+ map<string, string> aicoreSpec;
203+ map<string, string> intrinsics;
204+ map<string, string> socVersionInfos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
205+ InitRegbasePlatForm(platFormInfo, socInfos, aicoreSpec, intrinsics, socVersionInfos);
206+ 
207+ FusedCeMaxSumUtCompileInfo compileInfo;
208+ gert::StorageShape input0 = {{1024}, {1024}};
209+ gert::StorageShape input1 = {{1024}, {1024}};
210+ gert::StorageShape input2 = {{1024}, {1024}};
211+ gert::StorageShape input3 = {{1024}, {1024}};
212+ gert::StorageShape input4 = {{1024}, {1024}};
213+ gert::StorageShape input5 = {{1024, 3}, {1024, 3}};
214+ gert::StorageShape output0 = {{1024}, {1024}};
215+ gert::StorageShape output1 = {{1024, 3}, {1024, 3}};
216+ 
217+ auto param = gert::TilingData::CreateCap(4096);
218+ ASSERT_NE(param, nullptr);
219+ auto workspaceHolder = gert::ContinuousVector::Create<size_t>(4096);
220+ auto wsSize = reinterpret_cast<gert::ContinuousVector*>(workspaceHolder.get());
221+ auto holder = gert::TilingContextFaker()
222+ .NodeIoNum(6, 2)
223+ .IrInstanceNum({1, 1, 1, 1, 1, 1})
224+ .InputShapes({&input0, &input1, &input2, &input3, &input4, &input5})
225+ .OutputShapes({&output0, &output1})
226+ .CompileInfo(&compileInfo)
227+ .PlatformInfo(reinterpret_cast<char*>(&platFormInfo))
228+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
229+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
230+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
231+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
232+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
233+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
234+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
235+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
236+ .TilingData(param.get())
237+ .Workspace(wsSize)
238+ .Build();
239+ 
240+ FusedCrossEntropyLossWithMaxSumRegBaseTilingData tilingData;
241+ // V=3边界：G=16时G*v=48<VL=64，末组向量读越过tileElems，依赖vocab源缓冲的VL padding防越界
242+ // 每行96B(vocab双缓冲24+softmax双缓冲24+fp32副本12+行向量9buffer36)
243+ // = FloorAlign((253952-2048-2*64*(4+4))/96, 128) = FloorAlign(250112/96, 128) = FloorAlign(2605.3, 128) = 2560
244+ DoRegbaseTiling(holder, platFormInfo, socInfos, aicoreSpec, intrinsics, socVersionInfos, 0, 64, tilingData);
245+ EXPECT_EQ(tilingData.formerCoreNum, 0);
246+ EXPECT_EQ(tilingData.formerRows, 17);
247+ EXPECT_EQ(tilingData.latterRows, 16);
248+ EXPECT_EQ(tilingData.vPerLoop, 64);
249+ EXPECT_EQ(tilingData.vLen, 3);
250+ EXPECT_EQ(tilingData.vCores, 1);
251+ EXPECT_EQ(tilingData.vChunk, 3);
252+ EXPECT_EQ(tilingData.aPerLoop, 2560);
253+}
254+ 
140TEST_F(FusedCrossEntropyLossWithMaxSumArch3510Tiling, test_regbase_tiling_full_bf16_small_bt)255TEST_F(FusedCrossEntropyLossWithMaxSumArch3510Tiling, test_regbase_tiling_full_bf16_small_bt)
141{256{
142 fe::PlatFormInfos platFormInfo;257 fe::PlatFormInfos platFormInfo;