Star1333
1451
代码
代码
Issues114
Pull Requests256
Actions
项目讨论
Wiki
分析
Star1333
1451
已合并
TopkV2算子性能优化 #2821
cy_hw创建于 5月19日
TopkV2算子性能优化 #2821
已合并
cy_hw创建于 5月19日
从caoyan_huawei/ops-mathcy_br_518_xingneng合入到cann/ops-mathmaster
共 7 个文件变更+1841-394
Mmath/sort/op_kernel/arch35/sort_merge_intra_core.h+1-1
@@ -50,7 +50,7 @@ public:
50 const SortRegBaseTilingData* tilingData, TPipe* pipe);50 const SortRegBaseTilingData* tilingData, TPipe* pipe);
51 __aicore__ inline void Process();51 __aicore__ inline void Process();
52 52 
53-private:53+protected:
54 struct MergeListContext {54 struct MergeListContext {
55 uint32_t remains[MERGE_LIST_MAX_NUM] = {0};55 uint32_t remains[MERGE_LIST_MAX_NUM] = {0};
56 uint32_t gmOffsets[MERGE_LIST_MAX_NUM] = {0};56 uint32_t gmOffsets[MERGE_LIST_MAX_NUM] = {0};
Mmath/top_k_v2/op_host/arch35/top_k_v2_tiling_arch35.cpp+801-388
@@ -30,6 +30,10 @@ const uint32_t BIN_NUM = 256;
30const uint32_t TILE_SIZE_DECREASING_FACTOR = 32;30const uint32_t TILE_SIZE_DECREASING_FACTOR = 32;
31const uint32_t TMP_DATA_NUM = 7680; // 默认UB一次性能处理的非64位数据的个数，可根据场景动态调整31const uint32_t TMP_DATA_NUM = 7680; // 默认UB一次性能处理的非64位数据的个数，可根据场景动态调整
32const uint32_t TMP_DATA_NUM_B64 = 5120; // 默认UB一次性能处理的64位数据的个数，可根据场景动态调整32const uint32_t TMP_DATA_NUM_B64 = 5120; // 默认UB一次性能处理的64位数据的个数，可根据场景动态调整
33+const uint32_t SINGLE_BLOCK_DATA_NUM = 15360;
34+const uint32_t SINGLE_BLOCK_DATA_NUM_B64 = 10240;
35+const uint32_t SINGLE_CORE_DATA_NUM = 15360;
36+const uint32_t SINGLE_CORE_DATA_NUM_B64 = 10240;
33const uint64_t AGLIN_FACTOR = 32;37const uint64_t AGLIN_FACTOR = 32;
34const uint32_t SMALL_MAX_DATA_SZIE = 1024;38const uint32_t SMALL_MAX_DATA_SZIE = 1024;
35const uint32_t MERGE_SORT_TILING_OFFSET = 10000;39const uint32_t MERGE_SORT_TILING_OFFSET = 10000;
@@ -38,6 +42,8 @@ const uint32_t MULT_CORE_MODE = 2;
38const uint32_t MULT_CORE_OPTIM_MODE = 4;42const uint32_t MULT_CORE_OPTIM_MODE = 4;
39const uint32_t SINGLE_BLOCK_MODE = 3;43const uint32_t SINGLE_BLOCK_MODE = 3;
40const uint32_t SORT_AND_TOP_K_MODE = 5;44const uint32_t SORT_AND_TOP_K_MODE = 5;
45+const uint32_t FP32_MERGE_MORE_CORE_MODE = 6;
46+const uint32_t FP32_MERGE_INTRA_CORE_MODE = 7;
41const uint32_t INT64_BYTE = 8;47const uint32_t INT64_BYTE = 8;
42const uint32_t INT32_BYTE = 4;48const uint32_t INT32_BYTE = 4;
43// SortAndTopk的阈值，排序轴大于该阈值的场景，走sortAndTopK模板49// SortAndTopk的阈值，排序轴大于该阈值的场景，走sortAndTopK模板
@@ -45,6 +51,17 @@ const uint32_t SORT_AND_TOP_K_THRESHOLD = 10000000;
45const uint32_t CONST_SIMT_SPACE = 32768; // 获取到的UB大小需要预留32KB给simt51const uint32_t CONST_SIMT_SPACE = 32768; // 获取到的UB大小需要预留32KB给simt
46const uint32_t SUPPORT_SORT_MAX_BYTE_SIZE = 8000;52const uint32_t SUPPORT_SORT_MAX_BYTE_SIZE = 8000;
47const uint32_t SUPPORT_SORT_MAX_SIZE = 2000;53const uint32_t SUPPORT_SORT_MAX_SIZE = 2000;
54+const uint32_t TOPK_MERGE_SORT_MORE_CORE_TILING_KEY_FLOAT = 23003;
55+const uint32_t TOPK_MERGE_SORT_INTRA_CORE_TILING_KEY_FLOAT = 33003;
56+const double FP32_K_LAST_AXIS_LOWER_RATIO = 0.25;
57+const double FP32_K_LAST_AXIS_UPPER_RATIO = 0.50;
58+const uint32_t SORT_STRUCT_SIZE_FP32 = 8;
59+const uint32_t FP32_MERGE_SORT_MAX_SIZE = 4096;
60+const uint32_t MERGE_SORT_DISABLE_DOUBLE_BUFFER_SIZE = 2048;
61+const uint32_t MERGE_MORE_CORE_ONE_CORE_DATA_SIZE = 2048;
62+const uint32_t MERGE_MORE_CORE_LIST_MAX_NUM = 4;
63+const uint32_t MERGE_INTRA_CORE_SORT_ALIGN = 32;
64+const uint32_t MERGE_INTRA_CORE_MAX_BLOCKS = 256;
48const float LAST_LOOP_CORE_UTILIZATION = 0.7;65const float LAST_LOOP_CORE_UTILIZATION = 0.7;
49const uint32_t SMALL_LOOP_UPPER_NUM = 4;66const uint32_t SMALL_LOOP_UPPER_NUM = 4;
50const uint32_t SMALL_LOOP_LOWER_NUM = 2;67const uint32_t SMALL_LOOP_LOWER_NUM = 2;
@@ -70,6 +87,7 @@ struct TopkTileInfo {
70 uint32_t tailBatchNum = 0;87 uint32_t tailBatchNum = 0;
71 uint32_t tailTileNum = 0;88 uint32_t tailTileNum = 0;
72 int64_t topKOutLastAxisNum = 0;89 int64_t topKOutLastAxisNum = 0;
90+ bool multiCoreBigModel = false;
73};91};
74struct SortTileInfo {92struct SortTileInfo {
75 uint32_t coreNumNeed = 0;93 uint32_t coreNumNeed = 0;
@@ -102,7 +120,7 @@ struct SortTileInfo {
102 uint32_t blockTileNum = 0;120 uint32_t blockTileNum = 0;
103 uint32_t tailTileNum = 0;121 uint32_t tailTileNum = 0;
104};122};
105-struct TopkComputingNowTileSizeInfo {123+struct TopkComputeNowTileSizeInfo {
106 ge::DataType dataType;124 ge::DataType dataType;
107 ge::DataType indicesDType;125 ge::DataType indicesDType;
108 bool isLargest = true;126 bool isLargest = true;
@@ -113,6 +131,7 @@ struct TopkComputingNowTileSizeInfo {
113 uint32_t maxCoreNum = 0;131 uint32_t maxCoreNum = 0;
114 uint64_t ubSizePlatForm = 0;132 uint64_t ubSizePlatForm = 0;
115 uint64_t ubBlockAlignSize = 0; // ub的对齐数值，当前为32133 uint64_t ubBlockAlignSize = 0; // ub的对齐数值，当前为32
134+ uint32_t unsortedDimNum = 0;
116};135};
117static const std::map<ge::DataType, uint32_t> tilingDataTypeKeyMap = {136static const std::map<ge::DataType, uint32_t> tilingDataTypeKeyMap = {
118 {ge::DT_INT64, 1004}, {ge::DT_INT32, 1003}, {ge::DT_INT16, 1002}, {ge::DT_INT8, 1001},137 {ge::DT_INT64, 1004}, {ge::DT_INT32, 1003}, {ge::DT_INT16, 1002}, {ge::DT_INT8, 1001},
@@ -127,50 +146,65 @@ static const std::map<ge::DataType, uint32_t> optDataTypeBitMap = {
127static const std::map<ge::DataType, uint32_t> b64DataTypeBitMap = {{ge::DT_INT64, 8}, {ge::DT_UINT64, 8}};146static const std::map<ge::DataType, uint32_t> b64DataTypeBitMap = {{ge::DT_INT64, 8}, {ge::DT_UINT64, 8}};
128} // namespace topkV2DataInfo147} // namespace topkV2DataInfo
129 148 
130-uint32_t CeilAlignDiv(int64_t a, int64_t b)149+// ==================== Helper Functions ====================
150+ 
151+ 
152+inline uint32_t GetDataTypeSize(ge::DataType dataType)
R
RRuiWang_6月6日

这些为何不用GE的接口获取字节大小

likedislike
cy_hw
cy_hw
6月6日 评论：
131{153{
132- if (b == 0) {154+ return topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second;
133- return static_cast<uint32_t>(a);
134- }
135- return static_cast<uint32_t>((a + b - 1) / b);
136}155}
137 156 
138-template <typename T>157+inline bool IsDataType64Bit(ge::DataType dataType)
139-auto CeilAlignDivMul(int64_t a, int64_t b) -> T const
140{158{
141- if (b == 0) {159+ return topkV2DataInfo::b64DataTypeBitMap.count(dataType) != 0;
142- return static_cast<T>(a);160+}
143- }161+ 
144- return static_cast<T>(((a + b - 1) / b) * b);162+inline uint32_t GetDefaultTileDataSize(ge::DataType dataType)
163+{
164+ return IsDataType64Bit(dataType) ? topkV2DataInfo::TMP_DATA_NUM_B64 :
165+ topkV2DataInfo::TMP_DATA_NUM;
166+}
167+ 
168+inline uint32_t GetSingleBlockModelDefaultTileDataSize(ge::DataType dataType)
169+{
170+ return IsDataType64Bit(dataType) ? topkV2DataInfo::SINGLE_BLOCK_DATA_NUM_B64 :
171+ topkV2DataInfo::SINGLE_BLOCK_DATA_NUM;
172+}
173+ 
174+inline uint32_t GetSingleCoreModelDefaultTileDataSize(ge::DataType dataType)
175+{
176+ return IsDataType64Bit(dataType) ? topkV2DataInfo::SINGLE_CORE_DATA_NUM_B64 :
177+ topkV2DataInfo::SINGLE_CORE_DATA_NUM;
145}178}
146 179 
147ge::graphStatus GetTopkApiTmpBufferSize(180ge::graphStatus GetTopkApiTmpBufferSize(
148 gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData, uint32_t needDataNum, int64_t kValue,181 gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData, uint32_t needDataNum, int64_t kValue,
149 bool isLargest, ge::DataType dtype, bool isSort, uint32_t nowTileSize)182 bool isLargest, ge::DataType dtype, bool isSort, uint32_t nowTileSize)
150{183{
151- uint32_t maxBufferSize = 0;184+ uint32_t aglinInnerValue =
152- uint32_t minBufferSize = 0;185+ static_cast<uint32_t>(Ops::Base::CeilAlign(static_cast<uint64_t>(needDataNum), topkV2DataInfo::AGLIN_FACTOR));
153- uint32_t aglinInnerValue = ((needDataNum + topkV2DataInfo::AGLIN_FACTOR - 1) / topkV2DataInfo::AGLIN_FACTOR) *186+
154- topkV2DataInfo::AGLIN_FACTOR;187+ uint32_t aglinKValue = (topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) ?
155- uint32_t aglinKValue = 0;188+ std::min(static_cast<int64_t>(needDataNum), kValue) :
156- if (topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) {189+ std::min(static_cast<int64_t>(nowTileSize), kValue);
157- aglinKValue = std::min(static_cast<int64_t>(needDataNum), kValue);190+
158- } else {
159- aglinKValue = std::min(static_cast<int64_t>(nowTileSize), kValue);
160- }
161 AscendC::TopKConfig topkConfig;191 AscendC::TopKConfig topkConfig;
162 topkConfig.algo = AscendC::TopKAlgo::RADIX_SELECT;192 topkConfig.algo = AscendC::TopKAlgo::RADIX_SELECT;
163 topkConfig.order = AscendC::TopKOrder::UNSET;193 topkConfig.order = AscendC::TopKOrder::UNSET;
164 topkConfig.sorted = isSort;194 topkConfig.sorted = isSort;
195+
196+ uint32_t maxBufferSize = 0;
197+ uint32_t minBufferSize = 0;
165 bool isSuccess = AscendC::GetTopKMaxMinTmpSize(198 bool isSuccess = AscendC::GetTopKMaxMinTmpSize(
166- aglinInnerValue, 1, aglinKValue, false, false, AscendC::TopKMode::TOPK_NORMAL, isLargest, dtype, topkConfig,199+ aglinInnerValue, 1, aglinKValue, false, false, AscendC::TopKMode::TOPK_NORMAL,
167- maxBufferSize, minBufferSize);200+ isLargest, dtype, topkConfig, maxBufferSize, minBufferSize);
168- OP_LOGI("TopKV2TilingForAscendC", "Need tmp buffer %u byte for ac sort topk api", maxBufferSize);201+
169- OP_LOGI(202+ OP_LOGI("TopKV2TilingForAscendC", "TopK API buffer: kValue=%ld, alignedK=%u, alignedInner=%u, bufferSize=%u",
170- "TopKV2TilingForAscendC", "Init kValue=%ld aglinKValue=%u aglinInnerValue=%u", kValue, aglinKValue,203+ kValue, aglinKValue, aglinInnerValue, maxBufferSize);
171- aglinInnerValue);204+
172- OP_CHECK_IF(205+ OP_CHECK_IF(!isSuccess, OP_LOGE(context->GetNodeName(), "Failed to get TopK API buffer size"),
173- false == isSuccess, OP_LOGE(context->GetNodeName(), "Get topk api temp buffer fail"), return ge::GRAPH_FAILED);206+ return ge::GRAPH_FAILED);
207+
174 topkTilingData.set_topkAcApiTmpBufferSize(maxBufferSize);208 topkTilingData.set_topkAcApiTmpBufferSize(maxBufferSize);
175 return ge::GRAPH_SUCCESS;209 return ge::GRAPH_SUCCESS;
176}210}
@@ -180,68 +214,67 @@ uint64_t GetTopkMultiCoreRunTimeNeedSpace(
180 uint32_t indexDtypeSize, int64_t kValue)214 uint32_t indexDtypeSize, int64_t kValue)
181{215{
182 OP_CHECK_IF(tileData == 0, OP_LOGE("TopkV2", "tileData is 0"), return ge::GRAPH_FAILED);216 OP_CHECK_IF(tileData == 0, OP_LOGE("TopkV2", "tileData is 0"), return ge::GRAPH_FAILED);
183- OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return ge::GRAPH_FAILED);
184 217 
218+ uint64_t aglinFactor = topkV2DataInfo::AGLIN_FACTOR;
185 uint32_t lastDimTileNum = (static_cast<uint32_t>(lastAxisNum) + tileData - 1) / tileData;219 uint32_t lastDimTileNum = (static_cast<uint32_t>(lastAxisNum) + tileData - 1) / tileData;
186 uint32_t lastDimTileNumTimes = (lastDimTileNum + maxCoreNum - 1) / maxCoreNum;220 uint32_t lastDimTileNumTimes = (lastDimTileNum + maxCoreNum - 1) / maxCoreNum;
221+ uint64_t lastDimTileNumTimesAlign =
222+ Ops::Base::CeilAlign(static_cast<uint64_t>(sizeof(uint32_t) * lastDimTileNumTimes), aglinFactor);
187 uint64_t initUb = indexDtypeSize * topkV2DataInfo::BIN_NUM * (lastDimTileNumTimes + 1) +223 uint64_t initUb = indexDtypeSize * topkV2DataInfo::BIN_NUM * (lastDimTileNumTimes + 1) +
188- Ops::Base::CeilAlign(224+ lastDimTileNumTimesAlign * topkV2DataInfo::CONST_TWO;
189- static_cast<uint64_t>(sizeof(uint32_t) * lastDimTileNumTimes), topkV2DataInfo::AGLIN_FACTOR) *225+ 
190- topkV2DataInfo::CONST_TWO;
191 uint32_t factor = xDtypeSize * topkV2DataInfo::CONST_TWO + indexDtypeSize + indexToDtypeSize;226 uint32_t factor = xDtypeSize * topkV2DataInfo::CONST_TWO + indexDtypeSize + indexToDtypeSize;
192 227 
193 if (tileData < kValue) {228 if (tileData < kValue) {
194 factor += xDtypeSize + indexToDtypeSize + sizeof(int32_t);229 factor += xDtypeSize + indexToDtypeSize + sizeof(int32_t);
195 } else {230 } else {
196- initUb += Ops::Base::CeilAlign(static_cast<uint64_t>(kValue * sizeof(int32_t)), topkV2DataInfo::AGLIN_FACTOR) +231+ initUb += Ops::Base::CeilAlign(static_cast<uint64_t>(kValue * sizeof(int32_t)), aglinFactor) +
197- Ops::Base::CeilAlign(static_cast<uint64_t>(kValue * xDtypeSize), topkV2DataInfo::AGLIN_FACTOR) +232+ Ops::Base::CeilAlign(static_cast<uint64_t>(kValue * xDtypeSize), aglinFactor) +
198- Ops::Base::CeilAlign(static_cast<uint64_t>(kValue * indexToDtypeSize), topkV2DataInfo::AGLIN_FACTOR);233+ Ops::Base::CeilAlign(static_cast<uint64_t>(kValue * indexToDtypeSize), aglinFactor);
199 }234 }
200- OP_LOGI(235+ OP_LOGI("TopKV2TilingForAscendC", "tileData=%u, initUb=%u, factor = %u", tileData, initUb, factor);
201- "TopKV2TilingForAscendC", "GetTopkTempBuffer tileData=%u, initUb=%u, factor = %u", tileData, initUb, factor);
202 return initUb + factor * tileData;236 return initUb + factor * tileData;
203}237}
204 238 
205-uint32_t ComputeTopkTileData(239+uint32_t ComputeTopkRadixMoreCoreTileData(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
206- gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData, 240+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computeNowTileSizeInfo)
207- topkV2DataInfo::TopkComputingNowTileSizeInfo& computingNowTileSizeInfo)
208{241{
209- ge::DataType dataType = computingNowTileSizeInfo.dataType;242+ uint32_t xDtypeSize = GetDataTypeSize(computeNowTileSizeInfo.dataType);
210- ge::DataType indicesDType = computingNowTileSizeInfo.indicesDType;243+ uint32_t indexDtypeSize = GetDataTypeSize(computeNowTileSizeInfo.indicesDType);
211- bool isLargest = computingNowTileSizeInfo.isLargest;244+ uint32_t tileData = GetDefaultTileDataSize(computeNowTileSizeInfo.dataType);
212- bool isSort = computingNowTileSizeInfo.isSort;
213- int64_t lastAxisNum = computingNowTileSizeInfo.lastAxisNum;
214- int64_t kValue = computingNowTileSizeInfo.kValue;
215- uint32_t maxCoreNum = computingNowTileSizeInfo.maxCoreNum;
216- uint64_t ubSizePlatForm = computingNowTileSizeInfo.ubSizePlatForm;
217- 
218- uint32_t xDtypeSize = topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second;
219- uint32_t indexDtypeSize = topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second;
220-
221- uint32_t tileData = (topkV2DataInfo::b64DataTypeBitMap.count(dataType) != 0) ?
222- topkV2DataInfo::TMP_DATA_NUM_B64 : topkV2DataInfo::TMP_DATA_NUM;
223 245
224 uint64_t runTimeNeedSpace = GetTopkMultiCoreRunTimeNeedSpace(246 uint64_t runTimeNeedSpace = GetTopkMultiCoreRunTimeNeedSpace(
225- lastAxisNum, tileData, maxCoreNum, xDtypeSize, indexDtypeSize, indexDtypeSize, kValue);247+ computeNowTileSizeInfo.lastAxisNum, tileData, computeNowTileSizeInfo.maxCoreNum, xDtypeSize,
226- int64_t lastDimTileNum = CeilAlignDiv(lastAxisNum, static_cast<int64_t>(tileData));248+ indexDtypeSize, indexDtypeSize, computeNowTileSizeInfo.kValue);
227- int64_t maxInputK = std::max(lastDimTileNum, kValue);249+ 
250+ int64_t lastDimTileNum =
251+ Ops::Base::CeilDiv(static_cast<uint64_t>(computeNowTileSizeInfo.lastAxisNum), static_cast<uint64_t>(tileData));
252+ int64_t maxInputK = std::max(lastDimTileNum, computeNowTileSizeInfo.kValue);
228 253
229- GetTopkApiTmpBufferSize(context, topkTilingData, tileData, maxInputK, isLargest, dataType, isSort, tileData);254+ GetTopkApiTmpBufferSize(context, topkTilingData, tileData, maxInputK, computeNowTileSizeInfo.isLargest,
255+ computeNowTileSizeInfo.dataType, computeNowTileSizeInfo.isSort, tileData);
230 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();256 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
231 257
232- while (topkAcApiNeedBuffer + runTimeNeedSpace > ubSizePlatForm) {258+ while (topkAcApiNeedBuffer + runTimeNeedSpace > computeNowTileSizeInfo.ubSizePlatForm) {
233 tileData -= topkV2DataInfo::BIN_NUM;259 tileData -= topkV2DataInfo::BIN_NUM;
234- OP_CHECK_IF(tileData <= 0, OP_LOGE("TopkV2", "tileData is less than 0."), return ge::GRAPH_FAILED);260+ if (tileData <= 0) {
235- 261+ OP_LOGD("TopKV2TilingForAscendC", "tileData is less than or equal to 0.");
236- lastDimTileNum = CeilAlignDiv(lastAxisNum, static_cast<int64_t>(tileData));262+ return tileData;
237- maxInputK = std::max(lastDimTileNum, kValue);263+ }
238- GetTopkApiTmpBufferSize(context, topkTilingData, tileData, maxInputK, isLargest, dataType, isSort, tileData);264+ 
265+ lastDimTileNum =
266+ Ops::Base::CeilDiv(static_cast<uint64_t>(computeNowTileSizeInfo.lastAxisNum), static_cast<uint64_t>(tileData));
267+ maxInputK = std::max(lastDimTileNum, computeNowTileSizeInfo.kValue);
268+ GetTopkApiTmpBufferSize(context, topkTilingData, tileData, maxInputK, computeNowTileSizeInfo.isLargest,
269+ computeNowTileSizeInfo.dataType, computeNowTileSizeInfo.isSort, tileData);
270+ 
239 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();271 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
240 runTimeNeedSpace = GetTopkMultiCoreRunTimeNeedSpace(272 runTimeNeedSpace = GetTopkMultiCoreRunTimeNeedSpace(
241- lastAxisNum, tileData, maxCoreNum, xDtypeSize, indexDtypeSize, indexDtypeSize, kValue);273+ computeNowTileSizeInfo.lastAxisNum, tileData, computeNowTileSizeInfo.maxCoreNum, xDtypeSize,
274+ indexDtypeSize, indexDtypeSize, computeNowTileSizeInfo.kValue);
242 }275 }
243 276 
244- OP_LOGI("TopKV2TilingForAscendC", "tileData=%u, ApiTempBuffer=%u", tileData, topkAcApiNeedBuffer);277+ OP_LOGI("TopKV2TilingForAscendC", "Multi-core tile: data=%u, apiBuffer=%u", tileData, topkAcApiNeedBuffer);
245 return tileData;278 return tileData;
246}279}
247 280 
@@ -257,85 +290,132 @@ bool IsLastLoopCoreUtilizationSuccess(uint32_t unsortedDimNum, uint32_t tmpOneCo
257 return true;290 return true;
258 }291 }
259 // 最后一次loop剩余待处理的轴数量/每个核处理的dim要大于0.7，确保最后一个loop有超过一半的核在处理，尽可能提高利用率292 // 最后一次loop剩余待处理的轴数量/每个核处理的dim要大于0.7，确保最后一个loop有超过一半的核在处理，尽可能提高利用率
260- if (sortLoopTimes >= topkV2DataInfo::SMALL_LOOP_LOWER_NUM && sortLoopTimes <= topkV2DataInfo::SMALL_LOOP_UPPER_NUM && 293+ bool loopTimesCondition = sortLoopTimes >= topkV2DataInfo::SMALL_LOOP_LOWER_NUM &&
261- lastLoopDimNeedCoreNum < maxCoreNum * topkV2DataInfo::LAST_LOOP_CORE_UTILIZATION) {294+ sortLoopTimes <= topkV2DataInfo::SMALL_LOOP_UPPER_NUM;
295+ bool utilizationCondition = lastLoopDimNeedCoreNum < maxCoreNum * topkV2DataInfo::LAST_LOOP_CORE_UTILIZATION;
296+ if (loopTimesCondition && utilizationCondition) {
262 return false;297 return false;
263 }298 }
264 return true;299 return true;
265}300}
266 301 
267-uint32_t ComputeMergeSortTileData(302+uint32_t GetTileDataForMergeSort(uint32_t unsortedDimNum, uint32_t maxCoreNum,
268- TopKV2TilingDataSimd& topkTilingData, ge::DataType dataType, ge::DataType indicesDType, int64_t lastAxisNum,303+ uint32_t tileMaxData, uint32_t bufferNum, uint32_t aglinNum)
269- uint32_t maxCoreNum, uint32_t unsortedDimNum, uint64_t ubSizePlatForm)
270{304{
271- uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);305+ OP_CHECK_IF(bufferNum == 0, OP_LOGE("TopkV2", "mergeSort tiling bufferNum is invalid."),
272- uint32_t indexToDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second);306+ return topkV2DataInfo::SMALL_MAX_DATA_SZIE);
273- uint32_t convertTypeSize = (dataType == ge::DT_BF16) ?307+ OP_CHECK_IF(aglinNum == 0, OP_LOGE("TopkV2", "mergeSort tiling aglinNum is invalid."),
274- topkV2DataInfo::optDataTypeBitMap.find(ge::DT_FLOAT)->second :308+ return topkV2DataInfo::SMALL_MAX_DATA_SZIE);
275- topkV2DataInfo::optDataTypeBitMap.find(dataType)->second;309+ 
310+ uint32_t tileData = topkV2DataInfo::TMP_DATA_NUM;
311+ uint32_t oneCoreRowNum = (tileData / bufferNum) / aglinNum;
312+ // 按照每个核处理默认的 tileData 来计算一个核最多能处理多少行
313+ oneCoreRowNum = (oneCoreRowNum == 0) ? 1 : oneCoreRowNum;
314+ // virUnsortedDimNeedCoreNum: 默认最少需要多少核数
315+ uint32_t virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
316+
317+ // 默认最少需要的核数如果少于总核数, 说明如果按照之前的逻辑，会有核空闲，此时就应该将 unsortedDimNum
318+ // 均摊到所有的核上处理，然后返回tileData
319+ if (virUnsortedDimNeedCoreNum < maxCoreNum) {
320+ oneCoreRowNum = (unsortedDimNum + maxCoreNum - 1) / maxCoreNum;
321+ oneCoreRowNum = (oneCoreRowNum == 0) ? 1 : oneCoreRowNum;
322+ virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
323+ tileData = oneCoreRowNum * bufferNum * aglinNum;
324+ tileData = std::min(tileData, tileMaxData - topkV2DataInfo::BIN_NUM);
325+ return tileData;
326+ }
327+ 
328+ // 按照的默认tileData来计算需要的虚拟核数比总核数还多，说明tileData切分较小，没有最大限度利用最大处理tileMaxData数据能力
329+ // 在不大于tileMaxData的条件下需要适当增加
330+ while (virUnsortedDimNeedCoreNum >= maxCoreNum && topkV2DataInfo::BIN_NUM + tileData < tileMaxData) {
331+ tileData += topkV2DataInfo::BIN_NUM;
332+ oneCoreRowNum = (tileData / bufferNum) / aglinNum;
333+ oneCoreRowNum = (oneCoreRowNum == 0) ? 1 : oneCoreRowNum;
334+ virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
335+ }
336+
337+ uint32_t tmpTileData = tileData;
338+ // 在经过上述操作后的tileData, 可能导致在多loop中, 前几个loop核数处理数据比较理想，但是存在最后一个loop
339+ // 的极少数尾轴只有1个或者少数几个核在处理，导致最后一个loop将整个处理时间拉长,此时需要平衡最后一个loop,
340+ // 将tileData适当减少，使得处理尾loop的核数占总核数达到一定比例(0.7)，确保每个loop都是均匀处理数据
341+ while (!IsLastLoopCoreUtilizationSuccess(unsortedDimNum, oneCoreRowNum, maxCoreNum)) {
342+ if (tileData < topkV2DataInfo::BIN_NUM) {
343+ OP_LOGD("TopKV2TilingForAscendC", "tileData optimization =%u", tmpTileData);
344+ return tmpTileData;
345+ }
346+ tileData -= topkV2DataInfo::BIN_NUM;
347+ oneCoreRowNum = (tileData / bufferNum) / aglinNum;
348+ oneCoreRowNum = (oneCoreRowNum == 0) ? 1 : oneCoreRowNum;
349+ virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
350+ }
351+
352+ return tileData;
353+}
354+ 
355+uint32_t ComputeMergeSortTileData(
356+ TopKV2TilingDataSimd& topkTilingData, topkV2DataInfo::TopkComputeNowTileSizeInfo& computeTileSizeInfo)
357+{
358+ 
359+ ge::DataType dataType = computeTileSizeInfo.dataType;
360+ ge::DataType indicesDType = computeTileSizeInfo.indicesDType;
361+ int64_t lastAxisNum = computeTileSizeInfo.lastAxisNum;
362+ uint32_t maxCoreNum = computeTileSizeInfo.maxCoreNum;
363+ uint32_t unsortedDimNum= computeTileSizeInfo.unsortedDimNum;
364+ uint64_t ubSizePlatForm = computeTileSizeInfo.ubSizePlatForm;
365+ 
366+ uint32_t xDtypeSize = GetDataTypeSize(dataType);
367+ uint32_t indexToDtypeSize = GetDataTypeSize(indicesDType);
368+ uint32_t convertTypeSize = (dataType == ge::DT_BF16) ? GetDataTypeSize(ge::DT_FLOAT) :
369+ GetDataTypeSize(dataType);
276 370 
277 uint32_t mergeSortAcApiNeedBuffer = topkTilingData.get_mergSortAcApiNeedBufferSize();371 uint32_t mergeSortAcApiNeedBuffer = topkTilingData.get_mergSortAcApiNeedBufferSize();
278 uint32_t initUb = ubSizePlatForm - mergeSortAcApiNeedBuffer;372 uint32_t initUb = ubSizePlatForm - mergeSortAcApiNeedBuffer;
373+ OP_LOGD("TopKV2TilingForAscendC", "merge sort mergeSortAcApiNeedBuffer=%u, ubSizePlatForm=%u, "
374+ "convertTypeSize=%u", mergeSortAcApiNeedBuffer, ubSizePlatForm, convertTypeSize);
375+ 
279 uint32_t aglinNum = (static_cast<uint32_t>(lastAxisNum) + topkV2DataInfo::AGLIN_FACTOR - 1) /376 uint32_t aglinNum = (static_cast<uint32_t>(lastAxisNum) + topkV2DataInfo::AGLIN_FACTOR - 1) /
280 topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;377 topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;
281- uint32_t oneCoreRowNumSize = initUb - aglinNum * sizeof(uint32_t) -378+ uint32_t bufferNum = lastAxisNum >= topkV2DataInfo::MERGE_SORT_DISABLE_DOUBLE_BUFFER_SIZE ?
282- aglinNum * topkV2DataInfo::CONST_TWO * convertTypeSize * topkV2DataInfo::INT64_BYTE;379+ 1 : topkV2DataInfo::CONST_TWO;
283- uint32_t oneCoreRowNumMax =380+ uint32_t initSpace = aglinNum * sizeof(uint32_t) +
284- oneCoreRowNumSize / (aglinNum * topkV2DataInfo::CONST_TWO *381+ aglinNum * topkV2DataInfo::CONST_TWO * convertTypeSize * topkV2DataInfo::INT64_BYTE;
285- (topkV2DataInfo::CONST_TWO * xDtypeSize + indexToDtypeSize + convertTypeSize));
286- uint32_t tileMaxData = oneCoreRowNumMax * aglinNum * 2;
287- OP_LOGI("TopKV2TilingForAscendC", "tileMaxData=%u, maxCoreNum=%u", tileMaxData, maxCoreNum);
288 382 
289- 383+ if (initUb < initSpace) {
290- // 思路：1.占满核,均匀分核 2.循环数尽可能小384+ OP_LOGD("TopKV2TilingForAscendC", "Not enough remaining space, initUb=%u, tensorSpace=%u,"
291- uint32_t tileData = topkV2DataInfo::TMP_DATA_NUM;385+ " bufferNum=%u", initUb, initSpace, bufferNum);
292- uint32_t oneCoreRowNum = (tileData / topkV2DataInfo::CONST_TWO) / aglinNum;386+ // 空间不足返回默认最小tileData
293- oneCoreRowNum = (oneCoreRowNum == 0 ? 1 : oneCoreRowNum);387+ return topkV2DataInfo::SMALL_TILE_DATA_NUM;
294- uint32_t virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
295- if (virUnsortedDimNeedCoreNum < maxCoreNum) {
296- // 均匀分核
297- OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return ge::GRAPH_FAILED);
298- oneCoreRowNum = (unsortedDimNum + maxCoreNum - 1) / maxCoreNum;
299- oneCoreRowNum = (oneCoreRowNum == 0 ? 1 : oneCoreRowNum);
300- virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
301- tileData = oneCoreRowNum * topkV2DataInfo::CONST_TWO * aglinNum;
302- tileData = std::min(tileData, tileMaxData - topkV2DataInfo::BIN_NUM);
303- } else {
304- // 原先占满了核则增大tileData，但必须使核占满
305- while (virUnsortedDimNeedCoreNum >= maxCoreNum && tileData < tileMaxData - topkV2DataInfo::BIN_NUM) {
306- tileData += topkV2DataInfo::BIN_NUM;
307- oneCoreRowNum = (tileData / topkV2DataInfo::CONST_TWO) / aglinNum;
308- oneCoreRowNum = (oneCoreRowNum == 0 ? 1 : oneCoreRowNum);
309- virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
310- }
311- uint32_t tmpTileData = tileData;
312- while (!IsLastLoopCoreUtilizationSuccess(unsortedDimNum, oneCoreRowNum, maxCoreNum)) {
313- tileData -= topkV2DataInfo::BIN_NUM;
314- // 若自减到0,说明没有合适的tileData,放弃均匀分核,采用前值
315- if (tileData <= 0) {
316- OP_LOGD("TopKV2TilingForAscendC", "final tileData=%u", tmpTileData);
317- return tmpTileData;
318- }
319- oneCoreRowNum = (tileData / topkV2DataInfo::CONST_TWO) / aglinNum;
320- oneCoreRowNum = (oneCoreRowNum == 0 ? 1 : oneCoreRowNum);
321- virUnsortedDimNeedCoreNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
322- }
323 }388 }
324 389 
390+ uint32_t oneCoreRowNumSize = initUb - initSpace;
391+ // 每个核最多可以处理多少行
392+ uint32_t factor = topkV2DataInfo::CONST_TWO * xDtypeSize + indexToDtypeSize + convertTypeSize;
393+ uint32_t oneCoreNeedSpace = aglinNum * bufferNum * factor;
394+ uint32_t oneCoreRowNumMax = oneCoreRowNumSize / oneCoreNeedSpace;
395+ 
396+ uint32_t tileMaxData = oneCoreRowNumMax * aglinNum * bufferNum;
397+ OP_LOGD("TopKV2TilingForAscendC", "tileMaxData=%u, maxCoreNum=%u, oneCoreRowNumMax=%d, "
398+ "oneCoreRowNumSize=%u", tileMaxData, maxCoreNum, oneCoreRowNumMax, oneCoreRowNumSize);
399+ 
400+ uint32_t tileData =
401+ GetTileDataForMergeSort(unsortedDimNum, maxCoreNum, tileMaxData, bufferNum, aglinNum);
402+
325 return tileData;403 return tileData;
326}404}
327 405 
328void SetMergeSortTmpSize(406void SetMergeSortTmpSize(
329 gert::TilingContext* context, ge::DataType dataType, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData)407 gert::TilingContext* context, ge::DataType dataType, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData)
330{408{
331- uint32_t aglinDataSize = (static_cast<uint32_t>(lastAxisNum) + topkV2DataInfo::AGLIN_FACTOR - 1) /
332- topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;
333- uint32_t dataTypeSize = (dataType == ge::DT_BF16) ? topkV2DataInfo::optDataTypeBitMap.find(ge::DT_FLOAT)->second :
334- topkV2DataInfo::optDataTypeBitMap.find(dataType)->second;
335 auto platform_info = context->GetPlatformInfo();409 auto platform_info = context->GetPlatformInfo();
336 if (nullptr == platform_info) {410 if (nullptr == platform_info) {
337 OP_LOGE("TopKV2TilingForAscendC", "platform_info is nullptr.");411 OP_LOGE("TopKV2TilingForAscendC", "platform_info is nullptr.");
338 }412 }
413+ 
414+ uint32_t aglinDataSize = (static_cast<uint32_t>(lastAxisNum) + topkV2DataInfo::AGLIN_FACTOR - 1) /
415+ topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;
416+ uint32_t dataTypeSize = (dataType == ge::DT_BF16) ? GetDefaultTileDataSize(ge::DT_FLOAT) :
417+ GetDefaultTileDataSize(dataType);
418+ 
339 auto plat = platform_ascendc::PlatformAscendC(platform_info);419 auto plat = platform_ascendc::PlatformAscendC(platform_info);
340 uint32_t dataSizeNeed = AscendC::GetConcatTmpSize(plat, aglinDataSize, dataTypeSize);420 uint32_t dataSizeNeed = AscendC::GetConcatTmpSize(plat, aglinDataSize, dataTypeSize);
341 OP_LOGI("TopKV2TilingForAscendC", "Allocal buffer mergesort element len = %ld ac merge api", lastAxisNum);421 OP_LOGI("TopKV2TilingForAscendC", "Allocal buffer mergesort element len = %ld ac merge api", lastAxisNum);
@@ -343,41 +423,44 @@ void SetMergeSortTmpSize(
343 topkTilingData.set_mergSortAcApiNeedBufferSize(dataSizeNeed);423 topkTilingData.set_mergSortAcApiNeedBufferSize(dataSizeNeed);
344}424}
345 425 
346-void TileModeSmallSizeOptim(426+void TileModeSmallSizeOptim(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
347- uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,427+ topkV2DataInfo::TopkTileInfo& topkTileInfo, topkV2DataInfo::TopkComputeNowTileSizeInfo& computeTileSizeInfo)
348- topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)
349{428{
429+ uint32_t unsortedDimNum = computeTileSizeInfo.unsortedDimNum;
430+ uint32_t maxCoreNum = computeTileSizeInfo.maxCoreNum;
431+ int64_t lastAxisNum = computeTileSizeInfo.lastAxisNum;
432+ ge::DataType dataType = computeTileSizeInfo.dataType;
433+ 
434+ SetMergeSortTmpSize(context, dataType, lastAxisNum, topkTilingData);
435+ uint32_t nowTileSize = ComputeMergeSortTileData(topkTilingData, computeTileSizeInfo);
436+ 
350 uint32_t aglinNum = (static_cast<uint32_t>(lastAxisNum) + topkV2DataInfo::AGLIN_FACTOR - 1) /437 uint32_t aglinNum = (static_cast<uint32_t>(lastAxisNum) + topkV2DataInfo::AGLIN_FACTOR - 1) /
351 topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;438 topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;
352- uint32_t oneCoreRowNum = (nowTileSize / 2) / aglinNum;439+ uint32_t bufferNum = lastAxisNum >= topkV2DataInfo::MERGE_SORT_DISABLE_DOUBLE_BUFFER_SIZE ?
353- oneCoreRowNum = (oneCoreRowNum == 0 ? 1 : oneCoreRowNum);440+ 1 : topkV2DataInfo::CONST_TWO;
441+
442+ uint32_t oneCoreRowNum = std::max((nowTileSize / bufferNum) / aglinNum, 1U);
354 uint32_t virUnsortedDimNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;443 uint32_t virUnsortedDimNum = (unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum;
355- uint32_t coreNumNeed = 0;
356- OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return);
357 uint32_t sortLoopTimes = (virUnsortedDimNum + maxCoreNum - 1) / maxCoreNum;444 uint32_t sortLoopTimes = (virUnsortedDimNum + maxCoreNum - 1) / maxCoreNum;
358- if (sortLoopTimes == 1) {445+
359- uint32_t realCoreNum = virUnsortedDimNum % maxCoreNum;446+ uint32_t realCoreNum = virUnsortedDimNum % maxCoreNum;
360- if (realCoreNum == 0) {447+ uint32_t coreNumNeed = (sortLoopTimes == 1) ? (realCoreNum == 0 ? maxCoreNum : realCoreNum) : maxCoreNum;
361- realCoreNum = maxCoreNum;448+
362- }
363- coreNumNeed = realCoreNum;
364- } else {
365- coreNumNeed = maxCoreNum;
366- }
367 topkTilingData.set_sortLoopTimes(sortLoopTimes);449 topkTilingData.set_sortLoopTimes(sortLoopTimes);
368 topkTilingData.set_lastDimTileNum(1);450 topkTilingData.set_lastDimTileNum(1);
369 topkTilingData.set_unsortedDimParallel(coreNumNeed);451 topkTilingData.set_unsortedDimParallel(coreNumNeed);
370 topkTilingData.set_lastDimNeedCore(1);452 topkTilingData.set_lastDimNeedCore(1);
371 topkTilingData.set_numTileDataSize(lastAxisNum);453 topkTilingData.set_numTileDataSize(lastAxisNum);
454+ topkTilingData.set_keyParams4(bufferNum);
455+ 
372 topkTileInfo.ubRealLoadDataNum = lastAxisNum;456 topkTileInfo.ubRealLoadDataNum = lastAxisNum;
373 topkTileInfo.coreNumNeed = coreNumNeed;457 topkTileInfo.coreNumNeed = coreNumNeed;
374 topkTileInfo.lastDimTileNum = 1;458 topkTileInfo.lastDimTileNum = 1;
375 topkTileInfo.unsortedDimParallel = coreNumNeed;459 topkTileInfo.unsortedDimParallel = coreNumNeed;
376 topkTileInfo.oneCoreRowNum = oneCoreRowNum;460 topkTileInfo.oneCoreRowNum = oneCoreRowNum;
377- OP_LOGI("TopKV2TilingForAscendC", "Small size opt mode oneCoreRowNum=%u", oneCoreRowNum);461+ 
378- OP_LOGI(462+ OP_LOGI("TopKV2TilingForAscendC", "Small size opt mode coreNumNeed=%u sortLoopTimes=%u lastAxisNum=%ld, "
379- "TopKV2TilingForAscendC", "Small size opt mode coreNumNeed=%u sortLoopTimes=%u lastAxisNum=%ld", coreNumNeed,463+ "oneCoreRowNum=%u, nowTileSize=%u.", coreNumNeed, sortLoopTimes, lastAxisNum, oneCoreRowNum, nowTileSize);
380- sortLoopTimes, lastAxisNum);
381}464}
382 465 
383uint64_t GetSingleBlockTopkRunTimeNeedSpace(466uint64_t GetSingleBlockTopkRunTimeNeedSpace(
@@ -397,28 +480,31 @@ uint64_t GetSingleBlockTopkRunTimeNeedSpace(
397 return initUb;480 return initUb;
398}481}
399 482 
400-uint32_t ComputeSingleBlockTileData(483+uint32_t ComputeSingleBlockTileData(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
401- gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData, ge::DataType dataType,484+ ge::DataType dataType, ge::DataType indicesDType, bool isLargest, bool isSort, int64_t lastAxisNum,
402- ge::DataType indicesDType, bool isLargest, bool isSort, int64_t lastAxisNum, int64_t kValue,485+ int64_t kValue, uint64_t ubSizePlatForm)
403- uint64_t ubSizePlatForm)
404{486{
405- uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);487+ uint32_t xDtypeSize = GetDataTypeSize(dataType);
406- uint32_t indexToDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second);488+ uint32_t indexToDtypeSize = GetDataTypeSize(indicesDType);
407- uint32_t tileData = (topkV2DataInfo::b64DataTypeBitMap.count(dataType) != 0) ? topkV2DataInfo::TMP_DATA_NUM_B64 :489+ uint32_t tileData = GetSingleBlockModelDefaultTileDataSize(dataType);
408- topkV2DataInfo::TMP_DATA_NUM;490+
409 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);491 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);
410 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();492 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
411- uint64_t needSpace =493+ uint64_t needSpace = GetSingleBlockTopkRunTimeNeedSpace(lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue);
412- GetSingleBlockTopkRunTimeNeedSpace(lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue);494+
413 while (topkAcApiNeedBuffer + needSpace > ubSizePlatForm) {495 while (topkAcApiNeedBuffer + needSpace > ubSizePlatForm) {
414- tileData = tileData - topkV2DataInfo::BIN_NUM;496+ tileData -= topkV2DataInfo::BIN_NUM;
497+ if (tileData < lastAxisNum) {
498+ OP_LOGI("TopKV2TilingForAscendC", "tileData is less than lastAxisNum, tileData=%u", tileData);
499+ return 0U;
500+ }
415 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);501 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);
416 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();502 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
417 needSpace = GetSingleBlockTopkRunTimeNeedSpace(lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue);503 needSpace = GetSingleBlockTopkRunTimeNeedSpace(lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue);
418 }504 }
419- OP_LOGI(505+
420- "TopKV2TilingForAscendC", "single block model tileData=%u, TempBuffer=%lu, ApiTempBuffer=%u", tileData,506+ OP_LOGI("TopKV2TilingForAscendC", "single block model tileData=%u, TempBuffer=%lu, ApiTempBuffer=%u",
421- needSpace, topkAcApiNeedBuffer);507+ tileData, needSpace, topkAcApiNeedBuffer);
422 return tileData;508 return tileData;
423}509}
424 510 
@@ -443,45 +529,39 @@ uint64_t GetTopkMultiCoreOptimModeRunTimeNeedSpace(
443 529 
444bool IsMultiCoreOptimMode(530bool IsMultiCoreOptimMode(
445 gert::TilingContext* context, uint32_t& inputNowTileSize, TopKV2TilingDataSimd& topkTilingData,531 gert::TilingContext* context, uint32_t& inputNowTileSize, TopKV2TilingDataSimd& topkTilingData,
446- topkV2DataInfo::TopkComputingNowTileSizeInfo& computingNowTileSizeInfo)532+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computeNowTileSizeInfo)
447{533{
448- // get variable534+ uint32_t xDtypeSize = GetDataTypeSize(computeNowTileSizeInfo.dataType);
449- ge::DataType dataType = computingNowTileSizeInfo.dataType;535+ uint32_t indexToDtypeSize = GetDataTypeSize(computeNowTileSizeInfo.indicesDType);
450- ge::DataType indicesDType = computingNowTileSizeInfo.indicesDType;536+ int32_t tileData = GetDefaultTileDataSize(computeNowTileSizeInfo.dataType);
451- int64_t kValue = computingNowTileSizeInfo.kValue;537+
452- bool isLargest = computingNowTileSizeInfo.isLargest;538+ if (tileData < computeNowTileSizeInfo.kValue) {
453- bool isSort = computingNowTileSizeInfo.isSort;539+ OP_LOGD("TopKV2TilingForAscendC", "K value exceeds initial tileData");
454- int64_t lastAxisNum = computingNowTileSizeInfo.lastAxisNum;
455- uint64_t ubBlockAlignSize = computingNowTileSizeInfo.ubBlockAlignSize;
456- 
457- // computing runtime local tensor need space and topk api need space
458- uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);
459- uint32_t indexToDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second);
460- int32_t tileData = (topkV2DataInfo::b64DataTypeBitMap.count(dataType) != 0) ? topkV2DataInfo::TMP_DATA_NUM_B64 :
461- topkV2DataInfo::TMP_DATA_NUM;
462- if (tileData < kValue) {
463- OP_LOGD("TopKV2TilingForAscendC", "k is greater than init tileData.");
464 return false;540 return false;
465 }541 }
466- GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);542+ GetTopkApiTmpBufferSize(context, topkTilingData, tileData, computeNowTileSizeInfo.kValue,
543+ computeNowTileSizeInfo.isLargest, computeNowTileSizeInfo.dataType, computeNowTileSizeInfo.isSort, tileData);
467 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();544 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
468 uint64_t needSpace = GetTopkMultiCoreOptimModeRunTimeNeedSpace(545 uint64_t needSpace = GetTopkMultiCoreOptimModeRunTimeNeedSpace(
469- lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue, ubBlockAlignSize);546+ computeNowTileSizeInfo.lastAxisNum, tileData, xDtypeSize, indexToDtypeSize,
547+ computeNowTileSizeInfo.kValue, computeNowTileSizeInfo.ubBlockAlignSize);
470 OP_LOGD(548 OP_LOGD(
471 "TopKV2TilingForAscendC",549 "TopKV2TilingForAscendC",
472 "multi core optim model init tileData=%u, init tempBuffer=%lu, init apiTempBuffer=%u, xDtypeSize=%u, "550 "multi core optim model init tileData=%u, init tempBuffer=%lu, init apiTempBuffer=%u, xDtypeSize=%u, "
473 "indexToDtypeSize=%u",551 "indexToDtypeSize=%u",
474 tileData, needSpace, topkAcApiNeedBuffer, xDtypeSize, indexToDtypeSize);552 tileData, needSpace, topkAcApiNeedBuffer, xDtypeSize, indexToDtypeSize);
475- while (topkAcApiNeedBuffer + needSpace > computingNowTileSizeInfo.ubSizePlatForm) {553+ while (topkAcApiNeedBuffer + needSpace > computeNowTileSizeInfo.ubSizePlatForm) {
476- tileData = tileData - topkV2DataInfo::TILE_SIZE_DECREASING_FACTOR;554+ tileData -= topkV2DataInfo::TILE_SIZE_DECREASING_FACTOR;
477- if (tileData < kValue) {555+ if (tileData < computeNowTileSizeInfo.kValue) {
478- OP_LOGD("TopKV2TilingForAscendC", "k value is greater than tilingData.");556+ OP_LOGD("TopKV2TilingForAscendC", "K value exceeds adjusted tileData");
479 return false;557 return false;
480 }558 }
481- GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);559+ GetTopkApiTmpBufferSize(context, topkTilingData, tileData, computeNowTileSizeInfo.kValue,
560+ computeNowTileSizeInfo.isLargest, computeNowTileSizeInfo.dataType, computeNowTileSizeInfo.isSort, tileData);
482 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();561 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
483 needSpace = GetTopkMultiCoreOptimModeRunTimeNeedSpace(562 needSpace = GetTopkMultiCoreOptimModeRunTimeNeedSpace(
484- lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue, ubBlockAlignSize);563+ computeNowTileSizeInfo.lastAxisNum, tileData, xDtypeSize, indexToDtypeSize,
564+ computeNowTileSizeInfo.kValue, computeNowTileSizeInfo.ubBlockAlignSize);
485 OP_LOGD(565 OP_LOGD(
486 "TopKV2TilingForAscendC",566 "TopKV2TilingForAscendC",
487 "multi core optim model now tileData=%u, now tempBuffer=%lu, now apiTempBuffer=%u.", tileData, needSpace,567 "multi core optim model now tileData=%u, now tempBuffer=%lu, now apiTempBuffer=%u.", tileData, needSpace,
@@ -489,22 +569,25 @@ bool IsMultiCoreOptimMode(
489 }569 }
490 570 
491 // 在确定正确的tileData之后，必须确保尾轴是多核模式，否则会出现多核的tiling模式，走的是singleBlock的模板571 // 在确定正确的tileData之后，必须确保尾轴是多核模式，否则会出现多核的tiling模式，走的是singleBlock的模板
492- if (tileData >= lastAxisNum) {572+ if (tileData >= computeNowTileSizeInfo.lastAxisNum) {
493- OP_LOGD("TopKV2TilingForAscendC", "tileData is greater than lastAxisNum.");573+ OP_LOGD("TopKV2TilingForAscendC", "tileData >= lastAxisNum, not suitable for multi-core");
494 return false;574 return false;
495 }575 }
496 576 
497- // compute tileNum multiply by kValue577+ // Verify K * tileNum fits within tileData
498- OP_CHECK_IF(tileData == 0, OP_LOGE("TopkV2", "tileData is 0"), return ge::GRAPH_FAILED);578+ OP_CHECK_IF(tileData == 0, OP_LOGE("TopkV2", "tileData cannot be zero"), return false);
499- uint32_t lastDimTileNum = (lastAxisNum + tileData - 1) / tileData;579+
500- uint32_t inputTopkSize = kValue * lastDimTileNum;580+ uint32_t lastDimTileNum =
581+ Ops::Base::CeilDiv(static_cast<uint64_t>(computeNowTileSizeInfo.lastAxisNum), static_cast<uint64_t>(tileData));
582+ uint32_t inputTopkSize = computeNowTileSizeInfo.kValue * lastDimTileNum;
583+
501 if (inputTopkSize <= static_cast<uint32_t>(tileData)) {584 if (inputTopkSize <= static_cast<uint32_t>(tileData)) {
502 inputNowTileSize = tileData;585 inputNowTileSize = tileData;
503- OP_LOGI(586+ OP_LOGI("TopKV2TilingForAscendC", "Multi-core optim valid: tileData=%u, topkSize=%u",
504- "TopKV2TilingForAscendC", "multi core optim model final inputNowTileSize=%u, inputTopkSize=%u",587+ inputNowTileSize, inputTopkSize);
505- inputNowTileSize, inputTopkSize);
506 return true;588 return true;
507 }589 }
590+
508 return false;591 return false;
509}592}
510 593 
@@ -564,20 +647,18 @@ uint32_t ComputeSingleCoreTileData(
564 ge::DataType indicesDType, bool isLargest, bool isSort, int64_t lastAxisNum, int64_t kValue,647 ge::DataType indicesDType, bool isLargest, bool isSort, int64_t lastAxisNum, int64_t kValue,
565 uint64_t ubSizePlatForm)648 uint64_t ubSizePlatForm)
566{649{
567- uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);650+ uint32_t xDtypeSize = GetDataTypeSize(dataType);
568- uint32_t indexToDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second);651+ uint32_t indexToDtypeSize = GetDataTypeSize(indicesDType);
569- uint32_t tileData = (topkV2DataInfo::b64DataTypeBitMap.count(dataType) != 0) ? topkV2DataInfo::TMP_DATA_NUM_B64 :652+ uint32_t tileData = GetSingleCoreModelDefaultTileDataSize(dataType);
570- topkV2DataInfo::TMP_DATA_NUM;653+
571 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);654 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);
572 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();655 uint32_t topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
573 uint64_t needSpace =656 uint64_t needSpace =
574 GetSingleCoreTopkRunTimeNeedSpace(lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue, isSort);657 GetSingleCoreTopkRunTimeNeedSpace(lastAxisNum, tileData, xDtypeSize, indexToDtypeSize, kValue, isSort);
575- OP_LOGI(658+ 
576- "TopKV2TilingForAscendC", "single core model init tileData=%u, init TempBuffer=%lu, now ApiTempBuffer=%u",
577- tileData, needSpace, topkAcApiNeedBuffer);
578 while (topkAcApiNeedBuffer + needSpace > ubSizePlatForm) {659 while (topkAcApiNeedBuffer + needSpace > ubSizePlatForm) {
579- tileData = tileData - topkV2DataInfo::BIN_NUM;660+ tileData -= topkV2DataInfo::BIN_NUM;
580- OP_CHECK_IF(tileData == 0, OP_LOGE("TopkV2", "tileData is 0"), return ge::GRAPH_FAILED);661+ OP_CHECK_IF(tileData == 0, OP_LOGE("TopkV2", "tileData is 0"), return 0);
581 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);662 GetTopkApiTmpBufferSize(context, topkTilingData, tileData, kValue, isLargest, dataType, isSort, tileData);
582 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();663 topkAcApiNeedBuffer = topkTilingData.get_topkAcApiTmpBufferSize();
583 needSpace =664 needSpace =
@@ -604,7 +685,6 @@ void TileModeSmallSize(
604 uint32_t tailBatchNumTotal = unsortedDimNum % batchNumSingleLoop;685 uint32_t tailBatchNumTotal = unsortedDimNum % batchNumSingleLoop;
605 uint32_t tailBatchNumSingleCore = 0;686 uint32_t tailBatchNumSingleCore = 0;
606 uint32_t tailBatchNum = 0;687 uint32_t tailBatchNum = 0;
607- OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return);
608 uint32_t coreNumNeed = maxCoreNum;688 uint32_t coreNumNeed = maxCoreNum;
609 if (tailBatchNumTotal != 0) {689 if (tailBatchNumTotal != 0) {
610 tailBatchNumSingleCore = tailBatchNumTotal / maxCoreNum;690 tailBatchNumSingleCore = tailBatchNumTotal / maxCoreNum;
@@ -630,12 +710,11 @@ void TileModeSmallSize(
630}710}
631 711 
632void TileModeSingleCore(712void TileModeSingleCore(
633- uint32_t unsortedDimNum, uint32_t maxCoreNum, uint32_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,713+ uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,
634 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)714 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)
635{715{
636 OP_CHECK_IF(nowTileSize == 0, OP_LOGE("TopkV2", "nowTileSize is 0"), return);716 OP_CHECK_IF(nowTileSize == 0, OP_LOGE("TopkV2", "nowTileSize is 0"), return);
637 uint32_t lastDimTileNum = (lastAxisNum + nowTileSize - 1) / nowTileSize;717 uint32_t lastDimTileNum = (lastAxisNum + nowTileSize - 1) / nowTileSize;
638- OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return);
639 uint32_t sortLoopTimes = unsortedDimNum / maxCoreNum;718 uint32_t sortLoopTimes = unsortedDimNum / maxCoreNum;
640 uint32_t tailBatchNum = unsortedDimNum % maxCoreNum;719 uint32_t tailBatchNum = unsortedDimNum % maxCoreNum;
641 uint32_t tileNum = lastAxisNum / lastDimTileNum;720 uint32_t tileNum = lastAxisNum / lastDimTileNum;
@@ -665,9 +744,6 @@ void TileModeMediumSize(
665 uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,744 uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,
666 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)745 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)
667{746{
668- if (topkTileInfo.topKOutLastAxisNum > topkV2DataInfo::MAX_K_FOR_INT64) {
669- nowTileSize /= topkV2DataInfo::CONST_TWO;
670- }
671 OP_CHECK_IF(nowTileSize == 0, OP_LOGE("TopkV2", "nowTileSize is 0"), return);747 OP_CHECK_IF(nowTileSize == 0, OP_LOGE("TopkV2", "nowTileSize is 0"), return);
672 uint32_t lastDimTileNum = (static_cast<uint32_t>(lastAxisNum) + nowTileSize - 1) / nowTileSize;748 uint32_t lastDimTileNum = (static_cast<uint32_t>(lastAxisNum) + nowTileSize - 1) / nowTileSize;
673 uint32_t unsortedDimParallel = maxCoreNum / lastDimTileNum;749 uint32_t unsortedDimParallel = maxCoreNum / lastDimTileNum;
@@ -699,9 +775,6 @@ void TileModeBigSize(
699 uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,775 uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,
700 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)776 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)
701{777{
702- if (topkTileInfo.topKOutLastAxisNum > topkV2DataInfo::MAX_K_FOR_INT64) {
703- nowTileSize /= topkV2DataInfo::CONST_TWO;
704- }
705 OP_CHECK_IF(nowTileSize == 0, OP_LOGE("TopkV2", "nowTileSize is 0"), return);778 OP_CHECK_IF(nowTileSize == 0, OP_LOGE("TopkV2", "nowTileSize is 0"), return);
706 int64_t lastDimTileNum = (lastAxisNum + nowTileSize - 1) / nowTileSize;779 int64_t lastDimTileNum = (lastAxisNum + nowTileSize - 1) / nowTileSize;
707 uint32_t coreNumNeed = static_cast<uint32_t>(std::min(static_cast<int64_t>(maxCoreNum), lastDimTileNum));780 uint32_t coreNumNeed = static_cast<uint32_t>(std::min(static_cast<int64_t>(maxCoreNum), lastDimTileNum));
@@ -716,27 +789,68 @@ void TileModeBigSize(
716 topkTileInfo.coreNumNeed = maxCoreNum;789 topkTileInfo.coreNumNeed = maxCoreNum;
717 topkTileInfo.lastDimTileNum = lastDimTileNum;790 topkTileInfo.lastDimTileNum = lastDimTileNum;
718 topkTileInfo.unsortedDimParallel = 1;791 topkTileInfo.unsortedDimParallel = 1;
719- OP_LOGI(792+ topkTileInfo.multiCoreBigModel = true; // topk radix大规模多核标志
720- "TopKV2TilingForAscendC", "Big size mode coreNumNeed=%u sortLoopTimes=%u lastAxisNum=%ld", coreNumNeed,793+ OP_LOGI("TopKV2TilingForAscendC", "Big size mode coreNumNeed=%u sortLoopTimes=%u lastAxisNum=%ld", coreNumNeed,
721 unsortedDimNum, lastAxisNum);794 unsortedDimNum, lastAxisNum);
722- OP_LOGI(
723- "TopKV2TilingForAscendC", "Big size mode lastDimTileNum=%ld unsortedDimParallel=1 lastDimRealCore=%u",
724- lastDimTileNum, coreNumNeed);
725}795}
726 796 
797+/**
798+ * Topk自身多核基数模板
799+ * 进入条件:
800+ * 1. 能计算出正确的tileSize；
801+ */
802+void TileTopkMoreCoreMode(uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum,
803+ TopKV2TilingDataSimd& topkTilingData, topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)
804+{
805+ uint32_t sortedDimParallelData = (nowTileSize * maxCoreNum) / 2;
806+ if (lastAxisNum <= sortedDimParallelData) {
807+ TileModeMediumSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
808+ } else {
809+ if (topkTileInfo.topKOutLastAxisNum > topkV2DataInfo::MAX_K_FOR_INT64) {
810+ nowTileSize /= topkV2DataInfo::CONST_TWO;
811+ }
812+ TileModeBigSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
813+ }
814+}
815+ 
816+/**
817+ * Topk自身多核基数排序优化模板
818+ * 进入条件:
819+ * 1. k * (lastAxisNum / nowTileSize) <= nowTileSize, 假设 lastAxisNum 一共需要N个核处理，每个核计算出Topk的值；
820+ * 然后将前面每个核的Topk的值集中到一个核（前提是这个核的UB能装的下），再进行一次Topk处理;
821+ */
727void TileMultiCoreOptimSize(822void TileMultiCoreOptimSize(
728 uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,823 uint32_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, TopKV2TilingDataSimd& topkTilingData,
729 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)824 topkV2DataInfo::TopkTileInfo& topkTileInfo, uint32_t nowTileSize)
730{825{
731- const uint32_t sortedDimParallelData = (nowTileSize * maxCoreNum) / 2;826+ uint32_t sortedDimParallelData = (nowTileSize * maxCoreNum) / 2;
732- if (topkTileInfo.topKOutLastAxisNum > topkV2DataInfo::MAX_K_FOR_INT64) {
733- nowTileSize *= topkV2DataInfo::CONST_TWO;
734- }
735 if (lastAxisNum <= sortedDimParallelData) {827 if (lastAxisNum <= sortedDimParallelData) {
736 TileModeMediumSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);828 TileModeMediumSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
737 } else {829 } else {
738 TileModeBigSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);830 TileModeBigSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
739 }831 }
832+ // tiling和多核基数排序保持一致, modelType不同
833+ topkTilingData.set_modeType(topkV2DataInfo::MULT_CORE_OPTIM_MODE);
834+}
835+ 
836+/**
837+ * 进入条件：
838+ * 1. 尾轴不超过1000万;
839+ * 2. topk自有的基于多核radix基数排序算能够计算出正确的UB tileSize;
840+ */
841+bool IsTopkRadixMoreCoreMode(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
842+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computeNowTileSizeInfo, uint32_t &nowTileSize) {
843+ if (computeNowTileSizeInfo.lastAxisNum >= topkV2DataInfo::SORT_AND_TOP_K_THRESHOLD) {
844+ OP_LOGD("[TopKV2Tiling] lastAxisNum exceeds 10 million.");
845+ return false;
846+ }
847+ uint32_t tmpTileSize = ComputeTopkRadixMoreCoreTileData(context, topkTilingData, computeNowTileSizeInfo);
848+ if (tmpTileSize > 0) {
849+ nowTileSize = tmpTileSize;
850+ OP_LOGD("[TopKV2Tiling]", "radix more core final tileSize=%u", nowTileSize);
851+ return true;
852+ }
853+ return false;
740}854}
741 855 
742ge::graphStatus IsValidParam(gert::TilingContext* context)856ge::graphStatus IsValidParam(gert::TilingContext* context)
@@ -785,16 +899,227 @@ ge::graphStatus IsValidParam(gert::TilingContext* context)
785 return ge::GRAPH_SUCCESS;899 return ge::GRAPH_SUCCESS;
786}900}
787 901 
788-bool IsModeSingleCore(uint32_t unsortedDimNum, uint32_t maxCoreNum)902+/**
903+ * 判断是否为 Single Block 模式
904+ *
905+ * 执行条件：
906+ * 1. lastAxisNum 能一次性装入 UB 中，即 tileSize >= lastAxisNum
907+ * 这样一个核可以处理多个 lastAxisNum（batch），充分利用 UB 空间
908+ * 2. 能计算出有效的 tileSize（大于 0）
909+ */
910+bool IsSingleBlockMode(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
911+ uint32_t& tileSize, topkV2DataInfo::TopkComputeNowTileSizeInfo& computeTileSizeInfo)
912+{
913+ // 计算 Single Block 模式所需的 tileSize
914+ uint32_t computedTileSize = ComputeSingleBlockTileData(context, topkTilingData, computeTileSizeInfo.dataType,
915+ computeTileSizeInfo.indicesDType, computeTileSizeInfo.isLargest, computeTileSizeInfo.isSort,
916+ computeTileSizeInfo.lastAxisNum, computeTileSizeInfo.kValue, computeTileSizeInfo.ubSizePlatForm);
917+
918+ // 判断是否为 Single Block 模式
919+ // 条件1：计算出的 tileSize 必须大于 0（有效）
920+ // 条件2：tileSize >= lastAxisNum（能一次性装入 UB）
921+ bool isSingleBlockMode = (computedTileSize > 0) && (computedTileSize >= computeTileSizeInfo.lastAxisNum);
922+
923+ if (isSingleBlockMode) {
924+ tileSize = computedTileSize;
925+ OP_LOGI("[TopKV2Tiling]", "Single block mode enabled: tileSize=%u >= lastAxisNum=%ld",
926+ tileSize, computeTileSizeInfo.lastAxisNum);
927+ }
928+ return isSingleBlockMode;
929+}
930+ 
931+/**
932+ * 执行条件：
933+ * 1. 数据类型bf16, float16, float32, 尾轴小于等于1024;
934+ */
935+bool IsSmallSizeMergeSortMode(ge::DataType dataType, int64_t lastAxisNum) {
936+ bool isSmallMergeSort = lastAxisNum <= topkV2DataInfo::SMALL_MAX_DATA_SZIE &&
937+ topkV2DataInfo::optDataTypeBitMap.count(dataType) != 0;
938+ return isSmallMergeSort;
939+}
940+ 
941+/**
942+ * 执行条件：
943+ * 1. 非尾轴大于等于核数, B轴均匀分核确定是有性能收益, 均匀分核场景，需要考虑尾行处理的时间与核间同步时间的均衡；
944+ * 目前测试非均匀分核性能也有提升，故不区分是否均匀分核，直接返回true，后面如果有性能走这个模板有性能劣化可以考虑这一点;
945+ * 2. UB能找到合适的tileSize;
946+ */
947+bool IsSingleCoreMode(gert::TilingContext *context, TopKV2TilingDataSimd& topkTilingData, uint32_t &initTileSize,
948+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computTileInfo)
789{949{
790- // B轴小于核数，则不走该模板950+ if (computTileInfo.unsortedDimNum < computTileInfo.maxCoreNum) {
791- if (unsortedDimNum < maxCoreNum) {951+ OP_LOGD("[TopKV2Tiling], single core model unsortedDimNum is less than maxCoreNum.");
792 return false;952 return false;
793 }953 }
794- // B轴均匀分核确定是有性能收益；954+ uint32_t nowTileSizeTmp = ComputeSingleCoreTileData(context, topkTilingData, computTileInfo.dataType,
795- // 均匀分核场景，需要考虑尾行处理的时间与核间同步时间的均衡；955+ computTileInfo.indicesDType, computTileInfo.isLargest, computTileInfo.isSort, computTileInfo.lastAxisNum,
796- // 目前测试非均匀分核性能也有提升，故不区分是否均匀分核，直接返回true，后面如果有性能走这个模板有性能劣化可以考虑这一点956+ computTileInfo.kValue, computTileInfo.ubSizePlatForm);
797- return true;957+ if (nowTileSizeTmp > 0) {
958+ initTileSize = nowTileSizeTmp;
959+ OP_LOGD("[TopKV2Tiling]", "single core final tileSize=%u", initTileSize);
960+ return true;
961+ }
962+ return false;
963+}
964+ 
965+uint32_t AlignTopkMergeMoreCoreWorkspaceElems(int64_t elementNum)
966+{
967+ if (elementNum <= 0) {
968+ return 0;
969+ }
970+ return static_cast<uint32_t>(Ops::Base::CeilAlign(
971+ static_cast<uint64_t>(elementNum * topkV2DataInfo::SORT_STRUCT_SIZE_FP32),
972+ topkV2DataInfo::AGLIN_FACTOR) / topkV2DataInfo::SORT_STRUCT_SIZE_FP32);
973+}
974+ 
975+void GetTopkMergeMoreCoreFp32(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
976+ uint32_t maxCoreNum, uint32_t unsortedDimNum, int64_t lastAxisNum, int64_t outLastAxisNum,
977+ uint32_t onceMaxElements, uint64_t ubSizePlatForm)
978+{
979+ uint32_t splitCoreNum = Ops::Base::CeilDiv(static_cast<uint64_t>(lastAxisNum),
980+ static_cast<uint64_t>(topkV2DataInfo::MERGE_MORE_CORE_ONE_CORE_DATA_SIZE));
981+ uint32_t numTileDataSize = splitCoreNum == 0 ? 0 : static_cast<uint32_t>(lastAxisNum) / splitCoreNum;
982+ uint32_t coreNumNeed = unsortedDimNum * splitCoreNum;
983+ uint32_t onceMaxElementsAlign = (onceMaxElements / topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN) *
984+ topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN;
985+ 
986+ topkTilingData.set_modeType(topkV2DataInfo::FP32_MERGE_MORE_CORE_MODE);
987+ topkTilingData.set_sortLoopTimes(1);
988+ topkTilingData.set_unsortedDimParallel(unsortedDimNum);
989+ topkTilingData.set_lastDimNeedCore(splitCoreNum);
990+ topkTilingData.set_numTileDataSize(numTileDataSize);
991+ topkTilingData.set_lastDimTileNum(splitCoreNum);
992+ topkTilingData.set_oneCoreRowNum(1);
993+ topkTilingData.set_keyParams0(onceMaxElementsAlign);
994+ topkTilingData.set_lastAxisNum(lastAxisNum);
995+ topkTilingData.set_unsortedDimNum(unsortedDimNum);
996+ topkTilingData.set_topKRealValue(outLastAxisNum);
997+ topkTilingData.set_outputLastDimValue(outLastAxisNum);
998+ topkTilingData.set_lastDimTileNumTimes(Ops::Base::CeilDiv(static_cast<uint64_t>(splitCoreNum), static_cast<uint64_t>(maxCoreNum)));
999+ 
1000+ uint32_t alignInput = AlignTopkMergeMoreCoreWorkspaceElems(lastAxisNum);
1001+ size_t usrSize = static_cast<size_t>(unsortedDimNum) * alignInput *
1002+ topkV2DataInfo::SORT_STRUCT_SIZE_FP32 * topkV2DataInfo::CONST_TWO;
1003+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
1004+ userWorkSpaceSize[0] = usrSize + topkV2DataInfo::SYS_WORK_SPACE_SIZE;
1005+ context->SetTilingKey(topkV2DataInfo::TOPK_MERGE_SORT_MORE_CORE_TILING_KEY_FLOAT);
1006+ context->SetBlockDim(coreNumNeed);
1007+ context->SetLocalMemorySize(ubSizePlatForm);
1008+ context->SetScheduleMode(1);
1009+}
1010+ 
1011+uint32_t ComputeTopkMergeMoreCoreOnceMaxElements(uint64_t ubSizePlatForm, ge::DataType indicesDType)
1012+{
1013+ uint32_t indexBytes = GetDataTypeSize(indicesDType);
1014+ uint32_t bytesPerElem = topkV2DataInfo::MERGE_MORE_CORE_LIST_MAX_NUM *
1015+ topkV2DataInfo::SORT_STRUCT_SIZE_FP32 * topkV2DataInfo::CONST_TWO;
1016+ bytesPerElem += topkV2DataInfo::MERGE_MORE_CORE_LIST_MAX_NUM * static_cast<uint32_t>(sizeof(uint32_t));
1017+ bytesPerElem += topkV2DataInfo::MERGE_MORE_CORE_LIST_MAX_NUM * static_cast<uint32_t>(sizeof(float));
1018+ if (indexBytes == topkV2DataInfo::INT64_BYTE) {
1019+ bytesPerElem += topkV2DataInfo::MERGE_MORE_CORE_LIST_MAX_NUM * indexBytes;
1020+ }
1021+ return bytesPerElem == 0 ? 0 : static_cast<uint32_t>(ubSizePlatForm / bytesPerElem);
1022+}
1023+ 
1024+ge::graphStatus TileModeFp32MoreCoreSort(gert::TilingContext *context, TopKV2TilingDataSimd& topkTilingData,
1025+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computNowTileInfo) {
1026+ 
1027+ topkTilingData.set_isLargest(computNowTileInfo.isLargest);
1028+ topkTilingData.set_isSort(computNowTileInfo.isSort);
1029+ topkTilingData.set_isInInt32Range(computNowTileInfo.isInInt32Range);
1030+ topkTilingData.set_platformCoreNum(computNowTileInfo.maxCoreNum);
1031+ 
1032+ uint32_t mergeMoreCoreOnceMaxElements = ComputeTopkMergeMoreCoreOnceMaxElements(
1033+ computNowTileInfo.ubSizePlatForm, computNowTileInfo.indicesDType);
1034+ 
1035+ GetTopkMergeMoreCoreFp32(context, topkTilingData, static_cast<uint32_t>(computNowTileInfo.maxCoreNum),
1036+ computNowTileInfo.unsortedDimNum, computNowTileInfo.lastAxisNum, computNowTileInfo.kValue,
1037+ mergeMoreCoreOnceMaxElements, computNowTileInfo.ubSizePlatForm);
1038+ topkTilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
1039+ context->GetRawTilingData()->SetDataSize(topkTilingData.GetDataSize());
1040+ OP_LOGI("TopKV2TilingForAscendC", "TopKV2 fp32 merge more-core tiling end");
1041+ return ge::GRAPH_SUCCESS;
1042+}
1043+ 
1044+void GetTopkMergeIntraCoreFp32(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
1045+ uint32_t maxCoreNum, uint32_t unsortedDimNum, int64_t lastAxisNum, int64_t outLastAxisNum,
1046+ uint32_t blockSortSize, uint32_t extractChunkSize, uint64_t ubSizePlatForm)
1047+{
1048+ uint32_t blocksPerRow = Ops::Base::CeilDiv(static_cast<uint64_t>(lastAxisNum), static_cast<uint64_t>(blockSortSize));
1049+ uint32_t alignNum = blocksPerRow * blockSortSize;
1050+ uint32_t actualCoreNum = std::min(maxCoreNum, unsortedDimNum);
1051+ uint32_t batchPerCore = Ops::Base::CeilDiv(static_cast<uint64_t>(unsortedDimNum), static_cast<uint64_t>(actualCoreNum));
1052+ uint32_t sortLoopTimes = 1;
1053+ 
1054+ topkTilingData.set_modeType(topkV2DataInfo::FP32_MERGE_INTRA_CORE_MODE);
1055+ topkTilingData.set_sortLoopTimes(sortLoopTimes);
1056+ topkTilingData.set_unsortedDimParallel(actualCoreNum);
1057+ topkTilingData.set_lastDimNeedCore(actualCoreNum);
1058+ topkTilingData.set_numTileDataSize(blockSortSize);
1059+ topkTilingData.set_lastDimTileNum(blocksPerRow);
1060+ topkTilingData.set_oneCoreRowNum(batchPerCore);
1061+ topkTilingData.set_keyParams0(batchPerCore);
1062+ topkTilingData.set_keyParams1(blockSortSize * topkV2DataInfo::CONST_TWO);
1063+ topkTilingData.set_keyParams2(alignNum * topkV2DataInfo::CONST_TWO);
1064+ topkTilingData.set_keyParams3(alignNum);
1065+ topkTilingData.set_keyParams4(extractChunkSize);
1066+ topkTilingData.set_keyParams5(blockSortSize == 0 ? 0 :
1067+ static_cast<uint32_t>(std::numeric_limits<int32_t>::max() / blockSortSize));
1068+ topkTilingData.set_lastAxisNum(lastAxisNum);
1069+ topkTilingData.set_unsortedDimNum(unsortedDimNum);
1070+ topkTilingData.set_topKRealValue(outLastAxisNum);
1071+ topkTilingData.set_outputLastDimValue(outLastAxisNum);
1072+ topkTilingData.set_lastDimTileNumTimes(Ops::Base::CeilDiv(static_cast<uint64_t>(blocksPerRow), static_cast<uint64_t>(maxCoreNum)));
1073+ 
1074+ size_t perCoreWorkspace = static_cast<size_t>(alignNum) * topkV2DataInfo::SORT_STRUCT_SIZE_FP32 *
1075+ topkV2DataInfo::CONST_TWO;
1076+ size_t usrSize = perCoreWorkspace * actualCoreNum;
1077+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
1078+ userWorkSpaceSize[0] = usrSize + topkV2DataInfo::SYS_WORK_SPACE_SIZE;
1079+ context->SetTilingKey(topkV2DataInfo::TOPK_MERGE_SORT_INTRA_CORE_TILING_KEY_FLOAT);
1080+ context->SetBlockDim(actualCoreNum);
1081+ context->SetLocalMemorySize(ubSizePlatForm);
1082+ context->SetScheduleMode(1);
1083+}
1084+ 
1085+uint32_t ComputeTopkMergeIntraCoreBlockSortSize(uint64_t ubSizePlatForm)
1086+{
1087+ constexpr uint32_t phase2BytesPerElem = topkV2DataInfo::CONST_TWO * topkV2DataInfo::SORT_STRUCT_SIZE_FP32 *
1088+ topkV2DataInfo::CONST_TWO * topkV2DataInfo::CONST_TWO;
1089+ uint32_t blockSortSize = static_cast<uint32_t>(ubSizePlatForm / phase2BytesPerElem);
1090+ return (blockSortSize / topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN) *
1091+ topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN;
1092+}
1093+ 
1094+uint32_t ComputeTopkMergeIntraCoreExtractChunkSize(uint64_t ubSizePlatForm, ge::DataType indicesDType)
1095+{
1096+ uint32_t indexBytes = GetDataTypeSize(indicesDType);
1097+ uint32_t bytesPerElem = (topkV2DataInfo::SORT_STRUCT_SIZE_FP32 + sizeof(float) +
1098+ sizeof(int32_t) + indexBytes) * topkV2DataInfo::CONST_TWO;
1099+
1100+ uint32_t extractChunkSize = static_cast<uint32_t>(ubSizePlatForm / bytesPerElem);
1101+ return (extractChunkSize / topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN) *
1102+ topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN;
1103+}
1104+ 
1105+ge::graphStatus TileModeFp32IntraCoreSort(gert::TilingContext *context, TopKV2TilingDataSimd& topkTilingData,
1106+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computNowTileInfo) {
1107+ topkTilingData.set_isLargest(computNowTileInfo.isLargest);
1108+ topkTilingData.set_isSort(computNowTileInfo.isSort);
1109+ topkTilingData.set_isInInt32Range(computNowTileInfo.isInInt32Range);
1110+ topkTilingData.set_platformCoreNum(computNowTileInfo.maxCoreNum);
1111+ 
1112+ uint32_t mergeIntraCoreBlockSortSize = ComputeTopkMergeIntraCoreBlockSortSize(computNowTileInfo.ubSizePlatForm);
1113+ uint32_t mergeIntraCoreExtractChunkSize = ComputeTopkMergeIntraCoreExtractChunkSize(
1114+ computNowTileInfo.ubSizePlatForm, computNowTileInfo.indicesDType);
1115+ 
1116+ GetTopkMergeIntraCoreFp32(context, topkTilingData, static_cast<uint32_t>(computNowTileInfo.maxCoreNum),
1117+ computNowTileInfo.unsortedDimNum, computNowTileInfo.lastAxisNum, computNowTileInfo.kValue,
1118+ mergeIntraCoreBlockSortSize, mergeIntraCoreExtractChunkSize, computNowTileInfo.ubSizePlatForm);
1119+ topkTilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
1120+ context->GetRawTilingData()->SetDataSize(topkTilingData.GetDataSize());
1121+ OP_LOGI("TopKV2TilingForAscendC", "TopKV2 fp32 merge intra-core tiling end");
1122+ return ge::GRAPH_SUCCESS;
798}1123}
799 1124 
800// sort核间模板tiling计算相关函数1125// sort核间模板tiling计算相关函数
@@ -833,9 +1158,12 @@ void AdjTmpUb(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t tileData, uin
833void ComputeTileDataOne(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t lastDimTileNum, uint32_t ubExtra,1158void ComputeTileDataOne(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t lastDimTileNum, uint32_t ubExtra,
834 uint32_t &tileData, uint32_t tileFactor)1159 uint32_t &tileData, uint32_t tileFactor)
835{1160{
836- uint32_t allCore = CeilAlignDivMul<uint32_t>(int64_t(lastDimTileNum), int64_t(sortTileInfo.maxCoreNum));1161+ uint32_t allCore = static_cast<uint32_t>(
837- uint32_t newTileData = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(allCore));1162+ Ops::Base::CeilAlign(static_cast<uint64_t>(lastDimTileNum), static_cast<uint64_t>(sortTileInfo.maxCoreNum)));
838- tileData = CeilAlignDivMul<uint32_t>(int64_t(newTileData), int64_t(topkV2DataInfo::BIN_NUM));1163+ uint32_t newTileData = static_cast<uint32_t>(
1164+ Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.sortAxisNum), static_cast<uint64_t>(allCore)));
1165+ tileData = static_cast<uint32_t>(
1166+ Ops::Base::CeilAlign(static_cast<uint64_t>(newTileData), static_cast<uint64_t>(topkV2DataInfo::BIN_NUM)));
839 tileData = std::max(tileData, topkV2DataInfo::SMALL_TILE_DATA_NUM);1167 tileData = std::max(tileData, topkV2DataInfo::SMALL_TILE_DATA_NUM);
840 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);1168 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
841 AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);1169 AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
@@ -847,8 +1175,8 @@ bool NeedAdjTileData(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t &tileD
847{1175{
848 if (sortTileInfo.unSortDimNum == int64_t(1) && lastDimTileNum == uint32_t(1)) {1176 if (sortTileInfo.unSortDimNum == int64_t(1) && lastDimTileNum == uint32_t(1)) {
849 OP_LOGI("RadixSortTiling", "unSortDimNum and lastDimTileNum is 1");1177 OP_LOGI("RadixSortTiling", "unSortDimNum and lastDimTileNum is 1");
850- uint32_t newTileData = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(sortTileInfo.maxCoreNum));1178+ uint32_t newTileData = static_cast<uint32_t>(Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.sortAxisNum), static_cast<uint64_t>(sortTileInfo.maxCoreNum)));
851- newTileData = CeilAlignDivMul<uint32_t>(int64_t(newTileData), int64_t(topkV2DataInfo::BIN_NUM));1179+ newTileData = static_cast<uint32_t>(Ops::Base::CeilAlign(static_cast<uint64_t>(newTileData), static_cast<uint64_t>(topkV2DataInfo::BIN_NUM)));
852 tileData = std::max(newTileData, topkV2DataInfo::SMALL_TILE_DATA_NUM);1180 tileData = std::max(newTileData, topkV2DataInfo::SMALL_TILE_DATA_NUM);
853 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);1181 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
854 AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);1182 AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
@@ -865,7 +1193,7 @@ bool NeedAdjTileData(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t &tileD
865 OP_LOGI("RadixSortTiling", "unSortDimNum greater than 1,and unSortDimNum small and lastDimTileNum is one");1193 OP_LOGI("RadixSortTiling", "unSortDimNum greater than 1,and unSortDimNum small and lastDimTileNum is one");
866 uint32_t hCore = sortTileInfo.maxCoreNum / static_cast<uint32_t>(sortTileInfo.unSortDimNum);1194 uint32_t hCore = sortTileInfo.maxCoreNum / static_cast<uint32_t>(sortTileInfo.unSortDimNum);
867 uint32_t hTileData = static_cast<uint32_t>(sortTileInfo.sortAxisNum) / hCore;1195 uint32_t hTileData = static_cast<uint32_t>(sortTileInfo.sortAxisNum) / hCore;
868- tileData = CeilAlignDivMul<uint32_t>(int64_t(hTileData), int64_t(topkV2DataInfo::BIN_NUM));1196+ tileData = static_cast<uint32_t>(Ops::Base::CeilAlign(static_cast<uint64_t>(hTileData), static_cast<uint64_t>(topkV2DataInfo::BIN_NUM)));
869 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);1197 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
870 AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);1198 AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
871 return tileData;1199 return tileData;
@@ -874,15 +1202,15 @@ bool NeedAdjTileData(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t &tileD
874 // b大于1且h轴循环次数小于总核数，也就是b轴核数大于11202 // b大于1且h轴循环次数小于总核数，也就是b轴核数大于1
875 OP_LOGI("RadixSortTiling", "unSortDimNum is one, lastDimTileNum greater than one");1203 OP_LOGI("RadixSortTiling", "unSortDimNum is one, lastDimTileNum greater than one");
876 int64_t newTileData = sortTileInfo.sortAxisNum / int64_t(lastDimTileNum);1204 int64_t newTileData = sortTileInfo.sortAxisNum / int64_t(lastDimTileNum);
877- tileData = CeilAlignDivMul<uint32_t>(newTileData, int64_t(topkV2DataInfo::BIN_NUM));1205+ tileData = static_cast<uint32_t>(Ops::Base::CeilAlign(static_cast<uint64_t>(newTileData), static_cast<uint64_t>(topkV2DataInfo::BIN_NUM)));
878- lastDimTileNum = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(tileData));1206+ lastDimTileNum = static_cast<uint32_t>(Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.sortAxisNum), static_cast<uint64_t>(tileData)));
879 uint32_t bCore = lastDimTileNum == 0 ? sortTileInfo.maxCoreNum : sortTileInfo.maxCoreNum / lastDimTileNum;1207 uint32_t bCore = lastDimTileNum == 0 ? sortTileInfo.maxCoreNum : sortTileInfo.maxCoreNum / lastDimTileNum;
880 if (lastDimTileNum < sortTileInfo.maxCoreNum && sortTileInfo.unSortDimNum < int64_t(sortTileInfo.maxCoreNum)) {1208 if (lastDimTileNum < sortTileInfo.maxCoreNum && sortTileInfo.unSortDimNum < int64_t(sortTileInfo.maxCoreNum)) {
881 if (sortTileInfo.unSortDimNum < int64_t(bCore)) {1209 if (sortTileInfo.unSortDimNum < int64_t(bCore)) {
882 bCore = static_cast<uint32_t>(sortTileInfo.unSortDimNum);1210 bCore = static_cast<uint32_t>(sortTileInfo.unSortDimNum);
883 uint32_t hCore = sortTileInfo.maxCoreNum / bCore;1211 uint32_t hCore = sortTileInfo.maxCoreNum / bCore;
884- uint32_t tileDataNew = CeilAlignDiv(int64_t(sortTileInfo.sortAxisNum), int64_t(hCore));1212+ uint32_t tileDataNew = static_cast<uint32_t>(Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.sortAxisNum), static_cast<uint64_t>(hCore)));
885- tileData = CeilAlignDivMul<uint32_t>(int64_t(tileDataNew), int64_t(topkV2DataInfo::BIN_NUM));1213+ tileData = static_cast<uint32_t>(Ops::Base::CeilAlign(static_cast<uint64_t>(tileDataNew), static_cast<uint64_t>(topkV2DataInfo::BIN_NUM)));
886 }1214 }
887 }1215 }
888 if (bCore == static_cast<uint32_t>(1) && lastDimTileNum < sortTileInfo.maxCoreNum) {1216 if (bCore == static_cast<uint32_t>(1) && lastDimTileNum < sortTileInfo.maxCoreNum) {
@@ -920,7 +1248,7 @@ uint32_t ComputeTileData(topkV2DataInfo::SortTileInfo &sortTileInfo)
920 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);1248 SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
921 tmpUbSize = sortTileInfo.tmpUbSize;1249 tmpUbSize = sortTileInfo.tmpUbSize;
922 }1250 }
923- uint32_t lastDimTileNum = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(tileData));1251+ uint32_t lastDimTileNum = Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.sortAxisNum), static_cast<uint64_t>(tileData));
924 OP_LOGI("RadixSortTiling", "tileData %u, lastDimTileNum %u, tmpUbSize %u", tileData, lastDimTileNum, tmpUbSize);1252 OP_LOGI("RadixSortTiling", "tileData %u, lastDimTileNum %u, tmpUbSize %u", tileData, lastDimTileNum, tmpUbSize);
925 bool smallTile = (sortTileInfo.sortAxisNum <= static_cast<int64_t>(topkV2DataInfo::SMALL_TILE_DATA_NUM)) &&1253 bool smallTile = (sortTileInfo.sortAxisNum <= static_cast<int64_t>(topkV2DataInfo::SMALL_TILE_DATA_NUM)) &&
926 lastDimTileNum == uint32_t(1);1254 lastDimTileNum == uint32_t(1);
@@ -943,29 +1271,35 @@ void ComputeWorkSpace(gert::TilingContext *context, topkV2DataInfo::SortTileInfo
943 dtypeSizeWk = static_cast<uint32_t>(sizeof(int64_t));1271 dtypeSizeWk = static_cast<uint32_t>(sizeof(int64_t));
944 }1272 }
945 size_t excusiveBinsGmWkSize = static_cast<size_t>(sortTileInfo.keyParams1) * sortTileInfo.keyParams4 * dtypeSizeWk;1273 size_t excusiveBinsGmWkSize = static_cast<size_t>(sortTileInfo.keyParams1) * sortTileInfo.keyParams4 * dtypeSizeWk;
946- excusiveBinsGmWkSize = CeilAlignDivMul<size_t>(int64_t(excusiveBinsGmWkSize), int64_t(sortTileInfo.blockUbSize));1274+ excusiveBinsGmWkSize = static_cast<size_t>(
1275+ Ops::Base::CeilAlign(static_cast<uint64_t>(excusiveBinsGmWkSize), static_cast<uint64_t>(sortTileInfo.blockUbSize)));
947 1276 
948- size_t globalHistGmWkSize =1277+ size_t globalHistGmWkSize =
949- static_cast<size_t>(sortTileInfo.keyParams3) * sortTileInfo.keyParams2 * sortTileInfo.keyParams0 * dtypeSizeWk;1278+ static_cast<size_t>(sortTileInfo.keyParams3) * sortTileInfo.keyParams2 * sortTileInfo.keyParams0 * dtypeSizeWk;
950- globalHistGmWkSize = CeilAlignDivMul<size_t>(int64_t(globalHistGmWkSize), int64_t(sortTileInfo.blockUbSize));1279+ globalHistGmWkSize = static_cast<size_t>(
1280+ Ops::Base::CeilAlign(static_cast<uint64_t>(globalHistGmWkSize), static_cast<uint64_t>(sortTileInfo.blockUbSize)));
951 1281 
952 size_t outIdxDbWK = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel * dtypeSizeWk;1282 size_t outIdxDbWK = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel * dtypeSizeWk;
953- outIdxDbWK = CeilAlignDivMul<size_t>(int64_t(outIdxDbWK), int64_t(sortTileInfo.blockUbSize));1283+ outIdxDbWK = static_cast<size_t>(
1284+ Ops::Base::CeilAlign(static_cast<uint64_t>(outIdxDbWK), static_cast<uint64_t>(sortTileInfo.blockUbSize)));
954 1285 
955- size_t sortOutIdxGMWK = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel *1286+ size_t sortOutIdxGMWK = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel *
956- sortTileInfo.y2DtypeSize;1287+ sortTileInfo.y2DtypeSize;
957- sortOutIdxGMWK = CeilAlignDivMul<size_t>(int64_t(sortOutIdxGMWK), int64_t(sortTileInfo.blockUbSize));1288+ sortOutIdxGMWK = static_cast<size_t>(
1289+ Ops::Base::CeilAlign(static_cast<uint64_t>(sortOutIdxGMWK), static_cast<uint64_t>(sortTileInfo.blockUbSize)));
958 1290 
959 size_t histTileGmWk = static_cast<size_t>(sortTileInfo.lastDimTileNum) * topkV2DataInfo::BIN_NUM *1291 size_t histTileGmWk = static_cast<size_t>(sortTileInfo.lastDimTileNum) * topkV2DataInfo::BIN_NUM *
960 sortTileInfo.unsortedDimParallel * sizeof(int16_t) * topkV2DataInfo::CONST_2;1292 sortTileInfo.unsortedDimParallel * sizeof(int16_t) * topkV2DataInfo::CONST_2;
961 1293 
962 size_t xB8GmWkSize = static_cast<size_t>(sortTileInfo.lastDimTileNum) * sortTileInfo.numTileDataSize *1294 size_t xB8GmWkSize = static_cast<size_t>(sortTileInfo.lastDimTileNum) * sortTileInfo.numTileDataSize *
963 sortTileInfo.unsortedDimParallel;1295 sortTileInfo.unsortedDimParallel;
964- xB8GmWkSize = CeilAlignDivMul<size_t>(int64_t(xB8GmWkSize), int64_t(sortTileInfo.blockUbSize));1296+ xB8GmWkSize = static_cast<size_t>(
1297+ Ops::Base::CeilAlign(static_cast<uint64_t>(xB8GmWkSize), static_cast<uint64_t>(sortTileInfo.blockUbSize)));
965 1298 
966 size_t outValueDbWKSize = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel *1299 size_t outValueDbWKSize = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel *
967 sortTileInfo.dtypeSize *topkV2DataInfo::CONST_2;1300 sortTileInfo.dtypeSize *topkV2DataInfo::CONST_2;
968- outValueDbWKSize = CeilAlignDivMul<size_t>(int64_t(outValueDbWKSize), int64_t(sortTileInfo.blockUbSize));1301+ outValueDbWKSize = static_cast<size_t>(
1302+ Ops::Base::CeilAlign(static_cast<uint64_t>(outValueDbWKSize), static_cast<uint64_t>(sortTileInfo.blockUbSize)));
969 1303 
970 OP_LOGI("RadixSortTiling",1304 OP_LOGI("RadixSortTiling",
971 "excusiveBinsGmWkSize %lu, globalHistGmWkSize %lu, outIdxDbWK %lu, sortOutIdxGMWK %lu, histTileGmWk %lu,"1305 "excusiveBinsGmWkSize %lu, globalHistGmWkSize %lu, outIdxDbWK %lu, sortOutIdxGMWK %lu, histTileGmWk %lu,"
@@ -982,7 +1316,8 @@ ge::graphStatus GetRadixSortMoreCore(gert::TilingContext *context, topkV2DataInf
982{1316{
983 sortTileInfo.ubSize = sortTileInfo.ubSize - topkV2DataInfo::SIMT_UB;1317 sortTileInfo.ubSize = sortTileInfo.ubSize - topkV2DataInfo::SIMT_UB;
984 uint32_t tileData = ComputeTileData(sortTileInfo);1318 uint32_t tileData = ComputeTileData(sortTileInfo);
985- uint32_t lastDimTileNum = CeilAlignDiv(int64_t(sortTileInfo.sortAxisNum), int64_t(tileData));1319+ uint32_t lastDimTileNum =
1320+ Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.sortAxisNum), static_cast<uint64_t>(tileData));
986 if (sortTileInfo.maxCoreNum <= lastDimTileNum) {1321 if (sortTileInfo.maxCoreNum <= lastDimTileNum) {
987 sortTileInfo.unsortedDimParallel = static_cast<uint32_t>(1);1322 sortTileInfo.unsortedDimParallel = static_cast<uint32_t>(1);
988 } else {1323 } else {
@@ -993,7 +1328,8 @@ ge::graphStatus GetRadixSortMoreCore(gert::TilingContext *context, topkV2DataInf
993 }1328 }
994 }1329 }
995 sortTileInfo.numTileDataSize = tileData;1330 sortTileInfo.numTileDataSize = tileData;
996- sortTileInfo.sortLoopTimes = CeilAlignDiv(int64_t(sortTileInfo.unSortDimNum), int64_t(sortTileInfo.unsortedDimParallel));1331+ sortTileInfo.sortLoopTimes =
1332+ Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.unSortDimNum), static_cast<uint64_t>(sortTileInfo.unsortedDimParallel));
997 sortTileInfo.lastDimNeedCore = std::min(sortTileInfo.maxCoreNum, lastDimTileNum);1333 sortTileInfo.lastDimNeedCore = std::min(sortTileInfo.maxCoreNum, lastDimTileNum);
998 sortTileInfo.coreNumNeed = sortTileInfo.unsortedDimParallel * sortTileInfo.lastDimNeedCore;1334 sortTileInfo.coreNumNeed = sortTileInfo.unsortedDimParallel * sortTileInfo.lastDimNeedCore;
999 sortTileInfo.lastDimTileNum = lastDimTileNum;1335 sortTileInfo.lastDimTileNum = lastDimTileNum;
@@ -1005,18 +1341,23 @@ ge::graphStatus GetRadixSortMoreCore(gert::TilingContext *context, topkV2DataInf
1005 uint32_t allNumGloblHist = topkV2DataInfo::BIN_NUM * lastDimTileNum * sortTileInfo.dtypeSize *1341 uint32_t allNumGloblHist = topkV2DataInfo::BIN_NUM * lastDimTileNum * sortTileInfo.dtypeSize *
1006 sortTileInfo.unsortedDimParallel;1342 sortTileInfo.unsortedDimParallel;
1007 uint32_t allNumExcusiveBin = topkV2DataInfo::BIN_NUM * sortTileInfo.dtypeSize * sortTileInfo.unsortedDimParallel;1343 uint32_t allNumExcusiveBin = topkV2DataInfo::BIN_NUM * sortTileInfo.dtypeSize * sortTileInfo.unsortedDimParallel;
1008- uint32_t oneCoreSize = CeilAlignDiv(int64_t(allNumGloblHist), int64_t(sortTileInfo.coreNumNeed));1344+ uint32_t oneCoreSize =
1345+ Ops::Base::CeilDiv(static_cast<uint64_t>(allNumGloblHist), static_cast<uint64_t>(sortTileInfo.coreNumNeed));
1009 sortTileInfo.keyParams5 =1346 sortTileInfo.keyParams5 =
1010 std::max(static_cast<int64_t>(oneCoreSize), static_cast<int64_t>(sortTileInfo.blockUbSize));1347 std::max(static_cast<int64_t>(oneCoreSize), static_cast<int64_t>(sortTileInfo.blockUbSize));
1011- sortTileInfo.keyParams0 = CeilAlignDiv(int64_t(allNumGloblHist), int64_t(sortTileInfo.keyParams5));1348+ sortTileInfo.keyParams0 =
1012- sortTileInfo.keyParams3 = CeilAlignDiv(int64_t(sortTileInfo.keyParams5), int64_t(ubSizeNum));1349+ Ops::Base::CeilDiv(static_cast<uint64_t>(allNumGloblHist), static_cast<uint64_t>(sortTileInfo.keyParams5));
1350+ sortTileInfo.keyParams3 =
1351+ Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.keyParams5), static_cast<uint64_t>(ubSizeNum));
1013 sortTileInfo.keyParams2 = sortTileInfo.keyParams5 > ubSizeNum ? ubSizeNum : sortTileInfo.keyParams5;1352 sortTileInfo.keyParams2 = sortTileInfo.keyParams5 > ubSizeNum ? ubSizeNum : sortTileInfo.keyParams5;
1014 1353 
1015- uint32_t oneCoreSize1 = CeilAlignDiv(int64_t(allNumExcusiveBin), int64_t(sortTileInfo.coreNumNeed));1354+ uint32_t oneCoreSize1 =
1355+ Ops::Base::CeilDiv(static_cast<uint64_t>(allNumExcusiveBin), static_cast<uint64_t>(sortTileInfo.coreNumNeed));
1016 sortTileInfo.keyParams4 =1356 sortTileInfo.keyParams4 =
1017- std::max(static_cast<int64_t>(oneCoreSize1), static_cast<int64_t>(sortTileInfo.blockUbSize));1357+ std::max(static_cast<uint64_t>(oneCoreSize1), static_cast<uint64_t>(sortTileInfo.blockUbSize));
1018 1358 
1019- sortTileInfo.keyParams1 = CeilAlignDiv(int64_t(allNumExcusiveBin), int64_t(sortTileInfo.keyParams4));1359+ sortTileInfo.keyParams1 =
1360+ Ops::Base::CeilDiv(static_cast<uint64_t>(allNumExcusiveBin), static_cast<uint64_t>(sortTileInfo.keyParams4));
1020 1361 
1021 // 取前k个结果相关流程的tile计算1362 // 取前k个结果相关流程的tile计算
1022 uint32_t avilableUbSize = (sortTileInfo.ubSize - 1) / topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;1363 uint32_t avilableUbSize = (sortTileInfo.ubSize - 1) / topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;
@@ -1024,8 +1365,9 @@ ge::graphStatus GetRadixSortMoreCore(gert::TilingContext *context, topkV2DataInf
1024 OP_LOGE("TopKV2", "sortAndTopK Tiling avilableUbSize is zero"), return ge::GRAPH_FAILED);1365 OP_LOGE("TopKV2", "sortAndTopK Tiling avilableUbSize is zero"), return ge::GRAPH_FAILED);
1025 auto dataType = context->GetInputDesc(0)->GetDataType();1366 auto dataType = context->GetInputDesc(0)->GetDataType();
1026 auto indicesDType = context->GetOutputDesc(1)->GetDataType();1367 auto indicesDType = context->GetOutputDesc(1)->GetDataType();
1027- uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);1368+
1028- uint32_t indexToDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second);1369+ uint32_t xDtypeSize = GetDataTypeSize(dataType);
1370+ uint32_t indexToDtypeSize = GetDataTypeSize(indicesDType);
1029 uint32_t kGetDtypeSize = std::max(xDtypeSize, indexToDtypeSize);1371 uint32_t kGetDtypeSize = std::max(xDtypeSize, indexToDtypeSize);
1030 OP_CHECK_IF(kGetDtypeSize == 0, OP_LOGE("GetRadixSortMoreCore", "kGetDtypeSize is zero"), return ge::GRAPH_FAILED);1372 OP_CHECK_IF(kGetDtypeSize == 0, OP_LOGE("GetRadixSortMoreCore", "kGetDtypeSize is zero"), return ge::GRAPH_FAILED);
1031 sortTileInfo.tileDataSize = avilableUbSize / kGetDtypeSize;1373 sortTileInfo.tileDataSize = avilableUbSize / kGetDtypeSize;
@@ -1154,6 +1496,32 @@ void PrintTilindDataSort(gert::TilingContext *context, topkV2DataInfo::SortTileI
1154 return;1496 return;
1155}1497}
1156 1498 
1499+ge::graphStatus TileModeSortAndTopK(gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData,
1500+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computeNowTileSizeInfo) {
1501+ topkV2DataInfo::SortTileInfo sortTileInfo;
1502+ 
1503+ OP_CHECK_IF(SortCheckParams(context, sortTileInfo) != ge::GRAPH_SUCCESS,
1504+ OP_LOGE(context->GetNodeName(), "sort and topk check params failed"), return ge::GRAPH_FAILED);
1505+ topkTilingData.set_modeType(topkV2DataInfo::SORT_AND_TOP_K_MODE);
1506+ OP_LOGI("[TopKV2Tiling]", "topkTilingData.set_modeType is: %u, SORT_AND_TOP_K_MODE: %u",
1507+ topkTilingData.get_modeType() , topkV2DataInfo::SORT_AND_TOP_K_MODE);
1508+ sortTileInfo.maxCoreNum = static_cast<uint32_t>(computeNowTileSizeInfo.maxCoreNum);
1509+ sortTileInfo.isDescend = static_cast<bool>(computeNowTileSizeInfo.isLargest);
1510+ sortTileInfo.isInt32 = static_cast<uint32_t>(computeNowTileSizeInfo.lastAxisNum <= topkV2DataInfo::INT32_MAX_RANGE_VALUE_FOR_SORT);
1511+ sortTileInfo.topKRealValue = computeNowTileSizeInfo.kValue;
1512+ OP_CHECK_IF(GetRadixSortMoreCore(context, sortTileInfo) != ge::GRAPH_SUCCESS,
1513+ OP_LOGE(context->GetNodeName(), "Get RadixSortMoreCore tiling failed"), return ge::GRAPH_FAILED);
1514+ auto dataTypeKey = topkV2DataInfo::tilingDataTypeKeyMap.find(computeNowTileSizeInfo.dataType)->second;
1515+ context->SetTilingKey(dataTypeKey);
1516+ context->SetBlockDim(sortTileInfo.coreNumNeed);
1517+ context->SetLocalMemorySize(sortTileInfo.ubSize);
1518+ FillTilingDataSort(context, sortTileInfo, topkTilingData);
1519+ PrintTilindDataSort(context, sortTileInfo);
1520+ // sortAndTopK模板核心是Sort，不需要后续Topk相关的tiling计算过程
1521+ OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling end");
1522+ return ge::GRAPH_SUCCESS;
1523+}
1524+ 
1157bool needSortWithIndex(TopKV2TilingDataSimd& topkTilingData, bool isSorted, ge::DataType dataType)1525bool needSortWithIndex(TopKV2TilingDataSimd& topkTilingData, bool isSorted, ge::DataType dataType)
1158{1526{
1159 if (isSorted && topkTilingData.get_modeType() == topkV2DataInfo::MULT_CORE_MODE) {1527 if (isSorted && topkTilingData.get_modeType() == topkV2DataInfo::MULT_CORE_MODE) {
@@ -1173,19 +1541,106 @@ bool needSortWithIndex(TopKV2TilingDataSimd& topkTilingData, bool isSorted, ge::
1173 return false;1541 return false;
1174}1542}
1175 1543 
1544+/**
1545+ * 执行条件：
1546+ * 1. 仅处理float32数据类型, 4096 <= lastAxisNum <= 32768;
1547+ * 2. splitCoreNum = ceil(lastAxisNum / 2048) > 1 且非尾轴数量 unsortedDimNum * splitCoreNum <= maxCoreNum 且 K > 0;
1548+ */
1549+bool IsTopkMergeSortMoreCoreFp32Mode(topkV2DataInfo::TopkComputeNowTileSizeInfo& computNowTileInfo)
1550+{
1551+ if (computNowTileInfo.dataType != ge::DT_FLOAT || computNowTileInfo.kValue <= 0) {
1552+ return false;
1553+ }
1554+ double ratio = static_cast<double>(computNowTileInfo.kValue) / computNowTileInfo.lastAxisNum;
1555+ if (ratio < topkV2DataInfo::FP32_K_LAST_AXIS_LOWER_RATIO || topkV2DataInfo::FP32_K_LAST_AXIS_UPPER_RATIO <= ratio) {
1556+ return false;
1557+ }
1558+ 
1559+ uint32_t onceMaxElements =
1560+ ComputeTopkMergeMoreCoreOnceMaxElements(computNowTileInfo.ubSizePlatForm, computNowTileInfo.indicesDType);
1561+ if (computNowTileInfo.maxCoreNum == 0 || computNowTileInfo.unsortedDimNum == 0 ||
1562+ onceMaxElements < topkV2DataInfo::MERGE_INTRA_CORE_SORT_ALIGN) {
1563+ return false;
1564+ }
1565+ uint32_t splitCoreNum = Ops::Base::CeilDiv(static_cast<uint64_t>(computNowTileInfo.lastAxisNum),
1566+ static_cast<uint64_t>(topkV2DataInfo::MERGE_MORE_CORE_ONE_CORE_DATA_SIZE));
1567+ if (splitCoreNum <= 1 || splitCoreNum > computNowTileInfo.maxCoreNum) {
1568+ return false;
1569+ }
1570+ if (static_cast<uint64_t>(computNowTileInfo.unsortedDimNum) * splitCoreNum <= computNowTileInfo.maxCoreNum) {
1571+ OP_LOGD("TopKV2Tiling", "float32 more core model onceMaxElements = %d.", onceMaxElements);
1572+ return true;
1573+ }
1574+ return false;
1575+}
1576+ 
1577+/**
1578+ * 执行条件：
1579+ * 1. MoreCore 路由未命中, 既 IsTopkMergeSortMoreCoreFp32Mode 条件不成立;
1580+ * 2. dtype == fp32 且 尾轴 > 4096 且 blocksPerRow = ceil(N / blockSortSize) <= 256 且 K > 0
1581+ *
1582+ */
1583+bool IsTopkMergeSortIntraCoreFp32Mode(topkV2DataInfo::TopkComputeNowTileSizeInfo& computTileInfo)
1584+{
1585+ if (computTileInfo.dataType != ge::DT_FLOAT ||
1586+ computTileInfo.unsortedDimNum < computTileInfo.maxCoreNum / 2) {
1587+ return false;
1588+ }
1589+ if (computTileInfo.kValue <= 0) {
1590+ return false;
1591+ }
1592+ double ratio = static_cast<double>(computTileInfo.kValue) / computTileInfo.lastAxisNum;
1593+ if (ratio < topkV2DataInfo::FP32_K_LAST_AXIS_LOWER_RATIO ||
1594+ topkV2DataInfo::FP32_K_LAST_AXIS_UPPER_RATIO <= ratio) {
1595+ return false;
1596+ }
1597+ if (!computTileInfo.isSort && computTileInfo.kValue <= topkV2DataInfo::SUPPORT_SORT_MAX_SIZE) {
1598+ return false;
1599+ }
1600+ 
1601+ uint32_t blockSortSize = ComputeTopkMergeIntraCoreBlockSortSize(computTileInfo.ubSizePlatForm);
1602+ uint32_t extractChunkSize = ComputeTopkMergeIntraCoreExtractChunkSize(computTileInfo.ubSizePlatForm,
1603+ computTileInfo.indicesDType);
1604+ if (blockSortSize == 0 || extractChunkSize == 0) {
1605+ return false;
1606+ }
1607+ OP_LOGI("IsTopkMergeIntraCoreFp32", "blockSortSize = %d, extractChunkSize=%d, dataType=%d, kValue=%d.",
1608+ blockSortSize, extractChunkSize, computTileInfo.kValue);
1609+ uint32_t blocksPerRow = Ops::Base::CeilDiv(static_cast<uint64_t>(computTileInfo.lastAxisNum), static_cast<uint64_t>(blockSortSize));
1610+ return blocksPerRow > 1 && blocksPerRow <= topkV2DataInfo::MERGE_INTRA_CORE_MAX_BLOCKS;
1611+}
1612+ 
1613+ge::graphStatus TileModeFp32MergeSort(gert::TilingContext *context, TopKV2TilingDataSimd& topkTilingData,
1614+ topkV2DataInfo::TopkComputeNowTileSizeInfo& computNowTileInfo) {
1615+ if (IsTopkMergeSortMoreCoreFp32Mode(computNowTileInfo)) {
1616+ return TileModeFp32MoreCoreSort(context, topkTilingData, computNowTileInfo);
1617+ } else {
1618+ return TileModeFp32IntraCoreSort(context, topkTilingData, computNowTileInfo);
1619+ }
1620+}
1621+ 
1622+bool IsFp32MergeSortMode(topkV2DataInfo::TopkComputeNowTileSizeInfo& computNowTileInfo) {
1623+ bool isMoreCoreModel = IsTopkMergeSortMoreCoreFp32Mode(computNowTileInfo);
1624+ bool isIntraCoreModel = IsTopkMergeSortIntraCoreFp32Mode(computNowTileInfo);
1625+ return isMoreCoreModel || isIntraCoreModel;
1626+}
1627+ 
1176ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)1628ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
1177{1629{
1178 OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling start");1630 OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling start");
1179 TopKV2TilingDataSimd topkTilingData;1631 TopKV2TilingDataSimd topkTilingData;
1180- OP_CHECK_IF(1632+ OP_CHECK_IF(IsValidParam(context) == ge::GRAPH_FAILED, OP_LOGE("TopkV2", "Input param is invalid"),
1181- IsValidParam(context) == ge::GRAPH_FAILED, OP_LOGE("TopkV2", "Input param is invalid"),
1182 return ge::GRAPH_FAILED);1633 return ge::GRAPH_FAILED);
1634+ OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return ge::GRAPH_FAILED);
1635+
1183 const gert::Shape inputShape = context->GetInputShape(0)->GetStorageShape();1636 const gert::Shape inputShape = context->GetInputShape(0)->GetStorageShape();
1184 auto dataType = context->GetInputDesc(0)->GetDataType();1637 auto dataType = context->GetInputDesc(0)->GetDataType();
1185 const gert::Shape outShape = context->GetOutputShape(0)->GetStorageShape();1638 const gert::Shape outShape = context->GetOutputShape(0)->GetStorageShape();
1186 auto dataTypeKey = topkV2DataInfo::tilingDataTypeKeyMap.find(dataType)->second;1639 auto dataTypeKey = topkV2DataInfo::tilingDataTypeKeyMap.find(dataType)->second;
1187 auto indicesDType = context->GetOutputDesc(1)->GetDataType();1640 auto indicesDType = context->GetOutputDesc(1)->GetDataType();
1188 std::string opType(context->GetNodeType());1641 std::string opType(context->GetNodeType());
1642+ 
1643+ // check property
1189 auto const attrs = context->GetAttrs();1644 auto const attrs = context->GetAttrs();
1190 OP_CHECK_NULL_WITH_CONTEXT(context, attrs);1645 OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
1191 const bool* isSorted = attrs->GetAttrPointer<bool>(0);1646 const bool* isSorted = attrs->GetAttrPointer<bool>(0);
@@ -1197,19 +1652,26 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
1197 const bool* isLargest = attrs->GetAttrPointer<bool>(2);1652 const bool* isLargest = attrs->GetAttrPointer<bool>(2);
1198 OP_CHECK_NULL_WITH_CONTEXT(context, isLargest);1653 OP_CHECK_NULL_WITH_CONTEXT(context, isLargest);
1199 OP_LOGI(context->GetNodeName(), "isLargest=%d", *isLargest);1654 OP_LOGI(context->GetNodeName(), "isLargest=%d", *isLargest);
1655+ 
1200 // check the indices_dtype attr and actual value of indices output1656 // check the indices_dtype attr and actual value of indices output
1201 const int* indicesDTypeValuePtr = attrs->GetAttrPointer<int>(3);1657 const int* indicesDTypeValuePtr = attrs->GetAttrPointer<int>(3);
1202 OP_CHECK_NULL_WITH_CONTEXT(context, indicesDTypeValuePtr);1658 OP_CHECK_NULL_WITH_CONTEXT(context, indicesDTypeValuePtr);
1203 OP_LOGI(context->GetNodeName(), "indicesDTypeValuePtr=%d, outPutIndexType=%ld.", *indicesDTypeValuePtr, 1659 OP_LOGI(context->GetNodeName(), "indicesDTypeValuePtr=%d, outPutIndexType=%ld.", *indicesDTypeValuePtr,
1204 static_cast<int64_t>(indicesDType));1660 static_cast<int64_t>(indicesDType));
1205 1661 
1662+ // 获取输入张量的维度数量
1206 size_t inputDimNum = inputShape.GetDimNum();1663 size_t inputDimNum = inputShape.GetDimNum();
1207- OP_CHECK_IF(1664+
1208- *dimValuePtr < static_cast<int32_t>(-inputDimNum) || *dimValuePtr >= static_cast<int32_t>(inputDimNum),1665+ // 验证 dim 参数的有效性（必须在 [-inputDimNum, inputDimNum-1] 范围内）
1209- OP_LOGE_WITH_INVALID_ATTR(context->GetNodeName(), "dim",1666+ int32_t dimMin = -static_cast<int32_t>(inputDimNum);
1210- std::to_string(*dimValuePtr).c_str(),1667+ int32_t dimMax = static_cast<int32_t>(inputDimNum) - 1;
1211- (std::string("range [") + std::to_string(-static_cast<int64_t>(inputDimNum)) + ", " + std::to_string(static_cast<int64_t>(inputDimNum) - 1) + "]").c_str()),1668+ int32_t dimValue = *dimValuePtr;
1212- return ge::GRAPH_FAILED);1669+
1670+ OP_CHECK_IF(dimValue < dimMin || dimValue > dimMax,
1671+ OP_LOGE_WITH_INVALID_ATTR(context->GetNodeName(), "dim",
1672+ std::to_string(dimValue).c_str(),
1673+ (std::string("range [") + std::to_string(dimMin) + ", " + std::to_string(dimMax) + "]").c_str()),
1674+ return ge::GRAPH_FAILED);
1213 int64_t lastAxisNum = inputShape.GetDim(inputDimNum - 1);1675 int64_t lastAxisNum = inputShape.GetDim(inputDimNum - 1);
1214 uint32_t unsortedDimNum = 1;1676 uint32_t unsortedDimNum = 1;
1215 for (uint32_t i = 0; i < (inputDimNum - 1); i++) {1677 for (uint32_t i = 0; i < (inputDimNum - 1); i++) {
@@ -1221,107 +1683,72 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
1221 uint32_t isInInt32Range = static_cast<uint32_t>(lastAxisNum <= int32Max);1683 uint32_t isInInt32Range = static_cast<uint32_t>(lastAxisNum <= int32Max);
1222 // Get Platform Info1684 // Get Platform Info
1223 uint64_t ubSizePlatForm = 0;1685 uint64_t ubSizePlatForm = 0;
1686+ uint64_t originUbSizePlatForm = 0;
1224 auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());1687 auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
1225 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);1688 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
1226 1689 
1227 uint64_t ubBlock = static_cast<uint64_t>(Ops::Base::GetUbBlockSize(context));1690 uint64_t ubBlock = static_cast<uint64_t>(Ops::Base::GetUbBlockSize(context));
1228- OP_LOGI("[TopKV2Tiling]", "ubBlock size = : %u", ubBlock);1691+ originUbSizePlatForm = ubSizePlatForm;
1692+ OP_LOGI("[TopKV2Tiling]", "ubBlock size = : %u, originUbSizePlatForm=%d", ubBlock, originUbSizePlatForm);
1229 1693 
1230- // 校验ubSizePlatForm1694+ // 预留空间给SIMT
1231 ubSizePlatForm -= topkV2DataInfo::CONST_SIMT_SPACE;1695 ubSizePlatForm -= topkV2DataInfo::CONST_SIMT_SPACE;
1232- 1696+ // 用于模板tilingSize计算
1233- // sortAndTopk模板tiling处理流程，和topk其他模板没有关联性1697+ topkV2DataInfo::TopkComputeNowTileSizeInfo computeNowTileSizeInfo;
1234- if (lastAxisNum > topkV2DataInfo::SORT_AND_TOP_K_THRESHOLD) {1698+ computeNowTileSizeInfo.isLargest = *isLargest;
1235- topkV2DataInfo::SortTileInfo sortTileInfo;1699+ computeNowTileSizeInfo.isSort = *isSorted;
1236- OP_CHECK_IF(SortCheckParams(context, sortTileInfo) != ge::GRAPH_SUCCESS,1700+ computeNowTileSizeInfo.isInInt32Range = lastAxisNum <= int32Max;
1237- OP_LOGE(context->GetNodeName(), "sort and topk check params failed"), return ge::GRAPH_FAILED);1701+ computeNowTileSizeInfo.lastAxisNum = lastAxisNum;
1238- topkTilingData.set_modeType(topkV2DataInfo::SORT_AND_TOP_K_MODE);1702+ computeNowTileSizeInfo.kValue = outLastAxisNum;
1239- OP_LOGI("[TopKV2Tiling]", "topkTilingData.set_modeType is: %u, SORT_AND_TOP_K_MODE: %u",1703+ computeNowTileSizeInfo.maxCoreNum = maxCoreNum;
1240- topkTilingData.get_modeType() , topkV2DataInfo::SORT_AND_TOP_K_MODE);1704+ computeNowTileSizeInfo.ubSizePlatForm = ubSizePlatForm;
1241- sortTileInfo.maxCoreNum = static_cast<uint32_t>(maxCoreNum);1705+ computeNowTileSizeInfo.dataType = dataType;
1242- sortTileInfo.isDescend = static_cast<bool>(*isLargest);1706+ computeNowTileSizeInfo.indicesDType = indicesDType;
1243- sortTileInfo.isInt32 = static_cast<uint32_t>(lastAxisNum <= topkV2DataInfo::INT32_MAX_RANGE_VALUE_FOR_SORT);1707+ computeNowTileSizeInfo.ubBlockAlignSize = ubBlock;
1244- sortTileInfo.topKRealValue = outLastAxisNum;1708+ computeNowTileSizeInfo.unsortedDimNum = unsortedDimNum;
1245- OP_CHECK_IF(GetRadixSortMoreCore(context, sortTileInfo) != ge::GRAPH_SUCCESS,
1246- OP_LOGE(context->GetNodeName(), "Get RadixSortMoreCore tiling failed"), return ge::GRAPH_FAILED);
1247- context->SetTilingKey(dataTypeKey);
1248- context->SetBlockDim(sortTileInfo.coreNumNeed);
1249- context->SetLocalMemorySize(sortTileInfo.ubSize);
1250- FillTilingDataSort(context, sortTileInfo, topkTilingData);
1251- PrintTilindDataSort(context, sortTileInfo);
1252- // sortAndTopK模板核心是Sort，不需要后续Topk相关的tiling计算过程
1253- OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling end");
1254- return ge::GRAPH_SUCCESS;
1255- }
1256- 
1257- // 用于核间优化模板 tilingSize计算流程
1258- topkV2DataInfo::TopkComputingNowTileSizeInfo computingNowTileSizeInfo;
1259- computingNowTileSizeInfo.isLargest = *isLargest;
1260- computingNowTileSizeInfo.isSort = *isSorted;
1261- computingNowTileSizeInfo.isInInt32Range = lastAxisNum <= int32Max;
1262- computingNowTileSizeInfo.lastAxisNum = lastAxisNum;
1263- computingNowTileSizeInfo.kValue = outLastAxisNum;
1264- computingNowTileSizeInfo.maxCoreNum = maxCoreNum;
1265- computingNowTileSizeInfo.ubSizePlatForm = ubSizePlatForm;
1266- computingNowTileSizeInfo.dataType = dataType;
1267- computingNowTileSizeInfo.indicesDType = indicesDType;
1268- computingNowTileSizeInfo.ubBlockAlignSize = ubBlock;
1269 OP_LOGI(1709 OP_LOGI(
1270 "[TopKV2Tiling]",1710 "[TopKV2Tiling]",
1271- "computingNowTileSizeInfo isLargest: %u, isSort: %u, isInInt32Range: %u, lastAxisNum: %u, kValue: %u, "1711+ "computeNowTileSizeInfo isLargest: %u, isSort: %u, isInInt32Range: %u, lastAxisNum: %u, kValue: %u, "
1272 "maxCoreNum: %u, ubSizePlatForm: %u",1712 "maxCoreNum: %u, ubSizePlatForm: %u",
1273- computingNowTileSizeInfo.isLargest, computingNowTileSizeInfo.isSort, computingNowTileSizeInfo.isInInt32Range,1713+ computeNowTileSizeInfo.isLargest, computeNowTileSizeInfo.isSort, computeNowTileSizeInfo.isInInt32Range,
1274- computingNowTileSizeInfo.lastAxisNum, computingNowTileSizeInfo.kValue, computingNowTileSizeInfo.maxCoreNum,1714+ computeNowTileSizeInfo.lastAxisNum, computeNowTileSizeInfo.kValue, computeNowTileSizeInfo.maxCoreNum,
1275- computingNowTileSizeInfo.ubSizePlatForm);1715+ computeNowTileSizeInfo.ubSizePlatForm);
1276 1716 
1277- uint32_t nowTileSize = ComputeTopkTileData(context, topkTilingData, computingNowTileSizeInfo);
1278 topkV2DataInfo::TopkTileInfo topkTileInfo;1717 topkV2DataInfo::TopkTileInfo topkTileInfo;
1279 topkTileInfo.topKOutLastAxisNum = outLastAxisNum;1718 topkTileInfo.topKOutLastAxisNum = outLastAxisNum;
1280- 1719+ uint32_t nowTileSize = topkV2DataInfo::TMP_DATA_NUM;
1281- const uint32_t sortedDimParallelData = (nowTileSize * maxCoreNum) / 2;1720+ if (IsSmallSizeMergeSortMode(dataType, lastAxisNum)) {
1282- if (lastAxisNum <= topkV2DataInfo::SMALL_MAX_DATA_SZIE && topkV2DataInfo::optDataTypeBitMap.count(dataType) != 0) {1721+ TileModeSmallSizeOptim(context, topkTilingData, topkTileInfo, computeNowTileSizeInfo);
1283- SetMergeSortTmpSize(context, dataType, lastAxisNum, topkTilingData);
1284- uint32_t nowTileSizeTmp = ComputeMergeSortTileData(
1285- topkTilingData, dataType, indicesDType, lastAxisNum, maxCoreNum, unsortedDimNum, ubSizePlatForm);
1286- TileModeSmallSizeOptim(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSizeTmp);
1287 dataTypeKey += topkV2DataInfo::MERGE_SORT_TILING_OFFSET;1722 dataTypeKey += topkV2DataInfo::MERGE_SORT_TILING_OFFSET;
1288- } else if (lastAxisNum <= nowTileSize) {1723+ } else if (IsSingleBlockMode(context, topkTilingData, nowTileSize, computeNowTileSizeInfo)) {
1289- uint32_t nowTileSizeTmp = ComputeSingleBlockTileData(1724+ TileModeSmallSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
1290- context, topkTilingData, dataType, indicesDType, *isLargest, *isSorted, lastAxisNum, outLastAxisNum,1725+ } else if (IsFp32MergeSortMode(computeNowTileSizeInfo)) {
1291- ubSizePlatForm);1726+ return TileModeFp32MergeSort(context, topkTilingData, computeNowTileSizeInfo);
1292- TileModeSmallSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSizeTmp);1727+ } else if (IsSingleCoreMode(context, topkTilingData, nowTileSize, computeNowTileSizeInfo)) {
1293- } else if (IsModeSingleCore(unsortedDimNum, maxCoreNum)) {1728+ TileModeSingleCore(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
1294- uint32_t nowTileSizeTmp = ComputeSingleCoreTileData(1729+ } else if (IsMultiCoreOptimMode(context, nowTileSize, topkTilingData, computeNowTileSizeInfo)) {
1295- context, topkTilingData, dataType, indicesDType, *isLargest, *isSorted, lastAxisNum, outLastAxisNum,1730+ // topk核间模板优化，修改modeType, tilling不变
1296- ubSizePlatForm);
1297- OP_CHECK_IF(nowTileSizeTmp == 0, OP_LOGE("TopkV2", "nowTileSizeTmp is 0"), return ge::GRAPH_FAILED);
1298- TileModeSingleCore(
1299- unsortedDimNum, maxCoreNum, static_cast<uint32_t>(lastAxisNum), topkTilingData, topkTileInfo,
1300- nowTileSizeTmp);
1301- } else if (IsMultiCoreOptimMode(context, nowTileSize, topkTilingData, computingNowTileSizeInfo)) {
1302- // topk核间模板优化，修改modelType, tilling不变
1303 TileMultiCoreOptimSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);1731 TileMultiCoreOptimSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
1304- topkTilingData.set_modeType(topkV2DataInfo::MULT_CORE_OPTIM_MODE);1732+ } else if (IsTopkRadixMoreCoreMode(context, topkTilingData, computeNowTileSizeInfo, nowTileSize)) {
1305- } else if (lastAxisNum <= sortedDimParallelData) {1733+ TileTopkMoreCoreMode(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
1306- TileModeMediumSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize);
1307 } else {1734 } else {
1308- TileModeBigSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSize); // 核间模板1735+ // sortAndTopk模板处理超大尾轴
1736+ return TileModeSortAndTopK(context, topkTilingData, computeNowTileSizeInfo);
1309 }1737 }
1310- int64_t maxInputK = std::max(outLastAxisNum, topkTileInfo.lastDimTileNum);
1311 1738 
1312 if (topkTilingData.get_modeType() == topkV2DataInfo::MULT_CORE_MODE &&1739 if (topkTilingData.get_modeType() == topkV2DataInfo::MULT_CORE_MODE &&
1313- topkTileInfo.topKOutLastAxisNum > topkV2DataInfo::MAX_K_FOR_INT64) {1740+ topkTileInfo.topKOutLastAxisNum > topkV2DataInfo::MAX_K_FOR_INT64 && topkTileInfo.multiCoreBigModel) {
1741+ OP_LOGD("[TopKV2Tiling] radix more core tileSize value halved.");
1314 nowTileSize /= topkV2DataInfo::CONST_TWO;1742 nowTileSize /= topkV2DataInfo::CONST_TWO;
1315 }1743 }
1316 1744 
1317- GetTopkApiTmpBufferSize(1745+ int64_t maxInputK = std::max(outLastAxisNum, topkTileInfo.lastDimTileNum);
1318- context, topkTilingData, topkTileInfo.ubRealLoadDataNum, maxInputK, *isLargest, dataType, *isSorted,1746+ 
1319- nowTileSize);1747+ GetTopkApiTmpBufferSize(context, topkTilingData, topkTileInfo.ubRealLoadDataNum, maxInputK, *isLargest,
1320- OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TopkV2", "maxCoreNum is 0"), return ge::GRAPH_FAILED);1748+ dataType, *isSorted, nowTileSize);
1321 int64_t lastDimTileNumTimesValue = (topkTileInfo.lastDimTileNum + maxCoreNum - 1) / maxCoreNum;1749 int64_t lastDimTileNumTimesValue = (topkTileInfo.lastDimTileNum + maxCoreNum - 1) / maxCoreNum;
1322- OP_LOGI(1750+ OP_LOGD("[TopKV2Tiling]", "lastAxisNum: %ld, int32Max: %ld, isInInt32Range: %u, nowTileSize: %u", lastAxisNum,
1323- "[TopKV2Tiling]", "lastAxisNum: %ld, int32Max: %ld, isInInt32Range: %u, nowTileSize: %u", lastAxisNum, int32Max,1751+ int32Max, isInInt32Range, nowTileSize);
1324- isInInt32Range, nowTileSize);
1325 1752 
1326 // fill the topkTilingData1753 // fill the topkTilingData
1327 context->SetTilingKey(dataTypeKey);1754 context->SetTilingKey(dataTypeKey);
@@ -1357,57 +1784,43 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
1357 topkTilingData.get_tailLoopBatchNum(), topkTilingData.get_tailBatchNum(), topkTilingData.get_tailTileNum(),1784 topkTilingData.get_tailLoopBatchNum(), topkTilingData.get_tailBatchNum(), topkTilingData.get_tailTileNum(),
1358 topkTileInfo.coreNumNeed);1785 topkTileInfo.coreNumNeed);
1359 1786 
1360- // TopKV2 Workspace计算流程1787+ // TopKV2 Workspace 计算流程
1361 size_t usrSize = 0;1788 size_t usrSize = 0;
1362- OP_LOGI("[TopKV2Tiling]", "begin to calc TopKV2 Workspace size.");1789+
1363- if (isInInt32Range) {1790+ uint64_t alginFactor = topkV2DataInfo::AGLIN_FACTOR;
1364- if (topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) {1791+ uint32_t modeType = topkTilingData.get_modeType();
1365- usrSize = topkTileInfo.lastDimTileNum * topkV2DataInfo::BIN_NUM * topkTileInfo.unsortedDimParallel *1792+ uint32_t lastDimTileNum = topkTileInfo.lastDimTileNum;
1366- sizeof(int32_t);1793+ uint32_t unsortedDimParallel = topkTileInfo.unsortedDimParallel;
1367- OP_LOGI("[TopKV2Tiling]", "TopK V2 Workspace size is : %u", usrSize);1794+
1368- } else if (topkTilingData.get_modeType() == topkV2DataInfo::MULT_CORE_OPTIM_MODE) {1795+ // 提取公共变量：索引类型大小（根据数据范围选择 int32 或 int64）
1369- uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);1796+ size_t indexTypeSize = isInInt32Range ? sizeof(int32_t) : sizeof(int64_t);
1370- size_t tempSortResultNeedSize = Ops::Base::CeilAlign(1797+
1371- static_cast<uint64_t>(1798+ // 根据 modeType 计算 Workspace 大小
1372- xDtypeSize * topkTileInfo.lastDimTileNum * topkTileInfo.unsortedDimParallel *1799+ if (modeType == topkV2DataInfo::SINGLE_CORE_MODE) {
1373- topkTilingData.get_topKRealValue()),1800+ usrSize = lastDimTileNum * topkV2DataInfo::BIN_NUM * unsortedDimParallel * indexTypeSize;
1374- topkV2DataInfo::AGLIN_FACTOR);1801+ } else if (modeType == topkV2DataInfo::MULT_CORE_OPTIM_MODE && isInInt32Range) {
1375- size_t tempSortIndexNeedSize = Ops::Base::CeilAlign(1802+ uint32_t xDtypeSize = GetDataTypeSize(dataType);
1376- static_cast<uint64_t>(1803+ uint32_t xDtypeSizeFactor = lastDimTileNum * unsortedDimParallel;
1377- sizeof(int32_t) * topkTileInfo.lastDimTileNum * topkTileInfo.unsortedDimParallel *1804+ // 计算临时排序结果空间（对齐）
1378- topkTilingData.get_topKRealValue()),1805+ size_t tempSortResultNeedSize = Ops::Base::CeilAlign(
1379- topkV2DataInfo::AGLIN_FACTOR);1806+ static_cast<uint64_t>(xDtypeSize * xDtypeSizeFactor * outLastAxisNum), alginFactor);
1380- usrSize = tempSortResultNeedSize + tempSortIndexNeedSize;1807+ // 计算临时排序索引空间（对齐）
1381- OP_LOGI("[TopKV2Tiling]", "Workspace size for Radix Sort perf template within Int32Range is : %u", usrSize);1808+ size_t tempSortIndexNeedSize = Ops::Base::CeilAlign(
1382- } else {1809+ static_cast<uint64_t>(sizeof(int32_t) * xDtypeSizeFactor * outLastAxisNum), alginFactor);
1383- size_t dataSetCumSumHistNeedSize =1810+ usrSize = tempSortResultNeedSize + tempSortIndexNeedSize;
1384- topkV2DataInfo::BIN_NUM * sizeof(int32_t) * topkTileInfo.unsortedDimParallel;
1385- size_t dataSetTileTopkNeedSize = Ops::Base::CeilAlign(
1386- static_cast<uint64_t>(topkTileInfo.lastDimTileNum * sizeof(int32_t) * topkTileInfo.unsortedDimParallel),
1387- topkV2DataInfo::AGLIN_FACTOR);
1388- usrSize = dataSetCumSumHistNeedSize + dataSetTileTopkNeedSize * topkV2DataInfo::CONST_TWO;
1389- OP_LOGI("[TopKV2Tiling]", "Workspace size for Radix Sort more core within Int32Range is : %u", usrSize);
1390- }
1391 } else {1811 } else {
1392- if (topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) {1812+ size_t dataSetCumSumHistNeedSize = topkV2DataInfo::BIN_NUM * indexTypeSize * unsortedDimParallel;
1393- usrSize = topkTileInfo.lastDimTileNum * topkV2DataInfo::BIN_NUM * topkTileInfo.unsortedDimParallel *1813+ size_t dataSetTileTopkNeedSize = Ops::Base::CeilAlign(
1394- sizeof(int64_t);1814+ static_cast<uint64_t>(lastDimTileNum * indexTypeSize * unsortedDimParallel), alginFactor);
1395- OP_LOGI("[TopKV2Tiling]", "TopK V2 Workspace size is : %u", usrSize);1815+ usrSize = dataSetCumSumHistNeedSize + dataSetTileTopkNeedSize * topkV2DataInfo::CONST_TWO;
1396- } else {
1397- size_t dataSetCumSumHistNeedSize =
1398- topkV2DataInfo::BIN_NUM * sizeof(int64_t) * topkTileInfo.unsortedDimParallel;
1399- size_t dataSetTileTopkNeedSize = Ops::Base::CeilAlign(
1400- topkTileInfo.lastDimTileNum * sizeof(int64_t) * topkTileInfo.unsortedDimParallel,
1401- topkV2DataInfo::AGLIN_FACTOR);
1402- usrSize = dataSetCumSumHistNeedSize + dataSetTileTopkNeedSize * topkV2DataInfo::CONST_TWO;
1403- OP_LOGI("[TopKV2Tiling]", "Workspace size for Radix Sort more core beyond Int32Range is : %u", usrSize);
1404- }
1405 }1816 }
1406 1817 
1407 // sortWithIndex tiling&workspace计算流程，1818 // sortWithIndex tiling&workspace计算流程，
1408 if (needSortWithIndex(topkTilingData, *isSorted, dataType)) {1819 if (needSortWithIndex(topkTilingData, *isSorted, dataType)) {
1409- OP_CHECK_IF(sortWithIndex::RadixSortTilingOfIdx(context, topkTilingData, maxCoreNum, &usrSize) != ge::GRAPH_SUCCESS,1820+ OP_CHECK_IF(sortWithIndex::RadixSortTilingOfIdx(context, topkTilingData, maxCoreNum, &usrSize) !=
1410- OP_LOGE(context->GetNodeName(), "SortWithIndex Tiling Simt calc failed"), return ge::GRAPH_FAILED);1821+ ge::GRAPH_SUCCESS,
1822+ OP_LOGE(context->GetNodeName(), "SortWithIndex Tiling Simt calc failed"),
1823+ return ge::GRAPH_FAILED);
1411 }1824 }
1412 1825 
1413 // save tilingdata1826 // save tilingdata
Mmath/top_k_v2/op_kernel/arch35/top_k_merge_sort.h+5-3
@@ -71,6 +71,7 @@ public:
71 uint32_t sortLoopTimes_ = 0;71 uint32_t sortLoopTimes_ = 0;
72 uint32_t platformCoreNum_ = 0;72 uint32_t platformCoreNum_ = 0;
73 uint32_t outputLastDimValue_ = 0;73 uint32_t outputLastDimValue_ = 0;
74+ uint32_t bufferNum_ = DOUBLE_BUFFER;
74 // merge sort kernel75 // merge sort kernel
75 topkV2::KernelVbsMergeSort<T, CONVERT_TYPE, IS_LARGEST> vbsSort;76 topkV2::KernelVbsMergeSort<T, CONVERT_TYPE, IS_LARGEST> vbsSort;
76};77};
@@ -100,6 +101,7 @@ __aicore__ inline void MergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST,
100 sortLoopTimes_ = tilingData->sortLoopTimes;101 sortLoopTimes_ = tilingData->sortLoopTimes;
101 unsortedDimParallel_ = tilingData->unsortedDimParallel;102 unsortedDimParallel_ = tilingData->unsortedDimParallel;
102 oneCoreRowNum_ = tilingData->oneCoreRowNum;103 oneCoreRowNum_ = tilingData->oneCoreRowNum;
104+ bufferNum_ = tilingData->keyParams4;
103 // 高阶API需要的临时空间大小105 // 高阶API需要的临时空间大小
104 mergSortAcApiNeedBufferSize_ = tilingData->mergSortAcApiNeedBufferSize;106 mergSortAcApiNeedBufferSize_ = tilingData->mergSortAcApiNeedBufferSize;
105}107}
@@ -113,10 +115,10 @@ __aicore__ inline void MergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST,
113 outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));115 outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));
114 // init queue116 // init queue
115 uint64_t realNum = ROUND_UP_AGLIN(numTileData_) * oneCoreRowNum_;117 uint64_t realNum = ROUND_UP_AGLIN(numTileData_) * oneCoreRowNum_;
116- pipe->InitBuffer(inQueueX_, DOUBLE_BUFFER, ROUND_UP_AGLIN(realNum) * sizeof(T));118+ pipe->InitBuffer(inQueueX_, bufferNum_, ROUND_UP_AGLIN(realNum) * sizeof(T));
117- pipe->InitBuffer(outValueQueue_, DOUBLE_BUFFER, ROUND_UP_AGLIN(realNum * sizeof(T)));119+ pipe->InitBuffer(outValueQueue_, bufferNum_, ROUND_UP_AGLIN(realNum * sizeof(T)));
118 // 改为实际来支持int64_t120 // 改为实际来支持int64_t
119- pipe->InitBuffer(outIndexQueue_, DOUBLE_BUFFER, ROUND_UP_AGLIN(realNum * sizeof(INDEX_TYPE)));121+ pipe->InitBuffer(outIndexQueue_, bufferNum_, ROUND_UP_AGLIN(realNum * sizeof(INDEX_TYPE)));
120}122}
121 123 
122template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>124template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
Amath/top_k_v2/op_kernel/arch35/top_k_merge_sort_intra_core.h+357-0
@@ -0,0 +1,357 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file top_k_merge_sort_intra_core.h
13+ * \brief Intra-core block merge topk for fp32 (inherits from SortMergeIntraCore)
14+ * \details Inherits from Sort::SortMergeIntraCore and overrides 4 key functions to limit output to topKValue.
15+ * Uses this-> to access base class members (cleaner than using declarations).
16+ * Reuses 95% of base class code through inheritance.
17+ */
18+ 
19+#ifndef TOP_K_MERGE_SORT_INTRA_CORE_H
20+#define TOP_K_MERGE_SORT_INTRA_CORE_H
21+ 
22+#include <cmath>
23+#include "kernel_operator.h"
24+#include "op_kernel/math_util.h"
25+#include "op_kernel/platform_util.h"
26+#include "kernel_tiling/kernel_tiling.h"
27+#include "../../sort/arch35/sort_merge_intra_core.h"
28+#include "../../sort/arch35/sort_tiling_data.h"
29+ 
30+namespace topkV2 {
31+using namespace AscendC;
32+ 
33+template <typename ValueType, typename IndexType, bool IsDescend>
34+class TopKMergeSortIntraCore : public Sort::SortMergeIntraCore<ValueType, IndexType, IsDescend> {
35+public:
36+ using Base = Sort::SortMergeIntraCore<ValueType, IndexType, IsDescend>;
37+
38+ __aicore__ inline TopKMergeSortIntraCore() : Base() {}
39+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace,
40+ const TopKV2TilingDataSimd* tilingData, TPipe* pipe);
41+ __aicore__ inline void Process();
42+ 
43+private:
44+ uint32_t topKValue_ = 0;
45+ 
46+ __aicore__ inline uint32_t MergeSingleBatch();
47+ __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, typename Base::MergeListContext& ctx);
48+ __aicore__ inline void MergeOneGroup(uint32_t groupStart, uint32_t groupBlockCount,
49+ uint32_t fullBlockElemCount, uint32_t fullBlockSortLen, uint32_t lastBlockElemCount,
50+ uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset,
51+ uint32_t& mergedGroupElemCount);
52+ __aicore__ inline void CopyRemainingList(typename Base::MergeListContext& ctx,
53+ int64_t dstOffsetBase, uint32_t& dstCumulativeOffset, uint32_t& dstElemCount);
54+ __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion);
55+};
56+ 
57+template <typename ValueType, typename IndexType, bool IsDescend>
58+__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::Init(
59+ GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace,
60+ const TopKV2TilingDataSimd* tilingData, TPipe* pipe)
61+{
62+ if (tilingData == nullptr || pipe == nullptr) {
63+ return;
64+ }
65+
66+ topKValue_ = static_cast<uint32_t>(tilingData->topKRealValue);
67+
68+ SortRegBaseTilingData sortTilingData;
69+ sortTilingData.unsortedDimNum = tilingData->unsortedDimNum;
70+ sortTilingData.lastAxisNum = tilingData->lastAxisNum;
71+ sortTilingData.numTileDataSize = tilingData->numTileDataSize;
72+ sortTilingData.lastDimTileNum = tilingData->lastDimTileNum;
73+ sortTilingData.lastDimNeedCore = tilingData->lastDimNeedCore;
74+ sortTilingData.keyParams0 = tilingData->keyParams0;
75+ sortTilingData.keyParams3 = tilingData->keyParams3;
76+ sortTilingData.keyParams4 = tilingData->keyParams4;
77+ sortTilingData.keyParams5 = tilingData->keyParams5;
78+
79+ Base::Init(x, value, indices, workspace, &sortTilingData, pipe);
80+}
81+ 
82+template <typename ValueType, typename IndexType, bool IsDescend>
83+__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::Process()
84+{
85+ // Calculate batch range for this core
86+ int64_t startBatch = static_cast<int64_t>(this->blockIdx_) * this->batchPerCore_;
87+ int64_t endBatch = (startBatch + this->batchPerCore_ < this->batchNum_) ?
88+ (startBatch + this->batchPerCore_) : this->batchNum_;
89+ 
90+ if (startBatch >= this->batchNum_) {
91+ return; // This core has no work
92+ }
93+ 
94+ // Process each batch through all 3 phases before moving to the next batch.
95+ // This allows cache (workspace) to be reused across batches within the same core,
96+
97+ // Precompute buffer sizes (loop-invariant)
98+ uint32_t mergeBufferSize = Sort::MERGE_LIST_MAX_NUM * this->sortBufferSize_;
99+ uint32_t extractInSize = AscendC::GetSortLen<ValueType>(this->extractChunkSize_) * sizeof(ValueType);
100+ 
101+ for (int64_t batchIdx = startBatch; batchIdx < endBatch; batchIdx++) {
102+ // ========== Phase 1: Sort blocks in UB ==========
103+ this->pipe_->InitBuffer(this->inQueueX_, DOUBLE_BUFFER, this->blockSortSize_ * sizeof(ValueType));
104+ this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_);
105+ this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_);
106+ this->pipe_->InitBuffer(this->sortedOutQueue_, DOUBLE_BUFFER, this->sortBufferSize_);
107+ this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t));
108+ 
109+ this->SortSingleBatchInUb(this->inputXGm_, batchIdx * this->sortAxisNum_);
110+ 
111+ event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_MTE2));
112+ SetFlag<HardEvent::MTE3_MTE2>(eventId);
113+ WaitFlag<HardEvent::MTE3_MTE2>(eventId);
114+ 
115+ this->pipe_->Reset();
116+ 
117+ // ========== Phase 2: Merge sorted blocks (4-way) ==========
118+ this->pipe_->InitBuffer(this->mergeInQueue_, 1, mergeBufferSize);
119+ this->pipe_->InitBuffer(this->mergeOutQueue_, 1, mergeBufferSize);
120+ 
121+ uint32_t resultRegion = MergeSingleBatch();
122+ 
123+ eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_MTE2));
124+ SetFlag<HardEvent::MTE3_MTE2>(eventId);
125+ WaitFlag<HardEvent::MTE3_MTE2>(eventId);
126+ 
127+ this->pipe_->Reset();
128+ 
129+ // ========== Phase 3: Extract and copy output ==========
130+ this->pipe_->InitBuffer(this->extractInQueue_, DOUBLE_BUFFER, extractInSize);
131+ this->pipe_->InitBuffer(this->outValueQueue_, DOUBLE_BUFFER, this->extractChunkSize_ * sizeof(ValueType));
132+ this->pipe_->InitBuffer(this->outIdxQueue_, DOUBLE_BUFFER, this->extractChunkSize_ * sizeof(uint32_t));
133+ if constexpr (IsSameType<int64_t, IndexType>::value) {
134+ this->pipe_->InitBuffer(this->outIdxInt64Queue_, DOUBLE_BUFFER, this->extractChunkSize_ * sizeof(int64_t));
135+ }
136+ 
137+ ExtractAndCopyOut(batchIdx, resultRegion);
138+ 
139+ this->pipe_->Reset();
140+ }
141+}
142+ 
143+template <typename ValueType, typename IndexType, bool IsDescend>
144+__aicore__ inline uint32_t TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeSingleBatch()
145+{
146+ uint32_t numBlocks = this->blocksPerRow_;
147+ uint32_t fullBlockElemCount = this->blockSortSize_;
148+ uint32_t fullBlockSortLen = this->blockSortLen_;
149+ uint32_t lastBlockElemCount = this->lastBlockSize_;
150+ 
151+ uint32_t pingPongFlag = 0;
152+ uint32_t mergeRounds = 0;
153+ 
154+ while (numBlocks > 1) {
155+ uint32_t cumulativeOffset = 0;
156+ uint32_t newNumBlocks = 0;
157+ uint32_t newFullBlockElemCount = 0;
158+ uint32_t newLastBlockElemCount = 0;
159+ 
160+ for (uint32_t i = 0; i < numBlocks; i += Sort::MERGE_LIST_MAX_NUM) {
161+ uint32_t groupBlockCount = (i + Sort::MERGE_LIST_MAX_NUM <= numBlocks) ?
162+ Sort::MERGE_LIST_MAX_NUM : (numBlocks - i);
163+ 
164+ uint32_t mergedGroupElemCount = 0;
165+ MergeOneGroup(i, groupBlockCount, fullBlockElemCount, fullBlockSortLen, lastBlockElemCount,
166+ numBlocks, pingPongFlag, cumulativeOffset, mergedGroupElemCount);
167+ if (newNumBlocks == 0) {
168+ newFullBlockElemCount = mergedGroupElemCount;
169+ }
170+ newLastBlockElemCount = mergedGroupElemCount;
171+ newNumBlocks++;
172+ }
173+
174+ if (newNumBlocks == 0 || newNumBlocks >= numBlocks) {
175+ break;
176+ }
177+
178+ numBlocks = newNumBlocks;
179+ if (numBlocks == 1) {
180+ fullBlockElemCount = newLastBlockElemCount;
181+ } else {
182+ fullBlockElemCount = newFullBlockElemCount;
183+ }
184+ fullBlockSortLen = AscendC::GetSortLen<ValueType>(fullBlockElemCount);
185+ lastBlockElemCount = newLastBlockElemCount;
186+ pingPongFlag = 1 - pingPongFlag;
187+ mergeRounds++;
188+ event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_MTE2));
189+ SetFlag<HardEvent::MTE3_MTE2>(eventId);
190+ WaitFlag<HardEvent::MTE3_MTE2>(eventId);
191+ }
192+ 
193+ return (mergeRounds == 0) ? 0 : pingPongFlag;
194+}
195+ 
196+template <typename ValueType, typename IndexType, bool IsDescend>
197+__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeOneGroup(
198+ uint32_t groupStart, uint32_t groupBlockCount,
199+ uint32_t fullBlockElemCount, uint32_t fullBlockSortLen, uint32_t lastBlockElemCount,
200+ uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset,
201+ uint32_t& mergedGroupElemCount)
202+{
203+ typename Base::MergeListContext ctx;
204+ ctx.listCount = groupBlockCount;
205+
206+ int64_t srcRegionOffset = (pingPongFlag == 0) ? 0 : this->batchSortLen_;
207+ int64_t dstRegionOffset = (pingPongFlag == 0) ? this->batchSortLen_ : 0;
208+ 
209+ for (uint32_t j = 0; j < groupBlockCount; j++) {
210+ uint32_t blockIdx = groupStart + j;
211+ ctx.srcOffsets[j] = srcRegionOffset + blockIdx * fullBlockSortLen;
212+ ctx.elemCounts[j] = (blockIdx < numBlocks - 1) ? fullBlockElemCount : lastBlockElemCount;
213+ }
214+ 
215+ if (groupBlockCount > 1) {
216+ DoIncrementalMerge(dstRegionOffset + cumulativeOffset, ctx);
217+ 
218+ uint32_t totalElem = 0;
219+ for (uint32_t j = 0; j < groupBlockCount; j++) {
220+ totalElem += ctx.elemCounts[j];
221+ }
222+ mergedGroupElemCount = (totalElem > topKValue_) ? topKValue_ : totalElem; // TopK limitation
223+ cumulativeOffset += AscendC::GetSortLen<ValueType>(mergedGroupElemCount);
224+ } else if (groupBlockCount == 1) {
225+ mergedGroupElemCount = (ctx.elemCounts[0] > topKValue_) ? topKValue_ : ctx.elemCounts[0]; // TopK limitation
226+ if (ctx.elemCounts[0] > 0) {
227+ int64_t srcOffset = ctx.srcOffsets[0];
228+ uint32_t remainElems = mergedGroupElemCount;
229+ uint32_t srcChunkOffset = 0;
230+ while (remainElems > 0) {
231+ uint32_t chunkSize = (remainElems > this->blockSortSize_) ? this->blockSortSize_ : remainElems;
232+ if (chunkSize == 0) break;
233+ this->CopyBlockChunk(srcOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset),
234+ dstRegionOffset + cumulativeOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), chunkSize);
235+ remainElems -= chunkSize;
236+ srcChunkOffset += chunkSize;
237+ }
238+ cumulativeOffset += AscendC::GetSortLen<ValueType>(mergedGroupElemCount);
239+ }
240+ }
241+}
242+ 
243+template <typename ValueType, typename IndexType, bool IsDescend>
244+__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::DoIncrementalMerge(
245+ int64_t dstOffsetBase, typename Base::MergeListContext& ctx)
246+{
247+ if (ctx.listCount > Sort::MERGE_LIST_MAX_NUM) return;
248+ 
249+ for (uint32_t i = 0; i < ctx.listCount; i++) {
250+ ctx.remains[i] = ctx.elemCounts[i];
251+ }
252+
253+ uint32_t dstCumulativeOffset = 0, dstElemCount = 0, loopGuard = 0;
254+ while (dstElemCount < topKValue_) { // TopK: early termination
255+ uint32_t activeLists = 0;
256+ for (uint32_t i = 0; i < ctx.listCount; i++) {
257+ if (ctx.remains[i] > 0) activeLists++;
258+ }
259+ if (activeLists <= 1) break;
260+ 
261+ LocalTensor<ValueType> ubMainInput = this->mergeInQueue_.template AllocTensor<ValueType>();
262+ uint16_t elementCountList[Sort::MERGE_LIST_MAX_NUM] = {0, 0, 0, 0};
263+ uint32_t remainListNum = this->LoadListsToUb(ubMainInput, elementCountList, ctx);
264+ this->mergeInQueue_.EnQue(ubMainInput);
265+ 
266+ LocalTensor<ValueType> ubMainInputCalc = this->mergeInQueue_.template DeQue<ValueType>();
267+ LocalTensor<ValueType> dstLocal = this->mergeOutQueue_.template AllocTensor<ValueType>();
268+ 
269+ uint32_t listSortedNums[Sort::MERGE_LIST_MAX_NUM] = {0, 0, 0, 0};
270+ uint32_t mergedCount = this->ExecuteMrgSort(dstLocal, ubMainInputCalc,
271+ elementCountList, listSortedNums, remainListNum);
272+ if (mergedCount == 0) {
273+ this->mergeInQueue_.FreeTensor(ubMainInputCalc);
274+ this->mergeOutQueue_.FreeTensor(dstLocal);
275+ break;
276+ }
277+ 
278+ this->mergeOutQueue_.EnQue(dstLocal);
279+ this->mergeInQueue_.FreeTensor(ubMainInputCalc);
280+ 
281+ LocalTensor<ValueType> dstLocalOut = this->mergeOutQueue_.template DeQue<ValueType>();
282+ uint32_t copyCount = mergedCount;
283+ if (dstElemCount + copyCount > topKValue_) { // TopK: limit output
284+ copyCount = topKValue_ - dstElemCount;
285+ }
286+ DataCopyExtParams outCopyParams{1, static_cast<uint32_t>(AscendC::GetSortLen<ValueType>(copyCount) * sizeof(ValueType)), 0, 0, 0};
287+ DataCopyPad(this->cacheGm_[dstOffsetBase + dstCumulativeOffset], dstLocalOut, outCopyParams);
288+ this->mergeOutQueue_.FreeTensor(dstLocalOut);
289+ 
290+ uint32_t j = 0;
291+ for (uint32_t i = 0; i < ctx.listCount; i++) {
292+ if (ctx.remains[i] > 0) {
293+ ctx.gmOffsets[i] += AscendC::GetSortLen<ValueType>(listSortedNums[j]);
294+ ctx.remains[i] -= listSortedNums[j];
295+ j++;
296+ }
297+ }
298+ dstCumulativeOffset += AscendC::GetSortLen<ValueType>(copyCount);
299+ dstElemCount += copyCount;
300+ if (++loopGuard > this->maxMergeIterations_) break;
301+ }
302+ 
303+ CopyRemainingList(ctx, dstOffsetBase, dstCumulativeOffset, dstElemCount);
304+}
305+ 
306+template <typename ValueType, typename IndexType, bool IsDescend>
307+__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::CopyRemainingList(
308+ typename Base::MergeListContext& ctx, int64_t dstOffsetBase, uint32_t& dstCumulativeOffset, uint32_t& dstElemCount)
309+{
310+ for (uint32_t listIdx = 0; listIdx < ctx.listCount && dstElemCount < topKValue_; listIdx++) { // TopK: early stop
311+ while (ctx.remains[listIdx] > 0) {
312+ if (dstElemCount >= topKValue_) {
313+ break;
314+ }
315+ uint32_t loadCount = (ctx.remains[listIdx] > this->blockSortSize_) ? this->blockSortSize_ : ctx.remains[listIdx];
316+ uint32_t remainingTopK = topKValue_ - dstElemCount;
317+ loadCount = (loadCount > remainingTopK) ? remainingTopK : loadCount; // TopK: limit chunk
318+ if (loadCount == 0) break;
319+
320+ this->CopyBlockChunk(ctx.srcOffsets[listIdx] + ctx.gmOffsets[listIdx],
321+ dstOffsetBase + dstCumulativeOffset, loadCount);
322+ 
323+ ctx.gmOffsets[listIdx] += AscendC::GetSortLen<ValueType>(loadCount);
324+ ctx.remains[listIdx] -= loadCount;
325+ dstCumulativeOffset += AscendC::GetSortLen<ValueType>(loadCount);
326+ dstElemCount += loadCount;
327+ }
328+ }
329+}
330+ 
331+template <typename ValueType, typename IndexType, bool IsDescend>
332+__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut(
333+ int64_t batchIdx, uint32_t resultRegion)
334+{
335+ int64_t outputOffset = batchIdx * static_cast<int64_t>(topKValue_); // TopK: output offset
336+ 
337+ int64_t cacheBatchOffset = (resultRegion == 1) ? this->batchSortLen_ : 0;
338+ 
339+ uint32_t totalElem = topKValue_; // TopK: limit output count
340+ uint32_t elemProcessed = 0;
341+ uint32_t cacheOffset = 0;
342+ 
343+ while (elemProcessed < totalElem) {
344+ uint32_t elemCount = (elemProcessed + this->extractChunkSize_ <= totalElem) ?
345+ this->extractChunkSize_ : (totalElem - elemProcessed);
346+ if (elemCount == 0) break;
347+ 
348+ this->ExtractAndCopyChunk(cacheBatchOffset, cacheOffset, outputOffset, elemProcessed, elemCount);
349+ 
350+ elemProcessed += elemCount;
351+ cacheOffset += AscendC::GetSortLen<ValueType>(elemCount);
352+ }
353+}
354+ 
355+} // namespace topkV2
356+ 
357+#endif // TOP_K_MERGE_SORT_INTRA_CORE_H
Amath/top_k_v2/op_kernel/arch35/top_k_merge_sort_more_core.h+629-0
@@ -0,0 +1,629 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file top_k_merge_sort_more_core.h
13+ * \brief topk merge sort more-core kernel entry
14+ */
15+#ifndef TOP_K_MERGE_SORT_MORE_CORE_H
16+#define TOP_K_MERGE_SORT_MORE_CORE_H
17+#include <cmath>
18+#include "op_kernel/platform_util.h"
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "top_k_util_type_simd.h"
22+ 
23+namespace topkV2 {
24+using namespace AscendC;
25+constexpr uint32_t MORE_CORE_FP32_DTYPE = 4;
26+constexpr uint32_t MORE_CORE_DOUBLE_BUFFER = 1;
27+constexpr uint32_t MORE_CORE_BLOCK_UB = Ops::Base::GetUbBlockSize();
28+constexpr uint32_t MORE_CORE_MERGE_SORT_LIST_MAX_NUM = 4;
29+constexpr uint32_t MORE_CORE_DEALING_CONCAT_NUM_ONCE = 16;
30+constexpr uint32_t MORE_CORE_DEALING_SORT_NUM_ONCE = 32;
31+constexpr uint32_t MORE_CORE_DEALING_EXTRACT_NUM_ONCE = 32;
32+constexpr int32_t MORE_CORE_XOR_OP_VALUE_FP = 0x80000000;
33+constexpr int16_t MORE_CORE_XOR_OP_VALUE_HALF = 0x8000;
34+ 
35+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
36+struct 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);
39+ __aicore__ inline void Process();
40+ __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 InitIndexLocal(uint32_t tileNum, LocalTensor<uint32_t> sortedValueIndexLocal, int64_t offsetPerCore);
43+ __aicore__ inline void DoSort(uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal);
44+ __aicore__ inline void DoSortBf16(uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal);
45+ __aicore__ inline void CopyOutWorkSpace(uint32_t tileNum, int64_t offsetPerCore, LocalTensor<CONVERT_TYPE> sortedValueLocal);
46+ __aicore__ inline void CopyOutMultiCore();
47+ __aicore__ inline void SortInMultiCore();
48+ __aicore__ inline void FinalSortAndCopy();
49+ __aicore__ inline void MultiCoreInit();
50+ __aicore__ inline void FinalSortInit();
51+ __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t aglinTileNum);
52+ __aicore__ inline void CopyInMultiCore();
53+ __aicore__ inline void UpdateMrgParam();
54+ __aicore__ inline void DealingMergeSort();
55+ __aicore__ inline void UpdateSortInfo();
56+ __aicore__ inline void ExtractAndCopyOut();
57+ __aicore__ inline void ClearCache();
58+ __aicore__ inline uint32_t GetRemainingTopK() const;
59+public:
60+ TQue<QuePosition::VECIN, MORE_CORE_DOUBLE_BUFFER> inputQueue_;
61+ TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> outValueQueue_;
62+ TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> outIndexQueue_;
63+ TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> sortedQueue_;
64+ TQue<QuePosition::VECIN, MORE_CORE_DOUBLE_BUFFER> copyInQueue_;
65+ TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> castValueQueue_;
66+ TQue<QuePosition::VECOUT, MORE_CORE_DOUBLE_BUFFER> castIndexQueue_;
67+ // input value
68+ GlobalTensor<T> inputValueGm_;
69+ // output value
70+ GlobalTensor<T> outValueGm_;
71+ // output index
72+ GlobalTensor<INDEX_TYPE> outIndexGm_;
73+ 
74+ GlobalTensor<CONVERT_TYPE> workspaceGm_[2];
75+ GlobalTensor<CONVERT_TYPE> workspaceInput_;
76+ GlobalTensor<CONVERT_TYPE> workspaceOutput_;
77+ 
78+ TBuf<QuePosition::VECCALC> sortedValueUb_;
79+ TBuf<QuePosition::VECCALC> sortedValueIndexUb_;
80+ TBuf<QuePosition::VECCALC> concatTempBuf_;
81+ TBuf<QuePosition::VECCALC> sortTempBuf_;
82+ TBuf<QuePosition::VECCALC> inputValueTempBuf_;
83+ TBuf<QuePosition::VECCALC> sortedValueLocalCastTbuf_;
84+ 
85+ TPipe *pipe_;
86+ uint32_t blockIdx_ = 0;
87+ uint32_t numTileData_ = 0;
88+ uint32_t sortLoopRound_ = 0;
89+ uint32_t platformCoreNum_ = 0;
90+ uint32_t inputLastDimValue_ = 0;
91+ uint32_t outputTopKValue_ = 0;
92+ uint32_t frontCoreNum_ = 0;
93+ uint32_t rowIdx_ = 0;
94+ uint32_t rowCoreIdx_ = 0;
95+ uint32_t vfLenFp32_ = Ops::Base::GetVRegSize() / MORE_CORE_FP32_DTYPE;
96+ int64_t rowDataOffset_ = 0;
97+ int64_t rowWorkspaceOffset_ = 0;
98+ int64_t rowOutputOffset_ = 0;
99+ 
100+ // SortMultiCore
101+ int64_t listNum_{0};
102+ int64_t flag_ = 0;
103+ int64_t remainListNum_{0};
104+ int64_t outOffset_{0};
105+ int64_t offsets_[4] = {0};
106+ int64_t listRemainElements_[4] = {0};
107+ int64_t currentElements_{0};
108+ int64_t currentTailElements_{0};
109+ int64_t dealLengths_[4] = {0};
110+ int64_t allRemainElements_{0};
111+ int64_t curLoopSortedNum_{0};
112+ int64_t onceMaxElements_{0};
113+ uint16_t elementCountList_[4] = {0};
114+ uint16_t validBitTail_;
115+ uint32_t listSortedNums_[4] = {0};
116+ uint32_t workSpaceFlag_ = 0;
117+ LocalTensor<CONVERT_TYPE> ubInputs_[4];
118+ LocalTensor<CONVERT_TYPE> ubMainInput_;
119+};
120+ 
121+__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
122+{
123+ if (bytes == 0) {
124+ return 0;
125+ }
126+ return (elementNum * bytes + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB / bytes;
127+}
128+ 
129+template <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,
131+ GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe)
132+{
133+ blockIdx_ = GetBlockIdx();
134+ pipe_ = pipe;
135+ inputLastDimValue_ = static_cast<uint32_t>(tilingData->lastAxisNum);
136+ outputTopKValue_ = static_cast<uint32_t>(tilingData->topKRealValue);
137+ numTileData_ = tilingData->numTileDataSize;
138+ frontCoreNum_ = tilingData->lastDimNeedCore;
139+ uint32_t sortBufferSize = 8;
140+ rowIdx_ = frontCoreNum_ == 0 ? 0 : blockIdx_ / frontCoreNum_;
141+ rowCoreIdx_ = frontCoreNum_ == 0 ? 0 : blockIdx_ % frontCoreNum_;
142+ rowDataOffset_ = static_cast<int64_t>(rowIdx_) * static_cast<int64_t>(inputLastDimValue_);
143+ 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 per
145+ // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.
146+ rowWorkspaceOffset_ =
147+ static_cast<int64_t>(rowIdx_) * Align(inputLastDimValue_, sortBufferSize) * sortBufferSize /
148+ sizeof(CONVERT_TYPE) * 2;
149+ onceMaxElements_ = tilingData->keyParams0 / MORE_CORE_DEALING_SORT_NUM_ONCE * MORE_CORE_DEALING_SORT_NUM_ONCE;
150+ 
151+ inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));
152+ outValueGm_.SetGlobalBuffer((__gm__ T*)(value));
153+ outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));
154+ workspaceGm_[0].SetGlobalBuffer(
155+ (__gm__ CONVERT_TYPE*)(workSpace) + rowWorkspaceOffset_,
156+ Align(inputLastDimValue_, sortBufferSize) * sortBufferSize / sizeof(CONVERT_TYPE));
157+ workspaceGm_[1].SetGlobalBuffer(
158+ (__gm__ CONVERT_TYPE*)(workSpace) + rowWorkspaceOffset_ +
159+ Align(inputLastDimValue_, sortBufferSize) * sortBufferSize / sizeof(CONVERT_TYPE));
160+ 
161+ 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;
163+ pipe_->InitBuffer(inputQueue_, MORE_CORE_DOUBLE_BUFFER, alignTile * sizeof(T));
164+ 
165+ pipe_->InitBuffer(sortedValueUb_, alignTile * sortBufferSize);
166+ pipe_->InitBuffer(sortedValueIndexUb_, alignTile * sizeof(uint32_t));
167+ pipe_->InitBuffer(concatTempBuf_, alignTile * sortBufferSize);
168+ pipe_->InitBuffer(sortTempBuf_, alignTile * sortBufferSize);
169+ pipe_->InitBuffer(inputValueTempBuf_, alignTile * sizeof(CONVERT_TYPE));
170+ pipe_->InitBuffer(sortedValueLocalCastTbuf_, alignTile * sortBufferSize);
171+}
172+ 
173+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
174+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::Process()
175+{
176+ int64_t offsetPerCore = 0;
177+ if (rowCoreIdx_ < frontCoreNum_ - 1) {
178+ offsetPerCore = numTileData_ * rowCoreIdx_;
179+ SortInSingleCore(numTileData_, offsetPerCore);
180+ } else if (rowCoreIdx_ == frontCoreNum_ - 1) {
181+ uint32_t tailNum = inputLastDimValue_ - numTileData_ * (frontCoreNum_ - 1);
182+ offsetPerCore = numTileData_ * rowCoreIdx_;
183+ SortInSingleCore(tailNum, offsetPerCore);
184+ }
185+ SyncAll();
186+ pipe_->Reset();
187+ uint32_t sortBufferSize = 8;
188+ pipe_->InitBuffer(sortedQueue_, MORE_CORE_DOUBLE_BUFFER,
189+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sortBufferSize);
190+ pipe_->InitBuffer(copyInQueue_, MORE_CORE_DOUBLE_BUFFER,
191+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sortBufferSize);
192+ pipe_->InitBuffer(castValueQueue_, MORE_CORE_DOUBLE_BUFFER,
193+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(CONVERT_TYPE));
194+ pipe_->InitBuffer(castIndexQueue_, MORE_CORE_DOUBLE_BUFFER,
195+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(uint32_t));
196+ if constexpr (std::is_same<bfloat16_t, T>::value) {
197+ pipe_->InitBuffer(outValueQueue_, MORE_CORE_DOUBLE_BUFFER,
198+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(T));
199+ }
200+ if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {
201+ pipe_->InitBuffer(outIndexQueue_, MORE_CORE_DOUBLE_BUFFER,
202+ MORE_CORE_MERGE_SORT_LIST_MAX_NUM * onceMaxElements_ * sizeof(INDEX_TYPE));
203+ }
204+ 
205+ listNum_ = frontCoreNum_;
206+ currentElements_ = numTileData_;
207+ currentTailElements_ = inputLastDimValue_ - numTileData_ * (frontCoreNum_ - 1);
208+ uint32_t currentCoreNum;
209+ uint32_t remainListNum;
210+ while(listNum_ > MORE_CORE_MERGE_SORT_LIST_MAX_NUM) {
211+ workspaceInput_ = workspaceGm_[workSpaceFlag_];
212+ workspaceOutput_ = workspaceGm_[1 - workSpaceFlag_];
213+ SortInMultiCore();
214+ currentCoreNum = (listNum_ + MORE_CORE_MERGE_SORT_LIST_MAX_NUM - 1) / MORE_CORE_MERGE_SORT_LIST_MAX_NUM;
215+ remainListNum = listNum_ - (currentCoreNum - 1) * MORE_CORE_MERGE_SORT_LIST_MAX_NUM;
216+ currentTailElements_ = currentElements_ * (remainListNum - 1) + currentTailElements_;
217+ listNum_ = currentCoreNum;
218+ currentElements_ = currentElements_ * MORE_CORE_MERGE_SORT_LIST_MAX_NUM;
219+ workSpaceFlag_ = (workSpaceFlag_ + 1) % 2;
220+ }
221+ workspaceInput_ = workspaceGm_[workSpaceFlag_];
222+ workspaceOutput_ = workspaceGm_[1 - workSpaceFlag_];
223+ FinalSortAndCopy();
224+}
225+ 
226+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
227+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::SortInSingleCore(
228+ uint32_t tileNum, int64_t offsetPerCore)
229+{
230+ CopyInData(tileNum, offsetPerCore);
231+ LocalTensor<T> inputLocal = inputQueue_.DeQue<T>();
232+ LocalTensor<CONVERT_TYPE> sortedValueLocal = sortedValueUb_.Get<CONVERT_TYPE>();
233+ LocalTensor<uint32_t> sortedValueIndexLocal = sortedValueIndexUb_.Get<uint32_t>();
234+ InitIndexLocal(tileNum, sortedValueIndexLocal, offsetPerCore);
235+ if constexpr (std::is_same<bfloat16_t, T>::value) {
236+ DoSortBf16(tileNum, inputLocal, sortedValueLocal, sortedValueIndexLocal);
237+ } else {
238+ DoSort(tileNum, inputLocal, sortedValueLocal, sortedValueIndexLocal);
239+ }
240+ CopyOutWorkSpace(tileNum, offsetPerCore, sortedValueLocal);
241+ inputQueue_.FreeTensor(inputLocal);
242+}
243+ 
244+template <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,
246+ int64_t offsetPerCore)
247+{
248+ LocalTensor<T> inputLocal = inputQueue_.AllocTensor<T>();
249+ T defaultValue = IS_DESCEND ? static_cast<T>(-INFINITY) : static_cast<T>(NAN);
250+ uint32_t alignTile = (tileNum + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB;
251+ Duplicate(inputLocal, defaultValue, alignTile);
252+ 
253+ event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
254+ SetFlag<HardEvent::V_MTE2>(eventId);
255+ WaitFlag<HardEvent::V_MTE2>(eventId);
256+ 
257+ uint32_t currTileSizeAlign = Align(tileNum, sizeof(T));
258+ DataCopyExtParams copyParams;
259+ copyParams.blockCount = 1;
260+ copyParams.blockLen = tileNum * sizeof(T);
261+ copyParams.srcStride = 0;
262+ copyParams.dstStride = 0;
263+ DataCopyPadExtParams<T> padParams;
264+ padParams.isPad = true;
265+ padParams.rightPadding = currTileSizeAlign - tileNum;
266+ padParams.paddingValue = static_cast<T>(defaultValue);
267+ 
268+ AscendC::DataCopyPad(inputLocal, inputValueGm_[rowDataOffset_ + offsetPerCore], copyParams, padParams);
269+ inputQueue_.EnQue(inputLocal);
270+}
271+ 
272+template <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,
274+ LocalTensor<uint32_t> sortedValueIndexLocal, int64_t offsetPerCore)
275+{
276+ PipeBarrier<PIPE_ALL>();
277+ LocalTensor<int32_t> tempIndexLocal = sortedValueIndexLocal.ReinterpretCast<int32_t>();
278+ ArithProgression<int32_t>(tempIndexLocal, offsetPerCore, 1, tileNum);
279+ PipeBarrier<PIPE_ALL>();
280+}
281+ 
282+template <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,
284+ LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal)
285+{
286+ AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = sortTempBuf_.Get<CONVERT_TYPE>();
287+ AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = concatTempBuf_.Get<CONVERT_TYPE>();
288+ AscendC::LocalTensor<CONVERT_TYPE> inputValueTempLocal = inputValueTempBuf_.Get<CONVERT_TYPE>();
289+ AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = sortedValueLocalCastTbuf_.Get<CONVERT_TYPE>();
290+ 
291+ 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;
293+ uint32_t concatRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_CONCAT_NUM_ONCE - 1) / MORE_CORE_DEALING_CONCAT_NUM_ONCE;
294+ 
295+ AscendC::Cast(inputValueTempLocal, inputLocal, AscendC::RoundMode::CAST_NONE, aglinTileNum);
296+ if constexpr (!IS_DESCEND) {
297+ FlipSignBit(inputValueTempLocal, aglinTileNum);
298+ }
299+ AscendC::LocalTensor<CONVERT_TYPE> concatLocal;
300+ AscendC::Concat(concatLocal, inputValueTempLocal, concatTempLocal, concatRepeatTimes);
301+ AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal, sortRepeatTimes);
302+}
303+ 
304+template <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,
306+ LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal, LocalTensor<uint32_t> sortedValueIndexLocal)
307+{
308+ AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = sortTempBuf_.Get<CONVERT_TYPE>();
309+ AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = concatTempBuf_.Get<CONVERT_TYPE>();
310+ AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = sortedValueLocalCastTbuf_.Get<CONVERT_TYPE>();
311+ 
312+ 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;
314+ uint32_t concatRepeatTimes = (aglinTileNum + MORE_CORE_DEALING_CONCAT_NUM_ONCE - 1) / MORE_CORE_DEALING_CONCAT_NUM_ONCE;
315+ 
316+ if constexpr (!IS_DESCEND) {
317+ FlipSignBit(inputLocal, aglinTileNum);
318+ }
319+ AscendC::LocalTensor<CONVERT_TYPE> concatLocal;
320+ AscendC::Concat(concatLocal, inputLocal, concatTempLocal, concatRepeatTimes);
321+ AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal, sortRepeatTimes);
322+}
323+ 
324+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
325+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::FlipSignBit(
326+ LocalTensor<CONVERT_TYPE> xLocal, uint32_t aglinTileNum)
327+{
328+ if constexpr (std::is_same<float, CONVERT_TYPE>::value) {
329+ AscendC::LocalTensor<int32_t> castTensor = xLocal.template ReinterpretCast<int32_t>();
330+ AscendC::Adds(castTensor, castTensor, MORE_CORE_XOR_OP_VALUE_FP, aglinTileNum);
331+ } else if constexpr (std::is_same<half, CONVERT_TYPE>::value ){
332+ AscendC::LocalTensor<int16_t> castTensor = xLocal.template ReinterpretCast<int16_t>();
333+ AscendC::Adds(castTensor, castTensor, MORE_CORE_XOR_OP_VALUE_HALF, aglinTileNum);
334+ }
335+}
336+ 
337+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
338+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::CopyOutWorkSpace(
339+ uint32_t tileNum, int64_t offsetPerCore, LocalTensor<CONVERT_TYPE> sortedValueLocal)
340+{
341+ event_t eventIdVToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
342+ SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
343+ WaitFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
344+ 
345+ DataCopyExtParams copyParams;
346+ copyParams.blockCount = 1;
347+ copyParams.blockLen = GetSortLen<CONVERT_TYPE>(tileNum) * sizeof(CONVERT_TYPE);
348+ copyParams.srcStride = 0;
349+ copyParams.dstStride = 0;
350+ DataCopyPad(workspaceGm_[0][GetSortLen<CONVERT_TYPE>(offsetPerCore)], sortedValueLocal, copyParams);
351+}
352+ 
353+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
354+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::CopyOutMultiCore()
355+{
356+ LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.DeQue<CONVERT_TYPE>();
357+ 
358+ uint32_t len = curLoopSortedNum_;
359+ DataCopyExtParams copyParams;
360+ copyParams.blockCount = 1;
361+ copyParams.blockLen = GetSortLen<CONVERT_TYPE>(len) * sizeof(CONVERT_TYPE);
362+ copyParams.srcStride = 0;
363+ copyParams.dstStride = 0;
364+ DataCopyPad(workspaceOutput_[outOffset_], sortTempBuffer, copyParams);
365+ outOffset_ += GetSortLen<CONVERT_TYPE>(len);
366+ sortedQueue_.FreeTensor<CONVERT_TYPE>(sortTempBuffer);
367+}
368+ 
369+ 
370+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
371+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::SortInMultiCore()
372+{
373+ uint32_t needCoreNum = (listNum_ + MORE_CORE_MERGE_SORT_LIST_MAX_NUM - 1) / MORE_CORE_MERGE_SORT_LIST_MAX_NUM;
374+ if (rowCoreIdx_ < needCoreNum) {
375+ MultiCoreInit();
376+ for (; allRemainElements_ > 0;) {
377+ CopyInMultiCore();
378+ UpdateMrgParam();
379+ DealingMergeSort();
380+ UpdateSortInfo();
381+ CopyOutMultiCore();
382+ }
383+ ClearCache();
384+ }
385+ SyncAll();
386+}
387+ 
388+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
389+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::CopyInMultiCore()
390+{
391+ LocalTensor<CONVERT_TYPE> ubMainInput = copyInQueue_.AllocTensor<CONVERT_TYPE>();
392+ remainListNum_ = 0;
393+ for (int64_t i = 0, j = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
394+ dealLengths_[i] = (onceMaxElements_ > listRemainElements_[i] ? listRemainElements_[i] : onceMaxElements_);
395+ if (dealLengths_[i] > 0) {
396+ DataCopyExtParams copyParams;
397+ copyParams.blockCount = 1;
398+ copyParams.blockLen = GetSortLen<CONVERT_TYPE>(dealLengths_[i]) * sizeof(CONVERT_TYPE);
399+ copyParams.srcStride = 0;
400+ copyParams.dstStride = 0;
401+ DataCopyPadExtParams<CONVERT_TYPE> padParams{false, 0, 0, 0};
402+ DataCopyPad(ubMainInput[GetSortLen<CONVERT_TYPE>(onceMaxElements_) * i], workspaceInput_[offsets_[i]], copyParams, padParams);
403+ elementCountList_[j] = dealLengths_[i];
404+ remainListNum_ += 1;
405+ j++;
406+ }
407+ }
408+ copyInQueue_.EnQue(ubMainInput);
409+}
410+ 
411+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
412+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::DealingMergeSort() {
413+ LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.AllocTensor<CONVERT_TYPE>();
414+ LocalTensor<CONVERT_TYPE> ubMainInput = copyInQueue_.DeQue<CONVERT_TYPE>();
415+ LocalTensor<CONVERT_TYPE> tmpUbInputs[4];
416+ for (int64_t i = 0, j = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
417+ if (dealLengths_[i] > 0) {
418+ tmpUbInputs[j] = ubMainInput[GetSortLen<CONVERT_TYPE>(onceMaxElements_) * i];
419+ j++;
420+ }
421+ }
422+ if (remainListNum_ == 2) {
423+ MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);
424+ MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);
425+ } else if (remainListNum_ == 3) {
426+ MrgSortSrcList sortListTail =
427+ MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[0]);
428+ MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);
429+ } else if (remainListNum_ == 4) {
430+ MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[3]);
431+ MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, elementCountList_, listSortedNums_, validBitTail_, 1);
432+ } else {
433+ AscendC::Copy(sortTempBuffer, tmpUbInputs[0], Align(GetSortLen<CONVERT_TYPE>(elementCountList_[0]), sizeof(CONVERT_TYPE)));
434+ listSortedNums_[0] = elementCountList_[0];
435+ }
436+ sortedQueue_.EnQue(sortTempBuffer);
437+ copyInQueue_.FreeTensor(ubMainInput);
438+}
439+ 
440+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
441+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::UpdateSortInfo()
442+{
443+ curLoopSortedNum_ = 0;
444+ for (int64_t i = 0, j = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
445+ if (dealLengths_[i] > 0) {
446+ // update remain size
447+ listRemainElements_[i] -= listSortedNums_[j];
448+ allRemainElements_ -= listSortedNums_[j];
449+ // update offset
450+ offsets_[i] += GetSortOffset<CONVERT_TYPE>(listSortedNums_[j]);
451+ // update current loop sorted nums
452+ curLoopSortedNum_ += listSortedNums_[j];
453+ j++;
454+ }
455+ }
456+}
457+ 
458+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
459+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::MultiCoreInit()
460+{
461+ outOffset_ = GetSortLen<CONVERT_TYPE>(rowCoreIdx_ * MORE_CORE_MERGE_SORT_LIST_MAX_NUM * currentElements_);
462+ 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;
464+ if (blockNum < listNum_ - 1) {
465+ listRemainElements_[i] = currentElements_;
466+ offsets_[i] = GetSortOffset<CONVERT_TYPE>(blockNum * currentElements_);
467+ allRemainElements_ += listRemainElements_[i];
468+ } else if (blockNum == listNum_ - 1) {
469+ listRemainElements_[i] = currentTailElements_;
470+ offsets_[i] = GetSortOffset<CONVERT_TYPE>(blockNum * currentElements_);
471+ allRemainElements_ += currentTailElements_;
472+ } else {
473+ listRemainElements_[i] = 0;
474+ }
475+ }
476+}
477+ 
478+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
479+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::UpdateMrgParam()
480+{
481+ if (remainListNum_ == 2) {
482+ elementCountList_[2] = 0;
483+ elementCountList_[3] = 0;
484+ validBitTail_ = 0b0011;
485+ } else if (remainListNum_ == 3) {
486+ elementCountList_[3] = 0;
487+ validBitTail_ = 0b0111;
488+ } else if (remainListNum_ == 4) {
489+ validBitTail_ = 0b1111;
490+ } else {
491+ elementCountList_[1] = 0;
492+ elementCountList_[2] = 0;
493+ elementCountList_[3] = 0;
494+ validBitTail_ = 0b0001;
495+ }
496+}
497+ 
498+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
499+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::FinalSortAndCopy()
500+{
501+ if (rowCoreIdx_ == 0) {
502+ FinalSortInit();
503+ for (; allRemainElements_ > 0 && outOffset_ < outputTopKValue_;) {
504+ CopyInMultiCore();
505+ UpdateMrgParam();
506+ DealingMergeSort();
507+ UpdateSortInfo();
508+ ExtractAndCopyOut();
509+ }
510+ ClearCache();
511+ }
512+ SyncAll();
513+}
514+ 
515+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
516+__aicore__ inline uint32_t TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::GetRemainingTopK() const
517+{
518+ if (outOffset_ >= static_cast<int64_t>(outputTopKValue_)) {
519+ return 0;
520+ }
521+ return static_cast<uint32_t>(static_cast<int64_t>(outputTopKValue_) - outOffset_);
522+}
523+ 
524+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
525+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::ExtractAndCopyOut()
526+{
527+ uint32_t copyCount = curLoopSortedNum_;
528+ uint32_t remainingTopK = GetRemainingTopK();
529+ if (remainingTopK == 0) {
530+ LocalTensor<CONVERT_TYPE> sortTempBufferDrop = sortedQueue_.DeQue<CONVERT_TYPE>();
531+ sortedQueue_.FreeTensor(sortTempBufferDrop);
532+ return;
533+ }
534+ if (copyCount > remainingTopK) {
535+ copyCount = remainingTopK;
536+ }
537+ LocalTensor<T> ubOutput1;
538+ if constexpr (std::is_same<bfloat16_t, T>::value) {
539+ ubOutput1 = outValueQueue_.AllocTensor<T>();
540+ }
541+ LocalTensor<INDEX_TYPE> ubOutput2;
542+ if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {
543+ ubOutput2 = outIndexQueue_.AllocTensor<INDEX_TYPE>();
544+ }
545+ LocalTensor<CONVERT_TYPE> sortTempBuffer = sortedQueue_.DeQue<CONVERT_TYPE>();
546+ LocalTensor<CONVERT_TYPE> castValue = castValueQueue_.AllocTensor<CONVERT_TYPE>();
547+ 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));
549+ if constexpr (!IS_DESCEND) {
550+ FlipSignBit(castValue, ((curLoopSortedNum_ + MORE_CORE_BLOCK_UB - 1) / MORE_CORE_BLOCK_UB * MORE_CORE_BLOCK_UB));
551+ }
552+ DataCopyExtParams copyParamsValue;
553+ copyParamsValue.blockCount = 1;
554+ copyParamsValue.blockLen = copyCount * sizeof(T);
555+ copyParamsValue.srcStride = 0;
556+ copyParamsValue.dstStride = 0;
557+ 
558+ DataCopyExtParams copyParamsIndex;
559+ copyParamsIndex.blockCount = 1;
560+ copyParamsIndex.blockLen = copyCount * sizeof(INDEX_TYPE);
561+ copyParamsIndex.srcStride = 0;
562+ copyParamsIndex.dstStride = 0;
563+ 
564+ uint32_t sortedValueAlign = Align(curLoopSortedNum_, sizeof(CONVERT_TYPE));
565+ if constexpr (std::is_same<bfloat16_t, T>::value) {
566+ AscendC::Cast(ubOutput1, castValue, AscendC::RoundMode::CAST_RINT, sortedValueAlign);
567+ outValueQueue_.EnQue(ubOutput1);
568+ ubOutput1 = outValueQueue_.DeQue<T>();
569+ DataCopyPad(outValueGm_[rowOutputOffset_ + outOffset_], ubOutput1, copyParamsValue);
570+ outValueQueue_.FreeTensor(ubOutput1);
571+ } else {
572+ castValueQueue_.EnQue(castValue);
573+ castValue = castValueQueue_.DeQue<T>();
574+ DataCopyPad(outValueGm_[rowOutputOffset_ + outOffset_], castValue, copyParamsValue);
575+ }
576+ castValueQueue_.FreeTensor(castValue);
577+ 
578+ uint32_t sortedIndexAlign = Align(curLoopSortedNum_, sizeof(uint32_t));
579+ LocalTensor<int32_t> castIndexTemp = castIndex.template ReinterpretCast<int32_t>();
580+ if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {
581+ AscendC::Cast(ubOutput2, castIndexTemp, AscendC::RoundMode::CAST_NONE, sortedIndexAlign);
582+ outIndexQueue_.EnQue(ubOutput2);
583+ ubOutput2 = outIndexQueue_.DeQue<INDEX_TYPE>();
584+ DataCopyPad(outIndexGm_[rowOutputOffset_ + outOffset_], ubOutput2, copyParamsIndex);
585+ outIndexQueue_.FreeTensor(ubOutput2);
586+ } else {
587+ castIndexQueue_.EnQue(castIndex);
588+ castIndex = castIndexQueue_.DeQue<uint32_t>();
589+ castIndexTemp = castIndex.template ReinterpretCast<int32_t>();
590+ DataCopyPad(outIndexGm_[rowOutputOffset_ + outOffset_], castIndexTemp, copyParamsIndex);
591+ }
592+ castIndexQueue_.FreeTensor(castIndex);
593+ sortedQueue_.FreeTensor(sortTempBuffer);
594+ outOffset_ += copyCount;
595+}
596+ 
597+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
598+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::FinalSortInit()
599+{
600+ outOffset_ = 0;
601+ for (int64_t i = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
602+ if (i < listNum_ - 1) {
603+ listRemainElements_[i] = currentElements_;
604+ offsets_[i] = GetSortOffset<CONVERT_TYPE>(i * currentElements_);
605+ allRemainElements_ += listRemainElements_[i];
606+ } else if (i == listNum_ - 1) {
607+ listRemainElements_[i] = currentTailElements_;
608+ offsets_[i] = GetSortOffset<CONVERT_TYPE>(i * currentElements_);
609+ allRemainElements_ += currentTailElements_;
610+ } else {
611+ listRemainElements_[i] = 0;
612+ }
613+ }
614+}
615+ 
616+template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
617+__aicore__ inline void TopKMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::ClearCache()
618+{
619+ allRemainElements_ = 0;
620+ outOffset_ = 0;
621+ remainListNum_ = 0;
622+ for (int64_t i = 0; i < MORE_CORE_MERGE_SORT_LIST_MAX_NUM; i++) {
623+ offsets_[i] = 0;
624+ listRemainElements_[i] = 0;
625+ elementCountList_[i] = 0;
626+ }
627+}
628+} // namespace topkV2
629+#endif //TOP_K_MERGE_SORT_MORE_CORE_H
Mmath/top_k_v2/op_kernel/top_k_v2_apt.cpp+47-1
@@ -19,6 +19,8 @@
19#include "arch35/radix_sort_top_k.h"19#include "arch35/radix_sort_top_k.h"
20#include "arch35/radix_topk_constant.h"20#include "arch35/radix_topk_constant.h"
21#include "arch35/top_k_merge_sort.h"21#include "arch35/top_k_merge_sort.h"
22+#include "arch35/top_k_merge_sort_more_core.h"
23+#include "arch35/top_k_merge_sort_intra_core.h"
22#include "arch35/radix_sort_top_k_single_block.h"24#include "arch35/radix_sort_top_k_single_block.h"
23#include "arch35/radix_sort_top_k_single_core.h"25#include "arch35/radix_sort_top_k_single_core.h"
24#include "arch35/radix_sort_top_k_inter_core_template_optimization.h"26#include "arch35/radix_sort_top_k_inter_core_template_optimization.h"
@@ -41,6 +43,8 @@ using namespace SortAndTopK;
41#define TOPK_MERGE_SORT_TILING_KEY_FLOAT 1300343#define TOPK_MERGE_SORT_TILING_KEY_FLOAT 13003
42#define TOPK_MERGE_SORT_TILING_KEY_FLOAT16 1300244#define TOPK_MERGE_SORT_TILING_KEY_FLOAT16 13002
43#define TOPK_MERGE_SORT_TILING_KEY_BF16 1400245#define TOPK_MERGE_SORT_TILING_KEY_BF16 14002
46+#define TOPK_MERGE_SORT_MORE_CORE_TILING_KEY_FLOAT 23003
47+#define TOPK_MERGE_SORT_INTRA_CORE_TILING_KEY_FLOAT 33003
44 48 
45const uint32_t SINGLE_CORE_MODE = 1;49const uint32_t SINGLE_CORE_MODE = 1;
46const uint32_t MULT_CORE_OPTIM_MODE = 4;50const uint32_t MULT_CORE_OPTIM_MODE = 4;
@@ -280,6 +284,42 @@ __aicore__ inline void generateMergeTopKObject(
280 }284 }
281}285}
282 286 
287+template <typename T, typename CONVERT_TYPE, typename INDEX_DTYPE>
288+__aicore__ inline void generateMergeTopKMoreCoreObject(
289+ GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR globalWorkGm, GM_ADDR tiling)
290+{
291+ GET_TILING_DATA(tilingData, tiling);
292+ bool isLargest = (tilingData.isLargest > 0) ? true : false;
293+ TPipe pipe;
294+ if (isLargest) {
295+ topkV2::TopKMergeSortMoreCore<T, CONVERT_TYPE, true, INDEX_DTYPE> mergeSort;
296+ mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);
297+ mergeSort.Process();
298+ } else {
299+ topkV2::TopKMergeSortMoreCore<T, CONVERT_TYPE, false, INDEX_DTYPE> mergeSort;
300+ mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);
301+ mergeSort.Process();
302+ }
303+}
304+ 
305+template <typename T, typename CONVERT_TYPE, typename INDEX_DTYPE>
306+__aicore__ inline void generateMergeTopKIntraCoreObject(
307+ GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR globalWorkGm, GM_ADDR tiling)
308+{
309+ GET_TILING_DATA(tilingData, tiling);
310+ bool isLargest = (tilingData.isLargest > 0) ? true : false;
311+ TPipe pipe;
312+ if (isLargest) {
313+ topkV2::TopKMergeSortIntraCore<T, INDEX_DTYPE, true> mergeSort;
314+ mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);
315+ mergeSort.Process();
316+ } else {
317+ topkV2::TopKMergeSortIntraCore<T, INDEX_DTYPE, false> mergeSort;
318+ mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);
319+ mergeSort.Process();
320+ }
321+}
322+ 
283extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling)323extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling)
284{324{
285 if (workspace == nullptr) {325 if (workspace == nullptr) {
@@ -350,11 +390,17 @@ extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR val
350 #if ORIG_DTYPE_X == DT_FLOAT390 #if ORIG_DTYPE_X == DT_FLOAT
351 TILING_KEY_IS(TOPK_COMMON_TILING_KEY_FLOAT);391 TILING_KEY_IS(TOPK_COMMON_TILING_KEY_FLOAT);
352 TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_FLOAT);392 TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_FLOAT);
393+ TILING_KEY_IS(TOPK_MERGE_SORT_MORE_CORE_TILING_KEY_FLOAT);
394+ TILING_KEY_IS(TOPK_MERGE_SORT_INTRA_CORE_TILING_KEY_FLOAT);
353 395 
354 #if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_FLOAT396 #if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_FLOAT
355 generateOpObject<float, uint32_t, topkV2::B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);397 generateOpObject<float, uint32_t, topkV2::B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
356 #elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_FLOAT398 #elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_FLOAT
357 generateMergeTopKObject<float, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);399 generateMergeTopKObject<float, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);
400+ #elif TILING_KEY_VAR == TOPK_MERGE_SORT_MORE_CORE_TILING_KEY_FLOAT
401+ generateMergeTopKMoreCoreObject<float, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);
402+ #elif TILING_KEY_VAR == TOPK_MERGE_SORT_INTRA_CORE_TILING_KEY_FLOAT
403+ generateMergeTopKIntraCoreObject<float, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);
358 #endif404 #endif
359 #endif405 #endif
360 406 
@@ -380,4 +426,4 @@ extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR val
380 #endif426 #endif
381 #endif427 #endif
382}428}
383-#endif // TOP_K_V2_APT_H429+#endif // TOP_K_V2_APT_H
Mmath/top_k_v2/tests/ut/op_host/arch35/test_top_k_v2_tiling.cpp+1-1
@@ -52,7 +52,7 @@ TEST_F(TopKV2Tiling, test_tiling_small_merge_sort) {
52 },52 },
53 &compileInfo);53 &compileInfo);
54 uint64_t expectTilingKey = 13003;54 uint64_t expectTilingKey = 13003;
55- string expectTilingData = "4294967297 42949672961 4294967306 274877906976 0 4294967297 0 0 1 32 8 1 8 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ";55+ string expectTilingData = "4294967297 42949672961 4294967306 274877906976 0 4294967297 0 0 1 32 8 1 8 1 0 0 0 0 0 0 0 0 0 0 8589934592 0 0 0 ";
56 std::vector<size_t> expectWorkspaces = {16787584};56 std::vector<size_t> expectWorkspaces = {16787584};
57 ExecuteTestCase(57 ExecuteTestCase(
58 tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingData, expectWorkspaces);58 tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingData, expectWorkspaces);