Star1726
1643
Code
Code
Issues480
Pull Requests539
Actions
Discussion
Wiki
Analysis
Star1726
1643
Merged
TopkV2算子代码优化 #6015
cy_hwAdd on Sep 28
TopkV2算子代码优化 #6015
Merged
cy_hwAdd on Sep 28
Merged fromcaoyan_huawei/ops-mathcy_9.2.0_shangfen_clearCode_928intocann/ops-math9.2.0
Total: 10 files changed+902-850
Mmath/sort_with_index/op_host/arch35/sort_with_index_tiling.cpp+2-2
@@ -57,9 +57,9 @@ struct SortTileInfo {
57 uint32_t dtypeSize = 0;57 uint32_t dtypeSize = 0;
58 uint32_t y2DtypeSize = 0;58 uint32_t y2DtypeSize = 0;
59 uint32_t maxCoreNum = 0;59 uint32_t maxCoreNum = 0;
60+ uint32_t lastDimNeedCore = 0;
BruceLP
BruceLPBruceLPSep 28

关联新issue

likedislike
60 uint32_t numTileDataSize = 0;61 uint32_t numTileDataSize = 0;
61 uint32_t sortLoopTimes = 0;62 uint32_t sortLoopTimes = 0;
62- uint32_t lastDimNeedCore = 0;
63 uint32_t keyParams0 = 0;63 uint32_t keyParams0 = 0;
64 uint32_t keyParams1 = 0;64 uint32_t keyParams1 = 0;
65 uint32_t keyParams2 = 0;65 uint32_t keyParams2 = 0;
@@ -69,8 +69,8 @@ struct SortTileInfo {
69 uint32_t tmpUbSize = 0;69 uint32_t tmpUbSize = 0;
70 bool isDescend = false;70 bool isDescend = false;
71 ge::DataType dataType = ge::DT_UINT8;71 ge::DataType dataType = ge::DT_UINT8;
72- uint32_t isInt32 = 0;
73 int32_t xDimNum = 0;72 int32_t xDimNum = 0;
73+ uint32_t isInt32 = 0;
74 int64_t sortAxisNum = 1;74 int64_t sortAxisNum = 1;
75 int64_t unSortDimNum = 1;75 int64_t unSortDimNum = 1;
76};76};
Mmath/top_k_v2/op_host/arch35/sort_with_index_tiling.h+196-196
@@ -27,15 +27,11 @@ namespace sortWithIndex {
27constexpr size_t WORK_SPACE_SIZE = 16777216; // 16 * 1024 * 1024;27constexpr size_t WORK_SPACE_SIZE = 16777216; // 16 * 1024 * 1024;
28const uint32_t BIN_NUM = 256;28const uint32_t BIN_NUM = 256;
29const uint32_t TILE_DATA_NUM = 4096;29const uint32_t TILE_DATA_NUM = 4096;
30+const uint32_t UB_CONST_INT32 = 4096; // 输出idx为int32时kernel侧需要的固定ub大小
31+const uint32_t UB_CONST_INT64 = 7168; // 输出idx为int64时kernel侧需要的固定ub大小
30const uint32_t MEDIUM_TILE_DATA_NUM = 2048;32const uint32_t MEDIUM_TILE_DATA_NUM = 2048;
31const uint32_t SMALL_TILE_DATA_NUM = 1024;33const uint32_t SMALL_TILE_DATA_NUM = 1024;
32const uint32_t TILE_DATA_NUM_B64 = 2048;34const uint32_t TILE_DATA_NUM_B64 = 2048;
33-const uint32_t TMP_UB = 1024; // 暂预留给sort高级api的大小
34-const uint32_t CONST_10 = 10; // int32索引时, 计算各种tensor的乘法因子
35-const uint32_t CONST_14 = 14; // int64索引时, 计算各种tensor的乘法因子
36-const uint32_t CONST_6 = 6; // int64索引时, 计算各种tensor的乘法因子
37-const uint32_t CONST_1 = 1;
38-const uint32_t CONST_2 = 2;
39const uint32_t INT64_BYTE = 8;35const uint32_t INT64_BYTE = 8;
40const uint32_t INT32_BYTE = 4;36const uint32_t INT32_BYTE = 4;
41const uint32_t SIMT_UB = 32768; // 预留了32k给simt使用37const uint32_t SIMT_UB = 32768; // 预留了32k给simt使用
Mmath/top_k_v2/op_host/arch35/top_k_v2_tiling_arch35.cpp+1-1
@@ -1389,7 +1389,7 @@ static void InitTopKSmallAxisTileInfo(TopKSmallAxisRouteInfo& info,
1389 */1389 */
1390static bool SelectTopKSmallAxisRoute(const TopKSmallAxisRouteInfo& info, SmallAxisRoutePlan& plan, bool& isTwoStage)1390static bool SelectTopKSmallAxisRoute(const TopKSmallAxisRouteInfo& info, SmallAxisRoutePlan& plan, bool& isTwoStage)
1391{1391{
1392- bool selected = info.isNonLastAxis ? SelectNonLastSmallAxisRoute(info, plan) : SelectSmallAxisRoute(info, plan);1392+ bool selected = info.isNonLastAxis ? PickTopkNonLastSmallAxisRoute(info, plan) : PickTopkSmallAxisRoute(info, plan);
1393 if (!selected || (plan.kind != SmallAxisRouteKind::INSERTION && plan.kind != SmallAxisRouteKind::TWO_STAGE)) {1393 if (!selected || (plan.kind != SmallAxisRouteKind::INSERTION && plan.kind != SmallAxisRouteKind::TWO_STAGE)) {
1394 return false;1394 return false;
1395 }1395 }
Mmath/top_k_v2/op_host/arch35/top_k_v2_tiling_base.h+5-5
@@ -130,14 +130,14 @@ struct SortTileInfo {
130 uint32_t maxCoreNum = 0;130 uint32_t maxCoreNum = 0;
131 uint32_t numTileDataSize = 0;131 uint32_t numTileDataSize = 0;
132 uint64_t sortLoopTimes = 0;132 uint64_t sortLoopTimes = 0;
133- uint32_t lastDimNeedCore = 0;133+ uint32_t tmpUbSize = 0;
134 uint32_t keyParams0 = 0;134 uint32_t keyParams0 = 0;
135 uint32_t keyParams1 = 0;135 uint32_t keyParams1 = 0;
136 uint32_t keyParams2 = 0;136 uint32_t keyParams2 = 0;
137 uint32_t keyParams3 = 0;137 uint32_t keyParams3 = 0;
138 uint32_t keyParams4 = 0;138 uint32_t keyParams4 = 0;
139 uint32_t keyParams5 = 0;139 uint32_t keyParams5 = 0;
140- uint32_t tmpUbSize = 0;140+ uint32_t lastDimNeedCore = 0;
141 bool isDescend = false;141 bool isDescend = false;
142 ge::DataType dataType = ge::DT_UINT8;142 ge::DataType dataType = ge::DT_UINT8;
143 uint32_t isInt32 = 0;143 uint32_t isInt32 = 0;
@@ -209,8 +209,8 @@ struct TopKSmallAxisRouteInfo {
209 int64_t innerSize = 1;209 int64_t innerSize = 1;
210};210};
211 211 
212-bool SelectSmallAxisRoute(const TopKSmallAxisRouteInfo& info, SmallAxisRoutePlan& plan);212+bool PickTopkSmallAxisRoute(const TopKSmallAxisRouteInfo& info, SmallAxisRoutePlan& routePlan);
213-bool SelectNonLastSmallAxisRoute(const TopKSmallAxisRouteInfo& info, SmallAxisRoutePlan& plan);213+bool PickTopkNonLastSmallAxisRoute(const TopKSmallAxisRouteInfo& info, SmallAxisRoutePlan& routePlan);
214 214 
215struct TopkNonLastSmallAxisTileInfo {215struct TopkNonLastSmallAxisTileInfo {
216 int64_t rank = 0;216 int64_t rank = 0;
@@ -253,7 +253,7 @@ uint32_t GetSingleCoreModelDefaultTileDataSize(ge::DataType dataType);
253 253 
254// ==================== Common Align Helpers ====================254// ==================== Common Align Helpers ====================
255 255 
256-bool CeilAlignUint32(uint64_t rawSize, uint32_t alignSize, uint32_t& alignedSize);256+bool TopkCeilAlignUint32(uint64_t sizeToAlign, uint32_t alignment, uint32_t& alignedOut);
257 257 
258// ==================== FP32 MergeSort Helpers ====================258// ==================== FP32 MergeSort Helpers ====================
259 259 
Mmath/top_k_v2/op_kernel/arch35/top_k_merge_sort.h+20-19
@@ -183,8 +183,8 @@ __aicore__ inline void MergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST,
183 outValueQueue_.EnQue<T>(sortedValueLocal);183 outValueQueue_.EnQue<T>(sortedValueLocal);
184 inQueueX_.FreeTensor(xLocal);184 inQueueX_.FreeTensor(xLocal);
185 // copy result out185 // copy result out
186- uint64_t gmOffset = sortLoopRound_ * unsortedDimParallel_ * outputLastDimValue_ * oneCoreRowNum_;
187 uint64_t answerTileOffset = tileId * outputLastDimValue_ * oneCoreRowNum_;186 uint64_t answerTileOffset = tileId * outputLastDimValue_ * oneCoreRowNum_;
187+ uint64_t gmOffset = sortLoopRound_ * unsortedDimParallel_ * outputLastDimValue_ * oneCoreRowNum_;
188 188 
189 CopyValue2Gm(gmOffset, answerTileOffset, outputLastDimValue_, nowCoreRealRowNum);189 CopyValue2Gm(gmOffset, answerTileOffset, outputLastDimValue_, nowCoreRealRowNum);
190}190}
@@ -257,33 +257,34 @@ __aicore__ inline void MergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST,
257 IS_BITONIC_SORT>::CopyValue2Gm(uint64_t gmOffset, uint64_t tileOffset,257 IS_BITONIC_SORT>::CopyValue2Gm(uint64_t gmOffset, uint64_t tileOffset,
258 uint32_t outputLastDimValue, uint32_t oneCoreRowNum)258 uint32_t outputLastDimValue, uint32_t oneCoreRowNum)
259{259{
260- uint32_t aglinOneRowTileSize = ROUND_UP_AGLIN(numTileData_);260+ uint32_t alignedRowTileSize = ROUND_UP_AGLIN(numTileData_);
261 // value stride261 // value stride
262- uint32_t currTileSizeAlign = ROUND_UP_AGLIN(outputLastDimValue * sizeof(T)) / sizeof(T);262+ uint32_t alignedValueRowSize = ROUND_UP_AGLIN(outputLastDimValue * sizeof(T)) / sizeof(T);
263- uint32_t ubStrideValue = ((aglinOneRowTileSize - currTileSizeAlign) * sizeof(T)) / UB_AGLIN_VALUE;263+ uint32_t ubValueRowGapStride = ((alignedRowTileSize - alignedValueRowSize) * sizeof(T)) / UB_AGLIN_VALUE;
264 // index stride264 // index stride
265- uint32_t currSrcTileIndexSizeAlign = ROUND_UP_AGLIN(outputLastDimValue * sizeof(int32_t)) / sizeof(int32_t);265+ uint32_t alignedInt32IndexRowSize = ROUND_UP_AGLIN(outputLastDimValue * sizeof(int32_t)) / sizeof(int32_t);
266- uint32_t ubSrcStrideIndex = ((aglinOneRowTileSize - currSrcTileIndexSizeAlign) * sizeof(int32_t)) / UB_AGLIN_VALUE;266+ uint32_t ubInt32IndexRowGapStride = ((alignedRowTileSize - alignedInt32IndexRowSize) * sizeof(int32_t)) /
267+ UB_AGLIN_VALUE;
267 268 
268- uint32_t currTileIndexSizeAlign = ROUND_UP_AGLIN(outputLastDimValue * sizeof(INDEX_TYPE)) / sizeof(INDEX_TYPE);269+ uint32_t alignedIndexRowSize = ROUND_UP_AGLIN(outputLastDimValue * sizeof(INDEX_TYPE)) / sizeof(INDEX_TYPE);
269- uint32_t ubStrideIndex = ((aglinOneRowTileSize - currTileIndexSizeAlign) * sizeof(INDEX_TYPE)) / UB_AGLIN_VALUE;270+ uint32_t ubIndexRowGapStride = ((alignedRowTileSize - alignedIndexRowSize) * sizeof(INDEX_TYPE)) / UB_AGLIN_VALUE;
270 // copy result out271 // copy result out
271- AscendC::LocalTensor<T> outValueLocal = outValueQueue_.DeQue<T>();272+ AscendC::LocalTensor<T> sortedValueLocal = outValueQueue_.DeQue<T>();
272- AscendC::LocalTensor<INDEX_TYPE> outIndexLocal = outIndexQueue_.DeQue<INDEX_TYPE>();273+ AscendC::LocalTensor<INDEX_TYPE> sortedIndexLocal = outIndexQueue_.DeQue<INDEX_TYPE>();
273- AscendC::DataCopyExtParams dataCopyParamValue{static_cast<uint16_t>(oneCoreRowNum),274+ AscendC::DataCopyExtParams valueDataCopyParam{static_cast<uint16_t>(oneCoreRowNum),
274- static_cast<uint32_t>(outputLastDimValue * sizeof(T)), ubStrideValue,275+ static_cast<uint32_t>(outputLastDimValue * sizeof(T)),
275- 0, 0};276+ ubValueRowGapStride, 0, 0};
276- AscendC::DataCopyPad(outValueGm_[gmOffset + tileOffset], outValueLocal, dataCopyParamValue);277+ AscendC::DataCopyPad(outValueGm_[gmOffset + tileOffset], sortedValueLocal, valueDataCopyParam);
277 278 
278- AscendC::DataCopyExtParams dataCopyParamIndex{279+ AscendC::DataCopyExtParams indexDataCopyParam{
279 static_cast<uint16_t>(oneCoreRowNum), // 连续数据块的个数280 static_cast<uint16_t>(oneCoreRowNum), // 连续数据块的个数
280 static_cast<uint32_t>(outputLastDimValue * sizeof(INDEX_TYPE)), // 每个连续传输数据块的长度，长度为Byte281 static_cast<uint32_t>(outputLastDimValue * sizeof(INDEX_TYPE)), // 每个连续传输数据块的长度，长度为Byte
281- ubStrideIndex, // 源操作数，相邻连续数据块的间隔282+ ubIndexRowGapStride, // 源操作数，相邻连续数据块的间隔
282 0, // 目的操作数，相邻连续数据块的间隔283 0, // 目的操作数，相邻连续数据块的间隔
283 0};284 0};
284- AscendC::DataCopyPad(outIndexGm_[gmOffset + tileOffset], outIndexLocal, dataCopyParamIndex);285+ AscendC::DataCopyPad(outIndexGm_[gmOffset + tileOffset], sortedIndexLocal, indexDataCopyParam);
285- outIndexQueue_.FreeTensor(outIndexLocal);286+ outIndexQueue_.FreeTensor(sortedIndexLocal);
286- outValueQueue_.FreeTensor(outValueLocal);287+ outValueQueue_.FreeTensor(sortedValueLocal);
287}288}
288} // namespace topkV2289} // namespace topkV2
289#endif290#endif
Mmath/top_k_v2/op_kernel/arch35/top_k_merge_sort_intra_core.h+74-70
@@ -45,10 +45,10 @@ private:
45 45 
46 __aicore__ inline uint32_t MergeSingleBatch();46 __aicore__ inline uint32_t MergeSingleBatch();
47 __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, typename Base::MergeListContext& ctx);47 __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, typename Base::MergeListContext& ctx);
48- __aicore__ inline void MergeOneGroup(uint32_t groupStart, uint32_t groupBlockCount, uint32_t fullBlockElemCount,48+ __aicore__ inline void MergeOneGroup(uint32_t groupStartIdx, uint32_t groupSegCount, uint32_t fullSegElemNum,
49- uint32_t fullBlockSortLen, uint32_t lastBlockElemCount, uint32_t numBlocks,49+ uint32_t fullSegSortLen, uint32_t tailSegElemNum, uint32_t pendingSegNum,
50- uint32_t pingPongFlag, uint32_t& cumulativeOffset,50+ uint32_t pingPongRegionSel, uint32_t& passDstOffsetAcc,
51- uint32_t& mergedGroupElemCount);51+ uint32_t& groupMergedElemNum);
52 __aicore__ inline void CopyRemainingList(typename Base::MergeListContext& ctx, int64_t dstOffsetBase,52 __aicore__ inline void CopyRemainingList(typename Base::MergeListContext& ctx, int64_t dstOffsetBase,
53 uint32_t& dstCumulativeOffset, uint32_t& dstElemCount);53 uint32_t& dstCumulativeOffset, uint32_t& dstElemCount);
54 __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion);54 __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion);
@@ -141,100 +141,104 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::
141template <typename ValueType, typename IndexType, bool IsDescend>141template <typename ValueType, typename IndexType, bool IsDescend>
142__aicore__ inline uint32_t TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeSingleBatch()142__aicore__ inline uint32_t TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeSingleBatch()
143{143{
144- uint32_t numBlocks = this->blocksPerRow_;144+ uint32_t pendingSegNum = this->blocksPerRow_;
145- uint32_t fullBlockElemCount = this->blockSortSize_;145+ uint32_t fullSegElemNum = this->blockSortSize_;
146- uint32_t fullBlockSortLen = this->blockSortLen_;146+ uint32_t tailSegElemNum = this->lastBlockSize_;
147- uint32_t lastBlockElemCount = this->lastBlockSize_;147+ uint32_t fullSegSortLen = this->blockSortLen_;
148 148 
149- uint32_t pingPongFlag = 0;149+ uint32_t pingPongRegionSel = 0;
150- uint32_t mergeRounds = 0;150+ uint32_t mergePassNum = 0;
151 151 
152- while (numBlocks > 1) {152+ while (pendingSegNum > 1) {
153- uint32_t cumulativeOffset = 0;153+ uint32_t passDstOffsetAcc = 0;
154- uint32_t newNumBlocks = 0;154+ uint32_t nextSegNum = 0;
155- uint32_t newFullBlockElemCount = 0;155+ uint32_t nextTailSegElemNum = 0;
156- uint32_t newLastBlockElemCount = 0;156+ uint32_t nextFullSegElemNum = 0;
157 157 
158- for (uint32_t i = 0; i < numBlocks; i += Sort::MERGE_LIST_MAX_NUM) {158+ for (uint32_t groupStartIdx = 0; groupStartIdx < pendingSegNum; groupStartIdx += Sort::MERGE_LIST_MAX_NUM) {
159- uint32_t groupBlockCount = (i + Sort::MERGE_LIST_MAX_NUM <= numBlocks) ? Sort::MERGE_LIST_MAX_NUM :159+ uint32_t groupSegCount = (groupStartIdx + Sort::MERGE_LIST_MAX_NUM <= pendingSegNum) ?
160- (numBlocks - i);160+ Sort::MERGE_LIST_MAX_NUM :
161+ (pendingSegNum - groupStartIdx);
161 162 
162- uint32_t mergedGroupElemCount = 0;163+ uint32_t groupMergedElemNum = 0;
163- MergeOneGroup(i, groupBlockCount, fullBlockElemCount, fullBlockSortLen, lastBlockElemCount, numBlocks,164+ MergeOneGroup(groupStartIdx, groupSegCount, fullSegElemNum, fullSegSortLen, tailSegElemNum, pendingSegNum,
164- pingPongFlag, cumulativeOffset, mergedGroupElemCount);165+ pingPongRegionSel, passDstOffsetAcc, groupMergedElemNum);
165- if (newNumBlocks == 0) {166+ if (nextSegNum == 0) {
166- newFullBlockElemCount = mergedGroupElemCount;167+ nextFullSegElemNum = groupMergedElemNum;
167 }168 }
168- newLastBlockElemCount = mergedGroupElemCount;169+ nextTailSegElemNum = groupMergedElemNum;
169- newNumBlocks++;170+ nextSegNum++;
170 }171 }
171 172 
172- if (newNumBlocks == 0 || newNumBlocks >= numBlocks) {173+ if (nextSegNum == 0 || nextSegNum >= pendingSegNum) {
173 break;174 break;
174 }175 }
175 176 
176- numBlocks = newNumBlocks;177+ pendingSegNum = nextSegNum;
177- if (numBlocks == 1) {178+ if (pendingSegNum == 1) {
178- fullBlockElemCount = newLastBlockElemCount;179+ fullSegElemNum = nextTailSegElemNum;
179 } else {180 } else {
180- fullBlockElemCount = newFullBlockElemCount;181+ fullSegElemNum = nextFullSegElemNum;
181 }182 }
182- fullBlockSortLen = AscendC::GetSortLen<ValueType>(fullBlockElemCount);183+ fullSegSortLen = AscendC::GetSortLen<ValueType>(fullSegElemNum);
183- lastBlockElemCount = newLastBlockElemCount;184+ tailSegElemNum = nextTailSegElemNum;
184- pingPongFlag = 1 - pingPongFlag;185+ pingPongRegionSel = 1 - pingPongRegionSel;
185- mergeRounds++;186+ mergePassNum++;
186- event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_MTE2));187+ event_t barrierEvtId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_MTE2));
187- SetFlag<HardEvent::MTE3_MTE2>(eventId);188+ SetFlag<HardEvent::MTE3_MTE2>(barrierEvtId);
188- WaitFlag<HardEvent::MTE3_MTE2>(eventId);189+ WaitFlag<HardEvent::MTE3_MTE2>(barrierEvtId);
189 }190 }
190 191 
191- return (mergeRounds == 0) ? 0 : pingPongFlag;192+ return (mergePassNum == 0) ? 0 : pingPongRegionSel;
192}193}
193 194 
194template <typename ValueType, typename IndexType, bool IsDescend>195template <typename ValueType, typename IndexType, bool IsDescend>
195__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeOneGroup(196__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeOneGroup(
196- uint32_t groupStart, uint32_t groupBlockCount, uint32_t fullBlockElemCount, uint32_t fullBlockSortLen,197+ uint32_t groupStartIdx, uint32_t groupSegCount, uint32_t fullSegElemNum, uint32_t fullSegSortLen,
197- uint32_t lastBlockElemCount, uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset,198+ uint32_t tailSegElemNum, uint32_t pendingSegNum, uint32_t pingPongRegionSel, uint32_t& passDstOffsetAcc,
198- uint32_t& mergedGroupElemCount)199+ uint32_t& groupMergedElemNum)
199{200{
200- typename Base::MergeListContext ctx;201+ typename Base::MergeListContext segListCtx;
201- ctx.listCount = groupBlockCount;202+ segListCtx.listCount = groupSegCount;
202 203 
203- int64_t srcRegionOffset = (pingPongFlag == 0) ? 0 : this->batchSortLen_;204+ int64_t srcRegionBase = (pingPongRegionSel == 0) ? 0 : this->batchSortLen_;
204- int64_t dstRegionOffset = (pingPongFlag == 0) ? this->batchSortLen_ : 0;205+ int64_t dstRegionBase = (pingPongRegionSel == 0) ? this->batchSortLen_ : 0;
205 206 
206- for (uint32_t j = 0; j < groupBlockCount; j++) {207+ for (uint32_t segIdxInGroup = 0; segIdxInGroup < groupSegCount; segIdxInGroup++) {
207- uint32_t blockIdx = groupStart + j;208+ uint32_t pendingSegIdx = groupStartIdx + segIdxInGroup;
208- ctx.srcOffsets[j] = srcRegionOffset + blockIdx * fullBlockSortLen;209+ segListCtx.srcOffsets[segIdxInGroup] = srcRegionBase + pendingSegIdx * fullSegSortLen;
209- ctx.elemCounts[j] = (blockIdx < numBlocks - 1) ? fullBlockElemCount : lastBlockElemCount;210+ segListCtx.elemCounts[segIdxInGroup] = (pendingSegIdx < pendingSegNum - 1) ? fullSegElemNum : tailSegElemNum;
210 }211 }
211 212 
212- if (groupBlockCount > 1) {213+ if (groupSegCount > 1) {
213- DoIncrementalMerge(dstRegionOffset + cumulativeOffset, ctx);214+ DoIncrementalMerge(dstRegionBase + passDstOffsetAcc, segListCtx);
214 215 
215- uint32_t totalElem = 0;216+ uint32_t groupTotalElemNum = 0;
216- for (uint32_t j = 0; j < groupBlockCount; j++) {217+ for (uint32_t segIdxInGroup = 0; segIdxInGroup < groupSegCount; segIdxInGroup++) {
217- totalElem += ctx.elemCounts[j];218+ groupTotalElemNum += segListCtx.elemCounts[segIdxInGroup];
218 }219 }
219- mergedGroupElemCount = (totalElem > topKValue_) ? topKValue_ : totalElem; // TopK limitation220+ groupMergedElemNum = (groupTotalElemNum > topKValue_) ? topKValue_ : groupTotalElemNum; // TopK limitation
220- cumulativeOffset += AscendC::GetSortLen<ValueType>(mergedGroupElemCount);221+ passDstOffsetAcc += AscendC::GetSortLen<ValueType>(groupMergedElemNum);
221- } else if (groupBlockCount == 1) {222+ } else if (groupSegCount == 1) {
222- mergedGroupElemCount = (ctx.elemCounts[0] > topKValue_) ? topKValue_ : ctx.elemCounts[0]; // TopK limitation223+ groupMergedElemNum = (segListCtx.elemCounts[0] > topKValue_) ? topKValue_ :
223- if (ctx.elemCounts[0] > 0) {224+ segListCtx.elemCounts[0]; // TopK limitation
224- int64_t srcOffset = ctx.srcOffsets[0];225+ if (segListCtx.elemCounts[0] > 0) {
225- uint32_t remainElems = mergedGroupElemCount;226+ int64_t singleSegGmOffset = segListCtx.srcOffsets[0];
226- uint32_t srcChunkOffset = 0;227+ uint32_t remainCopyElemNum = groupMergedElemNum;
227- while (remainElems > 0) {228+ uint32_t srcChunkElemOffset = 0;
228- uint32_t chunkSize = (remainElems > this->blockSortSize_) ? this->blockSortSize_ : remainElems;229+ while (remainCopyElemNum > 0) {
229- if (chunkSize == 0)230+ uint32_t copyChunkElemNum = (remainCopyElemNum > this->blockSortSize_) ? this->blockSortSize_ :
231+ remainCopyElemNum;
232+ if (copyChunkElemNum == 0)
230 break;233 break;
231 this->CopyBlockChunk(234 this->CopyBlockChunk(
232- srcOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset),235+ singleSegGmOffset + AscendC::GetSortLen<ValueType>(srcChunkElemOffset),
233- dstRegionOffset + cumulativeOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), chunkSize);236+ dstRegionBase + passDstOffsetAcc + AscendC::GetSortLen<ValueType>(srcChunkElemOffset),
234- remainElems -= chunkSize;237+ copyChunkElemNum);
235- srcChunkOffset += chunkSize;238+ remainCopyElemNum -= copyChunkElemNum;
239+ srcChunkElemOffset += copyChunkElemNum;
236 }240 }
237- cumulativeOffset += AscendC::GetSortLen<ValueType>(mergedGroupElemCount);241+ passDstOffsetAcc += AscendC::GetSortLen<ValueType>(groupMergedElemNum);
238 }242 }
239 }243 }
240}244}
Mmath/top_k_v2/op_kernel/arch35/top_k_merge_sort_more_core.h+70-54
@@ -35,14 +35,21 @@ constexpr int16_t MORE_CORE_XOR_OP_VALUE_HALF = 0x8000;
35template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>35template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
36struct TopKMergeSortMoreCore {36struct TopKMergeSortMoreCore {
37 __aicore__ inline TopKMergeSortMoreCore() {}37 __aicore__ inline TopKMergeSortMoreCore() {}
38- __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe);38+ __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace,
39+ const TopKV2TilingDataSimd* tilingData, TPipe* pipe);
39 __aicore__ inline void Process();40 __aicore__ inline void Process();
40 __aicore__ inline void SortInSingleCore(uint32_t tileNum, int64_t offsetPerCore);41 __aicore__ inline void SortInSingleCore(uint32_t tileNum, int64_t offsetPerCore);
41 __aicore__ inline void CopyInData(uint32_t tileNum, int64_t offsetPerCore);42 __aicore__ inline void CopyInData(uint32_t tileNum, int64_t offsetPerCore);
42- __aicore__ inline void InitIndexLocal(uint32_t tileNum, LocalTensor<uint32_t> sortedValueIndexLocal, int64_t offsetPerCore);43+ __aicore__ inline void InitIndexLocal(uint32_t tileNum, LocalTensor<uint32_t> sortedValueIndexLocal,
43- __aicore__ inline void DoSort(uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal);44+ int64_t offsetPerCore);
44- __aicore__ inline void DoSortBf16(uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal);45+ __aicore__ inline void DoSort(uint32_t tileNum, LocalTensor<T> inputLocal,
45- __aicore__ inline void CopyOutWorkSpace(uint32_t tileNum, int64_t offsetPerCore, LocalTensor<CONVERT_TYPE> sortedValueLocal);46+ LocalTensor<CONVERT_TYPE> sortedValueLocal,
47+ LocalTensor<uint32_t> sortedValueIndexLocal);
48+ __aicore__ inline void DoSortBf16(uint32_t tileNum, LocalTensor<T> inputLocal,
49+ LocalTensor<CONVERT_TYPE> sortedValueLocal,
50+ LocalTensor<uint32_t> sortedValueIndexLocal);
51+ __aicore__ inline void CopyOutWorkSpace(uint32_t tileNum, int64_t offsetPerCore,
52+ LocalTensor<CONVERT_TYPE> sortedValueLocal);
46 __aicore__ inline void CopyOutMultiCore();53 __aicore__ inline void CopyOutMultiCore();
47 __aicore__ inline void SortInMultiCore();54 __aicore__ inline void SortInMultiCore();
48 __aicore__ inline void FinalSortAndCopy();55 __aicore__ inline void FinalSortAndCopy();
@@ -56,6 +63,7 @@ struct TopKMergeSortMoreCore {
56 __aicore__ inline void ExtractAndCopyOut();63 __aicore__ inline void ExtractAndCopyOut();
57 __aicore__ inline void ClearCache();64 __aicore__ inline void ClearCache();
58 __aicore__ inline uint32_t GetRemainingTopK() const;65 __aicore__ inline uint32_t GetRemainingTopK() const;
66+ 
59public:67public:
60 TQue<QuePosition::VECIN, MORE_CORE_DOUBLE_BUFFER> inputQueue_;68 TQue<QuePosition::VECIN, MORE_CORE_DOUBLE_BUFFER> inputQueue_;
61 TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> outValueQueue_;69 TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> outValueQueue_;
@@ -82,7 +90,7 @@ public:
82 TBuf<QuePosition::VECCALC> inputValueTempBuf_;90 TBuf<QuePosition::VECCALC> inputValueTempBuf_;
83 TBuf<QuePosition::VECCALC> sortedValueLocalCastTbuf_;91 TBuf<QuePosition::VECCALC> sortedValueLocalCastTbuf_;
84 92 
85- TPipe *pipe_;93+ TPipe* pipe_;
86 uint32_t blockIdx_ = 0;94 uint32_t blockIdx_ = 0;
87 uint32_t numTileData_ = 0;95 uint32_t numTileData_ = 0;
88 uint32_t sortLoopRound_ = 0;96 uint32_t sortLoopRound_ = 0;
@@ -103,11 +111,11 @@ public:
103 int64_t remainListNum_{0};111 int64_t remainListNum_{0};
104 int64_t outOffset_{0};112 int64_t outOffset_{0};
105 int64_t offsets_[4] = {0};113 int64_t offsets_[4] = {0};
106- int64_t listRemainElements_[4] = {0};
107- int64_t currentElements_{0};
108 int64_t currentTailElements_{0};114 int64_t currentTailElements_{0};
109 int64_t dealLengths_[4] = {0};115 int64_t dealLengths_[4] = {0};
110 int64_t allRemainElements_{0};116 int64_t allRemainElements_{0};
117+ int64_t listRemainElements_[4] = {0};
118+ int64_t currentElements_{0};
111 int64_t curLoopSortedNum_{0};119 int64_t curLoopSortedNum_{0};
112 int64_t onceMaxElements_{0};120 int64_t onceMaxElements_{0};
113 uint16_t elementCountList_[4] = {0};121 uint16_t elementCountList_[4] = {0};
@@ -127,8 +135,9 @@ __aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
127}135}
128 136 
129template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>137template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
130-__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::Init(GM_ADDR inputValue, 138+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::Init(
131- GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe)139+ GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const TopKV2TilingDataSimd* tilingData,
140+ TPipe* pipe)
132{141{
133 blockIdx_ = GetBlockIdx();142 blockIdx_ = GetBlockIdx();
134 pipe_ = pipe;143 pipe_ = pipe;
@@ -143,20 +152,17 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
143 rowOutputOffset_ = static_cast<int64_t>(rowIdx_) * static_cast<int64_t>(outputTopKValue_);152 rowOutputOffset_ = static_cast<int64_t>(rowIdx_) * static_cast<int64_t>(outputTopKValue_);
144 // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per153 // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per
145 // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.154 // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.
146- rowWorkspaceOffset_ =155+ rowWorkspaceOffset_ = static_cast<int64_t>(rowIdx_) * Align(inputLastDimValue_, sortBufferSize) * sortBufferSize /
147- static_cast<int64_t>(rowIdx_) * Align(inputLastDimValue_, sortBufferSize) * sortBufferSize /156+ sizeof(CONVERT_TYPE) * 2;
148- sizeof(CONVERT_TYPE) * 2;
149 onceMaxElements_ = tilingData->keyParams0 / MORE_CORE_DEALING_SORT_NUM_ONCE * MORE_CORE_DEALING_SORT_NUM_ONCE;157 onceMaxElements_ = tilingData->keyParams0 / MORE_CORE_DEALING_SORT_NUM_ONCE * MORE_CORE_DEALING_SORT_NUM_ONCE;
150 158 
151 inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));159 inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));
152 outValueGm_.SetGlobalBuffer((__gm__ T*)(value));160 outValueGm_.SetGlobalBuffer((__gm__ T*)(value));
153 outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));161 outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));
154- workspaceGm_[0].SetGlobalBuffer(162+ workspaceGm_[0].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + rowWorkspaceOffset_,
155- (__gm__ CONVERT_TYPE*)(workSpace) + rowWorkspaceOffset_,163+ Align(inputLastDimValue_, sortBufferSize) * sortBufferSize / sizeof(CONVERT_TYPE));
156- Align(inputLastDimValue_, sortBufferSize) * sortBufferSize / sizeof(CONVERT_TYPE));164+ workspaceGm_[1].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + rowWorkspaceOffset_ +
157- workspaceGm_[1].SetGlobalBuffer(165+ Align(inputLastDimValue_, sortBufferSize) * sortBufferSize / sizeof(CONVERT_TYPE));
158- (__gm__ CONVERT_TYPE*)(workSpace) + rowWorkspaceOffset_ +
159- Align(inputLastDimValue_, sortBufferSize) * sortBufferSize / sizeof(CONVERT_TYPE));
160 166 
161 uint32_t tailNum = inputLastDimValue_ - (frontCoreNum_ - 1) * numTileData_;167 uint32_t tailNum = inputLastDimValue_ - (frontCoreNum_ - 1) * numTileData_;
162 uint32_t alignTile = (tailNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB;168 uint32_t alignTile = (tailNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB;
@@ -186,20 +192,20 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
186 pipe_->Reset();192 pipe_->Reset();
187 uint32_t sortBufferSize = 8;193 uint32_t sortBufferSize = 8;
188 pipe_->InitBuffer(sortedQueue_, MORE_CORE_DOUBLE_BUFFER,194 pipe_->InitBuffer(sortedQueue_, MORE_CORE_DOUBLE_BUFFER,
189- MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sortBufferSize);195+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sortBufferSize);
190 pipe_->InitBuffer(copyInQueue_, MORE_CORE_DOUBLE_BUFFER,196 pipe_->InitBuffer(copyInQueue_, MORE_CORE_DOUBLE_BUFFER,
191- MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sortBufferSize);197+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sortBufferSize);
192 pipe_->InitBuffer(castValueQueue_, MORE_CORE_DOUBLE_BUFFER,198 pipe_->InitBuffer(castValueQueue_, MORE_CORE_DOUBLE_BUFFER,
193- MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(CONVERT_TYPE));199+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(CONVERT_TYPE));
194 pipe_->InitBuffer(castIndexQueue_, MORE_CORE_DOUBLE_BUFFER,200 pipe_->InitBuffer(castIndexQueue_, MORE_CORE_DOUBLE_BUFFER,
195- MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(uint32_t));201+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(uint32_t));
196 if constexpr (std::is_same<bfloat16_t, T>::value) {202 if constexpr (std::is_same<bfloat16_t, T>::value) {
197 pipe_->InitBuffer(outValueQueue_, MORE_CORE_DOUBLE_BUFFER,203 pipe_->InitBuffer(outValueQueue_, MORE_CORE_DOUBLE_BUFFER,
198- MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(T));204+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(T));
199 }205 }
200 if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {206 if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {
201 pipe_->InitBuffer(outIndexQueue_, MORE_CORE_DOUBLE_BUFFER,207 pipe_->InitBuffer(outIndexQueue_, MORE_CORE_DOUBLE_BUFFER,
202- MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(INDEX_TYPE));208+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(INDEX_TYPE));
203 }209 }
204 210 
205 listNum_ = frontCoreNum_;211 listNum_ = frontCoreNum_;
@@ -207,7 +213,7 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
207 currentTailElements_ = inputLastDimValue_ - numTileData_ * (frontCoreNum_ - 1);213 currentTailElements_ = inputLastDimValue_ - numTileData_ * (frontCoreNum_ - 1);
208 uint32_t currentCoreNum;214 uint32_t currentCoreNum;
209 uint32_t remainListNum;215 uint32_t remainListNum;
210- while(listNum_ > MORE_CORE_MERGE_SORT_LIST_MAX_NUM) {216+ while (listNum_ > MORE_CORE_MERGE_SORT_LIST_MAX_NUM) {
211 workspaceInput_ = workspaceGm_[workSpaceFlag_];217 workspaceInput_ = workspaceGm_[workSpaceFlag_];
212 workspaceOutput_ = workspaceGm_[1 - workSpaceFlag_];218 workspaceOutput_ = workspaceGm_[1 - workSpaceFlag_];
213 SortInMultiCore();219 SortInMultiCore();
@@ -242,8 +248,8 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
242}248}
243 249 
244template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>250template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
245-__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::CopyInData(uint32_t tileNum, 251+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::CopyInData(uint32_t tileNum,
246- int64_t offsetPerCore)252+ int64_t offsetPerCore)
247{253{
248 LocalTensor<T> inputLocal = inputQueue_.AllocTensor<T>();254 LocalTensor<T> inputLocal = inputQueue_.AllocTensor<T>();
249 T defaultValue = IS_DESCEND ? static_cast<T>(-INFINITY) : static_cast<T>(NAN);255 T defaultValue = IS_DESCEND ? static_cast<T>(-INFINITY) : static_cast<T>(NAN);
@@ -270,8 +276,8 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
270}276}
271 277 
272template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>278template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
273-__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::InitIndexLocal(uint32_t tileNum, 279+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::InitIndexLocal(
274- LocalTensor<uint32_t> sortedValueIndexLocal, int64_t offsetPerCore)280+ uint32_t tileNum, LocalTensor<uint32_t> sortedValueIndexLocal, int64_t offsetPerCore)
275{281{
276 PipeBarrier<PIPE_ALL>();282 PipeBarrier<PIPE_ALL>();
277 LocalTensor<int32_t> tempIndexLocal = sortedValueIndexLocal.ReinterpretCast<int32_t>();283 LocalTensor<int32_t> tempIndexLocal = sortedValueIndexLocal.ReinterpretCast<int32_t>();
@@ -280,17 +286,19 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
280}286}
281 287 
282template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>288template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
283-__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DoSortBf16(uint32_t tileNum, 289+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DoSortBf16(
284- LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal)290+ uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal,
291+ LocalTensor<uint32_t> sortedValueIndexLocal)
285{292{
286- AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = sortTempBuf_.Get<CONVERT_TYPE>();293+ AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = sortTempBuf_.Get<CONVERT_TYPE>();
287 AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = concatTempBuf_.Get<CONVERT_TYPE>();294 AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = concatTempBuf_.Get<CONVERT_TYPE>();
288 AscendC::LocalTensor<CONVERT_TYPE> inputValueTempLocal = inputValueTempBuf_.Get<CONVERT_TYPE>();295 AscendC::LocalTensor<CONVERT_TYPE> inputValueTempLocal = inputValueTempBuf_.Get<CONVERT_TYPE>();
289 AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = sortedValueLocalCastTbuf_.Get<CONVERT_TYPE>();296 AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = sortedValueLocalCastTbuf_.Get<CONVERT_TYPE>();
290 297 
291 uint32_t aglinTileNum = ((tileNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB) * MORE_CORE_BLOCK_UB;298 uint32_t aglinTileNum = ((tileNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB) * MORE_CORE_BLOCK_UB;
292 uint32_t sortRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_SORT_NUM_ONCE - 1) / MORE_CORE_DEALING_SORT_NUM_ONCE;299 uint32_t sortRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_SORT_NUM_ONCE - 1) / MORE_CORE_DEALING_SORT_NUM_ONCE;
293- uint32_t concatRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_CONCAT_NUM_ONCE - 1) / MORE_CORE_DEALING_CONCAT_NUM_ONCE;300+ uint32_t concatRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_CONCAT_NUM_ONCE - 1) /
301+ MORE_CORE_DEALING_CONCAT_NUM_ONCE;
294 302 
295 AscendC::Cast(inputValueTempLocal, inputLocal, AscendC::RoundMode::CAST_NONE, aglinTileNum);303 AscendC::Cast(inputValueTempLocal, inputLocal, AscendC::RoundMode::CAST_NONE, aglinTileNum);
296 if constexpr (!IS_DESCEND) {304 if constexpr (!IS_DESCEND) {
@@ -298,27 +306,31 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
298 }306 }
299 AscendC::LocalTensor<CONVERT_TYPE> concatLocal;307 AscendC::LocalTensor<CONVERT_TYPE> concatLocal;
300 AscendC::Concat(concatLocal, inputValueTempLocal, concatTempLocal, concatRepeatTimes);308 AscendC::Concat(concatLocal, inputValueTempLocal, concatTempLocal, concatRepeatTimes);
301- AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal, sortRepeatTimes);309+ AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal,
310+ sortRepeatTimes);
302}311}
303 312 
304template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>313template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
305-__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DoSort(uint32_t tileNum, 314+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DoSort(
306- LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal)315+ uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal,
316+ LocalTensor<uint32_t> sortedValueIndexLocal)
307{317{
308- AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = sortTempBuf_.Get<CONVERT_TYPE>();318+ AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = sortTempBuf_.Get<CONVERT_TYPE>();
309 AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = concatTempBuf_.Get<CONVERT_TYPE>();319 AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = concatTempBuf_.Get<CONVERT_TYPE>();
310 AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = sortedValueLocalCastTbuf_.Get<CONVERT_TYPE>();320 AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = sortedValueLocalCastTbuf_.Get<CONVERT_TYPE>();
311 321 
312 uint32_t aglinTileNum = ((tileNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB) * MORE_CORE_BLOCK_UB;322 uint32_t aglinTileNum = ((tileNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB) * MORE_CORE_BLOCK_UB;
313 uint32_t sortRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_SORT_NUM_ONCE - 1) / MORE_CORE_DEALING_SORT_NUM_ONCE;323 uint32_t sortRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_SORT_NUM_ONCE - 1) / MORE_CORE_DEALING_SORT_NUM_ONCE;
314- uint32_t concatRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_CONCAT_NUM_ONCE - 1) / MORE_CORE_DEALING_CONCAT_NUM_ONCE;324+ uint32_t concatRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_CONCAT_NUM_ONCE - 1) /
325+ MORE_CORE_DEALING_CONCAT_NUM_ONCE;
315 326 
316 if constexpr (!IS_DESCEND) {327 if constexpr (!IS_DESCEND) {
317 FlipSignBit(inputLocal, aglinTileNum);328 FlipSignBit(inputLocal, aglinTileNum);
318 }329 }
319 AscendC::LocalTensor<CONVERT_TYPE> concatLocal;330 AscendC::LocalTensor<CONVERT_TYPE> concatLocal;
320 AscendC::Concat(concatLocal, inputLocal, concatTempLocal, concatRepeatTimes);331 AscendC::Concat(concatLocal, inputLocal, concatTempLocal, concatRepeatTimes);
321- AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal, sortRepeatTimes);332+ AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal,
333+ sortRepeatTimes);
322}334}
323 335 
324template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>336template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
@@ -328,7 +340,7 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
328 if constexpr (std::is_same<float, CONVERT_TYPE>::value) {340 if constexpr (std::is_same<float, CONVERT_TYPE>::value) {
329 AscendC::LocalTensor<int32_t> castTensor = xLocal.template ReinterpretCast<int32_t>();341 AscendC::LocalTensor<int32_t> castTensor = xLocal.template ReinterpretCast<int32_t>();
330 AscendC::Adds(castTensor, castTensor, MORE_CORE_XOR_OP_VALUE_FP, aglinTileNum);342 AscendC::Adds(castTensor, castTensor, MORE_CORE_XOR_OP_VALUE_FP, aglinTileNum);
331- } else if constexpr (std::is_same<half, CONVERT_TYPE>::value ){343+ } else if constexpr (std::is_same<half, CONVERT_TYPE>::value) {
332 AscendC::LocalTensor<int16_t> castTensor = xLocal.template ReinterpretCast<int16_t>();344 AscendC::LocalTensor<int16_t> castTensor = xLocal.template ReinterpretCast<int16_t>();
333 AscendC::Adds(castTensor, castTensor, MORE_CORE_XOR_OP_VALUE_HALF, aglinTileNum);345 AscendC::Adds(castTensor, castTensor, MORE_CORE_XOR_OP_VALUE_HALF, aglinTileNum);
334 }346 }
@@ -366,7 +378,6 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
366 sortedQueue_.FreeTensor<CONVERT_TYPE>(sortTempBuffer);378 sortedQueue_.FreeTensor<CONVERT_TYPE>(sortTempBuffer);
367}379}
368 380 
369- 
370template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>381template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
371__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::SortInMultiCore()382__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::SortInMultiCore()
372{383{
@@ -399,7 +410,8 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
399 copyParams.srcStride = 0;410 copyParams.srcStride = 0;
400 copyParams.dstStride = 0;411 copyParams.dstStride = 0;
401 DataCopyPadExtParams<CONVERT_TYPE> padParams{false, 0, 0, 0};412 DataCopyPadExtParams<CONVERT_TYPE> padParams{false, 0, 0, 0};
402- DataCopyPad(ubMainInput[GetSortLen<CONVERT_TYPE>(onceMaxElements_) * i], workspaceInput_[offsets_[i]], copyParams, padParams);413+ DataCopyPad(ubMainInput[GetSortLen<CONVERT_TYPE>(onceMaxElements_) * i], workspaceInput_[offsets_[i]],
414+ copyParams, padParams);
403 elementCountList_[j] = dealLengths_[i];415 elementCountList_[j] = dealLengths_[i];
404 remainListNum_ += 1;416 remainListNum_ += 1;
405 j++;417 j++;
@@ -409,7 +421,8 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
409}421}
410 422 
411template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>423template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
412-__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DealingMergeSort() {424+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DealingMergeSort()
425+{
413 LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.AllocTensor<CONVERT_TYPE>();426 LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.AllocTensor<CONVERT_TYPE>();
414 LocalTensor<CONVERT_TYPE> ubMainInput = copyInQueue_.DeQue<CONVERT_TYPE>();427 LocalTensor<CONVERT_TYPE> ubMainInput = copyInQueue_.DeQue<CONVERT_TYPE>();
415 LocalTensor<CONVERT_TYPE> tmpUbInputs[4];428 LocalTensor<CONVERT_TYPE> tmpUbInputs[4];
@@ -423,14 +436,14 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
423 MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);436 MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);
424 MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);437 MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);
425 } else if (remainListNum_ == 3) {438 } else if (remainListNum_ == 3) {
426- MrgSortSrcList sortListTail =439+ MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[0]);
427- MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[0]);
428 MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);440 MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);
429 } else if (remainListNum_ == 4) {441 } else if (remainListNum_ == 4) {
430 MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[3]);442 MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[3]);
431 MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);443 MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);
432 } else {444 } else {
433- AscendC::Copy(sortTempBuffer, tmpUbInputs[0], Align(GetSortLen<CONVERT_TYPE>(elementCountList_[0]), sizeof(CONVERT_TYPE)));445+ AscendC::Copy(sortTempBuffer, tmpUbInputs[0],
446+ Align(GetSortLen<CONVERT_TYPE>(elementCountList_[0]), sizeof(CONVERT_TYPE)));
434 listSortedNums_[0] = elementCountList_[0];447 listSortedNums_[0] = elementCountList_[0];
435 }448 }
436 sortedQueue_.EnQue(sortTempBuffer);449 sortedQueue_.EnQue(sortTempBuffer);
@@ -462,11 +475,11 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
462 for (int64_t i = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {475 for (int64_t i = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
463 uint32_t blockNum = rowCoreIdx_ * MORE_CORE_MERGE_SORT_LIST_MAX_NUM + i;476 uint32_t blockNum = rowCoreIdx_ * MORE_CORE_MERGE_SORT_LIST_MAX_NUM + i;
464 if (blockNum < listNum_ - 1) {477 if (blockNum < listNum_ - 1) {
465- listRemainElements_[i] = currentElements_;478+ listRemainElements_[i] = currentElements_;
466 offsets_[i] = GetSortOffset<CONVERT_TYPE>(blockNum * currentElements_);479 offsets_[i] = GetSortOffset<CONVERT_TYPE>(blockNum * currentElements_);
467 allRemainElements_ += listRemainElements_[i];480 allRemainElements_ += listRemainElements_[i];
468 } else if (blockNum == listNum_ - 1) {481 } else if (blockNum == listNum_ - 1) {
469- listRemainElements_[i] = currentTailElements_;482+ listRemainElements_[i] = currentTailElements_;
470 offsets_[i] = GetSortOffset<CONVERT_TYPE>(blockNum * currentElements_);483 offsets_[i] = GetSortOffset<CONVERT_TYPE>(blockNum * currentElements_);
471 allRemainElements_ += currentTailElements_;484 allRemainElements_ += currentTailElements_;
472 } else {485 } else {
@@ -545,9 +558,12 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
545 LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.DeQue<CONVERT_TYPE>();558 LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.DeQue<CONVERT_TYPE>();
546 LocalTensor<CONVERT_TYPE> castValue = castValueQueue_.AllocTensor<CONVERT_TYPE>();559 LocalTensor<CONVERT_TYPE> castValue = castValueQueue_.AllocTensor<CONVERT_TYPE>();
547 LocalTensor<uint32_t> castIndex = castIndexQueue_.AllocTensor<uint32_t>();560 LocalTensor<uint32_t> castIndex = castIndexQueue_.AllocTensor<uint32_t>();
548- AscendC::Extract(castValue, castIndex, sortTempBuffer, ((curLoopSortedNum_ + MORE_CORE_DEALING_EXTRACT_NUM_ONCE - 1) / MORE_CORE_DEALING_EXTRACT_NUM_ONCE));561+ AscendC::Extract(
562+ castValue, castIndex, sortTempBuffer,
563+ ((curLoopSortedNum_ + MORE_CORE_DEALING_EXTRACT_NUM_ONCE - 1) / MORE_CORE_DEALING_EXTRACT_NUM_ONCE));
549 if constexpr (!IS_DESCEND) {564 if constexpr (!IS_DESCEND) {
550- FlipSignBit(castValue, ((curLoopSortedNum_ + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB));565+ FlipSignBit(castValue,
566+ ((curLoopSortedNum_ + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB));
551 }567 }
552 DataCopyExtParams copyParamsValue;568 DataCopyExtParams copyParamsValue;
553 copyParamsValue.blockCount = 1;569 copyParamsValue.blockCount = 1;
@@ -600,11 +616,11 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
600 outOffset_ = 0;616 outOffset_ = 0;
601 for (int64_t i = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {617 for (int64_t i = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
602 if (i < listNum_ - 1) {618 if (i < listNum_ - 1) {
603- listRemainElements_[i] = currentElements_;619+ listRemainElements_[i] = currentElements_;
604 offsets_[i] = GetSortOffset<CONVERT_TYPE>(i * currentElements_);620 offsets_[i] = GetSortOffset<CONVERT_TYPE>(i * currentElements_);
605 allRemainElements_ += listRemainElements_[i];621 allRemainElements_ += listRemainElements_[i];
606 } else if (i == listNum_ - 1) {622 } else if (i == listNum_ - 1) {
607- listRemainElements_[i] = currentTailElements_;623+ listRemainElements_[i] = currentTailElements_;
608 offsets_[i] = GetSortOffset<CONVERT_TYPE>(i * currentElements_);624 offsets_[i] = GetSortOffset<CONVERT_TYPE>(i * currentElements_);
609 allRemainElements_ += currentTailElements_;625 allRemainElements_ += currentTailElements_;
610 } else {626 } else {
@@ -626,4 +642,4 @@ __aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_
626 }642 }
627}643}
628} // namespace topkV2644} // namespace topkV2
629-#endif //TOP_K_MERGE_SORT_MORE_CORE_H645+#endif // TOP_K_MERGE_SORT_MORE_CORE_H
Mmath/top_k_v2/op_kernel/arch35/top_k_small_axis_insertion.h+62-58
@@ -34,17 +34,18 @@ using namespace AscendC;
34 */34 */
35template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T>35template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T>
36__simt_vf__ LAUNCH_BOUND(TopKSmallAxis::INSERTION_THREAD_NUM) __aicore__36__simt_vf__ LAUNCH_BOUND(TopKSmallAxis::INSERTION_THREAD_NUM) __aicore__
37- void SimtStoreTopKInsertionBatch(uint32_t validSegs, uint32_t kValue, uint32_t valueRowElems,37+ void SimtStoreTopKInsertionBatch(uint32_t validSegCount, uint32_t kValue, uint32_t valueRowSpan,
38- uint32_t indexRowElems, uint64_t outputStart, __ubuf__ CONVERT_TYPE* values,38+ uint32_t indexRowSpan, uint64_t outputStart, __ubuf__ CONVERT_TYPE* values,
39 __ubuf__ uint32_t* indices, __gm__ volatile T* outputValue,39 __ubuf__ uint32_t* indices, __gm__ volatile T* outputValue,
40 __gm__ volatile OUT_IDX_T* outputIndex)40 __gm__ volatile OUT_IDX_T* outputIndex)
41{41{
42- uint32_t total = validSegs * kValue;42+ uint32_t total = validSegCount * kValue;
43- for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += TopKSmallAxis::INSERTION_THREAD_NUM) {43+ for (uint32_t itemId = static_cast<uint32_t>(threadIdx.x); itemId < total;
44- uint32_t seg = idx / kValue;44+ itemId += TopKSmallAxis::INSERTION_THREAD_NUM) {
45- uint32_t rank = idx - seg * kValue;45+ uint32_t segIndex = itemId / kValue;
46- outputValue[outputStart + idx] = static_cast<T>(values[seg * valueRowElems + rank]);46+ uint32_t rankInSeg = itemId - segIndex * kValue;
47- outputIndex[outputStart + idx] = static_cast<OUT_IDX_T>(indices[seg * indexRowElems + rank]);47+ outputValue[outputStart + itemId] = static_cast<T>(values[segIndex * valueRowSpan + rankInSeg]);
48+ outputIndex[outputStart + itemId] = static_cast<OUT_IDX_T>(indices[segIndex * indexRowSpan + rankInSeg]);
48 }49 }
49}50}
50 51 
@@ -53,19 +54,21 @@ __simt_vf__ LAUNCH_BOUND(TopKSmallAxis::INSERTION_THREAD_NUM) __aicore__
53 */54 */
54template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T>55template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T>
55__simt_vf__ LAUNCH_BOUND(TopKSmallAxis::INSERTION_THREAD_NUM) __aicore__56__simt_vf__ LAUNCH_BOUND(TopKSmallAxis::INSERTION_THREAD_NUM) __aicore__
56- void SimtStoreTopKNonLastInsertionBatch(uint32_t validSegs, uint32_t kValue, uint32_t valueRowElems,57+ void SimtStoreTopKNonLastInsertionBatch(uint32_t validSegCount, uint32_t kValue, uint32_t valueRowSpan,
57- uint32_t indexRowElems, uint64_t outerBaseOffset, uint64_t innerStart,58+ uint32_t indexRowSpan, uint64_t outerSliceOffset, uint64_t innerTileStart,
58- uint64_t innerSize, __ubuf__ CONVERT_TYPE* values,59+ uint64_t innerDimSize, __ubuf__ CONVERT_TYPE* values,
59 __ubuf__ uint32_t* indices, __gm__ volatile T* outputValue,60 __ubuf__ uint32_t* indices, __gm__ volatile T* outputValue,
60 __gm__ volatile OUT_IDX_T* outputIndex)61 __gm__ volatile OUT_IDX_T* outputIndex)
61{62{
62- uint32_t total = validSegs * kValue;63+ uint32_t total = validSegCount * kValue;
63- for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += TopKSmallAxis::INSERTION_THREAD_NUM) {64+ for (uint32_t itemId = static_cast<uint32_t>(threadIdx.x); itemId < total;
64- uint32_t rank = idx / validSegs;65+ itemId += TopKSmallAxis::INSERTION_THREAD_NUM) {
65- uint32_t seg = idx - rank * validSegs;66+ uint32_t rankInSeg = itemId / validSegCount;
66- uint64_t gmOffset = outerBaseOffset + static_cast<uint64_t>(rank) * innerSize + innerStart + seg;67+ uint32_t segIndex = itemId - rankInSeg * validSegCount;
67- outputValue[gmOffset] = static_cast<T>(values[seg * valueRowElems + rank]);68+ uint64_t gmOffset = outerSliceOffset + static_cast<uint64_t>(rankInSeg) * innerDimSize + innerTileStart +
68- outputIndex[gmOffset] = static_cast<OUT_IDX_T>(indices[seg * indexRowElems + rank]);69+ segIndex;
70+ outputValue[gmOffset] = static_cast<T>(values[segIndex * valueRowSpan + rankInSeg]);
71+ outputIndex[gmOffset] = static_cast<OUT_IDX_T>(indices[segIndex * indexRowSpan + rankInSeg]);
69 }72 }
70}73}
71 74 
@@ -81,7 +84,7 @@ class TopKSmallAxisInsertion
81public:84public:
82 __aicore__ inline TopKSmallAxisInsertion() {}85 __aicore__ inline TopKSmallAxisInsertion() {}
83 __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const TopKV2TilingDataSimd* tilingData,86 __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const TopKV2TilingDataSimd* tilingData,
84- TPipe* pipe);87+ TPipe* pipePtr);
85 __aicore__ inline void Process() { Base::Process(); }88 __aicore__ inline void Process() { Base::Process(); }
86 89 
87 friend Base;90 friend Base;
@@ -102,11 +105,11 @@ private:
102 105 
103 // Main pipeline functions (in call order)106 // Main pipeline functions (in call order)
104 __aicore__ inline bool IsProcessInvalid() const;107 __aicore__ inline bool IsProcessInvalid() const;
105- __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchId) const;108+ __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchIdx) const;
106- __aicore__ inline void ProcessBatch(uint32_t batchId, uint32_t validSegs);109+ __aicore__ inline void ProcessBatch(uint32_t batchIdx, uint32_t validSegCount);
107- __aicore__ inline void RunBitonicFinalize(uint32_t validSegs);110+ __aicore__ inline void RunBitonicFinalize(uint32_t validSegCount);
108- __aicore__ inline void StoreTopK(int64_t segStart, uint32_t validSegs);111+ __aicore__ inline void StoreTopK(int64_t segStart, uint32_t validSegCount);
109- __aicore__ inline void StoreNonLastTopK(uint64_t outerId, uint64_t innerStart, uint32_t validSegs);112+ __aicore__ inline void StoreNonLastTopK(uint64_t outerId, uint64_t innerTileStart, uint32_t validSegCount);
110 113 
111 GlobalTensor<T> inputGm_;114 GlobalTensor<T> inputGm_;
112 GlobalTensor<T> outValueGm_;115 GlobalTensor<T> outValueGm_;
@@ -123,9 +126,9 @@ private:
123 126 
124template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>127template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
125__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::Init(128__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::Init(
126- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const TopKV2TilingDataSimd* tilingData, TPipe* pipe)129+ GM_ADDR x, GM_ADDR values, GM_ADDR indices, const TopKV2TilingDataSimd* tilingData, TPipe* pipePtr)
127{130{
128- if (tilingData == nullptr || pipe == nullptr) {131+ if (tilingData == nullptr || pipePtr == nullptr) {
129 return;132 return;
130 }133 }
131 blockIdx_ = GetBlockIdx();134 blockIdx_ = GetBlockIdx();
@@ -138,7 +141,7 @@ __aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
138 isNonLastAxis_ = tilingData->keyParams3 != 0U;141 isNonLastAxis_ = tilingData->keyParams3 != 0U;
139 innerLoopNum_ = tilingData->keyParams2;142 innerLoopNum_ = tilingData->keyParams2;
140 if (isNonLastAxis_) {143 if (isNonLastAxis_) {
141- // Mode-8 convention: unsortedDimNum is innerSize and oneCoreRowNum is outerSize.144+ // Mode-8 convention: unsortedDimNum is innerSize_ and oneCoreRowNum is outerSize.
142 outerSize_ = static_cast<int64_t>(tilingData->oneCoreRowNum);145 outerSize_ = static_cast<int64_t>(tilingData->oneCoreRowNum);
143 innerSize_ = static_cast<int64_t>(tilingData->unsortedDimNum);146 innerSize_ = static_cast<int64_t>(tilingData->unsortedDimNum);
144 } else {147 } else {
@@ -154,7 +157,7 @@ __aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
154 return;157 return;
155 }158 }
156 159 
157- Base::InitInsertionBuffers(pipe);160+ Base::InitInsertionBuffers(pipePtr);
158}161}
159 162 
160template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>163template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
@@ -166,21 +169,21 @@ __aicore__ inline bool TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
166 169 
167template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>170template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
168__aicore__ inline uint32_t TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::ComputeValidSegs(171__aicore__ inline uint32_t TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::ComputeValidSegs(
169- uint32_t batchId) const172+ uint32_t batchIdx) const
170{173{
171 if (isNonLastAxis_) {174 if (isNonLastAxis_) {
172 // Non-last batching iterates inner tiles inside each outer slice; the tail tile175 // Non-last batching iterates inner tiles inside each outer slice; the tail tile
173 // may contain fewer segments than the configured batch size.176 // may contain fewer segments than the configured batch size.
174- uint32_t innerTileId = batchId % innerLoopNum_;177+ uint32_t innerTileId = batchIdx % innerLoopNum_;
175- int64_t innerStart = static_cast<int64_t>(innerTileId) * static_cast<int64_t>(segmentsPerBatch_);178+ int64_t innerTileStart = static_cast<int64_t>(innerTileId) * static_cast<int64_t>(segmentsPerBatch_);
176- int64_t innerRemain = innerSize_ - innerStart;179+ int64_t innerRemain = innerSize_ - innerTileStart;
177 if (innerRemain <= 0) {180 if (innerRemain <= 0) {
178 return 0;181 return 0;
179 }182 }
180 return innerRemain >= static_cast<int64_t>(segmentsPerBatch_) ? segmentsPerBatch_ :183 return innerRemain >= static_cast<int64_t>(segmentsPerBatch_) ? segmentsPerBatch_ :
181 static_cast<uint32_t>(innerRemain);184 static_cast<uint32_t>(innerRemain);
182 }185 }
183- int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(segmentsPerBatch_);186+ int64_t segStart = static_cast<int64_t>(batchIdx) * static_cast<int64_t>(segmentsPerBatch_);
184 int64_t segRemain = totalSegs_ - segStart;187 int64_t segRemain = totalSegs_ - segStart;
185 if (segRemain <= 0) {188 if (segRemain <= 0) {
186 return 0;189 return 0;
@@ -192,54 +195,55 @@ __aicore__ inline uint32_t TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, Is
192}195}
193 196 
194template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>197template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
195-__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::ProcessBatch(uint32_t batchId,198+__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::ProcessBatch(
196- uint32_t validSegs)199+ uint32_t batchIdx, uint32_t validSegCount)
197{200{
198- int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(segmentsPerBatch_);201+ int64_t segStart = static_cast<int64_t>(batchIdx) * static_cast<int64_t>(segmentsPerBatch_);
199 uint64_t outerId = 0;202 uint64_t outerId = 0;
200- uint64_t innerStart = 0;203+ uint64_t innerTileStart = 0;
201 if (isNonLastAxis_) {204 if (isNonLastAxis_) {
202- // batchId is linearized as [outerId, innerTileId] for non-last-axis work.205+ // batchIdx is linearized as [outerId, innerTileId] for non-last-axis work.
203- outerId = static_cast<uint64_t>(batchId / innerLoopNum_);206+ outerId = static_cast<uint64_t>(batchIdx / innerLoopNum_);
204- uint32_t innerTileId = batchId % innerLoopNum_;207+ uint32_t innerTileId = batchIdx % innerLoopNum_;
205- innerStart = static_cast<uint64_t>(innerTileId) * static_cast<uint64_t>(segmentsPerBatch_);208+ innerTileStart = static_cast<uint64_t>(innerTileId) * static_cast<uint64_t>(segmentsPerBatch_);
206- uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);209+ uint64_t outerSliceOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);
207- Base::LoadNonLastBatch(inputGm_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_), validSegs);210+ Base::LoadNonLastBatch(inputGm_, outerSliceOffset, innerTileStart, static_cast<uint64_t>(innerSize_),
211+ validSegCount);
208 } else {212 } else {
209- Base::LoadContiguousBatch(inputGm_, segStart * static_cast<int64_t>(segmentLen_), validSegs, false);213+ Base::LoadContiguousBatch(inputGm_, segStart * static_cast<int64_t>(segmentLen_), validSegCount, false);
210 }214 }
211- Base::SortBatch(validSegs);215+ Base::SortBatch(validSegCount);
212- RunBitonicFinalize(validSegs);216+ RunBitonicFinalize(validSegCount);
213 if (isNonLastAxis_) {217 if (isNonLastAxis_) {
214- StoreNonLastTopK(outerId, innerStart, validSegs);218+ StoreNonLastTopK(outerId, innerTileStart, validSegCount);
215 } else {219 } else {
216- StoreTopK(segStart, validSegs);220+ StoreTopK(segStart, validSegCount);
217 }221 }
218}222}
219 223 
220template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>224template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
221__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::RunBitonicFinalize(225__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::RunBitonicFinalize(
222- uint32_t validSegs)226+ uint32_t validSegCount)
223{227{
224- if (!useBitonicFinalize_ || validSegs == 0U) {228+ if (!useBitonicFinalize_ || validSegCount == 0U) {
225 return;229 return;
226 }230 }
227- RunBitonicFinalizeSelectionRows<CONVERT_TYPE, uint32_t, IsDescend>(values_, indices_, kValue_, validSegs,231+ RunBitonicFinalizeSelectionRows<CONVERT_TYPE, uint32_t, IsDescend>(values_, indices_, kValue_, validSegCount,
228 valueRowStride_, indexRowStride_);232 valueRowStride_, indexRowStride_);
229}233}
230 234 
231template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>235template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
232__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::StoreTopK(int64_t segStart,236__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::StoreTopK(int64_t segStart,
233- uint32_t validSegs)237+ uint32_t validSegCount)
234{238{
235- if (validSegs == 0U) {239+ if (validSegCount == 0U) {
236 return;240 return;
237 }241 }
238 // Output layout is [totalSegs, k]; each sorted segment contributes its first k elements.242 // Output layout is [totalSegs, k]; each sorted segment contributes its first k elements.
239 uint64_t outputStart = static_cast<uint64_t>(segStart) * static_cast<uint64_t>(kValue_);243 uint64_t outputStart = static_cast<uint64_t>(segStart) * static_cast<uint64_t>(kValue_);
240 asc_vf_call<SimtStoreTopKInsertionBatch<T, CONVERT_TYPE, OUT_IDX_T>>(244 asc_vf_call<SimtStoreTopKInsertionBatch<T, CONVERT_TYPE, OUT_IDX_T>>(
241- dim3(TopKSmallAxis::INSERTION_THREAD_NUM), validSegs, kValue_, valueRowStride_, indexRowStride_, outputStart,245+ dim3(TopKSmallAxis::INSERTION_THREAD_NUM), validSegCount, kValue_, valueRowStride_, indexRowStride_,
242- (__ubuf__ CONVERT_TYPE*)values_.GetPhyAddr(), (__ubuf__ uint32_t*)indices_.GetPhyAddr(),246+ outputStart, (__ubuf__ CONVERT_TYPE*)values_.GetPhyAddr(), (__ubuf__ uint32_t*)indices_.GetPhyAddr(),
243 (__gm__ volatile T*)outValueGm_.GetPhyAddr(), (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());247 (__gm__ volatile T*)outValueGm_.GetPhyAddr(), (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());
244 event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));248 event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
245 SetFlag<HardEvent::V_S>(eventId);249 SetFlag<HardEvent::V_S>(eventId);
@@ -248,13 +252,13 @@ __aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
248 252 
249template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>253template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
250__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::StoreNonLastTopK(254__aicore__ inline void TopKSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::StoreNonLastTopK(
251- uint64_t outerId, uint64_t innerStart, uint32_t validSegs)255+ uint64_t outerId, uint64_t innerTileStart, uint32_t validSegCount)
252{256{
253 // Output layout is [outer, k, inner]; the topk axis replaces segmentLen with kValue_.257 // Output layout is [outer, k, inner]; the topk axis replaces segmentLen with kValue_.
254 uint64_t outputBase = outerId * static_cast<uint64_t>(kValue_) * static_cast<uint64_t>(innerSize_);258 uint64_t outputBase = outerId * static_cast<uint64_t>(kValue_) * static_cast<uint64_t>(innerSize_);
255 asc_vf_call<SimtStoreTopKNonLastInsertionBatch<T, CONVERT_TYPE, OUT_IDX_T>>(259 asc_vf_call<SimtStoreTopKNonLastInsertionBatch<T, CONVERT_TYPE, OUT_IDX_T>>(
256- dim3(TopKSmallAxis::INSERTION_THREAD_NUM), validSegs, kValue_, valueRowStride_, indexRowStride_, outputBase,260+ dim3(TopKSmallAxis::INSERTION_THREAD_NUM), validSegCount, kValue_, valueRowStride_, indexRowStride_, outputBase,
257- innerStart, static_cast<uint64_t>(innerSize_), (__ubuf__ CONVERT_TYPE*)values_.GetPhyAddr(),261+ innerTileStart, static_cast<uint64_t>(innerSize_), (__ubuf__ CONVERT_TYPE*)values_.GetPhyAddr(),
258 (__ubuf__ uint32_t*)indices_.GetPhyAddr(), (__gm__ volatile T*)outValueGm_.GetPhyAddr(),262 (__ubuf__ uint32_t*)indices_.GetPhyAddr(), (__gm__ volatile T*)outValueGm_.GetPhyAddr(),
259 (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());263 (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());
260 event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));264 event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
Mmath/top_k_v2/op_kernel/arch35/top_k_small_axis_two_stage.h+60-56
@@ -33,17 +33,18 @@ using namespace AscendC;
33 */33 */
34template <typename T, typename OUT_IDX_T>34template <typename T, typename OUT_IDX_T>
35__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__35__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__
36- void SimtStoreTopKTwoStageBatch(uint32_t validSegs, uint32_t segmentLen, uint32_t kValue, uint64_t outputStart,36+ void SimtStoreTopKTwoStageBatch(uint32_t validSegCount, uint32_t segLen, uint32_t kValue, uint64_t outputStart,
37 __ubuf__ T* finalValues, __ubuf__ uint32_t* finalIdx,37 __ubuf__ T* finalValues, __ubuf__ uint32_t* finalIdx,
38 __gm__ volatile T* outputValue, __gm__ volatile OUT_IDX_T* outputIndex)38 __gm__ volatile T* outputValue, __gm__ volatile OUT_IDX_T* outputIndex)
39{39{
40- uint32_t total = validSegs * kValue;40+ uint32_t total = validSegCount * kValue;
41- for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += SmallAxisCommon::TWO_STAGE_THREAD_NUM) {41+ for (uint32_t itemId = static_cast<uint32_t>(threadIdx.x); itemId < total;
42- uint32_t seg = idx / kValue;42+ itemId += SmallAxisCommon::TWO_STAGE_THREAD_NUM) {
43- uint32_t rank = idx - seg * kValue;43+ uint32_t segIndex = itemId / kValue;
44- uint32_t srcOffset = seg * segmentLen + rank;44+ uint32_t rankInSeg = itemId - segIndex * kValue;
45- outputValue[outputStart + idx] = finalValues[srcOffset];45+ uint32_t srcOffset = segIndex * segLen + rankInSeg;
46- outputIndex[outputStart + idx] = static_cast<OUT_IDX_T>(finalIdx[srcOffset]);46+ outputValue[outputStart + itemId] = finalValues[srcOffset];
47+ outputIndex[outputStart + itemId] = static_cast<OUT_IDX_T>(finalIdx[srcOffset]);
47 }48 }
48}49}
49 50 
@@ -52,17 +53,19 @@ __simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__
52 */53 */
53template <typename T, typename OUT_IDX_T>54template <typename T, typename OUT_IDX_T>
54__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__55__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__
55- void SimtStoreTopKNonLastTwoStageBatch(uint32_t validSegs, uint32_t segmentLen, uint32_t kValue,56+ void SimtStoreTopKNonLastTwoStageBatch(uint32_t validSegCount, uint32_t segLen, uint32_t kValue,
56- uint64_t outerBaseOffset, uint64_t innerStart, uint64_t innerSize,57+ uint64_t outerSliceOffset, uint64_t innerTileStart, uint64_t innerDimSize,
57 __ubuf__ T* finalValues, __ubuf__ uint32_t* finalIdx,58 __ubuf__ T* finalValues, __ubuf__ uint32_t* finalIdx,
58 __gm__ volatile T* outputValue, __gm__ volatile OUT_IDX_T* outputIndex)59 __gm__ volatile T* outputValue, __gm__ volatile OUT_IDX_T* outputIndex)
59{60{
60- uint32_t total = validSegs * kValue;61+ uint32_t total = validSegCount * kValue;
61- for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += SmallAxisCommon::TWO_STAGE_THREAD_NUM) {62+ for (uint32_t itemId = static_cast<uint32_t>(threadIdx.x); itemId < total;
62- uint32_t rank = idx / validSegs;63+ itemId += SmallAxisCommon::TWO_STAGE_THREAD_NUM) {
63- uint32_t seg = idx - rank * validSegs;64+ uint32_t rankInSeg = itemId / validSegCount;
64- uint32_t srcOffset = seg * segmentLen + rank;65+ uint32_t segIndex = itemId - rankInSeg * validSegCount;
65- uint64_t gmOffset = outerBaseOffset + static_cast<uint64_t>(rank) * innerSize + innerStart + seg;66+ uint32_t srcOffset = segIndex * segLen + rankInSeg;
67+ uint64_t gmOffset = outerSliceOffset + static_cast<uint64_t>(rankInSeg) * innerDimSize + innerTileStart +
68+ segIndex;
66 outputValue[gmOffset] = finalValues[srcOffset];69 outputValue[gmOffset] = finalValues[srcOffset];
67 outputIndex[gmOffset] = static_cast<OUT_IDX_T>(finalIdx[srcOffset]);70 outputIndex[gmOffset] = static_cast<OUT_IDX_T>(finalIdx[srcOffset]);
68 }71 }
@@ -78,7 +81,7 @@ class TopKSmallAxisTwoStage
78public:81public:
79 __aicore__ inline TopKSmallAxisTwoStage() {}82 __aicore__ inline TopKSmallAxisTwoStage() {}
80 __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const TopKV2TilingDataSimd* tilingData,83 __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const TopKV2TilingDataSimd* tilingData,
81- TPipe* pipe);84+ TPipe* pipePtr);
82 __aicore__ inline void Process() { Base::Process(); }85 __aicore__ inline void Process() { Base::Process(); }
83 86 
84 friend Base;87 friend Base;
@@ -96,11 +99,11 @@ private:
96 99 
97 // Main pipeline functions (in call order)100 // Main pipeline functions (in call order)
98 __aicore__ inline bool IsProcessInvalid() const;101 __aicore__ inline bool IsProcessInvalid() const;
99- __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchId) const;102+ __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchIdx) const;
100- __aicore__ inline void ProcessBatch(uint32_t batchId, uint32_t validSegs);103+ __aicore__ inline void ProcessBatch(uint32_t batchIdx, uint32_t validSegCount);
101- __aicore__ inline void RunBitonicFinalize(uint32_t validSegs);104+ __aicore__ inline void RunBitonicFinalize(uint32_t validSegCount);
102- __aicore__ inline void StoreTopK(int64_t segStart, uint32_t validSegs);105+ __aicore__ inline void StoreTopK(int64_t segStart, uint32_t validSegCount);
103- __aicore__ inline void StoreNonLastTopK(uint64_t outerId, uint64_t innerStart, uint32_t validSegs);106+ __aicore__ inline void StoreNonLastTopK(uint64_t outerId, uint64_t innerTileStart, uint32_t validSegCount);
104 107 
105 GlobalTensor<T> inputGm_;108 GlobalTensor<T> inputGm_;
106 GlobalTensor<T> outValueGm_;109 GlobalTensor<T> outValueGm_;
@@ -118,12 +121,12 @@ private:
118template <typename T, typename OUT_IDX_T, bool IsDescend>121template <typename T, typename OUT_IDX_T, bool IsDescend>
119__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,122__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
120 const TopKV2TilingDataSimd* tilingData,123 const TopKV2TilingDataSimd* tilingData,
121- TPipe* pipe)124+ TPipe* pipePtr)
122{125{
123- if (tilingData == nullptr || pipe == nullptr) {126+ if (tilingData == nullptr || pipePtr == nullptr) {
124 return;127 return;
125 }128 }
126- pipe_ = pipe;129+ pipe_ = pipePtr;
127 blockIdx_ = GetBlockIdx();130 blockIdx_ = GetBlockIdx();
128 blockDim_ = GetBlockNum();131 blockDim_ = GetBlockNum();
129 batchSize_ = tilingData->keyParams0;132 batchSize_ = tilingData->keyParams0;
@@ -151,7 +154,7 @@ __aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::Init(GM_A
151 }154 }
152 155 
153 constexpr uint32_t kAliasElemBytes = static_cast<uint32_t>(sizeof(uint32_t));156 constexpr uint32_t kAliasElemBytes = static_cast<uint32_t>(sizeof(uint32_t));
154- Base::InitSortBuffers(pipe, maxFlatElems_, tilingData->tmpUbSize, tilingData->keyParams2 != 0U, kAliasElemBytes);157+ Base::InitSortBuffers(pipePtr, maxFlatElems_, tilingData->tmpUbSize, tilingData->keyParams2 != 0U, kAliasElemBytes);
155}158}
156 159 
157template <typename T, typename OUT_IDX_T, bool IsDescend>160template <typename T, typename OUT_IDX_T, bool IsDescend>
@@ -162,12 +165,12 @@ __aicore__ inline bool TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::IsProcess
162}165}
163 166 
164template <typename T, typename OUT_IDX_T, bool IsDescend>167template <typename T, typename OUT_IDX_T, bool IsDescend>
165-__aicore__ inline uint32_t TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::ComputeValidSegs(uint32_t batchId) const168+__aicore__ inline uint32_t TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::ComputeValidSegs(uint32_t batchIdx) const
166{169{
167 if (isNonLastAxis_) {170 if (isNonLastAxis_) {
168 // Non-last batches are grouped by outer slice, then by inner tile.171 // Non-last batches are grouped by outer slice, then by inner tile.
169 // The last inner tile can be shorter than batchSize_.172 // The last inner tile can be shorter than batchSize_.
170- uint32_t tileInOuter = batchId % innerLoopNum_;173+ uint32_t tileInOuter = batchIdx % innerLoopNum_;
171 int64_t tileStart = static_cast<int64_t>(tileInOuter) * static_cast<int64_t>(batchSize_);174 int64_t tileStart = static_cast<int64_t>(tileInOuter) * static_cast<int64_t>(batchSize_);
172 int64_t remainingInTile = innerSize_ - tileStart;175 int64_t remainingInTile = innerSize_ - tileStart;
173 if (remainingInTile <= 0) {176 if (remainingInTile <= 0) {
@@ -176,7 +179,7 @@ __aicore__ inline uint32_t TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::Compu
176 return remainingInTile >= static_cast<int64_t>(batchSize_) ? batchSize_ :179 return remainingInTile >= static_cast<int64_t>(batchSize_) ? batchSize_ :
177 static_cast<uint32_t>(remainingInTile);180 static_cast<uint32_t>(remainingInTile);
178 }181 }
179- int64_t batchStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);182+ int64_t batchStart = static_cast<int64_t>(batchIdx) * static_cast<int64_t>(batchSize_);
180 int64_t remainingSegs = totalSegs_ - batchStart;183 int64_t remainingSegs = totalSegs_ - batchStart;
181 if (remainingSegs <= 0) {184 if (remainingSegs <= 0) {
182 return 0;185 return 0;
@@ -188,48 +191,49 @@ __aicore__ inline uint32_t TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::Compu
188}191}
189 192 
190template <typename T, typename OUT_IDX_T, bool IsDescend>193template <typename T, typename OUT_IDX_T, bool IsDescend>
191-__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::ProcessBatch(uint32_t batchId,194+__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::ProcessBatch(uint32_t batchIdx,
192- uint32_t validSegs)195+ uint32_t validSegCount)
193{196{
194- uint32_t totalElems = validSegs * segmentLen_;197+ uint32_t batchElemTotal = validSegCount * segmentLen_;
195- int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);198+ int64_t segStart = static_cast<int64_t>(batchIdx) * static_cast<int64_t>(batchSize_);
196 // Pipeline: load -> two-stage sort (Stage1 flat sort + Stage2 per-segment restore) -> topk store.199 // Pipeline: load -> two-stage sort (Stage1 flat sort + Stage2 per-segment restore) -> topk store.
197 uint64_t outerId = 0;200 uint64_t outerId = 0;
198- uint64_t innerStart = 0;201+ uint64_t innerTileStart = 0;
199 if (isNonLastAxis_) {202 if (isNonLastAxis_) {
200- // batchId is linearized as outerId * innerLoopNum_ + innerTileId.203+ // batchIdx is linearized as outerId * innerLoopNum_ + innerTileId.
201- outerId = static_cast<uint64_t>(batchId / innerLoopNum_);204+ outerId = static_cast<uint64_t>(batchIdx / innerLoopNum_);
202- uint32_t innerTileId = batchId % innerLoopNum_;205+ uint32_t innerTileId = batchIdx % innerLoopNum_;
203- innerStart = static_cast<uint64_t>(innerTileId) * static_cast<uint64_t>(batchSize_);206+ innerTileStart = static_cast<uint64_t>(innerTileId) * static_cast<uint64_t>(batchSize_);
204- uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);207+ uint64_t outerSliceOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);
205- Base::LoadNonLastBatch(inputGm_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_), validSegs,208+ Base::LoadNonLastBatch(inputGm_, outerSliceOffset, innerTileStart, static_cast<uint64_t>(innerSize_),
206- totalElems);209+ validSegCount, batchElemTotal);
207 } else {210 } else {
208- Base::LoadContiguousBatch(inputGm_, segStart * static_cast<int64_t>(segmentLen_), totalElems);211+ Base::LoadContiguousBatch(inputGm_, segStart * static_cast<int64_t>(segmentLen_), batchElemTotal);
209 }212 }
210- Base::RunTwoStageSort(totalElems);213+ Base::RunTwoStageSort(batchElemTotal);
211- RunBitonicFinalize(validSegs);214+ RunBitonicFinalize(validSegCount);
212 if (isNonLastAxis_) {215 if (isNonLastAxis_) {
213- StoreNonLastTopK(outerId, innerStart, validSegs);216+ StoreNonLastTopK(outerId, innerTileStart, validSegCount);
214 } else {217 } else {
215- StoreTopK(segStart, validSegs);218+ StoreTopK(segStart, validSegCount);
216 }219 }
217}220}
218 221 
219template <typename T, typename OUT_IDX_T, bool IsDescend>222template <typename T, typename OUT_IDX_T, bool IsDescend>
220-__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::RunBitonicFinalize(uint32_t validSegs)223+__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::RunBitonicFinalize(uint32_t validSegCount)
221{224{
222- if (!useBitonicFinalize_ || validSegs == 0U) {225+ if (!useBitonicFinalize_ || validSegCount == 0U) {
223 return;226 return;
224 }227 }
225- RunBitonicFinalizeSelectionRows<T, uint32_t, IsDescend>(finalValues_, finalIdx_, kValue_, validSegs, segmentLen_,228+ RunBitonicFinalizeSelectionRows<T, uint32_t, IsDescend>(finalValues_, finalIdx_, kValue_, validSegCount,
226- segmentLen_);229+ segmentLen_, segmentLen_);
227}230}
228 231 
229template <typename T, typename OUT_IDX_T, bool IsDescend>232template <typename T, typename OUT_IDX_T, bool IsDescend>
230-__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreTopK(int64_t segStart, uint32_t validSegs)233+__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreTopK(int64_t segStart,
234+ uint32_t validSegCount)
231{235{
232- if (validSegs == 0U) {236+ if (validSegCount == 0U) {
233 return;237 return;
234 }238 }
235 // Rank-inverse writes final buffers via SIMT VF; BuildOutputs writes them via vector APIs.239 // Rank-inverse writes final buffers via SIMT VF; BuildOutputs writes them via vector APIs.
@@ -241,7 +245,7 @@ __aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreTopK
241 // Output layout is [totalSegs, k]; each sorted segment contributes its first k elements.245 // Output layout is [totalSegs, k]; each sorted segment contributes its first k elements.
242 uint64_t outputStart = static_cast<uint64_t>(segStart) * static_cast<uint64_t>(kValue_);246 uint64_t outputStart = static_cast<uint64_t>(segStart) * static_cast<uint64_t>(kValue_);
243 asc_vf_call<SimtStoreTopKTwoStageBatch<T, OUT_IDX_T>>(247 asc_vf_call<SimtStoreTopKTwoStageBatch<T, OUT_IDX_T>>(
244- dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kValue_, outputStart,248+ dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegCount, segmentLen_, kValue_, outputStart,
245 (__ubuf__ T*)finalValues_.GetPhyAddr(), (__ubuf__ uint32_t*)finalIdx_.GetPhyAddr(),249 (__ubuf__ T*)finalValues_.GetPhyAddr(), (__ubuf__ uint32_t*)finalIdx_.GetPhyAddr(),
246 (__gm__ volatile T*)outValueGm_.GetPhyAddr(), (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());250 (__gm__ volatile T*)outValueGm_.GetPhyAddr(), (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());
247 eventIdVToS = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));251 eventIdVToS = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
@@ -251,8 +255,8 @@ __aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreTopK
251 255 
252template <typename T, typename OUT_IDX_T, bool IsDescend>256template <typename T, typename OUT_IDX_T, bool IsDescend>
253__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreNonLastTopK(uint64_t outerId,257__aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreNonLastTopK(uint64_t outerId,
254- uint64_t innerStart,258+ uint64_t innerTileStart,
255- uint32_t validSegs)259+ uint32_t validSegCount)
256{260{
257 // Output layout is [outer, k, inner]; the topk axis replaces segmentLen with kValue_.261 // Output layout is [outer, k, inner]; the topk axis replaces segmentLen with kValue_.
258 uint64_t outputBase = outerId * static_cast<uint64_t>(kValue_) * static_cast<uint64_t>(innerSize_);262 uint64_t outputBase = outerId * static_cast<uint64_t>(kValue_) * static_cast<uint64_t>(innerSize_);
@@ -260,7 +264,7 @@ __aicore__ inline void TopKSmallAxisTwoStage<T, OUT_IDX_T, IsDescend>::StoreNonL
260 SetFlag<HardEvent::V_S>(eventIdVToS);264 SetFlag<HardEvent::V_S>(eventIdVToS);
261 WaitFlag<HardEvent::V_S>(eventIdVToS);265 WaitFlag<HardEvent::V_S>(eventIdVToS);
262 asc_vf_call<SimtStoreTopKNonLastTwoStageBatch<T, OUT_IDX_T>>(266 asc_vf_call<SimtStoreTopKNonLastTwoStageBatch<T, OUT_IDX_T>>(
263- dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kValue_, outputBase, innerStart,267+ dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegCount, segmentLen_, kValue_, outputBase, innerTileStart,
264 static_cast<uint64_t>(innerSize_), (__ubuf__ T*)finalValues_.GetPhyAddr(),268 static_cast<uint64_t>(innerSize_), (__ubuf__ T*)finalValues_.GetPhyAddr(),
265 (__ubuf__ uint32_t*)finalIdx_.GetPhyAddr(), (__gm__ volatile T*)outValueGm_.GetPhyAddr(),269 (__ubuf__ uint32_t*)finalIdx_.GetPhyAddr(), (__gm__ volatile T*)outValueGm_.GetPhyAddr(),
266 (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());270 (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());