Topk算子支持超大排序轴处理-ops-math-AtomGit

Mmath/sort_with_index/op_kernel/arch35/constant_var_simd.h+2-0

	@@ -13,6 +13,7 @@
13	*/	13	*/
14	#ifndef SORT_WITH_INDEX_CONSTANT_VAR_SIMD_H	14	#ifndef SORT_WITH_INDEX_CONSTANT_VAR_SIMD_H
15	#define SORT_WITH_INDEX_CONSTANT_VAR_SIMD_H	15	#define SORT_WITH_INDEX_CONSTANT_VAR_SIMD_H
		16	+namespace SortWithIndex {
16	const uint32_t THREAD_DIM_NUM = 1024;	17	const uint32_t THREAD_DIM_NUM = 1024;
17	const int32_t ONE_TIMES_B64_NUM = 32;	18	const int32_t ONE_TIMES_B64_NUM = 32;
18	const int32_t ONE_TIMES_B32_NUM = 64;	19	const int32_t ONE_TIMES_B32_NUM = 64;
	@@ -64,4 +65,5 @@ const int16_t XOR_OP_VALUE_HALF = 0x8000;
64	const int32_t SMALL_SORT_MAX_DATA_SIZE = 128;	65	const int32_t SMALL_SORT_MAX_DATA_SIZE = 128;
65	const uint32_t CONCAT_AGLIN_VALUE = 16;	66	const uint32_t CONCAT_AGLIN_VALUE = 16;
66	const uint32_t DOUBLE_BUFFER = 2;	67	const uint32_t DOUBLE_BUFFER = 2;
		68	+}
67	#endif	69	#endif

Mmath/sort_with_index/op_kernel/arch35/merge_sort.h+4-0

	@@ -20,7 +20,10 @@
20	#include "../../sort/arch35/util_type_simd.h" // 引入使用 ROUND_UP_AGLIN	20	#include "../../sort/arch35/util_type_simd.h" // 引入使用 ROUND_UP_AGLIN
21	#include "util_type_simd.h"	21	#include "util_type_simd.h"
22		22
		23	+namespace SortWithIndex {
		24	+
23	using namespace AscendC;	25	using namespace AscendC;
		26	+
24	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>	27	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
25	struct MergeSort {	28	struct MergeSort {
26	__aicore__ inline MergeSort() {}	29	__aicore__ inline MergeSort() {}
	@@ -277,4 +280,5 @@ __aicore__ inline void MergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST,
277	outIndexQueue_.FreeTensor(outIndexLocal);	280	outIndexQueue_.FreeTensor(outIndexLocal);
278	outValueQueue_.FreeTensor(outValueLocal);	281	outValueQueue_.FreeTensor(outValueLocal);
279	}	282	}
		283	+}
280	#endif	284	#endif

Mmath/sort_with_index/op_kernel/arch35/merge_sort_simd.h+4-0

	@@ -18,7 +18,10 @@
18	#include "constant_var_simd.h"	18	#include "constant_var_simd.h"
19	#include "util_type_simd.h"	19	#include "util_type_simd.h"
20		20
		21	+namespace SortWithIndex {
		22	+
21	using namespace AscendC;	23	using namespace AscendC;
		24	+
22	template <typename T, typename CONVERT_TYPE, bool IS_DESCEND>	25	template <typename T, typename CONVERT_TYPE, bool IS_DESCEND>
23	struct KernelVbsMergeSort {	26	struct KernelVbsMergeSort {
24	public:	27	public:
	@@ -170,4 +173,5 @@ __aicore__ inline void KernelVbsMergeSort<T, CONVERT_TYPE, IS_DESCEND>::flipSign
170	AscendC::Adds(castTensor, castTensor, XOR_OP_VALUE_HALF, aglinTileSize);	173	AscendC::Adds(castTensor, castTensor, XOR_OP_VALUE_HALF, aglinTileSize);
171	}	174	}
172	}	175	}
		176	+}
173	#endif	177	#endif

Mmath/sort_with_index/op_kernel/arch35/merge_sort_with_index.h+3-1

	@@ -23,6 +23,8 @@
23	#include <algorithm>	23	#include <algorithm>
24	#include "util_type_simd.h"	24	#include "util_type_simd.h"
25		25
		26	+namespace SortWithIndex {
		27	+
26	using namespace AscendC;	28	using namespace AscendC;
27		29
28	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>	30	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
	@@ -275,5 +277,5 @@ __aicore__ inline void MergeSortWithIndex<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_
275	}	277	}
276	this->inQueueXInt64Index_.EnQue(indexLocal);	278	this->inQueueXInt64Index_.EnQue(indexLocal);
277	}	279	}
278	-	280	+}
279	#endif	281	#endif

Mmath/sort_with_index/op_kernel/arch35/merge_sort_with_index_simd.h+4-1

	@@ -17,7 +17,10 @@
17	#include "constant_var_simd.h"	17	#include "constant_var_simd.h"
18	#include "merge_sort_simd.h"	18	#include "merge_sort_simd.h"
19		19
		20	+namespace SortWithIndex {
		21	+
20	using namespace AscendC;	22	using namespace AscendC;
		23	+
21	template <typename T, typename CONVERT_TYPE, bool IS_DESCEND>	24	template <typename T, typename CONVERT_TYPE, bool IS_DESCEND>
22	struct KernelVbsMergeSortWithIndex : public KernelVbsMergeSort<T, CONVERT_TYPE, IS_DESCEND>{	25	struct KernelVbsMergeSortWithIndex : public KernelVbsMergeSort<T, CONVERT_TYPE, IS_DESCEND>{
23	public:	26	public:
	@@ -101,5 +104,5 @@ __aicore__ inline void KernelVbsMergeSortWithIndex<T, CONVERT_TYPE, IS_DESCEND>:
101	}	104	}
102	AscendC::Cast(sortedValueLocal, sortedValueLocalCast, AscendC::RoundMode::CAST_RINT, aglinTileSize * nowCoreRealRowNum);	105	AscendC::Cast(sortedValueLocal, sortedValueLocalCast, AscendC::RoundMode::CAST_RINT, aglinTileSize * nowCoreRealRowNum);
103	}	106	}
104	-	107	+}
105	#endif	108	#endif

Mmath/sort_with_index/op_kernel/arch35/radix_sort_with_index_multi_block.h+3-3

Mmath/sort_with_index/op_kernel/arch35/radix_sort_with_index_single_block.h+2-2

	@@ -29,13 +29,13 @@ public:
29	__aicore__ inline void Init(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,	29	__aicore__ inline void Init(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,
30	GM_ADDR workspace, const SortWithIndexTilingDataSimt* tilingData, TPipe* pipe);	30	GM_ADDR workspace, const SortWithIndexTilingDataSimt* tilingData, TPipe* pipe);
31	__aicore__ inline void Process();	31	__aicore__ inline void Process();
		32	+ __aicore__ inline void ProcessSingleBlock(GlobalTensor<XType> xGm, GlobalTensor<IndexType> indexGm, uint64_t loopRound);
32		33
33	private:	34	private:
34	__aicore__ inline void CopyIn(GlobalTensor<XType> inputX, uint64_t tileOffset, uint32_t tileData);	35	__aicore__ inline void CopyIn(GlobalTensor<XType> inputX, uint64_t tileOffset, uint32_t tileData);
35	__aicore__ inline void CopyIndexIn(GlobalTensor<IndexType> inputIndex, uint64_t tileOffset, uint32_t tileData);	36	__aicore__ inline void CopyIndexIn(GlobalTensor<IndexType> inputIndex, uint64_t tileOffset, uint32_t tileData);
36	- __aicore__ inline void ProcessSingleBlock(GlobalTensor<XType> xGm, GlobalTensor<IndexType> indexGm, uint64_t loopRound);
37		37
38	-private:	38	+protected:
39	// 输入GlobalTensor	39	// 输入GlobalTensor
40	GlobalTensor<XType> xGm_;	40	GlobalTensor<XType> xGm_;
41	GlobalTensor<IndexType> indexGm_;	41	GlobalTensor<IndexType> indexGm_;

Mmath/sort_with_index/op_kernel/arch35/util_type_simd.h+31-0

	@@ -15,6 +15,9 @@
15	#define SORT_WITH_INDEX_UTIL_TYPE_SIMD_H	15	#define SORT_WITH_INDEX_UTIL_TYPE_SIMD_H
16	#include "kernel_operator.h"	16	#include "kernel_operator.h"
17	#include "constant_var_simd.h"	17	#include "constant_var_simd.h"
		18	+
		19	+namespace SortWithIndex {
		20	+
18	template <typename Tp, Tp v>	21	template <typename Tp, Tp v>
19	struct integral_constant {	22	struct integral_constant {
20	static constexpr Tp value = v;	23	static constexpr Tp value = v;
	@@ -31,4 +34,32 @@ __aicore__ inline T SortGetMin(T left, T right)
31	{	34	{
32	return (left > right ? right : left);	35	return (left > right ? right : left);
33	}	36	}
		37	+
		38	+#ifdef TOP_K_V2_APT_H
		39	+struct SortWithIndexTilingDataSimt {
		40	+ int32_t isDescend;
		41	+ uint32_t sortLoopTimes;
		42	+ uint32_t unsortedDimParallel;
		43	+ uint64_t unsortedDimNum;
		44	+ uint32_t lastDimTileNum;
		45	+ uint32_t lastDimNeedCore;
		46	+ uint32_t numTileDataSize;
		47	+ uint32_t sortAcApiNeedBufferSize;
		48	+ uint32_t mergSortAcApiNeedBufferSize;
		49	+ uint32_t oneCoreRowNum;
		50	+ uint32_t outputLastDimValue;
		51	+ uint32_t isInInt32Range;
		52	+ int64_t lastAxisNum;
		53	+ uint32_t keyParams0;
		54	+ uint32_t keyParams1;
		55	+ uint32_t keyParams2;
		56	+ uint32_t keyParams3;
		57	+ uint32_t keyParams4;
		58	+ uint32_t keyParams5;
		59	+ uint32_t tmpUbSize;
		60	+ uint32_t modeType;
		61	+};
		62	+#endif
		63	+
		64	+}
34	#endif	65	#endif

Mmath/top_k_v2/CMakeLists.txt+2-3

	@@ -15,6 +15,5 @@ set(SUPPORT_TILING_DIR "arch35" "arch35")
15	add_all_modules_sources(OPTYPE top_k_v2 ACLNNTYPE aclnn_exclude	15	add_all_modules_sources(OPTYPE top_k_v2 ACLNNTYPE aclnn_exclude
16	COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT}	16	COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT}
17	TILING_DIR ${SUPPORT_TILING_DIR}	17	TILING_DIR ${SUPPORT_TILING_DIR}
18	- DISABLE_IN_OPP TRUE)	18	+ DISABLE_IN_OPP TRUE
19	-	19	+ DEPENDENCIES sort sort_with_index)
20	-

Amath/top_k_v2/op_host/arch35/sort_with_index_tiling.h+566-0

@@ -0,0 +1,566 @@
	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 sort_with_index_tiling.h
	13	+* \brief sort_with_index ac tiling impl
	14	+*/
	15	+
	16	+#include "top_k_v2_tiling_arch35.h"
	17	+#include "log/log.h"
	18	+#include "tiling/tiling_api.h"
	19	+#include "op_host/tiling_util.h"
	20	+#include "util/platform_util.h"
	21	+
	22	+namespace optiling {
	23	+namespace sortWithIndex {
	24	+constexpr size_t WORK_SPACE_SIZE = 16777216; // 16 * 1024 * 1024;
	25	+const uint32_t BIN_NUM = 256;
	26	+const uint32_t TILE_DATA_NUM = 4096;
	27	+const uint32_t MEDIUM_TILE_DATA_NUM = 2048;
	28	+const uint32_t SMALL_TILE_DATA_NUM = 1024;
	29	+const uint32_t TILE_DATA_NUM_B64 = 2048;
	30	+const uint32_t TMP_UB = 1024; // 暂预留给sort高级api的大小
	31	+const uint32_t CONST_10 = 10; // int32索引时, 计算各种tensor的乘法因子
	32	+const uint32_t CONST_14 = 14; // int64索引时, 计算各种tensor的乘法因子
	33	+const uint32_t CONST_6 = 6; // int64索引时, 计算各种tensor的乘法因子
	34	+const uint32_t CONST_1 = 1;
	35	+const uint32_t CONST_2 = 2;
	36	+const uint32_t INT64_BYTE = 8;
	37	+const uint32_t INT32_BYTE = 4;
	38	+const uint32_t SIMT_UB = 32768; // 预留了32k给simt使用
	39	+const uint32_t NEED_UB_SIZE_BYTE = 221184; // 预留了32k给simt使用
	40	+const uint32_t SMALL_SORT_MAX_DATA_SIZE = 512;
	41	+const uint32_t AGLIN_VALUE = 32;
	42	+const uint32_t MERGE_SORT_TILING_OFFSET = 10000;
	43	+const uint32_t UB_CONST_INT32 = 4096; // 输出idx为int32时kernel侧需要的固定ub大小
	44	+const uint32_t UB_CONST_INT64 = 7168; // 输出idx为int64时kernel侧需要的固定ub大小
	45	+// 排序轴在int32范围内的最大值, 超过这个值, cutsum，前缀和就要用int64数据范围表示
	46	+const uint32_t INT32_MAX_RANGE_VALUE = 1073741823;
	47	+const uint32_t SMALL_SIZE_OPTIM_MODE = 0;
	48	+const uint32_t SMALL_SIZE_MODE = 1;
	49	+const uint32_t MULT_CORE_MODE = 2;
	50	+struct SortTileInfo {
	51	+ uint32_t coreNumNeed = 0;
	52	+ uint32_t lastDimTileNum = 0;
	53	+ uint32_t unsortedDimParallel = 1;
	54	+ uint32_t oneCoreRowNum = 1;
	55	+ uint32_t ubSize = 0;
	56	+ uint32_t blockUbSize = 0;
	57	+ uint32_t dtypeSize = 0;
	58	+ uint32_t y2DtypeSize = 0;
	59	+ uint32_t maxCoreNum = 0;
	60	+ uint32_t numTileDataSize = 0;
	61	+ uint32_t sortLoopTimes = 0;
	62	+ uint32_t lastDimNeedCore = 0;
	63	+ uint32_t keyParams0 = 0;
	64	+ uint32_t keyParams1 = 0;
	65	+ uint32_t keyParams2 = 0;
	66	+ uint32_t keyParams3 = 0;
	67	+ uint32_t keyParams4 = 0;
	68	+ uint32_t keyParams5 = 0;
	69	+ uint32_t tmpUbSize = 0;
	70	+ bool isDescend = false;
	71	+ ge::DataType dataType = ge::DT_UINT8;
	72	+ uint32_t isInt32 = 0;
	73	+ int32_t xDimNum = 0;
	74	+ int64_t sortAxisNum = 1;
	75	+ int64_t unSortDimNum = 1;
	76	+};
	77	+
	78	+static const std::map<ge::DataType, uint32_t> tilingDataTypeKeyMap = {
	79	+ {ge::DT_INT64, 1004}, {ge::DT_INT32, 1003}, {ge::DT_INT16, 1002}, {ge::DT_INT8, 1001},
	80	+ {ge::DT_UINT64, 2004}, {ge::DT_UINT32, 2003}, {ge::DT_UINT16, 2002}, {ge::DT_UINT8, 2001},
	81	+ {ge::DT_FLOAT, 3003}, {ge::DT_FLOAT16, 3002}, {ge::DT_BF16, 4002}};
	82	+static const std::map<ge::DataType, uint32_t> tilingDataTypeBitMap = {
	G G Guoqh3月27日适当添加空行
	83	+ {ge::DT_INT64, 8}, {ge::DT_INT32, 4}, {ge::DT_INT16, 2}, {ge::DT_INT8, 1},
	84	+ {ge::DT_UINT64, 8}, {ge::DT_UINT32, 4}, {ge::DT_UINT16, 2}, {ge::DT_UINT8, 1},
	85	+ {ge::DT_FLOAT, 4}, {ge::DT_FLOAT16, 2}, {ge::DT_BF16, 2}};
	86	+static const std::map<ge::DataType, uint32_t> optDataTypeBitMap = {
	87	+ {ge::DT_FLOAT, 4}, {ge::DT_FLOAT16, 2}, {ge::DT_BF16, 2}};
	88	+
	89	+uint32_t CeilDiv1(int64_t a, int64_t b)
	G G Guoqh3月27日 CeilDiv
	90	+{
	91	+ if (b == 0) {
	92	+ return static_cast<uint32_t>(a);
	93	+ }
	94	+ return static_cast<uint32_t>((a + b - 1) / b);
	95	+}
	96	+
	97	+template <typename T>
	98	+auto CeilDivMul1(int64_t a, int64_t b) ->T const
	G G Guoqh3月27日 CeilDivMul
	99	+{
	100	+ if (b == 0) {
	101	+ return static_cast<T>(a);
	102	+ }
	103	+ return static_cast<T>(((a + b - 1) / b) * b);
	104	+}
	105	+
	106	+void SetSortTmpSizeOfIdx(ge::DataType dataType, int64_t lastAxisNum, uint32_t tileData, bool isDescend, bool hasIndex,
	G G Guoqh3月27日参数个数过多，不要超过五个
	107	+ TopKV2TilingDataSimd& topkTilingData, sortWithIndex::SortTileInfo& sortTileInfo)
	108	+{
	109	+ int64_t reanLen = std::min(lastAxisNum, static_cast<int64_t>(tileData));
	110	+ std::vector<int64_t> shapeVec = {reanLen};
	111	+ ge::Shape srcShape(shapeVec);
	112	+ AscendC::SortConfig config;
	113	+ config.type = AscendC::SortType::RADIX_SORT;
	114	+ config.isDescend = isDescend;
	115	+ // SortWithIndex and no need to cut axis
	116	+ config.hasSrcIndex = hasIndex && (reanLen == lastAxisNum);
	117	+ config.hasDstIndex = true;
	118	+ uint32_t maxValue = 0;
	119	+ uint32_t minValue = 0;
	120	+ AscendC::GetSortMaxMinTmpSize(srcShape, dataType, ge::DT_UINT32, false, config, maxValue, minValue);
	121	+ OP_LOGW("[SortWithIndexTilingForAscendC]", "Allocal buffer element len = %lu ac sort api", reanLen);
	122	+ OP_LOGW("[SortWithIndexTilingForAscendC]", "Need tmp buffer %u byte for ac sort api", maxValue);
	123	+ topkTilingData.set_sortAcApiNeedBufferSizeForSort(maxValue);
	124	+ sortTileInfo.tmpUbSize = maxValue;
	125	+}
	126	+
	127	+void SetMergeSortTmpSizeOfIdx(gert::TilingContext* context, ge::DataType dataType, int64_t lastAxisNum,
	128	+ TopKV2TilingDataSimd& topkTilingData)
	129	+{
	130	+ uint32_t reanLen = 0;
	131	+ if ((lastAxisNum <= sortWithIndex::SMALL_SORT_MAX_DATA_SIZE) && (sortWithIndex::optDataTypeBitMap.count(dataType) != 0)) {
	132	+ reanLen = std::min(static_cast<uint32_t>(lastAxisNum), sortWithIndex::SMALL_SORT_MAX_DATA_SIZE);
	133	+ }
	134	+ uint32_t aglinDataSize = static_cast<uint32_t>((reanLen + sortWithIndex::AGLIN_VALUE - 1) /
	135	+ sortWithIndex::AGLIN_VALUE * sortWithIndex::AGLIN_VALUE);
	136	+ uint32_t dataTypeSize = (dataType == ge::DT_BF16) ? sortWithIndex::optDataTypeBitMap.find(ge::DT_FLOAT)->second :
	137	+ sortWithIndex::optDataTypeBitMap.find(dataType)->second;
	138	+ auto platform_info = context->GetPlatformInfo();
	139	+ if (nullptr == platform_info) {
	140	+ OP_LOGE("[SortWithIndexTilingForAscendC]", "platform_info is nullptr.");
	141	+ }
	142	+ auto plat = platform_ascendc::PlatformAscendC(platform_info);
	143	+ uint32_t dataSizeNeed = AscendC::GetConcatTmpSize(plat, aglinDataSize, dataTypeSize);
	144	+ OP_LOGW("[SortWithIndexTilingForAscendC]", "Allocal buffer mergesort element len = %u ac sort api", reanLen);
	145	+ OP_LOGW("[SortWithIndexTilingForAscendC]", "Merge sort need tmp buffer %u byte for ac api", dataSizeNeed);
	146	+ topkTilingData.set_mergSortAcApiNeedBufferSizeForSort(dataSizeNeed);
	147	+}
	148	+
	149	+void TileModeSmallSizeOptimOfIdx(
	150	+ uint64_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum, uint32_t tileData, SortTileInfo& sortTileInfo)
	151	+{
	152	+ uint32_t aglinNum = static_cast<uint32_t>((lastAxisNum + AGLIN_VALUE - 1) / AGLIN_VALUE * AGLIN_VALUE);
	153	+ uint32_t oneCoreRowNum = static_cast<uint32_t>((tileData / 2) / aglinNum);
	154	+ oneCoreRowNum = static_cast<uint32_t>(oneCoreRowNum == 0 ? 1 : oneCoreRowNum);
	155	+ uint32_t virUnsortedDimNum = static_cast<uint32_t>((unsortedDimNum + oneCoreRowNum - 1) / oneCoreRowNum);
	156	+ uint32_t coreNumNeed = 0;
	157	+ OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TileModeSmallSizeOptimOfIdx", "maxCoreNum is zero"), return);
	158	+ uint32_t sortLoopTimes = static_cast<uint32_t>((virUnsortedDimNum + maxCoreNum - 1) / maxCoreNum);
	159	+ if (sortLoopTimes == 1u) {
	160	+ uint32_t realCoreNum = virUnsortedDimNum % maxCoreNum;
	161	+ if (realCoreNum == 0u) {
	162	+ realCoreNum = maxCoreNum;
	163	+ }
	164	+ coreNumNeed = realCoreNum;
	165	+ } else {
	166	+ coreNumNeed = maxCoreNum;
	167	+ }
	168	+ sortTileInfo.coreNumNeed = coreNumNeed;
	169	+ sortTileInfo.lastDimTileNum = 1U;
	170	+ sortTileInfo.unsortedDimParallel = coreNumNeed;
	171	+ sortTileInfo.oneCoreRowNum = oneCoreRowNum;
	172	+ sortTileInfo.lastDimNeedCore = 1;
	173	+ sortTileInfo.sortLoopTimes = sortLoopTimes;
	174	+ sortTileInfo.numTileDataSize = lastAxisNum;
	175	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "Small size opt mode coreNumNeed=%u, sortLoopTimes=%u, lastAxisNum=%ld, "
	176	+ "oneCoreRowNum=%ld, ubsize=%lu.", coreNumNeed, sortLoopTimes, lastAxisNum, oneCoreRowNum, sortTileInfo.ubSize);
	177	+}
	178	+
	179	+void TileModeSmallSizeOfIdx(uint64_t unsortedDimNum, uint32_t maxCoreNum, int64_t lastAxisNum,
	180	+ sortWithIndex::SortTileInfo& sortTileInfo)
	181	+{
	182	+ uint32_t coreNumNeed = 0;
	183	+ OP_CHECK_IF(maxCoreNum == 0, OP_LOGE("TileModeSmallSizeOfIdx", "maxCoreNum is zero"), return);
	184	+ uint32_t sortLoopTimes = static_cast<uint32_t>((unsortedDimNum + maxCoreNum - 1) / maxCoreNum);
	185	+ if (sortLoopTimes == 1u) {
	186	+ uint32_t realCoreNum = unsortedDimNum % maxCoreNum;
	187	+ if (realCoreNum == 0u) {
	188	+ realCoreNum = maxCoreNum;
	189	+ }
	190	+ coreNumNeed = realCoreNum;
	191	+ } else {
	192	+ coreNumNeed = maxCoreNum;
	193	+ }
	194	+ sortTileInfo.coreNumNeed = coreNumNeed;
	195	+ sortTileInfo.lastDimTileNum = static_cast<uint32_t>(1);
	196	+ sortTileInfo.unsortedDimParallel = coreNumNeed;
	197	+ sortTileInfo.lastDimNeedCore = 1;
	198	+ sortTileInfo.numTileDataSize = static_cast<uint32_t>(lastAxisNum);
	199	+ sortTileInfo.sortLoopTimes = sortLoopTimes;
	200	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "Small size mode coreNumNeed=%u sortLoopTimes=%u "
	201	+ "lastAxisNum=%ld", coreNumNeed, sortLoopTimes, lastAxisNum);
	202	+}
	203	+
	204	+void PrintTilingDataOfIdx(sortWithIndex::SortTileInfo& sortTileInfo, TopKV2TilingDataSimd &topkTilingData)
	205	+{
	206	+ OP_LOGI(
	207	+ "[Print SortWithIndexTilingForAscendC TilingData]",
	208	+ "coreNum is %u, lastAxisNum is %ld, isInInt32Range is %u, "
	209	+ "sortLoopTimes is %u, unsortedDimParallel is %u, unsortedDimNum is %u, "
	210	+ "lastDimTileNum is %u, lastDimNeedCore is %u, numTileDataSize is %u, "
	211	+ "sortAcApiNeedBufferSize is %u, mergSortAcApiNeedBufferSize is %u, "
	212	+ "oneCoreRowNum is %u, outputLastDimValue is %u, tmp ub size is %u, "
	213	+ "keyParams0 is %u, keyParams1 is %u, keyParams2 is %u, keyParams3 is %u, keyParams4 is %u, "
	214	+ "keyParams5 is %u, ub avalibal size=%lu, modeType=%u.",
	215	+ sortTileInfo.coreNumNeed, topkTilingData.get_lastAxisNumForSort(), topkTilingData.get_isInInt32RangeForSort(),
	216	+ topkTilingData.get_sortLoopTimesForSort(), topkTilingData.get_unsortedDimParallelForSort(),
	217	+ topkTilingData.get_unsortedDimNumForSort(), topkTilingData.get_lastDimTileNumForSort(),
	218	+ topkTilingData.get_lastDimNeedCoreForSort(),topkTilingData.get_numTileDataSizeForSort(),
	219	+ topkTilingData.get_sortAcApiNeedBufferSizeForSort(),topkTilingData.get_mergSortAcApiNeedBufferSizeForSort(),
	220	+ topkTilingData.get_oneCoreRowNumForSort(), topkTilingData.get_outputLastDimValueForSort(),
	221	+ topkTilingData.get_tmpUbSize(), topkTilingData.get_keyParams0(), topkTilingData.get_keyParams1(),
	222	+ topkTilingData.get_keyParams2(), topkTilingData.get_keyParams3(), topkTilingData.get_keyParams4(),
	223	+ topkTilingData.get_keyParams5(), sortTileInfo.ubSize, topkTilingData.get_modeTypeForSort());
	224	+ return;
	225	+}
	226	+
	227	+void FillRadixSortTilingDataSort(sortWithIndex::SortTileInfo &sortTileInfo, TopKV2TilingDataSimd &topkTilingData)
	228	+{
	229	+ topkTilingData.set_numTileDataSizeForSort(sortTileInfo.numTileDataSize);
	230	+ topkTilingData.set_unsortedDimParallelForSort(sortTileInfo.unsortedDimParallel);
	231	+ topkTilingData.set_lastDimTileNumForSort(sortTileInfo.lastDimTileNum);
	232	+ topkTilingData.set_sortLoopTimesForSort(sortTileInfo.sortLoopTimes);
	233	+ topkTilingData.set_lastDimNeedCoreForSort(sortTileInfo.lastDimNeedCore);
	234	+ topkTilingData.set_keyParams0(sortTileInfo.keyParams0);
	235	+ topkTilingData.set_keyParams1(sortTileInfo.keyParams1);
	236	+ topkTilingData.set_keyParams2(sortTileInfo.keyParams2);
	237	+ topkTilingData.set_keyParams3(sortTileInfo.keyParams3);
	238	+ topkTilingData.set_keyParams4(sortTileInfo.keyParams4);
	239	+ topkTilingData.set_keyParams5(sortTileInfo.keyParams5);
	240	+ topkTilingData.set_tmpUbSize(sortTileInfo.tmpUbSize);
	241	+ topkTilingData.set_lastAxisNumForSort(sortTileInfo.sortAxisNum);
	242	+ topkTilingData.set_unsortedDimNumForSort(sortTileInfo.unSortDimNum);
	243	+ return;
	244	+}
	245	+
	246	+void SetSortTmpSize1(ge::DataType dataType, uint32_t tileData, bool isDescend, SortTileInfo &sortTileInfo)
	247	+{
	248	+ int64_t realLen = std::min(sortTileInfo.sortAxisNum, static_cast<int64_t>(tileData));
	249	+ std::vector<int64_t> shapeVec = { realLen };
	250	+ ge::Shape srcShape(shapeVec);
	251	+ AscendC::SortConfig config;
	252	+ config.type = AscendC::SortType::RADIX_SORT;
	253	+ config.isDescend = isDescend;
	254	+ config.hasSrcIndex = false;
	255	+ config.hasDstIndex = true;
	256	+ uint32_t maxValue = 0, minValue = 0;
	257	+ AscendC::GetSortMaxMinTmpSize(srcShape, dataType, ge::DT_UINT32, false, config, maxValue, minValue);
	258	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "api of sort shape is %ld, maxUb is %u", realLen, maxValue);
	259	+ sortTileInfo.tmpUbSize = maxValue;
	260	+ return;
	261	+}
	262	+
	263	+uint32_t ComputeRemainUb1(SortTileInfo &sortTileInfo, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor)
	264	+{
	265	+ uint32_t tmpUb = sortTileInfo.ubSize - (ubExtra + tileFactor * tileData);
	266	+ OP_LOGD("[SortWithIndexTilingForAscendC]", "ComputeRemainUb1 ubSize=%u, ubExtra=%lu, tileFactor=%lu, "
	267	+ "tileData=%lu, tmpUb=%lu.", sortTileInfo.ubSize, ubExtra, tileFactor, tileData, tmpUb);
	268	+ return tmpUb;
	269	+}
	270	+
	271	+void AdjTmpUb1(SortTileInfo &sortTileInfo, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor)
	272	+{
	273	+ uint32_t remainUbNew = ComputeRemainUb1(sortTileInfo, tileData, ubExtra, tileFactor) - sortTileInfo.tmpUbSize;
	274	+ remainUbNew = remainUbNew > sortTileInfo.blockUbSize ? (remainUbNew - sortTileInfo.blockUbSize) : uint32_t(0);
	275	+ uint32_t alignUbSize = (remainUbNew / sortTileInfo.blockUbSize) * sortTileInfo.blockUbSize;
	276	+ OP_LOGD("[SortWithIndexTilingForAscendC]", "alignUbSize %u, sortTileInfo.tmpUbSize=%lu, "
	277	+ "sortTileInfo.blockUbSize=%lu.", alignUbSize, sortTileInfo.tmpUbSize, sortTileInfo.blockUbSize);
	278	+ sortTileInfo.tmpUbSize = sortTileInfo.tmpUbSize + alignUbSize; // 剩余的ub都给tmpUbsize
	279	+}
	280	+
	281	+void ComputeTileDataOne1(SortTileInfo &sortTileInfo, uint32_t lastDimTileNum, uint32_t ubExtra, uint32_t &tileData,
	282	+ uint32_t tileFactor)
	283	+{
	284	+ uint32_t allCore = CeilDivMul1<uint32_t>(int64_t(lastDimTileNum), int64_t(sortTileInfo.maxCoreNum));
	285	+ uint32_t newTileData = CeilDiv1(sortTileInfo.sortAxisNum, int64_t(allCore));
	286	+ tileData = CeilDivMul1<uint32_t>(int64_t(newTileData), int64_t(BIN_NUM));
	287	+ tileData = std::max(tileData, SMALL_TILE_DATA_NUM);
	288	+ SetSortTmpSize1(ge::DT_UINT8, tileData, false, sortTileInfo);
	289	+ AdjTmpUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	290	+ return;
	291	+}
	292	+
	293	+bool NeedAdjTileData1(SortTileInfo &sortTileInfo, uint32_t &tileData, uint32_t lastDimTileNum, uint32_t ubExtra,
	294	+ uint32_t tileFactor)
	295	+{
	296	+ if (sortTileInfo.unSortDimNum == 1L && lastDimTileNum == 1U) {
	297	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "unSortDimNum and lastDimTileNum is 1");
	298	+ uint32_t newTileData = CeilDiv1(sortTileInfo.sortAxisNum, int64_t(sortTileInfo.maxCoreNum));
	299	+ newTileData = CeilDivMul1<uint32_t>(int64_t(newTileData), int64_t(BIN_NUM));
	300	+ tileData = std::max(newTileData, SMALL_TILE_DATA_NUM);
	301	+ SetSortTmpSize1(ge::DT_UINT8, tileData, false, sortTileInfo);
	302	+ AdjTmpUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	303	+ return true;
	304	+ }
	305	+ if (sortTileInfo.unSortDimNum == 1L \|\| (lastDimTileNum >= sortTileInfo.maxCoreNum)) {
	306	+ // b为1时，尽量均匀分核，同时保证处理的最小的tile_data为1024
	307	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "unSortDimNum is 1 and lastDimTileNum greater than allCore");
	308	+ ComputeTileDataOne1(sortTileInfo, lastDimTileNum, ubExtra, tileData, tileFactor);
	309	+ return true;
	310	+ }
	311	+ if (sortTileInfo.unSortDimNum > 1L && sortTileInfo.unSortDimNum < int64_t(sortTileInfo.maxCoreNum) &&
	312	+ lastDimTileNum == 1U) {
	313	+ OP_LOGI("[SortWithIndexTilingForAscendC]",
	314	+ "unSortDimNum greater than 1,and unSortDimNum small and lastDimTileNum is one");
	315	+ uint32_t hCore = sortTileInfo.maxCoreNum / static_cast<uint32_t>(sortTileInfo.unSortDimNum);
	316	+ uint32_t hTileData = static_cast<uint32_t>(sortTileInfo.sortAxisNum) / hCore;
	317	+ tileData = CeilDivMul1<uint32_t>(int64_t(hTileData), int64_t(BIN_NUM));
	318	+ SetSortTmpSize1(ge::DT_UINT8, tileData, false, sortTileInfo);
	319	+ AdjTmpUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	320	+ return tileData;
	321	+ }
	322	+ if (sortTileInfo.unSortDimNum > 1L && lastDimTileNum > 1U) {
	323	+ // b大于1且h轴循环次数小于总核数，也就是b轴核数大于1
	324	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "unSortDimNum is one, lastDimTileNum greater than one");
	325	+ int64_t newTileData = sortTileInfo.sortAxisNum / int64_t(lastDimTileNum);
	J J JimmyLam23月27日建议使用static_cast代替c风格的类型转换
	326	+ tileData = CeilDivMul1<uint32_t>(newTileData, int64_t(BIN_NUM));
	327	+ lastDimTileNum = CeilDiv1(sortTileInfo.sortAxisNum, int64_t(tileData));
	328	+ uint32_t bCore = lastDimTileNum == 0 ? sortTileInfo.maxCoreNum : sortTileInfo.maxCoreNum / lastDimTileNum;
	329	+ if (lastDimTileNum < sortTileInfo.maxCoreNum && sortTileInfo.unSortDimNum < int64_t(sortTileInfo.maxCoreNum)) {
	330	+ if (sortTileInfo.unSortDimNum < int64_t(bCore)) {
	331	+ bCore = static_cast<uint32_t>(sortTileInfo.unSortDimNum);
	332	+ uint32_t hCore = sortTileInfo.maxCoreNum / bCore;
	333	+ uint32_t tileDataNew = CeilDiv1(int64_t(sortTileInfo.sortAxisNum), int64_t(hCore));
	334	+ tileData = CeilDivMul1<uint32_t>(int64_t(tileDataNew), int64_t(BIN_NUM));
	335	+ }
	336	+ }
	337	+ if (bCore == 1U && lastDimTileNum < sortTileInfo.maxCoreNum) {
	338	+ ComputeTileDataOne1(sortTileInfo, lastDimTileNum, ubExtra, tileData, tileFactor);
	339	+ }
	340	+ SetSortTmpSize1(ge::DT_UINT8, tileData, false, sortTileInfo);
	341	+ AdjTmpUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	342	+ return true;
	343	+ }
	344	+ return false;
	345	+}
	346	+
	347	+uint32_t ComputeTileData1(SortTileInfo &sortTileInfo)
	348	+{
	349	+ uint32_t ubExtra;
	350	+ uint32_t tileFactor;
	351	+ if (sortTileInfo.isInt32 == 0U) { // 数据范围超过int32, y2DtypeSize表示索引类型
	352	+ ubExtra = UB_CONST_INT64;
	353	+ tileFactor = CONST_6 + sortTileInfo.dtypeSize + sortTileInfo.y2DtypeSize;
	354	+ } else {
	355	+ ubExtra = UB_CONST_INT32;
	356	+ tileFactor = CONST_6 + sortTileInfo.dtypeSize + sortTileInfo.y2DtypeSize;
	357	+ }
	358	+
	359	+ uint32_t tileData = (sortTileInfo.ubSize - ubExtra) / tileFactor;
	360	+ tileData = (tileData / BIN_NUM) * BIN_NUM;
	361	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "ubExtra=%u, tileFactor=%u, dtypeSize=%u, y2DtypeSize=%lu, "
	362	+ "tileData=%lu.", ubExtra, tileFactor, sortTileInfo.dtypeSize, sortTileInfo.y2DtypeSize, tileData);
	363	+ uint32_t remainUb = ComputeRemainUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	364	+ SetSortTmpSize1(ge::DT_UINT8, tileData, false, sortTileInfo);
	365	+
	366	+ uint32_t tmpUbSize = sortTileInfo.tmpUbSize;
	367	+ while (tmpUbSize > remainUb) {
	368	+ tileData = tileData - BIN_NUM;
	369	+ remainUb = ComputeRemainUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	370	+ SetSortTmpSize1(ge::DT_UINT8, tileData, false, sortTileInfo);
	371	+ tmpUbSize = sortTileInfo.tmpUbSize;
	372	+ }
	373	+ uint32_t lastDimTileNum = CeilDiv1(sortTileInfo.sortAxisNum, int64_t(tileData));
	374	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "tileData %u, lastDimTileNum %u, tmpUbSize %u", tileData,
	375	+ lastDimTileNum, tmpUbSize);
	376	+ bool smallTile =
	377	+ (sortTileInfo.sortAxisNum <= static_cast<int64_t>(SMALL_TILE_DATA_NUM)) && lastDimTileNum == uint32_t(1);
	378	+ if ((lastDimTileNum % sortTileInfo.maxCoreNum == 0U) \|\| smallTile) {
	379	+ OP_LOGI("[SortWithIndexTilingForAscendC]", "lastDimTileNum align or smallTile");
	380	+ AdjTmpUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	381	+ return tileData;
	382	+ }
	383	+ if (NeedAdjTileData1(sortTileInfo, tileData, lastDimTileNum, ubExtra, tileFactor)) {
	384	+ return tileData;
	385	+ }
	386	+ AdjTmpUb1(sortTileInfo, tileData, ubExtra, tileFactor);
	387	+ return tileData;
	388	+}
	389	+
	390	+void ComputeWorkSpace1(sortWithIndex::SortTileInfo &sortTileInfo, size_t* usrSize)
	391	+{
	392	+ uint32_t dtypeSizeWk = static_cast<uint32_t>(sizeof(int32_t));
	393	+ if (sortTileInfo.isInt32 == 0U) {
	394	+ dtypeSizeWk = static_cast<uint32_t>(sizeof(int64_t));
	395	+ }
	396	+ size_t excusiveBinsGmWkSize = static_cast<size_t>(sortTileInfo.keyParams1) * sortTileInfo.keyParams4 * dtypeSizeWk;
	397	+ excusiveBinsGmWkSize = sortWithIndex::CeilDivMul1<size_t>(int64_t(excusiveBinsGmWkSize),
	398	+ int64_t(sortTileInfo.blockUbSize));
	399	+
	400	+ size_t globalHistGmWkSize =
	401	+ static_cast<size_t>(sortTileInfo.keyParams3) * sortTileInfo.keyParams2 * sortTileInfo.keyParams0 * dtypeSizeWk;
	402	+ globalHistGmWkSize = sortWithIndex::CeilDivMul1<size_t>(int64_t(globalHistGmWkSize), int64_t(sortTileInfo.blockUbSize));
	403	+
	404	+ size_t outIdxDbWK =
	405	+ static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel * sortTileInfo.y2DtypeSize;
	406	+ outIdxDbWK = sortWithIndex::CeilDivMul1<size_t>(int64_t(outIdxDbWK), int64_t(sortTileInfo.blockUbSize));
	407	+
	408	+ size_t histTileGmWk = static_cast<size_t>(sortTileInfo.lastDimTileNum) * sortWithIndex::BIN_NUM *
	409	+ sortTileInfo.unsortedDimParallel * sizeof(int16_t) * sortWithIndex::CONST_2;
	410	+
	411	+ size_t xB8GmWkSize = static_cast<size_t>(sortTileInfo.lastDimTileNum) * sortTileInfo.numTileDataSize *
	412	+ sortTileInfo.unsortedDimParallel;
	413	+ xB8GmWkSize = sortWithIndex::CeilDivMul1<size_t>(int64_t(xB8GmWkSize), int64_t(sortTileInfo.blockUbSize));
	414	+
	415	+ size_t outValueDbWKSize =
	416	+ static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel * sortTileInfo.dtypeSize;
	417	+ outValueDbWKSize = sortWithIndex::CeilDivMul1<size_t>(int64_t(outValueDbWKSize), int64_t(sortTileInfo.blockUbSize));
	418	+ *usrSize += excusiveBinsGmWkSize + globalHistGmWkSize + outIdxDbWK + histTileGmWk + xB8GmWkSize + outValueDbWKSize;
	419	+ OP_LOGD("[SortWithIndexTilingForAscendC]",
	420	+ "excusiveBinsGmWkSize=%lu, globalHistGmWkSize=%lu, histTileGmWk=%lu,"
	421	+ " xB8GmWkSize=%lu, outValueDbWKSize=%lu, outIdxDbWK=%lu, usrSize=%lu.",
	422	+ excusiveBinsGmWkSize, globalHistGmWkSize, histTileGmWk, xB8GmWkSize, outValueDbWKSize, outIdxDbWK, *usrSize);
	423	+ return;
	424	+}
	425	+
	426	+void TileMoreCoreModeOfIdx(sortWithIndex::SortTileInfo &sortTileInfo, size_t* usrSize)
	427	+{
	428	+ uint32_t tileData = sortWithIndex::ComputeTileData1(sortTileInfo);
	429	+ uint32_t lastDimTileNum = sortWithIndex::CeilDiv1(int64_t(sortTileInfo.sortAxisNum), int64_t(tileData));
	J J JimmyLam23月27日该函数大于50行，建议拆分
	430	+ if (sortTileInfo.maxCoreNum <= lastDimTileNum) {
	431	+ sortTileInfo.unsortedDimParallel = 1U;
	432	+ } else {
	433	+ sortTileInfo.unsortedDimParallel =
	434	+ lastDimTileNum == 0U ? sortTileInfo.maxCoreNum : sortTileInfo.maxCoreNum / lastDimTileNum;
	435	+ if (sortTileInfo.unSortDimNum < static_cast<int64_t>(sortTileInfo.unsortedDimParallel)) {
	436	+ sortTileInfo.unsortedDimParallel = static_cast<uint32_t>(sortTileInfo.unSortDimNum);
	437	+ }
	438	+ }
	439	+ sortTileInfo.numTileDataSize = tileData;
	440	+ sortTileInfo.sortLoopTimes =
	441	+ sortWithIndex::CeilDiv1(int64_t(sortTileInfo.unSortDimNum), int64_t(sortTileInfo.unsortedDimParallel));
	442	+ sortTileInfo.lastDimNeedCore = std::min(sortTileInfo.maxCoreNum, lastDimTileNum);
	443	+ sortTileInfo.coreNumNeed = sortTileInfo.unsortedDimParallel * sortTileInfo.lastDimNeedCore;
	444	+ sortTileInfo.lastDimTileNum = lastDimTileNum;
	445	+
	446	+ uint32_t ubSizeNum = sortTileInfo.tmpUbSize / static_cast<uint32_t>(sizeof(uint32_t));
	447	+ if (sortTileInfo.isInt32 == 0U) {
	448	+ ubSizeNum = sortTileInfo.tmpUbSize / static_cast<uint32_t>(sizeof(int64_t));
	449	+ }
	450	+ uint32_t allNumGloblHist = sortWithIndex::BIN_NUM * lastDimTileNum * sortTileInfo.dtypeSize *
	451	+ sortTileInfo.unsortedDimParallel;
	452	+ uint32_t allNumExcusiveBin = sortWithIndex::BIN_NUM * sortTileInfo.dtypeSize * sortTileInfo.unsortedDimParallel;
	453	+ uint32_t oneCoreSize = sortWithIndex::CeilDiv1(int64_t(allNumGloblHist), int64_t(sortTileInfo.coreNumNeed));
	454	+ sortTileInfo.keyParams5 =
	455	+ std::max(static_cast<int64_t>(oneCoreSize), static_cast<int64_t>(sortTileInfo.blockUbSize));
	456	+ sortTileInfo.keyParams0 = sortWithIndex::CeilDiv1(int64_t(allNumGloblHist), int64_t(sortTileInfo.keyParams5));
	457	+ sortTileInfo.keyParams3 = sortWithIndex::CeilDiv1(int64_t(sortTileInfo.keyParams5), int64_t(ubSizeNum));
	458	+ sortTileInfo.keyParams2 = sortTileInfo.keyParams5 > ubSizeNum ? ubSizeNum : sortTileInfo.keyParams5;
	459	+
	460	+ uint32_t oneCoreSize1 = sortWithIndex::CeilDiv1(int64_t(allNumExcusiveBin), int64_t(sortTileInfo.coreNumNeed));
	461	+ sortTileInfo.keyParams4 =
	462	+ std::max(static_cast<int64_t>(oneCoreSize1), static_cast<int64_t>(sortTileInfo.blockUbSize));
	463	+
	464	+ sortTileInfo.keyParams1 = sortWithIndex::CeilDiv1(int64_t(allNumExcusiveBin), int64_t(sortTileInfo.keyParams4));
	465	+ ComputeWorkSpace1(sortTileInfo, usrSize);
	466	+ return;
	467	+}
	468	+
	469	+ge::graphStatus RadixSortTilingOfIdx(gert::TilingContext* context, TopKV2TilingDataSimd &topkTilingData,
	470	+ int32_t maxCoreNum, size_t* usrSize)
	471	+{
	472	+ OP_LOGI(context->GetNodeName(), "SortWithIndexTIling for topk start");
	473	+ auto dataType = context->GetInputDesc(0)->GetDataType();
	474	+ const gert::Shape outShape = context->GetOutputShape(0)->GetStorageShape();
	475	+ auto y2DType = context->GetOutputDesc(1)->GetDataType();
	476	+ auto tilingKey = sortWithIndex::tilingDataTypeKeyMap.find(dataType)->second;
	477	+ std::string opType(context->GetNodeType());
	478	+
	479	+ auto platformInfo = context->GetPlatformInfo();
	480	+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
	481	+ uint64_t ubSize = 0;
	482	+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
	483	+ OP_CHECK_IF(ubSize <= static_cast<uint64_t>(sortWithIndex::SIMT_UB),
	484	+ OP_LOGE(context->GetNodeName(), "block total ub size must greater than simtUb, "
	485	+ "but is %lu", ubSize), return ge::GRAPH_FAILED);
	486	+ OP_LOGW(context->GetNodeName(), "Get op_type[%s]", opType.c_str());
	487	+
	488	+ auto const attrs = context->GetAttrs();
	489	+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
	490	+ uint32_t xDimNum = outShape.GetDimNum();
	491	+ const bool* isDescending = attrs->GetAttrPointer<bool>(2);
	492	+ OP_CHECK_NULL_WITH_CONTEXT(context, isDescending);
	493	+ OP_LOGI(context->GetNodeName(), "isDescending=%u", *isDescending);
	494	+ int64_t sortAxisNum = outShape.GetDim(xDimNum - 1);
	495	+ uint64_t unSortDimNum = 1;
	496	+ for (uint32_t i = 0u; i < static_cast<uint32_t>(xDimNum - 1); i++) {
	497	+ unSortDimNum *= outShape.GetDim(i);
	498	+ }
	499	+
	500	+ uint32_t isInInt32Range = static_cast<uint32_t>(sortAxisNum <= sortWithIndex::INT32_MAX_RANGE_VALUE);
	501	+ topkTilingData.set_isInInt32RangeForSort(isInInt32Range);
	502	+
	503	+ uint32_t tileData = sortWithIndex::TILE_DATA_NUM;
	504	+ if (dataType == ge::DT_UINT64 \|\| dataType == ge::DT_INT64) {
	505	+ tileData = sortWithIndex::TILE_DATA_NUM_B64;
	506	+ } else {
	507	+ tileData = sortWithIndex::TILE_DATA_NUM;
	508	+ }
	509	+
	510	+ sortWithIndex::SortTileInfo sortTileInfo;
	511	+ // 预留给SIMT使用
	512	+ sortTileInfo.ubSize = ubSize - sortWithIndex::SIMT_UB;
	513	+ uint32_t blockUbAglinSize = Ops::Base::GetUbBlockSize(context);
	514	+ sortTileInfo.blockUbSize = blockUbAglinSize;
	515	+ sortTileInfo.dtypeSize = sortWithIndex::tilingDataTypeBitMap.find(dataType)->second;
	516	+ sortTileInfo.y2DtypeSize = sortWithIndex::tilingDataTypeBitMap.find(y2DType)->second;
	517	+ sortTileInfo.maxCoreNum = maxCoreNum;
	518	+ sortTileInfo.dataType = dataType;
	519	+ sortTileInfo.xDimNum = xDimNum;
	520	+ sortTileInfo.sortAxisNum = sortAxisNum;
	521	+ sortTileInfo.unSortDimNum = unSortDimNum;
	522	+ sortTileInfo.isInt32 = isInInt32Range;
	523	+ sortTileInfo.numTileDataSize = tileData;
	524	+
	525	+ // 设置高级api tmpUbSize需要的空间
	526	+ SetSortTmpSizeOfIdx(dataType, sortAxisNum, tileData, *isDescending, true, topkTilingData, sortTileInfo);
	527	+ if (sortAxisNum <= sortWithIndex::SMALL_SORT_MAX_DATA_SIZE && sortWithIndex::optDataTypeBitMap.count(dataType) != 0) {
	528	+ topkTilingData.set_modeTypeForSort(sortWithIndex::SMALL_SIZE_OPTIM_MODE);
	529	+ uint32_t tileDataS = sortWithIndex::TILE_DATA_NUM;
	530	+ TileModeSmallSizeOptimOfIdx(unSortDimNum, maxCoreNum, sortAxisNum, tileDataS, sortTileInfo);
	531	+ SetMergeSortTmpSizeOfIdx(context, dataType, sortAxisNum, topkTilingData);
	532	+ tilingKey += sortWithIndex::MERGE_SORT_TILING_OFFSET;
	533	+ } else if (sortAxisNum <= static_cast<int64_t>(tileData)) {
	534	+ topkTilingData.set_modeTypeForSort(sortWithIndex::SMALL_SIZE_MODE);
	535	+ TileModeSmallSizeOfIdx(unSortDimNum, maxCoreNum, sortAxisNum, sortTileInfo);
	536	+ } else {
	537	+ // more core radix sort case
	538	+ topkTilingData.set_modeTypeForSort(sortWithIndex::MULT_CORE_MODE);
	539	+ TileMoreCoreModeOfIdx(sortTileInfo, usrSize);
	540	+ }
	541	+ OP_LOGI(context->GetNodeName(), "ubSize: %ld, ubAglinSize: %ld, dtypeSize: %u, y2DtypeSize=%u,"
	542	+ " sortTileInfo.ubSize=%u, maxCoreNum=%lu, usrSize=%d, modeType=%u.",
	543	+ ubSize, blockUbAglinSize, sortTileInfo.dtypeSize, sortTileInfo.y2DtypeSize, sortTileInfo.ubSize, maxCoreNum,
	544	+ *usrSize, topkTilingData.get_modeTypeForSort());
	545	+
	546	+ topkTilingData.set_tilingKeyForSort(tilingKey);
	547	+ topkTilingData.set_lastAxisNumForSort(sortAxisNum);
	548	+ topkTilingData.set_unsortedDimNumForSort(unSortDimNum);
	549	+ topkTilingData.set_oneCoreRowNumForSort(sortTileInfo.oneCoreRowNum);
	550	+ topkTilingData.set_outputLastDimValueForSort(sortAxisNum);
	551	+ FillRadixSortTilingDataSort(sortTileInfo, topkTilingData);
	552	+ PrintTilingDataOfIdx(sortTileInfo, topkTilingData);
	553	+
	554	+ // add sortwithindex workspace
	555	+ int64_t topkValuesGmSize = CeilDivMul1<int64_t>(int64_t(sortAxisNum * unSortDimNum * sortTileInfo.dtypeSize),
	556	+ int64_t(AGLIN_VALUE));
	557	+ int64_t topkIndicesGmSize = CeilDivMul1<int64_t>(int64_t(sortAxisNum * unSortDimNum * sortTileInfo.y2DtypeSize),
	558	+ int64_t(AGLIN_VALUE));
	559	+ usrSize = usrSize + topkValuesGmSize + topkIndicesGmSize;
	560	+ OP_LOGI(context->GetNodeName(),
	561	+ "RadixSortTilingOfIdx final usrSize=%d, topkValuesGmSize: %d, topkIndicesGmSize: %d.",
	562	+ *usrSize, topkValuesGmSize, topkIndicesGmSize);
	563	+ return ge::GRAPH_SUCCESS;
	564	+}
	565	+} // namespace sortWithIndex
	566	+} // namespace optiling

Mmath/top_k_v2/op_host/arch35/top_k_v2_tiling_arch35.cpp+475-13

	@@ -16,8 +16,11 @@
16	#include "log/log.h"	16	#include "log/log.h"
17	#include "util/math_util.h"	17	#include "util/math_util.h"
18	#include "util/platform_util.h"	18	#include "util/platform_util.h"
		19	+#include "atvoss/broadcast/broadcast_tiling.h"
		20	+#include "sort_with_index_tiling.h"
19		21
20	namespace optiling {	22	namespace optiling {
		23	+namespace topkV2 {
21	namespace topkV2DataInfo {	24	namespace topkV2DataInfo {
22	const uint32_t CONST_TWO = 2;	25	const uint32_t CONST_TWO = 2;
23	const uint32_t CONST_THREE = 3;	26	const uint32_t CONST_THREE = 3;
	@@ -33,15 +36,26 @@ const uint32_t SINGLE_CORE_MODE = 1;
33	const uint32_t MULT_CORE_MODE = 2;	36	const uint32_t MULT_CORE_MODE = 2;
34	const uint32_t MULT_CORE_OPTIM_MODE = 4;	37	const uint32_t MULT_CORE_OPTIM_MODE = 4;
35	const uint32_t SINGLE_BLOCK_MODE = 3;	38	const uint32_t SINGLE_BLOCK_MODE = 3;
		39	+const uint32_t SORT_AND_TOP_K_MODE = 5;
36	const uint32_t INT64_BYTE = 8;	40	const uint32_t INT64_BYTE = 8;
37	const uint32_t INT32_BYTE = 4;	41	const uint32_t INT32_BYTE = 4;
38	-const uint32_t SINGLE_CORE_THRESHOLD =	42	+// SortAndTopk的阈值，排序轴大于该阈值的场景，走sortAndTopK模板
39	- 10000000; // 网络case，测试结果发现singlecore时间为2600，而老模板性能为9000，因此调整走SingleCore的阈值	43	+const uint32_t SORT_AND_TOP_K_THRESHOLD = 10000000;
40	const uint32_t CONST_SIMT_SPACE = 32768; // 获取到的UB大小需要预留32KB给simt	44	const uint32_t CONST_SIMT_SPACE = 32768; // 获取到的UB大小需要预留32KB给simt
41	const uint32_t SUPPORT_SORT_MAX_BYTE_SIZE = 8000;	45	const uint32_t SUPPORT_SORT_MAX_BYTE_SIZE = 8000;
		46	+const uint32_t SUPPORT_SORT_MAX_SIZE = 2000;
42	const float LAST_LOOP_CORE_UTILIZATION = 0.7;	47	const float LAST_LOOP_CORE_UTILIZATION = 0.7;
43	const uint32_t SMALL_LOOP_UPPER_NUM = 4;	48	const uint32_t SMALL_LOOP_UPPER_NUM = 4;
44	const uint32_t SMALL_LOOP_LOWER_NUM = 2;	49	const uint32_t SMALL_LOOP_LOWER_NUM = 2;
		50	+const uint32_t SIMT_UB = 32768; // SortAndTopK模板需要预留32k给simt使用
		51	+const uint32_t UB_CONST_INT32 = 4096; // 输出idx为int32时kernel侧需要的固定ub大小
		52	+const uint32_t UB_CONST_INT64 = 7168; // 输出idx为int64时kernel侧需要的固定ub大小
		53	+const uint32_t CONST_10 = 10;
		54	+const uint32_t CONST_14 = 14;
		55	+const uint32_t CONST_2 = 2;
		56	+const uint32_t SMALL_TILE_DATA_NUM = 1024; // 测试数据得出一次至少处理1024，sort性能比较好
		57	+// 排序轴在int32范围内的最大值, 超过这个值, cutsum，前缀和就要用int64数据范围表示
		58	+const uint32_t INT32_MAX_RANGE_VALUE_FOR_SORT = 1073741823;
45		59
46	constexpr size_t SYS_WORK_SPACE_SIZE = static_cast<size_t>(16 * 1024 * 1024);	60	constexpr size_t SYS_WORK_SPACE_SIZE = static_cast<size_t>(16 * 1024 * 1024);
47	struct TopkTileInfo {	61	struct TopkTileInfo {
	@@ -56,6 +70,37 @@ struct TopkTileInfo {
56	uint32_t tailTileNum = 0;	70	uint32_t tailTileNum = 0;
57	int64_t topKOutLastAxisNum = 0;	71	int64_t topKOutLastAxisNum = 0;
58	};	72	};
		73	+struct SortTileInfo {
		74	+ uint32_t coreNumNeed = 0;
		75	+ uint32_t lastDimTileNum = 0;
		76	+ uint32_t unsortedDimParallel = 1;
		77	+ uint32_t oneCoreRowNum = 1;
		78	+ uint32_t ubSize = 0;
		79	+ uint32_t blockUbSize = 0;
		80	+ uint32_t dtypeSize = 0;
		81	+ uint32_t y2DtypeSize = 0;
		82	+ uint32_t maxCoreNum = 0;
		83	+ uint32_t numTileDataSize = 0;
		84	+ uint32_t sortLoopTimes = 0;
		85	+ uint32_t lastDimNeedCore = 0;
		86	+ uint32_t keyParams0 = 0;
		87	+ uint32_t keyParams1 = 0;
		88	+ uint32_t keyParams2 = 0;
		89	+ uint32_t keyParams3 = 0;
		90	+ uint32_t keyParams4 = 0;
		91	+ uint32_t keyParams5 = 0;
		92	+ uint32_t tmpUbSize = 0;
		93	+ bool isDescend = false;
		94	+ ge::DataType dataType = ge::DT_UINT8;
		95	+ uint32_t isInt32 = 0;
		96	+ int32_t xDimNum = 0;
		97	+ int64_t sortAxisNum = 1;
		98	+ int64_t unSortDimNum = 1;
		99	+ int64_t topKRealValue = 0;
		100	+ uint32_t tileDataSize = 0;
		101	+ uint32_t blockTileNum = 0;
		102	+ uint32_t tailTileNum = 0;
		103	+};
59	struct TopkComputingNowTileSizeInfo {	104	struct TopkComputingNowTileSizeInfo {
60	ge::DataType dataType;	105	ge::DataType dataType;
61	ge::DataType indicesDType;	106	ge::DataType indicesDType;
	@@ -81,6 +126,23 @@ static const std::map<ge::DataType, uint32_t> optDataTypeBitMap = {
81	static const std::map<ge::DataType, uint32_t> b64DataTypeBitMap = {{ge::DT_INT64, 8}, {ge::DT_UINT64, 8}};	126	static const std::map<ge::DataType, uint32_t> b64DataTypeBitMap = {{ge::DT_INT64, 8}, {ge::DT_UINT64, 8}};
82	} // namespace topkV2DataInfo	127	} // namespace topkV2DataInfo
83		128
		129	+uint32_t CeilAlignDiv(int64_t a, int64_t b)
		130	+{
		131	+ if (b == 0) {
		132	+ return static_cast<uint32_t>(a);
		133	+ }
		134	+ return static_cast<uint32_t>((a + b - 1) / b);
		135	+}
		136	+
		137	+template <typename T>
		138	+auto CeilAlignDivMul(int64_t a, int64_t b) -> T const
		139	+{
		140	+ if (b == 0) {
		141	+ return static_cast<T>(a);
		142	+ }
		143	+ return static_cast<T>(((a + b - 1) / b) * b);
		144	+}
		145	+
84	ge::graphStatus GetTopkApiTmpBufferSize(	146	ge::graphStatus GetTopkApiTmpBufferSize(
85	gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData, uint32_t needDataNum, int64_t kValue,	147	gert::TilingContext* context, TopKV2TilingDataSimd& topkTilingData, uint32_t needDataNum, int64_t kValue,
86	bool isLargest, ge::DataType dtype, bool isSort, uint32_t nowTileSize)	148	bool isLargest, ge::DataType dtype, bool isSort, uint32_t nowTileSize)
	@@ -611,7 +673,8 @@ void TileModeMediumSize(
611	topkTilingData.set_lastDimNeedCore(lastDimTileNum);	673	topkTilingData.set_lastDimNeedCore(lastDimTileNum);
612	topkTilingData.set_numTileDataSize(nowTileSize);	674	topkTilingData.set_numTileDataSize(nowTileSize);
613	topkTileInfo.ubRealLoadDataNum = nowTileSize;	675	topkTileInfo.ubRealLoadDataNum = nowTileSize;
614	- topkTileInfo.coreNumNeed = coreNumNeed;	676	+ // 为了适配TopK拼接sortwithindex，TopK MediumSize Mode需要满核运行
		677	+ topkTileInfo.coreNumNeed = maxCoreNum;
615	topkTileInfo.lastDimTileNum = lastDimTileNum;	678	topkTileInfo.lastDimTileNum = lastDimTileNum;
616	topkTileInfo.unsortedDimParallel = unsortedDimParallel;	679	topkTileInfo.unsortedDimParallel = unsortedDimParallel;
617	OP_LOGI(	680	OP_LOGI(
	@@ -638,7 +701,8 @@ void TileModeBigSize(
638	topkTilingData.set_lastDimNeedCore(coreNumNeed);	701	topkTilingData.set_lastDimNeedCore(coreNumNeed);
639	topkTilingData.set_numTileDataSize(nowTileSize);	702	topkTilingData.set_numTileDataSize(nowTileSize);
640	topkTileInfo.ubRealLoadDataNum = nowTileSize;	703	topkTileInfo.ubRealLoadDataNum = nowTileSize;
641	- topkTileInfo.coreNumNeed = coreNumNeed;	704	+ // 为了适配TopK拼接sortwithindex，TopK BigSize Mode需要满核运行
		705	+ topkTileInfo.coreNumNeed = maxCoreNum;
642	topkTileInfo.lastDimTileNum = lastDimTileNum;	706	topkTileInfo.lastDimTileNum = lastDimTileNum;
643	topkTileInfo.unsortedDimParallel = 1;	707	topkTileInfo.unsortedDimParallel = 1;
644	OP_LOGI(	708	OP_LOGI(
	@@ -702,12 +766,8 @@ ge::graphStatus IsValidParam(gert::TilingContext* context)
702	return ge::GRAPH_SUCCESS;	766	return ge::GRAPH_SUCCESS;
703	}	767	}
704		768
705	-bool IsModeSingleCore(uint32_t unsortedDimNum, uint32_t maxCoreNum, uint32_t lastAxisNum)	769	+bool IsModeSingleCore(uint32_t unsortedDimNum, uint32_t maxCoreNum)
706	{	770	{
707	- // 排序轴如果大于singlecore的阈值，则不走该模板
708	- if (lastAxisNum > topkV2DataInfo::SINGLE_CORE_THRESHOLD) {
709	- return false;
710	- }
711	// B轴小于核数，则不走该模板	771	// B轴小于核数，则不走该模板
712	if (unsortedDimNum < maxCoreNum) {	772	if (unsortedDimNum < maxCoreNum) {
713	return false;	773	return false;
	@@ -718,6 +778,372 @@ bool IsModeSingleCore(uint32_t unsortedDimNum, uint32_t maxCoreNum, uint32_t las
718	return true;	778	return true;
719	}	779	}
720		780
		781	+// sort核间模板tiling计算相关函数
		782	+uint32_t ComputeRemainUb(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor)
		783	+{
		784	+ uint32_t tmpUb = sortTileInfo.ubSize - (ubExtra + tileFactor * tileData);
		785	+ return tmpUb;
		786	+}
		787	+
		788	+void SetSortTmpSize(ge::DataType dataType, uint32_t tileData, bool isDescend, topkV2DataInfo::SortTileInfo &sortTileInfo)
		789	+{
		790	+ int64_t realLen = std::min(sortTileInfo.sortAxisNum, static_cast<int64_t>(tileData));
		791	+ std::vector<int64_t> shapeVec = { realLen };
		792	+ ge::Shape srcShape(shapeVec);
		793	+ AscendC::SortConfig config;
		794	+ config.type = AscendC::SortType::RADIX_SORT;
		795	+ config.isDescend = isDescend;
		796	+ config.hasSrcIndex = false;
		797	+ config.hasDstIndex = true;
		798	+ uint32_t maxValue = 0, minValue = 0;
		799	+ AscendC::GetSortMaxMinTmpSize(srcShape, dataType, ge::DT_UINT32, false, config, maxValue, minValue);
		800	+ OP_LOGI("RadixSortTiling", "api of sort shape is %ld, maxUb is %u", realLen, maxValue);
		801	+ sortTileInfo.tmpUbSize = maxValue;
		802	+ return;
		803	+}
		804	+
		805	+void AdjTmpUb(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor)
		806	+{
		807	+ uint32_t remainUbNew = ComputeRemainUb(sortTileInfo, tileData, ubExtra, tileFactor) - sortTileInfo.tmpUbSize;
		808	+ remainUbNew = remainUbNew > sortTileInfo.blockUbSize ? (remainUbNew - sortTileInfo.blockUbSize) : uint32_t(0);
		809	+ uint32_t alignUbSize = (remainUbNew / sortTileInfo.blockUbSize) * sortTileInfo.blockUbSize;
		810	+ OP_LOGI("RadixSortTiling", "alignUbSize %u", alignUbSize);
		811	+ sortTileInfo.tmpUbSize = sortTileInfo.tmpUbSize + alignUbSize; // 剩余的ub都给tmpUbsize
		812	+}
		813	+
		814	+void ComputeTileDataOne(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t lastDimTileNum, uint32_t ubExtra,
		815	+ uint32_t &tileData, uint32_t tileFactor)
		816	+{
		817	+ uint32_t allCore = CeilAlignDivMul<uint32_t>(int64_t(lastDimTileNum), int64_t(sortTileInfo.maxCoreNum));
		818	+ uint32_t newTileData = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(allCore));
		819	+ tileData = CeilAlignDivMul<uint32_t>(int64_t(newTileData), int64_t(topkV2DataInfo::BIN_NUM));
		820	+ tileData = std::max(tileData, topkV2DataInfo::SMALL_TILE_DATA_NUM);
		821	+ SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
		822	+ AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
		823	+ return;
		824	+}
		825	+
		826	+bool NeedAdjTileData(topkV2DataInfo::SortTileInfo &sortTileInfo, uint32_t &tileData, uint32_t lastDimTileNum,
		827	+ uint32_t ubExtra, uint32_t tileFactor)
		828	+{
		829	+ if (sortTileInfo.unSortDimNum == int64_t(1) && lastDimTileNum == uint32_t(1)) {
		830	+ OP_LOGI("RadixSortTiling", "unSortDimNum and lastDimTileNum is 1");
		831	+ uint32_t newTileData = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(sortTileInfo.maxCoreNum));
		832	+ newTileData = CeilAlignDivMul<uint32_t>(int64_t(newTileData), int64_t(topkV2DataInfo::BIN_NUM));
		833	+ tileData = std::max(newTileData, topkV2DataInfo::SMALL_TILE_DATA_NUM);
		834	+ SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
		835	+ AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
		836	+ return true;
		837	+ }
		838	+ if (sortTileInfo.unSortDimNum == int64_t(1) \|\| (lastDimTileNum >= sortTileInfo.maxCoreNum)) {
		839	+ // b为1时，尽量均匀分核，同时保证处理的最小的tile_data为1024
		840	+ OP_LOGI("RadixSortTiling", "unSortDimNum is 1 and lastDimTileNum greater than allCore");
		841	+ ComputeTileDataOne(sortTileInfo, lastDimTileNum, ubExtra, tileData, tileFactor);
		842	+ return true;
		843	+ }
		844	+ if (sortTileInfo.unSortDimNum > int64_t(1) && sortTileInfo.unSortDimNum < int64_t(sortTileInfo.maxCoreNum) &&
		845	+ lastDimTileNum == uint32_t(1)) {
		846	+ OP_LOGI("RadixSortTiling", "unSortDimNum greater than 1,and unSortDimNum small and lastDimTileNum is one");
		847	+ uint32_t hCore = sortTileInfo.maxCoreNum / static_cast<uint32_t>(sortTileInfo.unSortDimNum);
		848	+ uint32_t hTileData = static_cast<uint32_t>(sortTileInfo.sortAxisNum) / hCore;
		849	+ tileData = CeilAlignDivMul<uint32_t>(int64_t(hTileData), int64_t(topkV2DataInfo::BIN_NUM));
		850	+ SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
		851	+ AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
		852	+ return tileData;
		853	+ }
		854	+ if (sortTileInfo.unSortDimNum > int64_t(1) && lastDimTileNum > uint32_t(1)) {
		855	+ // b大于1且h轴循环次数小于总核数，也就是b轴核数大于1
		856	+ OP_LOGI("RadixSortTiling", "unSortDimNum is one, lastDimTileNum greater than one");
		857	+ int64_t newTileData = sortTileInfo.sortAxisNum / int64_t(lastDimTileNum);
		858	+ tileData = CeilAlignDivMul<uint32_t>(newTileData, int64_t(topkV2DataInfo::BIN_NUM));
		859	+ lastDimTileNum = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(tileData));
		860	+ uint32_t bCore = lastDimTileNum == 0 ? sortTileInfo.maxCoreNum : sortTileInfo.maxCoreNum / lastDimTileNum;
		861	+ if (lastDimTileNum < sortTileInfo.maxCoreNum && sortTileInfo.unSortDimNum < int64_t(sortTileInfo.maxCoreNum)) {
		862	+ if (sortTileInfo.unSortDimNum < int64_t(bCore)) {
		863	+ bCore = static_cast<uint32_t>(sortTileInfo.unSortDimNum);
		864	+ uint32_t hCore = sortTileInfo.maxCoreNum / bCore;
		865	+ uint32_t tileDataNew = CeilAlignDiv(int64_t(sortTileInfo.sortAxisNum), int64_t(hCore));
		866	+ tileData = CeilAlignDivMul<uint32_t>(int64_t(tileDataNew), int64_t(topkV2DataInfo::BIN_NUM));
		867	+ }
		868	+ }
		869	+ if (bCore == static_cast<uint32_t>(1) && lastDimTileNum < sortTileInfo.maxCoreNum) {
		870	+ ComputeTileDataOne(sortTileInfo, lastDimTileNum, ubExtra, tileData, tileFactor);
		871	+ }
		872	+ SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
		873	+ AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
		874	+ return true;
		875	+ }
		876	+ return false;
		877	+}
		878	+
		879	+uint32_t ComputeTileData(topkV2DataInfo::SortTileInfo &sortTileInfo)
		880	+{
		881	+ uint32_t ubExtra;
		882	+ uint32_t tileFactor;
		883	+ if (sortTileInfo.isInt32 == static_cast<uint32_t>(0)) {
		884	+ ubExtra = topkV2DataInfo::UB_CONST_INT64;
		885	+ tileFactor = topkV2DataInfo::CONST_14 + sortTileInfo.dtypeSize;
		886	+ } else {
		887	+ ubExtra = topkV2DataInfo::UB_CONST_INT32;
		888	+ tileFactor = topkV2DataInfo::CONST_10 + sortTileInfo.dtypeSize;
		889	+ }
		890	+
		891	+ uint32_t tileData = (sortTileInfo.ubSize - ubExtra) / tileFactor;
		892	+ tileData = (tileData / topkV2DataInfo::BIN_NUM) * topkV2DataInfo::BIN_NUM;
		893	+
		894	+ uint32_t remainUb = ComputeRemainUb(sortTileInfo, tileData, ubExtra, tileFactor);
		895	+ SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
		896	+
		897	+ uint32_t tmpUbSize = sortTileInfo.tmpUbSize;
		898	+ while (tmpUbSize > remainUb) {
		899	+ tileData = tileData - topkV2DataInfo::BIN_NUM;
		900	+ remainUb = ComputeRemainUb(sortTileInfo, tileData, ubExtra, tileFactor);
		901	+ SetSortTmpSize(ge::DT_UINT8, tileData, false, sortTileInfo);
		902	+ tmpUbSize = sortTileInfo.tmpUbSize;
		903	+ }
		904	+ uint32_t lastDimTileNum = CeilAlignDiv(sortTileInfo.sortAxisNum, int64_t(tileData));
		905	+ OP_LOGI("RadixSortTiling", "tileData %u, lastDimTileNum %u, tmpUbSize %u", tileData, lastDimTileNum, tmpUbSize);
		906	+ bool smallTile = (sortTileInfo.sortAxisNum <= static_cast<int64_t>(topkV2DataInfo::SMALL_TILE_DATA_NUM)) &&
		907	+ lastDimTileNum == uint32_t(1);
		908	+ if ((lastDimTileNum % sortTileInfo.maxCoreNum == static_cast<uint32_t>(0)) \|\| smallTile) {
		909	+ OP_LOGI("RadixSortTiling", "lastDimTileNum align or smallTile");
		910	+ AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
		911	+ return tileData;
		912	+ }
		913	+ if (NeedAdjTileData(sortTileInfo, tileData, lastDimTileNum, ubExtra, tileFactor)) {
		914	+ return tileData;
		915	+ }
		916	+ AdjTmpUb(sortTileInfo, tileData, ubExtra, tileFactor);
		917	+ return tileData;
		918	+}
		919	+
		920	+void ComputeWorkSpace(gert::TilingContext *context, topkV2DataInfo::SortTileInfo &sortTileInfo)
		921	+{
		922	+ uint32_t dtypeSizeWk = static_cast<uint32_t>(sizeof(int32_t));
		923	+ if (sortTileInfo.isInt32 == static_cast<uint32_t>(0)) {
		924	+ dtypeSizeWk = static_cast<uint32_t>(sizeof(int64_t));
		925	+ }
		926	+ size_t excusiveBinsGmWkSize = static_cast<size_t>(sortTileInfo.keyParams1) * sortTileInfo.keyParams4 * dtypeSizeWk;
		927	+ excusiveBinsGmWkSize = CeilAlignDivMul<size_t>(int64_t(excusiveBinsGmWkSize), int64_t(sortTileInfo.blockUbSize));
		928	+
		929	+ size_t globalHistGmWkSize =
		930	+ static_cast<size_t>(sortTileInfo.keyParams3) * sortTileInfo.keyParams2 * sortTileInfo.keyParams0 * dtypeSizeWk;
		931	+ globalHistGmWkSize = CeilAlignDivMul<size_t>(int64_t(globalHistGmWkSize), int64_t(sortTileInfo.blockUbSize));
		932	+
		933	+ size_t outIdxDbWK = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel * dtypeSizeWk;
		934	+ outIdxDbWK = CeilAlignDivMul<size_t>(int64_t(outIdxDbWK), int64_t(sortTileInfo.blockUbSize));
		935	+
		936	+ size_t sortOutIdxGMWK = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel *
		937	+ sortTileInfo.y2DtypeSize;
		938	+ sortOutIdxGMWK = CeilAlignDivMul<size_t>(int64_t(sortOutIdxGMWK), int64_t(sortTileInfo.blockUbSize));
		939	+
		940	+ size_t histTileGmWk = static_cast<size_t>(sortTileInfo.lastDimTileNum) * topkV2DataInfo::BIN_NUM *
		941	+ sortTileInfo.unsortedDimParallel * sizeof(int16_t) * topkV2DataInfo::CONST_2;
		942	+
		943	+ size_t xB8GmWkSize = static_cast<size_t>(sortTileInfo.lastDimTileNum) * sortTileInfo.numTileDataSize *
		944	+ sortTileInfo.unsortedDimParallel;
		945	+ xB8GmWkSize = CeilAlignDivMul<size_t>(int64_t(xB8GmWkSize), int64_t(sortTileInfo.blockUbSize));
		946	+
		947	+ size_t outValueDbWKSize = static_cast<size_t>(sortTileInfo.sortAxisNum) * sortTileInfo.unsortedDimParallel *
		948	+ sortTileInfo.dtypeSize *topkV2DataInfo::CONST_2;
		949	+ outValueDbWKSize = CeilAlignDivMul<size_t>(int64_t(outValueDbWKSize), int64_t(sortTileInfo.blockUbSize));
		950	+
		951	+ OP_LOGI("RadixSortTiling",
		952	+ "excusiveBinsGmWkSize %lu, globalHistGmWkSize %lu, outIdxDbWK %lu, sortOutIdxGMWK %lu, histTileGmWk %lu,"
		953	+ " xB8GmWkSize %lu, outValueDbWKSize %lu ",
		954	+ excusiveBinsGmWkSize, globalHistGmWkSize, outIdxDbWK, sortOutIdxGMWK, histTileGmWk, xB8GmWkSize, outValueDbWKSize);
		955	+ size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);
		956	+ size_t usrSize = excusiveBinsGmWkSize + globalHistGmWkSize + outIdxDbWK + sortOutIdxGMWK + histTileGmWk +
		957	+ xB8GmWkSize + outValueDbWKSize;
		958	+ userWorkSpaceSize[0] = usrSize + topkV2DataInfo::SYS_WORK_SPACE_SIZE;
		959	+ return;
		960	+}
		961	+
		962	+ge::graphStatus GetRadixSortMoreCore(gert::TilingContext *context, topkV2DataInfo::SortTileInfo &sortTileInfo)
		963	+{
		964	+ sortTileInfo.ubSize = sortTileInfo.ubSize - topkV2DataInfo::SIMT_UB;
		965	+ uint32_t tileData = ComputeTileData(sortTileInfo);
		966	+ uint32_t lastDimTileNum = CeilAlignDiv(int64_t(sortTileInfo.sortAxisNum), int64_t(tileData));
		967	+ if (sortTileInfo.maxCoreNum <= lastDimTileNum) {
		968	+ sortTileInfo.unsortedDimParallel = static_cast<uint32_t>(1);
		969	+ } else {
		970	+ sortTileInfo.unsortedDimParallel = lastDimTileNum == 0 ? sortTileInfo.maxCoreNum : sortTileInfo.maxCoreNum /
		971	+ lastDimTileNum;
		972	+ if (sortTileInfo.unSortDimNum < static_cast<int64_t>(sortTileInfo.unsortedDimParallel)) {
		973	+ sortTileInfo.unsortedDimParallel = static_cast<uint32_t>(sortTileInfo.unSortDimNum);
		974	+ }
		975	+ }
		976	+ sortTileInfo.numTileDataSize = tileData;
		977	+ sortTileInfo.sortLoopTimes = CeilAlignDiv(int64_t(sortTileInfo.unSortDimNum), int64_t(sortTileInfo.unsortedDimParallel));
		978	+ sortTileInfo.lastDimNeedCore = std::min(sortTileInfo.maxCoreNum, lastDimTileNum);
		979	+ sortTileInfo.coreNumNeed = sortTileInfo.unsortedDimParallel * sortTileInfo.lastDimNeedCore;
		980	+ sortTileInfo.lastDimTileNum = lastDimTileNum;
		981	+
		982	+ uint32_t ubSizeNum = sortTileInfo.tmpUbSize / static_cast<uint32_t>(sizeof(uint32_t));
		983	+ if (sortTileInfo.isInt32 == static_cast<uint32_t>(0)) {
		984	+ ubSizeNum = sortTileInfo.tmpUbSize / static_cast<uint32_t>(sizeof(int64_t));
		985	+ }
		986	+ uint32_t allNumGloblHist = topkV2DataInfo::BIN_NUM * lastDimTileNum * sortTileInfo.dtypeSize *
		987	+ sortTileInfo.unsortedDimParallel;
		988	+ uint32_t allNumExcusiveBin = topkV2DataInfo::BIN_NUM * sortTileInfo.dtypeSize * sortTileInfo.unsortedDimParallel;
		989	+ uint32_t oneCoreSize = CeilAlignDiv(int64_t(allNumGloblHist), int64_t(sortTileInfo.coreNumNeed));
		990	+ sortTileInfo.keyParams5 =
		991	+ std::max(static_cast<int64_t>(oneCoreSize), static_cast<int64_t>(sortTileInfo.blockUbSize));
		992	+ sortTileInfo.keyParams0 = CeilAlignDiv(int64_t(allNumGloblHist), int64_t(sortTileInfo.keyParams5));
		993	+ sortTileInfo.keyParams3 = CeilAlignDiv(int64_t(sortTileInfo.keyParams5), int64_t(ubSizeNum));
		994	+ sortTileInfo.keyParams2 = sortTileInfo.keyParams5 > ubSizeNum ? ubSizeNum : sortTileInfo.keyParams5;
		995	+
		996	+ uint32_t oneCoreSize1 = CeilAlignDiv(int64_t(allNumExcusiveBin), int64_t(sortTileInfo.coreNumNeed));
		997	+ sortTileInfo.keyParams4 =
		998	+ std::max(static_cast<int64_t>(oneCoreSize1), static_cast<int64_t>(sortTileInfo.blockUbSize));
		999	+
		1000	+ sortTileInfo.keyParams1 = CeilAlignDiv(int64_t(allNumExcusiveBin), int64_t(sortTileInfo.keyParams4));
		1001	+
		1002	+ // 取前k个结果相关流程的tile计算
		1003	+ uint32_t avilableUbSize = (sortTileInfo.ubSize - 1) / topkV2DataInfo::AGLIN_FACTOR * topkV2DataInfo::AGLIN_FACTOR;
		1004	+ OP_CHECK_IF(avilableUbSize == 0,
		1005	+ OP_LOGE("TopKV2", "sortAndTopK Tiling avilableUbSize is zero"), return ge::GRAPH_FAILED);
		1006	+ auto dataType = context->GetInputDesc(0)->GetDataType();
		1007	+ auto indicesDType = context->GetOutputDesc(1)->GetDataType();
		1008	+ uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);
		1009	+ uint32_t indexToDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(indicesDType)->second);
		1010	+ uint32_t kGetDtypeSize = std::max(xDtypeSize, indexToDtypeSize);
		1011	+ OP_CHECK_IF(kGetDtypeSize == 0, OP_LOGE("GetRadixSortMoreCore", "kGetDtypeSize is zero"), return ge::GRAPH_FAILED);
		1012	+ sortTileInfo.tileDataSize = avilableUbSize / kGetDtypeSize;
		1013	+ uint32_t totalTileNum = (sortTileInfo.topKRealValue + sortTileInfo.tileDataSize - 1) / sortTileInfo.tileDataSize;
		1014	+ sortTileInfo.blockTileNum = totalTileNum / sortTileInfo.maxCoreNum;
		1015	+ sortTileInfo.tailTileNum = totalTileNum % sortTileInfo.maxCoreNum;
		1016	+ OP_CHECK_IF(sortTileInfo.blockTileNum == 0 && sortTileInfo.tailTileNum == 0,
		1017	+ OP_LOGE("TopKV2", "sortAndTopK blockTileNum & tailTileNum is wrong!"), return ge::GRAPH_FAILED);
		1018	+
		1019	+ ComputeWorkSpace(context, sortTileInfo);
		1020	+ context->SetScheduleMode(1);
		1021	+ return ge::GRAPH_SUCCESS;
		1022	+}
		1023	+
		1024	+ge::graphStatus CheckInputAndOutput(gert::TilingContext *context, topkV2DataInfo::SortTileInfo &sortTileInfo)
		1025	+{
		1026	+ auto platformInfo = context->GetPlatformInfo();
		1027	+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
		1028	+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
		1029	+ uint64_t ubSize = 0;
		1030	+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
		1031	+ OP_CHECK_IF(ubSize <= static_cast<uint64_t>(topkV2DataInfo::SIMT_UB),
		1032	+ OP_LOGE(context->GetNodeName(), "allUb must greater than simtUb, but is %lu", ubSize),
		1033	+ return ge::GRAPH_FAILED);
		1034	+ sortTileInfo.blockUbSize = Ops::Base::GetUbBlockSize(context);
		1035	+ OP_LOGI(context->GetNodeName(), "ubSize is %ld, blockUbSize %u", ubSize, sortTileInfo.blockUbSize);
		1036	+ sortTileInfo.ubSize = ubSize;
		1037	+ auto inputShapePtr = context->GetInputShape(0);
		1038	+ OP_CHECK_NULL_WITH_CONTEXT(context, inputShapePtr);
		1039	+ const gert::Shape &inputShape = Ops::Base::EnsureNotScalar(inputShapePtr->GetStorageShape());
		1040	+ auto yStorage = context->GetOutputShape(0);
		1041	+ OP_CHECK_NULL_WITH_CONTEXT(context, yStorage);
		1042	+ const gert::Shape &outShape = Ops::Base::EnsureNotScalar(yStorage->GetStorageShape());
		1043	+ OP_CHECK_IF(inputShape.GetShapeSize() == 0 \|\| outShape.GetShapeSize() == 0,
		1044	+ OP_LOGE(context->GetNodeName(), "not support empty input or output"),
		1045	+ return ge::GRAPH_FAILED);
		1046	+ int32_t xDimNum = inputShape.GetDimNum();
		1047	+ sortTileInfo.xDimNum = xDimNum;
		1048	+ int64_t sortAxisNum = inputShape.GetDim(xDimNum - 1);
		1049	+ sortTileInfo.sortAxisNum = sortAxisNum;
		1050	+ int64_t unSortDimNum = static_cast<int64_t>(1);
		1051	+ for (uint32_t i = 0; i < static_cast<uint32_t>((xDimNum - 1)); i++) {
		1052	+ int64_t dimSize = static_cast<int64_t>(inputShape.GetDim(i));
		1053	+ unSortDimNum *= dimSize;
		1054	+ }
		1055	+ sortTileInfo.unSortDimNum = unSortDimNum;
		1056	+ return ge::GRAPH_SUCCESS;
		1057	+}
		1058	+
		1059	+ge::graphStatus SortCheckParams(gert::TilingContext *context, topkV2DataInfo::SortTileInfo &sortTileInfo)
		1060	+{
		1061	+ OP_CHECK_IF(CheckInputAndOutput(context, sortTileInfo) != ge::GRAPH_SUCCESS,
		1062	+ OP_LOGE(context->GetNodeName(), "CheckInputAndOutput failed"), return ge::GRAPH_FAILED);
		1063	+ auto inputDescPtr = context->GetInputDesc(0);
		1064	+ OP_CHECK_NULL_WITH_CONTEXT(context, inputDescPtr);
		1065	+ ge::DataType dataType = inputDescPtr->GetDataType();
		1066	+ OP_CHECK_IF(topkV2DataInfo::tilingDataTypeBitMap.count(dataType) == 0,
		1067	+ OP_LOGE(context->GetNodeName(), "Not support data type"), return ge::GRAPH_FAILED);
		1068	+ sortTileInfo.dataType = dataType;
		1069	+ sortTileInfo.dtypeSize = topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second;
		1070	+ auto outDescPtr = context->GetOutputDesc(1);
		1071	+ OP_CHECK_NULL_WITH_CONTEXT(context, outDescPtr);
		1072	+ auto y2DType = outDescPtr->GetDataType();
		1073	+ auto outDescPtr0 = context->GetOutputDesc(0);
		1074	+ OP_CHECK_NULL_WITH_CONTEXT(context, outDescPtr0);
		1075	+ auto y1DType = outDescPtr0->GetDataType();
		1076	+ OP_CHECK_IF((y2DType != ge::DT_INT64) && (y2DType != ge::DT_INT32),
		1077	+ OP_LOGE(context->GetNodeName(), "Not support y2 type"), return ge::GRAPH_FAILED);
		1078	+ OP_CHECK_IF(y1DType != dataType,
		1079	+ OP_LOGE(context->GetNodeName(), "input0 dtype must be same as output0 dtype"),
		1080	+ return ge::GRAPH_FAILED);
		1081	+ sortTileInfo.y2DtypeSize = topkV2DataInfo::tilingDataTypeBitMap.find(y2DType)->second;
		1082	+ return ge::GRAPH_SUCCESS;
		1083	+}
		1084	+
		1085	+void FillTilingDataSort(gert::TilingContext *context, topkV2DataInfo::SortTileInfo &sortTileInfo,
		1086	+ TopKV2TilingDataSimd &topkTilingData)
		1087	+{
		1088	+ topkTilingData.set_isLargest(sortTileInfo.isDescend);
		1089	+ topkTilingData.set_isInInt32Range(sortTileInfo.isInt32);
		1090	+ topkTilingData.set_numTileDataSize(sortTileInfo.numTileDataSize);
		1091	+ topkTilingData.set_unsortedDimParallel(sortTileInfo.unsortedDimParallel);
		1092	+ topkTilingData.set_lastDimTileNum(sortTileInfo.lastDimTileNum);
		1093	+ topkTilingData.set_sortLoopTimes(sortTileInfo.sortLoopTimes);
		1094	+ topkTilingData.set_lastDimNeedCore(sortTileInfo.lastDimNeedCore);
		1095	+ topkTilingData.set_keyParams0(sortTileInfo.keyParams0);
		1096	+ topkTilingData.set_keyParams1(sortTileInfo.keyParams1);
		1097	+ topkTilingData.set_keyParams2(sortTileInfo.keyParams2);
		1098	+ topkTilingData.set_keyParams3(sortTileInfo.keyParams3);
		1099	+ topkTilingData.set_keyParams4(sortTileInfo.keyParams4);
		1100	+ topkTilingData.set_keyParams5(sortTileInfo.keyParams5);
		1101	+ topkTilingData.set_tmpUbSize(sortTileInfo.tmpUbSize);
		1102	+ topkTilingData.set_lastAxisNum(sortTileInfo.sortAxisNum);
		1103	+ topkTilingData.set_unsortedDimNum(sortTileInfo.unSortDimNum);
		1104	+ topkTilingData.set_topKRealValue(sortTileInfo.topKRealValue);
		1105	+ topkTilingData.set_sortAndTopkTileDataSize(sortTileInfo.tileDataSize);
		1106	+ topkTilingData.set_sortAndTopkBlockTileNum(sortTileInfo.blockTileNum);
		1107	+ topkTilingData.set_sortAndTopkTailTileNum(sortTileInfo.tailTileNum);
		1108	+ topkTilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
		1109	+ context->GetRawTilingData()->SetDataSize(topkTilingData.GetDataSize());
		1110	+ return;
		1111	+}
		1112	+
		1113	+void PrintTilindDataSort(gert::TilingContext *context, topkV2DataInfo::SortTileInfo &sortTileInfo)
		1114	+{
		1115	+ OP_LOGI(context->GetNodeName(),
		1116	+ "realCoreNum %u, numTileDataSize %u, unsortedDimParallel %u, "
		1117	+ "lastDimTileNum %u, sortLoopTimes %u, lastDimNeedCore %u, keyParams0 %u, keyParams1 %u "
		1118	+ "keyParams2 %u, keyParams3 %u, keyParams4 %u, keyParams5 %u, tmpUbSize %u, "
		1119	+ "lastAxisNum %ld, unsortedDimNum %ld, topKRealValue %ld, tileDataSize %u, blockTileNum %u, tailTileNum %u",
		1120	+ sortTileInfo.coreNumNeed, sortTileInfo.numTileDataSize, sortTileInfo.unsortedDimParallel,
		1121	+ sortTileInfo.lastDimTileNum, sortTileInfo.sortLoopTimes, sortTileInfo.lastDimNeedCore, sortTileInfo.keyParams0,
		1122	+ sortTileInfo.keyParams1, sortTileInfo.keyParams2, sortTileInfo.keyParams3, sortTileInfo.keyParams4,
		1123	+ sortTileInfo.keyParams5, sortTileInfo.tmpUbSize, sortTileInfo.sortAxisNum, sortTileInfo.unSortDimNum,
		1124	+ sortTileInfo.topKRealValue, sortTileInfo.tileDataSize, sortTileInfo.blockTileNum, sortTileInfo.tailTileNum);
		1125	+ return;
		1126	+}
		1127	+
		1128	+bool needSortWithIndex(TopKV2TilingDataSimd& topkTilingData, bool isSorted, ge::DataType dataType)
		1129	+{
		1130	+ if (isSorted && topkTilingData.get_modeType() == topkV2DataInfo::MULT_CORE_MODE) {
		1131	+ if(topkTilingData.get_topKRealValue() <= topkV2DataInfo::SUPPORT_SORT_MAX_SIZE) {
		1132	+ return false;
		1133	+ }
		1134	+ return true;
		1135	+ }
		1136	+ uint32_t xDtypeSize = static_cast<uint32_t>(topkV2DataInfo::tilingDataTypeBitMap.find(dataType)->second);
		1137	+ if (isSorted && topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) {
		1138	+ if (topkTilingData.get_topKRealValue() <= topkV2DataInfo::SUPPORT_SORT_MAX_SIZE &&
		1139	+ topkTilingData.get_topKRealValue() * xDtypeSize <= topkV2DataInfo::SUPPORT_SORT_MAX_BYTE_SIZE) {
		1140	+ return false;
		1141	+ }
		1142	+ return true;
		1143	+ }
		1144	+ return false;
		1145	+}
		1146	+
721	ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)	1147	ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
722	{	1148	{
723	OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling start");	1149	OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling start");
	@@ -777,6 +1203,7 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
777	topkV2DataInfo::TopkTileInfo topkTileInfo;	1203	topkV2DataInfo::TopkTileInfo topkTileInfo;
778	topkTileInfo.topKOutLastAxisNum = outLastAxisNum;	1204	topkTileInfo.topKOutLastAxisNum = outLastAxisNum;
779		1205
		1206	+ // 用于核间优化模板 tilingSize计算流程
780	topkV2DataInfo::TopkComputingNowTileSizeInfo computingNowTileSizeInfo;	1207	topkV2DataInfo::TopkComputingNowTileSizeInfo computingNowTileSizeInfo;
781	computingNowTileSizeInfo.isLargest = *isLargest;	1208	computingNowTileSizeInfo.isLargest = *isLargest;
782	computingNowTileSizeInfo.isSort = *isSorted;	1209	computingNowTileSizeInfo.isSort = *isSorted;
	@@ -808,7 +1235,28 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
808	context, topkTilingData, dataType, indicesDType, isLargest, isSorted, lastAxisNum, outLastAxisNum,	1235	context, topkTilingData, dataType, indicesDType, isLargest, isSorted, lastAxisNum, outLastAxisNum,
809	ubSizePlatForm);	1236	ubSizePlatForm);
810	TileModeSmallSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSizeTmp);	1237	TileModeSmallSize(unsortedDimNum, maxCoreNum, lastAxisNum, topkTilingData, topkTileInfo, nowTileSizeTmp);
811	- } else if (IsModeSingleCore(unsortedDimNum, maxCoreNum, static_cast<uint32_t>(lastAxisNum))) {	1238	+ } else if (lastAxisNum > topkV2DataInfo::SORT_AND_TOP_K_THRESHOLD) {
		1239	+ topkV2DataInfo::SortTileInfo sortTileInfo;
		1240	+ OP_CHECK_IF(SortCheckParams(context, sortTileInfo) != ge::GRAPH_SUCCESS,
		1241	+ OP_LOGE(context->GetNodeName(), "sort and topk check params failed"), return ge::GRAPH_FAILED);
		1242	+ topkTilingData.set_modeType(topkV2DataInfo::SORT_AND_TOP_K_MODE);
		1243	+ OP_LOGI("[TopKV2Tiling]", "topkTilingData.set_modeType is: %u, SORT_AND_TOP_K_MODE: %u",
		1244	+ topkTilingData.get_modeType() , topkV2DataInfo::SORT_AND_TOP_K_MODE);
		1245	+ sortTileInfo.maxCoreNum = static_cast<uint32_t>(maxCoreNum);
		1246	+ sortTileInfo.isDescend = static_cast<bool>(isLargest);
		1247	+ sortTileInfo.isInt32 = static_cast<uint32_t>(lastAxisNum <= topkV2DataInfo::INT32_MAX_RANGE_VALUE_FOR_SORT);
		1248	+ sortTileInfo.topKRealValue = outLastAxisNum;
		1249	+ OP_CHECK_IF(GetRadixSortMoreCore(context, sortTileInfo) != ge::GRAPH_SUCCESS,
		1250	+ OP_LOGE(context->GetNodeName(), "Get RadixSortMoreCore tiling failed"), return ge::GRAPH_FAILED);
		1251	+ context->SetTilingKey(dataTypeKey);
		1252	+ context->SetBlockDim(sortTileInfo.coreNumNeed);
		1253	+ context->SetLocalMemorySize(sortTileInfo.ubSize);
		1254	+ FillTilingDataSort(context, sortTileInfo, topkTilingData);
		1255	+ PrintTilindDataSort(context, sortTileInfo);
		1256	+ // sortAndTopK模板核心是Sort，不需要后续Topk相关的tiling计算过程
		1257	+ OP_LOGI("TopKV2TilingForAscendC", "TopKV2Tiling end");
		1258	+ return ge::GRAPH_SUCCESS;
		1259	+ } else if (IsModeSingleCore(unsortedDimNum, maxCoreNum)) {
812	uint32_t nowTileSizeTmp = ComputeSingleCoreTileData(	1260	uint32_t nowTileSizeTmp = ComputeSingleCoreTileData(
813	context, topkTilingData, dataType, indicesDType, isLargest, isSorted, lastAxisNum, outLastAxisNum,	1261	context, topkTilingData, dataType, indicesDType, isLargest, isSorted, lastAxisNum, outLastAxisNum,
814	ubSizePlatForm);	1262	ubSizePlatForm);
	@@ -840,6 +1288,8 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
840	OP_LOGI(	1288	OP_LOGI(
841	"[TopKV2Tiling]", "lastAxisNum: %ld, int32Max: %ld, isInInt32Range: %u, nowTileSize: %u", lastAxisNum, int32Max,	1289	"[TopKV2Tiling]", "lastAxisNum: %ld, int32Max: %ld, isInInt32Range: %u, nowTileSize: %u", lastAxisNum, int32Max,
842	isInInt32Range, nowTileSize);	1290	isInInt32Range, nowTileSize);
		1291	+
		1292	+ // fill the topkTilingData
843	context->SetTilingKey(dataTypeKey);	1293	context->SetTilingKey(dataTypeKey);
844	context->SetBlockDim(topkTileInfo.coreNumNeed);	1294	context->SetBlockDim(topkTileInfo.coreNumNeed);
845	context->SetScheduleMode(1);	1295	context->SetScheduleMode(1);
	@@ -857,9 +1307,6 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
857	topkTilingData.set_tailLoopBatchNum(topkTileInfo.tailLoopBatchNum);	1307	topkTilingData.set_tailLoopBatchNum(topkTileInfo.tailLoopBatchNum);
858	topkTilingData.set_tailBatchNum(topkTileInfo.tailBatchNum);	1308	topkTilingData.set_tailBatchNum(topkTileInfo.tailBatchNum);
859	topkTilingData.set_tailTileNum(topkTileInfo.tailTileNum);	1309	topkTilingData.set_tailTileNum(topkTileInfo.tailTileNum);
860	- topkTilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
861	- context->GetRawTilingData()->SetDataSize(topkTilingData.GetDataSize());
862	-
863	OP_LOGI(	1310	OP_LOGI(
864	context->GetNodeName(),	1311	context->GetNodeName(),
865	"TopK V2 tilingData tilingKey is %u, isLargest is %u, modelType is %u,"	1312	"TopK V2 tilingData tilingKey is %u, isLargest is %u, modelType is %u,"
	@@ -876,7 +1323,9 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
876	topkTilingData.get_tailLoopBatchNum(), topkTilingData.get_tailBatchNum(), topkTilingData.get_tailTileNum(),	1323	topkTilingData.get_tailLoopBatchNum(), topkTilingData.get_tailBatchNum(), topkTilingData.get_tailTileNum(),
877	topkTileInfo.coreNumNeed);	1324	topkTileInfo.coreNumNeed);
878		1325
		1326	+ // TopKV2 Workspace计算流程
879	size_t usrSize = 0;	1327	size_t usrSize = 0;
		1328	+ OP_LOGI("[TopKV2Tiling]", "begin to calc TopKV2 Workspace size.");
880	if (isInInt32Range) {	1329	if (isInInt32Range) {
881	if (topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) {	1330	if (topkTilingData.get_modeType() == topkV2DataInfo::SINGLE_CORE_MODE) {
882	usrSize = topkTileInfo.lastDimTileNum * topkV2DataInfo::BIN_NUM * topkTileInfo.unsortedDimParallel *	1331	usrSize = topkTileInfo.lastDimTileNum * topkV2DataInfo::BIN_NUM * topkTileInfo.unsortedDimParallel *
	@@ -920,6 +1369,18 @@ ge::graphStatus TopKV2Tiling(gert::TilingContext* context, int32_t maxCoreNum)
920	OP_LOGI("[TopKV2Tiling]", "Workspace size for Radix Sort more core beyond Int32Range is : %u", usrSize);	1369	OP_LOGI("[TopKV2Tiling]", "Workspace size for Radix Sort more core beyond Int32Range is : %u", usrSize);
921	}	1370	}
922	}	1371	}
		1372	+
		1373	+ // sortWithIndex tiling&workspace计算流程，
		1374	+ if (needSortWithIndex(topkTilingData, *isSorted, dataType)) {
		1375	+ OP_CHECK_IF(sortWithIndex::RadixSortTilingOfIdx(context, topkTilingData, maxCoreNum, &usrSize) != ge::GRAPH_SUCCESS,
		1376	+ OP_LOGE(context->GetNodeName(), "SortWithIndex Tiling Simt calc failed"), return ge::GRAPH_FAILED);
		1377	+ }
		1378	+
		1379	+ // save tilingdata
		1380	+ topkTilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
		1381	+ context->GetRawTilingData()->SetDataSize(topkTilingData.GetDataSize());
		1382	+
		1383	+ // set userWorkSpaceSize
923	size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);	1384	size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
924	userWorkSpaceSize[0] = usrSize + topkV2DataInfo::SYS_WORK_SPACE_SIZE;	1385	userWorkSpaceSize[0] = usrSize + topkV2DataInfo::SYS_WORK_SPACE_SIZE;
925	OP_LOGI("[TopKV2Tiling]", "user & system WorkSpace Size is : %u", userWorkSpaceSize[0]);	1386	OP_LOGI("[TopKV2Tiling]", "user & system WorkSpace Size is : %u", userWorkSpaceSize[0]);
	@@ -955,4 +1416,5 @@ ge::graphStatus Tiling4TopKV2(gert::TilingContext* context)
955	}	1416	}
956		1417
957	IMPL_OP_OPTILING(TopKV2).Tiling(Tiling4TopKV2).TilingParse<TopKV2CompileInfo>(TilingPrepareForTopKV2);	1418	IMPL_OP_OPTILING(TopKV2).Tiling(Tiling4TopKV2).TilingParse<TopKV2CompileInfo>(TilingPrepareForTopKV2);
		1419	+} // namespace topkV2
958	} // namespace optiling	1420	} // namespace optiling

Mmath/top_k_v2/op_host/arch35/top_k_v2_tiling_arch35.h+26-0

	@@ -21,6 +21,7 @@
21	#include "tiling/tiling_api.h"	21	#include "tiling/tiling_api.h"
22	namespace optiling {	22	namespace optiling {
23	BEGIN_TILING_DATA_DEF(TopKV2TilingDataSimd)	23	BEGIN_TILING_DATA_DEF(TopKV2TilingDataSimd)
		24	+ // TopK field
24	TILING_DATA_FIELD_DEF(int32_t, isLargest);	25	TILING_DATA_FIELD_DEF(int32_t, isLargest);
25	TILING_DATA_FIELD_DEF(int32_t, isSort);	26	TILING_DATA_FIELD_DEF(int32_t, isSort);
26	TILING_DATA_FIELD_DEF(uint32_t, sortLoopTimes);	27	TILING_DATA_FIELD_DEF(uint32_t, sortLoopTimes);
	@@ -43,6 +44,31 @@ BEGIN_TILING_DATA_DEF(TopKV2TilingDataSimd)
43	TILING_DATA_FIELD_DEF(int64_t, lastDimTileNum);	44	TILING_DATA_FIELD_DEF(int64_t, lastDimTileNum);
44	TILING_DATA_FIELD_DEF(int64_t, outputLastDimValue);	45	TILING_DATA_FIELD_DEF(int64_t, outputLastDimValue);
45	TILING_DATA_FIELD_DEF(int64_t, lastDimTileNumTimes);	46	TILING_DATA_FIELD_DEF(int64_t, lastDimTileNumTimes);
		47	+ // sort field
		48	+ TILING_DATA_FIELD_DEF(uint32_t, tilingKeyForSort);
		49	+ TILING_DATA_FIELD_DEF(uint32_t, sortLoopTimesForSort);
		50	+ TILING_DATA_FIELD_DEF(uint32_t, unsortedDimParallelForSort);
		51	+ TILING_DATA_FIELD_DEF(uint64_t, unsortedDimNumForSort);
		52	+ TILING_DATA_FIELD_DEF(uint32_t, lastDimTileNumForSort);
		53	+ TILING_DATA_FIELD_DEF(uint32_t, lastDimNeedCoreForSort);
		54	+ TILING_DATA_FIELD_DEF(uint32_t, numTileDataSizeForSort);
		55	+ TILING_DATA_FIELD_DEF(uint32_t, sortAcApiNeedBufferSizeForSort);
		56	+ TILING_DATA_FIELD_DEF(uint32_t, mergSortAcApiNeedBufferSizeForSort);
		57	+ TILING_DATA_FIELD_DEF(uint32_t, oneCoreRowNumForSort);
		58	+ TILING_DATA_FIELD_DEF(uint32_t, outputLastDimValueForSort);
		59	+ TILING_DATA_FIELD_DEF(uint32_t, isInInt32RangeForSort);
		60	+ TILING_DATA_FIELD_DEF(int64_t, lastAxisNumForSort);
		61	+ TILING_DATA_FIELD_DEF(uint32_t, modeTypeForSort);
		62	+ TILING_DATA_FIELD_DEF(uint32_t, keyParams0);
		63	+ TILING_DATA_FIELD_DEF(uint32_t, keyParams1);
		64	+ TILING_DATA_FIELD_DEF(uint32_t, keyParams2);
		65	+ TILING_DATA_FIELD_DEF(uint32_t, keyParams3);
		66	+ TILING_DATA_FIELD_DEF(uint32_t, keyParams4);
		67	+ TILING_DATA_FIELD_DEF(uint32_t, keyParams5);
		68	+ TILING_DATA_FIELD_DEF(uint32_t, tmpUbSize);
		69	+ TILING_DATA_FIELD_DEF(uint32_t, sortAndTopkTileDataSize);
		70	+ TILING_DATA_FIELD_DEF(uint32_t, sortAndTopkBlockTileNum);
		71	+ TILING_DATA_FIELD_DEF(uint32_t, sortAndTopkTailTileNum);
46	END_TILING_DATA_DEF;	72	END_TILING_DATA_DEF;
47	REGISTER_TILING_DATA_CLASS(TopKV2, TopKV2TilingDataSimd)	73	REGISTER_TILING_DATA_CLASS(TopKV2, TopKV2TilingDataSimd)
48	struct TopKV2CompileInfo {	74	struct TopKV2CompileInfo {

Mmath/top_k_v2/op_kernel/arch35/radix_sort_top_k.h+213-129

	@@ -26,8 +26,10 @@
26	#include "top_k_radix_block_sort_b64.h"	26	#include "top_k_radix_block_sort_b64.h"
27	#include "top_k_util_type_simd.h"	27	#include "top_k_util_type_simd.h"
28	#include "radix_topk_util.h"	28	#include "radix_topk_util.h"
		29	+#include "sort_with_index_entry.h"
29		30
30	using namespace AscendC;	31	using namespace AscendC;
		32	+using namespace topkV2;
31	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	33	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
32	struct RadixSortTopK {	34	struct RadixSortTopK {
33	__aicore__ inline RadixSortTopK() {}	35	__aicore__ inline RadixSortTopK() {}
	@@ -163,7 +165,20 @@ public:
163	GlobalTensor<T_INDEX> cumSumBinsGm_;	165	GlobalTensor<T_INDEX> cumSumBinsGm_;
164	GlobalTensor<uint32_t> tileTopkValueGm_;	166	GlobalTensor<uint32_t> tileTopkValueGm_;
165	GlobalTensor<uint32_t> tileTopkRemainValueGm_;	167	GlobalTensor<uint32_t> tileTopkRemainValueGm_;
		168	+ // TopKV2 output value
		169	+ GlobalTensor<T> topkValuesGm_;
		170	+ // TopKV2 output index
		171	+ GlobalTensor<T_INDEX_TO> topkIndicesGm_;
		172	+
		173	+ // sortWithindex相关字段
166	GM_ADDR workspace_;	174	GM_ADDR workspace_;
		175	+ GM_ADDR sortWithIndexWorkspace_;
		176	+ GM_ADDR topkValuesGmAddr_;
		177	+ GM_ADDR topkIndicesGmAddr_;
		178	+ GM_ADDR valueAddr_;
		179	+ GM_ADDR indicesAddr_;
		180	+ const TopKV2TilingDataSimd* tilingDataPtr_;
		181	+
167	LocalTensor<UNSIGNED_TYPE> inputXCopy_;	182	LocalTensor<UNSIGNED_TYPE> inputXCopy_;
168	LocalTensor<uint32_t> tileTopkValue_;	183	LocalTensor<uint32_t> tileTopkValue_;
169	LocalTensor<uint32_t> remainTileTopkValue_;	184	LocalTensor<uint32_t> remainTileTopkValue_;
	@@ -182,6 +197,9 @@ public:
182	uint32_t topkAcApiTmpBufferSize_ = 0;	197	uint32_t topkAcApiTmpBufferSize_ = 0;
183	UNSIGNED_TYPE histDataMask = 0;	198	UNSIGNED_TYPE histDataMask = 0;
184	UNSIGNED_TYPE highBitMask = 0;	199	UNSIGNED_TYPE highBitMask = 0;
		200	+ bool needSortWithIndex_ = false;
		201	+ bool inBlockDimRange_ = true;
		202	+ uint64_t oneBlock_ = Ops::Base::GetUbBlockSize();
185	};	203	};
186		204
187	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	205	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
	@@ -197,24 +215,64 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
197	InitPara(inputValue, k, value, indices, workSpace, tilingData);	215	InitPara(inputValue, k, value, indices, workSpace, tilingData);
198	// cumsum gm	216	// cumsum gm
199	uint32_t workSpaceOffset = 0;	217	uint32_t workSpaceOffset = 0;
200	- uint32_t oneBlockNumB32 = UB_AGLIN_VALUE / static_cast<uint32_t>(sizeof(int32_t));	218	+ cumSumBinsGm_.SetGlobalBuffer((__gm__ T_INDEX)(workspace_ + workSpaceOffset), topkV2::RADIX_SORT_BIN_NUM unsortedDimParallel_);
201	- cumSumBinsGm_.SetGlobalBuffer((__gm__ T_INDEX*)workspace_,	219	+ workSpaceOffset += topkV2::RADIX_SORT_BIN_NUM * unsortedDimParallel_ * sizeof(T_INDEX);
202	- RADIX_SORT_BIN_NUM * unsortedDimParallel_);	220	+
203	- workSpaceOffset += RADIX_SORT_BIN_NUM * unsortedDimParallel_;	221	+ uint32_t tileTopkValueGmOffset = lastDimTileNum_ * unsortedDimParallel_;
204	if constexpr (sizeof(T_INDEX) == sizeof(int64_t)) {	222	if constexpr (sizeof(T_INDEX) == sizeof(int64_t)) {
205	- workSpaceOffset = workSpaceOffset * CONST_TWO;	223	+ tileTopkValueGmOffset = tileTopkValueGmOffset * topkV2::CONST_TWO;
206	}	224	}
207	- tileTopkValueGm_.SetGlobalBuffer((__gm__ uint32_t*)workspace_ + workSpaceOffset,	225	+ tileTopkValueGmOffset = CeilDivMul(tileTopkValueGmOffset * sizeof(uint32_t), int64_t(oneBlock_)) / sizeof(uint32_t);
208	- lastDimTileNum_ * unsortedDimParallel_);	226	+ tileTopkValueGm_.SetGlobalBuffer((__gm__ uint32_t*)(workspace_ + workSpaceOffset), tileTopkValueGmOffset);
209	- workSpaceOffset += CeilDivMul(lastDimTileNum_ * unsortedDimParallel_, oneBlockNumB32);	227	+ workSpaceOffset += tileTopkValueGmOffset * sizeof(uint32_t);
210	- tileTopkRemainValueGm_.SetGlobalBuffer((__gm__ uint32_t*)workspace_ + workSpaceOffset,	228	+
211	- lastDimTileNum_ * unsortedDimParallel_);	229	+ uint32_t tileTopkRemainValueGmOffset = lastDimTileNum_ * unsortedDimParallel_;
		230	+ if constexpr (sizeof(T_INDEX) == sizeof(int64_t)) {
		231	+ tileTopkRemainValueGmOffset = tileTopkRemainValueGmOffset * topkV2::CONST_TWO;
		232	+ }
		233	+ tileTopkRemainValueGmOffset = CeilDivMul(tileTopkRemainValueGmOffset * sizeof(uint32_t), int64_t(oneBlock_)) / sizeof(uint32_t);
		234	+ tileTopkRemainValueGm_.SetGlobalBuffer((__gm__ uint32_t*)(workspace_ + workSpaceOffset), tileTopkRemainValueGmOffset);
		235	+ workSpaceOffset += tileTopkRemainValueGmOffset * sizeof(uint32_t);
		236	+
		237	+ if (IS_SORT) {
		238	+ if (topkValueInput_ <= SUPPORT_SORT_MAX_SIZE) {
		239	+ needSortWithIndex_ = false;
		240	+ } else {
		J J JimmyLam23月27日是否简化为needSortWithIndex_ = topkValueInput_ > SUPPORT_SORT_MAX_SIZE BruceLP 3月31日评论： 3月31日评论：
		241	+ needSortWithIndex_ = true;
		242	+ }
		243	+ }
		244	+ if (needSortWithIndex_) {
		245	+ // sortWithIndex 尾轴的大小
		246	+ uint32_t lastAxisNumForSort = tilingData->lastAxisNumForSort;
		247	+ // sortWithIndex 外轴的大小
		248	+ uint32_t unsortedDimNumForSort = tilingData->unsortedDimNumForSort;
		249	+ uint64_t topkIndicesGmOffset = lastAxisNumForSort * unsortedDimNumForSort;
		250	+ topkIndicesGmOffset =
		251	+ CeilAlignDivMul<uint64_t>(int64_t(topkIndicesGmOffset * sizeof(T_INDEX_TO)), int64_t(oneBlock_)) /
		252	+ sizeof(T_INDEX_TO);
		253	+ topkIndicesGm_.SetGlobalBuffer((__gm__ T_INDEX_TO*)(workspace_ + workSpaceOffset), topkIndicesGmOffset);
		254	+ topkIndicesGmAddr_ = workspace_ + workSpaceOffset;
		255	+ workSpaceOffset += topkIndicesGmOffset * sizeof(T_INDEX_TO);
		256	+
		257	+ uint64_t topkValuesGmOffset = lastAxisNumForSort * unsortedDimNumForSort;
		258	+ topkValuesGmOffset = CeilAlignDivMul<uint64_t>(int64_t(topkValuesGmOffset * sizeof(T)), int64_t(oneBlock_)) /
		259	+ sizeof(T);
		260	+ topkValuesGm_.SetGlobalBuffer((__gm__ T*)(workspace_ + workSpaceOffset), topkValuesGmOffset);
		261	+ topkValuesGmAddr_ = workspace_ + workSpaceOffset;
		262	+ workSpaceOffset += topkValuesGmOffset * sizeof(T);
		263	+
		264	+ sortWithIndexWorkspace_ = workspace_ + workSpaceOffset;
		265	+ valueAddr_ = value;
		266	+ indicesAddr_ = indices;
		267	+ tilingDataPtr_ = tilingData;
		268	+ }
		269	+
212	// vec calc buffer	270	// vec calc buffer
213	- pipe.InitBuffer(blockCumSumTbuf_, ROUND_UP_AGLIN(RADIX_SORT_BIN_NUM * sizeof(T_INDEX) * lastDimTileNumTimes_));	271	+ pipe.InitBuffer(blockCumSumTbuf_, ROUND_UP_AGLIN(topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX) * lastDimTileNumTimes_));
214	pipe.InitBuffer(inputXCopyTbuf_, ROUND_UP_AGLIN(numTileData_ * sizeof(UNSIGNED_TYPE)));	272	pipe.InitBuffer(inputXCopyTbuf_, ROUND_UP_AGLIN(numTileData_ * sizeof(UNSIGNED_TYPE)));
215	pipe.InitBuffer(tileTopkValueTbuf_, ROUND_UP_AGLIN(lastDimTileNumTimes_ * sizeof(uint32_t)));	273	pipe.InitBuffer(tileTopkValueTbuf_, ROUND_UP_AGLIN(lastDimTileNumTimes_ * sizeof(uint32_t)));
216	pipe.InitBuffer(remainTileTopkValueTbuf_, ROUND_UP_AGLIN(lastDimTileNumTimes_ * sizeof(uint32_t)));	274	pipe.InitBuffer(remainTileTopkValueTbuf_, ROUND_UP_AGLIN(lastDimTileNumTimes_ * sizeof(uint32_t)));
217	- pipe.InitBuffer(dataSetCumSumTbuf_, ROUND_UP_AGLIN(RADIX_SORT_BIN_NUM * sizeof(T_INDEX)));	275	+ pipe.InitBuffer(dataSetCumSumTbuf_, ROUND_UP_AGLIN(topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX)));
218	pipe.InitBuffer(topkSrcIndexTbuf_, ROUND_UP_AGLIN(numTileData_) * sizeof(T_INDEX));	276	pipe.InitBuffer(topkSrcIndexTbuf_, ROUND_UP_AGLIN(numTileData_) * sizeof(T_INDEX));
219	pipe.InitBuffer(topkSrcIndexCopyTbuf_, ROUND_UP_AGLIN(numTileData_) * sizeof(T_INDEX_TO));	277	pipe.InitBuffer(topkSrcIndexCopyTbuf_, ROUND_UP_AGLIN(numTileData_) * sizeof(T_INDEX_TO));
220	pipe.InitBuffer(sortedShareMemTbuf_, ROUND_UP_AGLIN(topkAcApiTmpBufferSize_));	278	pipe.InitBuffer(sortedShareMemTbuf_, ROUND_UP_AGLIN(topkAcApiTmpBufferSize_));
	@@ -222,8 +280,8 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
222	tileTopkValue_ = tileTopkValueTbuf_.Get<uint32_t>();	280	tileTopkValue_ = tileTopkValueTbuf_.Get<uint32_t>();
223	remainTileTopkValue_ = remainTileTopkValueTbuf_.Get<uint32_t>();	281	remainTileTopkValue_ = remainTileTopkValueTbuf_.Get<uint32_t>();
224	// clear ub buffer	282	// clear ub buffer
225	- Duplicate(tileTopkValue_, CLEAR_UB_VALUE, lastDimTileNumTimes_);	283	+ Duplicate(tileTopkValue_, topkV2::CLEAR_UB_VALUE, lastDimTileNumTimes_);
226	- Duplicate(remainTileTopkValue_, CLEAR_UB_VALUE, lastDimTileNumTimes_);	284	+ Duplicate(remainTileTopkValue_, topkV2::CLEAR_UB_VALUE, lastDimTileNumTimes_);
227	// init queue	285	// init queue
228	pipe.InitBuffer(inQueueX_, 1, ROUND_UP_AGLIN(numTileData_) * sizeof(T));	286	pipe.InitBuffer(inQueueX_, 1, ROUND_UP_AGLIN(numTileData_) * sizeof(T));
229	uint32_t outQueueNum = TopkGetMin<uint32_t>(numTileData_, topkValueInitInput_);	287	uint32_t outQueueNum = TopkGetMin<uint32_t>(numTileData_, topkValueInitInput_);
	@@ -254,6 +312,7 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
254	topkAcApiTmpBufferSize_ = tilingData->topkAcApiTmpBufferSize;	312	topkAcApiTmpBufferSize_ = tilingData->topkAcApiTmpBufferSize;
255	topkValueInput_ = tilingData->topKRealValue;	313	topkValueInput_ = tilingData->topKRealValue;
256	topkValueInitInput_ = tilingData->topKRealValue;	314	topkValueInitInput_ = tilingData->topKRealValue;
		315	+ inBlockDimRange_ = GetBlockIdx() < lastDimTileNum_ * unsortedDimParallel_;
257	inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));	316	inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));
258	topkValueGm_.SetGlobalBuffer((__gm__ T*)(value));	317	topkValueGm_.SetGlobalBuffer((__gm__ T*)(value));
259	topkValueIndexGm_.SetGlobalBuffer((__gm__ T_INDEX_TO*)(indices));	318	topkValueIndexGm_.SetGlobalBuffer((__gm__ T_INDEX_TO*)(indices));
	@@ -268,6 +327,11 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
268	uint64_t loopOffset = i * unsortedDimParallel_ * totalDataNum_;	327	uint64_t loopOffset = i * unsortedDimParallel_ * totalDataNum_;
269	ProcessMultiBlockTopK(inputValueGm_[loopOffset]);	328	ProcessMultiBlockTopK(inputValueGm_[loopOffset]);
270	}	329	}
		330	+ if (needSortWithIndex_) {
		331	+ pipe.Reset();
		332	+ sortwithindexForTopK<T_INDEX_TO>(topkValuesGmAddr_, topkIndicesGmAddr_, valueAddr_, indicesAddr_,
		333	+ sortWithIndexWorkspace_, tilingDataPtr_, &pipe);
		334	+ }
271	}	335	}
272		336
273	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	337	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
	@@ -298,7 +362,7 @@ __aicore__ inline T_INDEX RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST,
298	WaitFlag<HardEvent::MTE2_V>(eventIdMte);	362	WaitFlag<HardEvent::MTE2_V>(eventIdMte);
299	// reduce sum	363	// reduce sum
300	uint32_t inputElementNum = (tileId + 1);	364	uint32_t inputElementNum = (tileId + 1);
301	- uint16_t repateTime = (inputElementNum + ONE_TIMES_B32_NUM - 1) / ONE_TIMES_B32_NUM;	365	+ uint16_t repateTime = (inputElementNum + topkV2::ONE_TIMES_B32_NUM - 1) / topkV2::ONE_TIMES_B32_NUM;
302	__local_mem__ uint32_t* lastDimTileTopKPtr = (__ubuf__ uint32_t*)lastDimTileTopKInfo.GetPhyAddr();	366	__local_mem__ uint32_t* lastDimTileTopKPtr = (__ubuf__ uint32_t*)lastDimTileTopKInfo.GetPhyAddr();
303	__local_mem__ uint32_t* lastDimTileTopKCopyPtr = lastDimTileTopKPtr;	367	__local_mem__ uint32_t* lastDimTileTopKCopyPtr = lastDimTileTopKPtr;
304	__VEC_SCOPE__ {	368	__VEC_SCOPE__ {
	@@ -310,7 +374,7 @@ __aicore__ inline T_INDEX RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST,
310	// mask value	374	// mask value
311	MicroAPI::MaskReg dataMask = MicroAPI::UpdateMask<uint32_t>(inputElementNum);	375	MicroAPI::MaskReg dataMask = MicroAPI::UpdateMask<uint32_t>(inputElementNum);
312	// load input	376	// load input
313	- MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(inputVectorOne, lastDimTileTopKPtr, ONE_TIMES_B32_NUM);	377	+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(inputVectorOne, lastDimTileTopKPtr, topkV2::ONE_TIMES_B32_NUM);
314	// reduce sum	378	// reduce sum
315	MicroAPI::RegTensor<uint32_t> reduceSumTensor;	379	MicroAPI::RegTensor<uint32_t> reduceSumTensor;
316	MicroAPI::ReduceSum(reduceSumTensor, inputVectorOne, dataMask);	380	MicroAPI::ReduceSum(reduceSumTensor, inputVectorOne, dataMask);
	@@ -373,16 +437,18 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
373	uint32_t unsortedAxisId = GetBlockIdx() / lastDimRealCore_;	437	uint32_t unsortedAxisId = GetBlockIdx() / lastDimRealCore_;
374	uint32_t tileTopkValueIndex = (tileId / platformCoreNum_);	438	uint32_t tileTopkValueIndex = (tileId / platformCoreNum_);
375	T_INDEX tilePrevCusumValue = ((boundaryBin >= 1) ? tileCusumBuffer(boundaryBin - 1) : 0);	439	T_INDEX tilePrevCusumValue = ((boundaryBin >= 1) ? tileCusumBuffer(boundaryBin - 1) : 0);
376	- CopyDataToGm(tileTopkValueGm_[oneRowGmOffset], tileTopkValue_[tileTopkValueIndex], tileId, 1);	440	+ if (inBlockDimRange_) {
377	- if (topkValueInput_ > 0) {	441	+ CopyDataToGm(tileTopkValueGm_[oneRowGmOffset], tileTopkValue_[tileTopkValueIndex], tileId, 1);
378	- remainTileTopkValue_(tileTopkValueIndex) =	442	+ if (topkValueInput_ > 0) {
379	- tileCusumBuffer(boundaryBin) - tilePrevCusumValue;	443	+ remainTileTopkValue_(tileTopkValueIndex) =
380	- CopyDataToGm(tileTopkRemainValueGm_[oneRowGmOffset], remainTileTopkValue_[tileTopkValueIndex], tileId, 1);	444	+ tileCusumBuffer(boundaryBin) - tilePrevCusumValue;
		445	+ CopyDataToGm(tileTopkRemainValueGm_[oneRowGmOffset], remainTileTopkValue_[tileTopkValueIndex], tileId, 1);
		446	+ }
381	}	447	}
382	// core sync	448	// core sync
383	PipeBarrier<PIPE_ALL>();	449	PipeBarrier<PIPE_ALL>();
384	SyncAll();	450	SyncAll();
385	- if (topkValueInput_ > 0) {	451	+ if (inBlockDimRange_ && topkValueInput_ > 0) {
386	// get tile offset	452	// get tile offset
387	// get cusum and store to tileTopkRemainValueGm_	453	// get cusum and store to tileTopkRemainValueGm_
388	uint32_t cumSumValue = GetTileTopkValueOffset(tileTopkRemainValueGm_, tileId, tileCount, oneRowGmOffset);	454	uint32_t cumSumValue = GetTileTopkValueOffset(tileTopkRemainValueGm_, tileId, tileCount, oneRowGmOffset);
	@@ -393,18 +459,20 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
393	// core sync	459	// core sync
394	PipeBarrier<PIPE_ALL>();	460	PipeBarrier<PIPE_ALL>();
395	SyncAll();	461	SyncAll();
396	- StoreFinalAnswer2Gm(inputLocalTensor, tileId, tileCount, oneRowGmOffset);	462	+ if (inBlockDimRange_) {
		463	+ StoreFinalAnswer2Gm(inputLocalTensor, tileId, tileCount, oneRowGmOffset);
		464	+ }
397	}	465	}
398		466
399	template <typename T_INDEX, typename T_INDEX_TO>	467	template <typename T_INDEX, typename T_INDEX_TO>
400	-__simt_vf__ LAUNCH_BOUND(RADIX_SORT_BIN_NUM)	468	+__simt_vf__ LAUNCH_BOUND(topkV2::RADIX_SORT_BIN_NUM)
401	__aicore__ inline void CopyCumSumToGmB64(__gm__ T_INDEX cumSumBinsGm_, __ubuf__ int32_t tileCusumBuffer,	469	__aicore__ inline void CopyCumSumToGmB64(__gm__ T_INDEX cumSumBinsGm_, __ubuf__ int32_t tileCusumBuffer,
402	uint32_t cumSumBinOffset, uint64_t tileTopkOffsetInUb)	470	uint32_t cumSumBinOffset, uint64_t tileTopkOffsetInUb)
403	{	471	{
404	- for (int i = Simt::GetThreadIdx(); i < RADIX_SORT_BIN_NUM / 2; i += RADIX_SORT_BIN_NUM) {	472	+ for (int i = Simt::GetThreadIdx(); i < topkV2::RADIX_SORT_BIN_NUM / 2; i += topkV2::RADIX_SORT_BIN_NUM) {
405	#pragma unroll	473	#pragma unroll
406	for (int j = 0; j < 2; j++) {	474	for (int j = 0; j < 2; j++) {
407	- uint32_t offset = i + j * RADIX_SORT_BIN_NUM / 2;	475	+ uint32_t offset = i + j * topkV2::RADIX_SORT_BIN_NUM / 2;
408	T_INDEX srcData = static_cast<T_INDEX>(tileCusumBuffer[tileTopkOffsetInUb + offset]);	476	T_INDEX srcData = static_cast<T_INDEX>(tileCusumBuffer[tileTopkOffsetInUb + offset]);
409	Simt::AtomicAdd<T_INDEX>(cumSumBinsGm_ + cumSumBinOffset + offset, srcData);	477	Simt::AtomicAdd<T_INDEX>(cumSumBinsGm_ + cumSumBinOffset + offset, srcData);
410	}	478	}
	@@ -412,11 +480,11 @@ __aicore__ inline void CopyCumSumToGmB64(__gm__ T_INDEX *cumSumBinsGm_, __ubuf__
412	}	480	}
413		481
414	template <typename T_INDEX, typename T_INDEX_TO>	482	template <typename T_INDEX, typename T_INDEX_TO>
415	-__simt_vf__ LAUNCH_BOUND(RADIX_SORT_BIN_NUM)	483	+__simt_vf__ LAUNCH_BOUND(topkV2::RADIX_SORT_BIN_NUM)
416	__aicore__ inline void CopyCumSumToGmB8B16B32(__gm__ T_INDEX cumSumBinsGm_, __ubuf__ int32_t tileCusumBuffer,	484	__aicore__ inline void CopyCumSumToGmB8B16B32(__gm__ T_INDEX cumSumBinsGm_, __ubuf__ int32_t tileCusumBuffer,
417	uint32_t cumSumBinOffset, uint64_t tileTopkOffsetInUb)	485	uint32_t cumSumBinOffset, uint64_t tileTopkOffsetInUb)
418	{	486	{
419	- for (int i = Simt::GetThreadIdx(); i < RADIX_SORT_BIN_NUM; i+= RADIX_SORT_BIN_NUM) {	487	+ for (int i = Simt::GetThreadIdx(); i < topkV2::RADIX_SORT_BIN_NUM; i+= topkV2::RADIX_SORT_BIN_NUM) {
420	uint32_t offset = i;	488	uint32_t offset = i;
421	T_INDEX srcData = static_cast<T_INDEX>(tileCusumBuffer[tileTopkOffsetInUb + offset]);	489	T_INDEX srcData = static_cast<T_INDEX>(tileCusumBuffer[tileTopkOffsetInUb + offset]);
422	Simt::AtomicAdd<T_INDEX>(cumSumBinsGm_ + cumSumBinOffset + offset, srcData);	490	Simt::AtomicAdd<T_INDEX>(cumSumBinsGm_ + cumSumBinOffset + offset, srcData);
	@@ -435,7 +503,7 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
435	}	503	}
436	// tile id	504	// tile id
437	uint32_t startTileId = GetBlockIdx() % lastDimRealCore_;	505	uint32_t startTileId = GetBlockIdx() % lastDimRealCore_;
438	- uint32_t cumSumBinOffset = unsortedAxisId * RADIX_SORT_BIN_NUM;	506	+ uint32_t cumSumBinOffset = unsortedAxisId * topkV2::RADIX_SORT_BIN_NUM;
439	uint32_t inputXUnsortedAxisOffset = unsortedAxisId * totalDataNum_;	507	uint32_t inputXUnsortedAxisOffset = unsortedAxisId * totalDataNum_;
440	// local buffer	508	// local buffer
441	LocalTensor<int32_t> tileCusumBuffer = blockCumSumTbuf_.Get<int32_t>();	509	LocalTensor<int32_t> tileCusumBuffer = blockCumSumTbuf_.Get<int32_t>();
	@@ -450,32 +518,32 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
450	T_INDEX boundaryBinPrevCuSum = -1;	518	T_INDEX boundaryBinPrevCuSum = -1;
451	T_INDEX boundaryBinCuSum = -1;	519	T_INDEX boundaryBinCuSum = -1;
452	// clear ub	520	// clear ub
453	- Duplicate(tileTopkValue_, CLEAR_UB_VALUE, lastDimTileNumTimes_);	521	+ Duplicate(tileTopkValue_, topkV2::CLEAR_UB_VALUE, lastDimTileNumTimes_);
454	// get global hist	522	// get global hist
455	for(int32_t round = (NUM_PASS - 1); round >= 0; round--) {	523	for(int32_t round = (NUM_PASS - 1); round >= 0; round--) {
456	- if (topkValueInput_ > 0) {	524	+ if (inBlockDimRange_ && topkValueInput_ > 0) {
457	// clear ub buffer	525	// clear ub buffer
458	- Duplicate(tileCusumBuffer, static_cast<int32_t>(CLEAR_UB_VALUE), RADIX_SORT_BIN_NUM * lastDimTileNumTimes_);	526	+ Duplicate(tileCusumBuffer, static_cast<int32_t>(topkV2::CLEAR_UB_VALUE), topkV2::RADIX_SORT_BIN_NUM * lastDimTileNumTimes_);
459	// clear gm buffer	527	// clear gm buffer
460	if (startTileId == 0) {	528	if (startTileId == 0) {
461	LocalTensor<T_INDEX> reuseBuf2InitCumsumTemp = dataSetCumSumTbuf_.Get<T_INDEX>();	529	LocalTensor<T_INDEX> reuseBuf2InitCumsumTemp = dataSetCumSumTbuf_.Get<T_INDEX>();
462	- Duplicate(reuseBuf2InitCumsumTemp, static_cast<T_INDEX>(CLEAR_UB_VALUE), RADIX_SORT_BIN_NUM);	530	+ Duplicate(reuseBuf2InitCumsumTemp, static_cast<T_INDEX>(topkV2::CLEAR_UB_VALUE), topkV2::RADIX_SORT_BIN_NUM);
463	event_t eventIdMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));	531	event_t eventIdMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
464	SetFlag<HardEvent::V_MTE3>(eventIdMte3);	532	SetFlag<HardEvent::V_MTE3>(eventIdMte3);
465	WaitFlag<HardEvent::V_MTE3>(eventIdMte3);	533	WaitFlag<HardEvent::V_MTE3>(eventIdMte3);
466	DataCopyExtParams copyParams {1, 1, 0, 0, 0};	534	DataCopyExtParams copyParams {1, 1, 0, 0, 0};
467	- copyParams.blockLen = RADIX_SORT_BIN_NUM * sizeof(T_INDEX);	535	+ copyParams.blockLen = topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX);
468	DataCopyPad(cumSumBinsGm_[cumSumBinOffset], reuseBuf2InitCumsumTemp, copyParams);	536	DataCopyPad(cumSumBinsGm_[cumSumBinOffset], reuseBuf2InitCumsumTemp, copyParams);
469	dataSetCumSumTbuf_.FreeTensor(reuseBuf2InitCumsumTemp);	537	dataSetCumSumTbuf_.FreeTensor(reuseBuf2InitCumsumTemp);
470	}	538	}
471	}	539	}
472	PipeBarrier<PIPE_ALL>();	540	PipeBarrier<PIPE_ALL>();
473	SyncAll();	541	SyncAll();
474	- if (topkValueInput_ > 0) {	542	+ if (inBlockDimRange_ && topkValueInput_ > 0) {
475	for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {	543	for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {
476	T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);	544	T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);
477	// tile top ub offset	545	// tile top ub offset
478	- uint64_t tileTopkOffsetInUb = tileTopkValueIndex * RADIX_SORT_BIN_NUM;	546	+ uint64_t tileTopkOffsetInUb = tileTopkValueIndex * topkV2::RADIX_SORT_BIN_NUM;
479	// offset	547	// offset
480	uint64_t tileOffset = tileId * numTileData_;	548	uint64_t tileOffset = tileId * numTileData_;
481	int32_t tileDataStart = tileId * numTileData_;	549	int32_t tileDataStart = tileId * numTileData_;
	@@ -517,14 +585,14 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
517	SetFlag<HardEvent::V_MTE3>(eventId);	585	SetFlag<HardEvent::V_MTE3>(eventId);
518	WaitFlag<HardEvent::V_MTE3>(eventId);	586	WaitFlag<HardEvent::V_MTE3>(eventId);
519	if constexpr (IsSameType<T_INDEX, int64_t>::value) {	587	if constexpr (IsSameType<T_INDEX, int64_t>::value) {
520	- if constexpr (NUM_PASS == B64_BITE_SIZE) {	588	+ if constexpr (NUM_PASS == topkV2::B64_BITE_SIZE) {
521	- Simt::VF_CALL<CopyCumSumToGmB64<T_INDEX, T_INDEX_TO>>(Simt::Dim3(RADIX_SORT_BIN_NUM),	589	+ Simt::VF_CALL<CopyCumSumToGmB64<T_INDEX, T_INDEX_TO>>(Simt::Dim3(topkV2::RADIX_SORT_BIN_NUM),
522	(__gm__ T_INDEX *)(cumSumBinsGm_.GetPhyAddr()),	590	(__gm__ T_INDEX *)(cumSumBinsGm_.GetPhyAddr()),
523	(__ubuf__ int32_t *)(tileCusumBuffer.GetPhyAddr()),	591	(__ubuf__ int32_t *)(tileCusumBuffer.GetPhyAddr()),
524	cumSumBinOffset,	592	cumSumBinOffset,
525	tileTopkOffsetInUb);	593	tileTopkOffsetInUb);
526	} else {	594	} else {
527	- Simt::VF_CALL<CopyCumSumToGmB8B16B32<T_INDEX, T_INDEX_TO>>(Simt::Dim3(RADIX_SORT_BIN_NUM),	595	+ Simt::VF_CALL<CopyCumSumToGmB8B16B32<T_INDEX, T_INDEX_TO>>(Simt::Dim3(topkV2::RADIX_SORT_BIN_NUM),
528	(__gm__ T_INDEX *)(cumSumBinsGm_.GetPhyAddr()),	596	(__gm__ T_INDEX *)(cumSumBinsGm_.GetPhyAddr()),
529	(__ubuf__ int32_t *)(tileCusumBuffer.GetPhyAddr()),	597	(__ubuf__ int32_t *)(tileCusumBuffer.GetPhyAddr()),
530	cumSumBinOffset,	598	cumSumBinOffset,
	@@ -535,7 +603,7 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
535	// copy ub to gm	603	// copy ub to gm
536	DataCopyExtParams dataCopyParam;	604	DataCopyExtParams dataCopyParam;
537	dataCopyParam.blockCount = 1;	605	dataCopyParam.blockCount = 1;
538	- dataCopyParam.blockLen = RADIX_SORT_BIN_NUM * sizeof(int32_t);	606	+ dataCopyParam.blockLen = topkV2::RADIX_SORT_BIN_NUM * sizeof(int32_t);
539	dataCopyParam.srcStride = 0;	607	dataCopyParam.srcStride = 0;
540	dataCopyParam.dstStride = 0;	608	dataCopyParam.dstStride = 0;
541	DataCopyPad(cumSumBinsGm_[cumSumBinOffset],	609	DataCopyPad(cumSumBinsGm_[cumSumBinOffset],
	@@ -550,15 +618,15 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
550	}	618	}
551	PipeBarrier<PIPE_ALL>();	619	PipeBarrier<PIPE_ALL>();
552	SyncAll();	620	SyncAll();
553	- if (topkValueInput_ > 0) {	621	+ if (inBlockDimRange_ && topkValueInput_ > 0) {
554	// find bucket boundary	622	// find bucket boundary
555	// scalar calc	623	// scalar calc
556	FindBoundary(boundaryBin, boundaryBinPrev, boundaryBinCuSum,	624	FindBoundary(boundaryBin, boundaryBinPrev, boundaryBinCuSum,
557	boundaryBinPrevCuSum, cumSumBinOffset);	625	boundaryBinPrevCuSum, cumSumBinOffset);
558	- UNSIGNED_TYPE oneRoundMask = static_cast<UNSIGNED_TYPE>(boundaryBin) << (round * SHIFT_BIT_NUM);	626	+ UNSIGNED_TYPE oneRoundMask = static_cast<UNSIGNED_TYPE>(boundaryBin) << (round * topkV2::SHIFT_BIT_NUM);
559	involvedDataMask += oneRoundMask;	627	involvedDataMask += oneRoundMask;
560	// update and mask	628	// update and mask
561	- UNSIGNED_TYPE shiftMask = static_cast<UNSIGNED_TYPE>(0xFF) << (round * SHIFT_BIT_NUM);	629	+ UNSIGNED_TYPE shiftMask = static_cast<UNSIGNED_TYPE>(0xFF) << (round * topkV2::SHIFT_BIT_NUM);
562	andDataMask += shiftMask;	630	andDataMask += shiftMask;
563	// update topk value	631	// update topk value
564	topkValueInput_ -= boundaryBinPrevCuSum;	632	topkValueInput_ -= boundaryBinPrevCuSum;
	@@ -567,7 +635,7 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
567	T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);	635	T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);
568	// update tile topk and topk	636	// update tile topk and topk
569	if (boundaryBinPrev >= 0) {	637	if (boundaryBinPrev >= 0) {
570	- tileTopkValue_(tileTopkValueIndex) += tileCusumBuffer[tileTopkValueIndex * RADIX_SORT_BIN_NUM](boundaryBinPrev);	638	+ tileTopkValue_(tileTopkValueIndex) += tileCusumBuffer[tileTopkValueIndex * topkV2::RADIX_SORT_BIN_NUM](boundaryBinPrev);
571	}	639	}
572	}	640	}
573	}	641	}
	@@ -579,7 +647,7 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
579	// medium mode	647	// medium mode
580	for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {	648	for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {
581	T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);	649	T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);
582	- uint64_t tileTopkOffsetInUb = tileTopkValueIndex * RADIX_SORT_BIN_NUM;	650	+ uint64_t tileTopkOffsetInUb = tileTopkValueIndex * topkV2::RADIX_SORT_BIN_NUM;
583	StoreAnswer2Gm(xLocal, tileCusumBuffer[tileTopkOffsetInUb], boundaryBin,	651	StoreAnswer2Gm(xLocal, tileCusumBuffer[tileTopkOffsetInUb], boundaryBin,
584	tileCount, tileId, oneRowTopKValueOffset);	652	tileCount, tileId, oneRowTopKValueOffset);
585	}	653	}
	@@ -589,14 +657,16 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
589	}	657	}
590	PipeBarrier<PIPE_ALL>();	658	PipeBarrier<PIPE_ALL>();
591	SyncAll();	659	SyncAll();
592	- if (IS_SORT && (startTileId == 0) && topkValueInitInput_ <= SUPPORT_SORT_MAX_SIZE) {	660	+ if (inBlockDimRange_) {
593	- if (NUM_PASS != B64_BITE_SIZE) {	661	+ if (IS_SORT && (startTileId == 0) && !needSortWithIndex_) {
594	- SortTopKRes(xLocal);	662	+ if (NUM_PASS != topkV2::B64_BITE_SIZE) {
595	- } else if (NUM_PASS == B64_BITE_SIZE && topkValueInitInput_ <= (SUPPORT_SORT_MAX_SIZE / 2)) {	663	+ SortTopKRes(xLocal);
596	- SortTopKRes(xLocal);	664	+ } else if (NUM_PASS == topkV2::B64_BITE_SIZE && topkValueInitInput_ <= (SUPPORT_SORT_MAX_SIZE / 2)) {
		665	+ SortTopKRes(xLocal);
		666	+ }
597	}	667	}
		668	+ inQueueX_.FreeTensor(xLocal);
598	}	669	}
599	- inQueueX_.FreeTensor(xLocal);
600	}	670	}
601		671
602	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	672	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
	@@ -609,23 +679,23 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
609	{	679	{
610	// load cumsum	680	// load cumsum
611	LocalTensor<T_INDEX> cumSumLocal = dataSetCumSumTbuf_.Get<T_INDEX>();	681	LocalTensor<T_INDEX> cumSumLocal = dataSetCumSumTbuf_.Get<T_INDEX>();
612	- uint32_t currTileSizeAlign = ROUND_UP_AGLIN(RADIX_SORT_BIN_NUM * sizeof(T_INDEX)) / sizeof(T_INDEX);	682	+ uint32_t currTileSizeAlign = ROUND_UP_AGLIN(topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX)) / sizeof(T_INDEX);
613	DataCopyPadExtParams<T_INDEX> padParams;	683	DataCopyPadExtParams<T_INDEX> padParams;
614	padParams.isPad = true;	684	padParams.isPad = true;
615	- padParams.rightPadding = currTileSizeAlign - RADIX_SORT_BIN_NUM;	685	+ padParams.rightPadding = currTileSizeAlign - topkV2::RADIX_SORT_BIN_NUM;
616	padParams.paddingValue = static_cast<T_INDEX>(0);	686	padParams.paddingValue = static_cast<T_INDEX>(0);
617	DataCopyExtParams dataCopyParam;	687	DataCopyExtParams dataCopyParam;
618	dataCopyParam.blockCount = 1;	688	dataCopyParam.blockCount = 1;
619	- dataCopyParam.blockLen = RADIX_SORT_BIN_NUM * sizeof(T_INDEX);	689	+ dataCopyParam.blockLen = topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX);
620	dataCopyParam.srcStride = 0;	690	dataCopyParam.srcStride = 0;
621	dataCopyParam.dstStride = 0;	691	dataCopyParam.dstStride = 0;
622	DataCopyPad(cumSumLocal, cumSumBinsGm_[cumSumBinOffset], dataCopyParam, padParams);	692	DataCopyPad(cumSumLocal, cumSumBinsGm_[cumSumBinOffset], dataCopyParam, padParams);
623	event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));	693	event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
624	SetFlag<HardEvent::MTE2_S>(eventId);	694	SetFlag<HardEvent::MTE2_S>(eventId);
625	WaitFlag<HardEvent::MTE2_S>(eventId);	695	WaitFlag<HardEvent::MTE2_S>(eventId);
626	- if (cumSumLocal(RADIX_SORT_BIN_NUM - 1) <= topkValueInput_) {	696	+ if (cumSumLocal(topkV2::RADIX_SORT_BIN_NUM - 1) <= topkValueInput_) {
627	boundaryBin = -1;	697	boundaryBin = -1;
628	- boundaryBinPrev = RADIX_SORT_BIN_NUM - 1;	698	+ boundaryBinPrev = topkV2::RADIX_SORT_BIN_NUM - 1;
629	boundaryBinCuSum = -1;	699	boundaryBinCuSum = -1;
630	boundaryBinPrevCuSum = cumSumLocal(boundaryBinPrev);	700	boundaryBinPrevCuSum = cumSumLocal(boundaryBinPrev);
631	return ;	701	return ;
	@@ -649,16 +719,16 @@ __aicore__ inline int32_t RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST,
649	LocalTensor<T_INDEX> cumSumLocal)	719	LocalTensor<T_INDEX> cumSumLocal)
650	{	720	{
651	int32_t left = 0;	721	int32_t left = 0;
652	- int32_t right = RADIX_SORT_BIN_NUM - 1;	722	+ int32_t right = topkV2::RADIX_SORT_BIN_NUM - 1;
653	while(left <= right) {	723	while(left <= right) {
654	int mid = (right + left) / 2;	724	int mid = (right + left) / 2;
655	if (cumSumLocal(mid) == topkValueInput_) {	725	if (cumSumLocal(mid) == topkValueInput_) {
656	- if ((mid + 1) < RADIX_SORT_BIN_NUM && cumSumLocal(mid + 1) > topkValueInput_) {	726	+ if ((mid + 1) < topkV2::RADIX_SORT_BIN_NUM && cumSumLocal(mid + 1) > topkValueInput_) {
657	return mid;	727	return mid;
658	} else {	728	} else {
659	left = mid + 1;	729	left = mid + 1;
660	}	730	}
661	- } else if ((mid + 1) < RADIX_SORT_BIN_NUM && cumSumLocal(mid + 1) > topkValueInput_ && cumSumLocal(mid) < topkValueInput_) {	731	+ } else if ((mid + 1) < topkV2::RADIX_SORT_BIN_NUM && cumSumLocal(mid + 1) > topkValueInput_ && cumSumLocal(mid) < topkValueInput_) {
662	return mid;	732	return mid;
663	} else if (cumSumLocal(mid) < topkValueInput_) {	733	} else if (cumSumLocal(mid) < topkValueInput_) {
664	left = mid + 1;	734	left = mid + 1;
	@@ -685,29 +755,31 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
685	}	755	}
686	// tile id	756	// tile id
687	uint32_t startTileId = GetBlockIdx() % lastDimRealCore_;	757	uint32_t startTileId = GetBlockIdx() % lastDimRealCore_;
688	- uint32_t cumSumBinOffset = unsortedAxisId * RADIX_SORT_BIN_NUM;	758	+ uint32_t cumSumBinOffset = unsortedAxisId * topkV2::RADIX_SORT_BIN_NUM;
689	uint32_t inputXUnsortedAxisOffset = unsortedAxisId * totalDataNum_;	759	uint32_t inputXUnsortedAxisOffset = unsortedAxisId * totalDataNum_;
690	uint32_t oneRowTopKValueOffset = unsortedAxisId * tileCount;	760	uint32_t oneRowTopKValueOffset = unsortedAxisId * tileCount;
691	UpdateTileRealTopK(startTileId, tileCount,	761	UpdateTileRealTopK(startTileId, tileCount,
692	tileCusumBuffer, boundaryBin,	762	tileCusumBuffer, boundaryBin,
693	oneRowTopKValueOffset);	763	oneRowTopKValueOffset);
694	- // copy data	764	+ if (inBlockDimRange_) {
695	- for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {	765	+ // copy data
696	- uint32_t tileTopkValueIndex = (tileId / platformCoreNum_);	766	+ for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {
697	- uint64_t tileTopkOffsetInUb = tileTopkValueIndex * RADIX_SORT_BIN_NUM;	767	+ uint32_t tileTopkValueIndex = (tileId / platformCoreNum_);
698	- // offset	768	+ uint64_t tileTopkOffsetInUb = tileTopkValueIndex * topkV2::RADIX_SORT_BIN_NUM;
699	- uint64_t tileOffset = tileId * numTileData_;	769	+ // offset
700	- int32_t tileDataStart = tileId * numTileData_;	770	+ uint64_t tileOffset = tileId * numTileData_;
701	- int32_t remainTileDataNum = totalDataNum_ - tileDataStart;	771	+ int32_t tileDataStart = tileId * numTileData_;
702	- if (remainTileDataNum < 0) {	772	+ int32_t remainTileDataNum = totalDataNum_ - tileDataStart;
703	- break;	773	+ if (remainTileDataNum < 0) {
		774	+ break;
		775	+ }
		776	+ int32_t currTileNum = TopkGetMin<int32_t>(remainTileDataNum, static_cast<int32_t>(numTileData_));
		777	+ // copy gm to ub
		778	+ CopyDataInWithReuseBuffer(inputX[inputXUnsortedAxisOffset], xLocal,
		779	+ tileOffset, currTileNum);
		780	+ xLocal = inQueueX_.DeQue<T>();
		781	+ StoreFinalAnswer2Gm(xLocal, tileId, tileCount, oneRowTopKValueOffset);
704	}	782	}
705	- int32_t currTileNum = TopkGetMin<int32_t>(remainTileDataNum, static_cast<int32_t>(numTileData_));
706	- // copy gm to ub
707	- CopyDataInWithReuseBuffer(inputX[inputXUnsortedAxisOffset], xLocal,
708	- tileOffset, currTileNum);
709	- xLocal = inQueueX_.DeQue<T>();
710	- StoreFinalAnswer2Gm(xLocal, tileId, tileCount, oneRowTopKValueOffset);
711	}	783	}
712	}	784	}
713		785
	@@ -719,45 +791,49 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
719	T_INDEX boundaryBin,	791	T_INDEX boundaryBin,
720	uint64_t oneRowGmOffset)	792	uint64_t oneRowGmOffset)
721	{	793	{
722	- LocalTensor<uint32_t> reuseBuffer2Copy = topkSrcIndexTbuf_.Get<uint32_t>();	794	+ if (inBlockDimRange_) {
723	- for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {	795	+ LocalTensor<uint32_t> reuseBuffer2Copy = topkSrcIndexTbuf_.Get<uint32_t>();
724	- T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);	796	+ for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {
725	- uint64_t tileTopkOffsetInUb = tileTopkValueIndex * RADIX_SORT_BIN_NUM;	797	+ T_INDEX tileTopkValueIndex = (tileId / platformCoreNum_);
726	- uint32_t tilePrevCusumValue = ((boundaryBin >= 1) ? tileCusumBuffer[tileTopkOffsetInUb](boundaryBin - 1) : 0);	798	+ uint64_t tileTopkOffsetInUb = tileTopkValueIndex * topkV2::RADIX_SORT_BIN_NUM;
727	- CopyUb2Ub(reuseBuffer2Copy, tileTopkValue_, tileTopkValueIndex);	799	+ uint32_t tilePrevCusumValue = ((boundaryBin >= 1) ? tileCusumBuffer[tileTopkOffsetInUb](boundaryBin - 1) : 0);
728	- event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));	800	+ CopyUb2Ub(reuseBuffer2Copy, tileTopkValue_, tileTopkValueIndex);
729	- SetFlag<HardEvent::V_MTE3>(eventId);	801	+ event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
730	- WaitFlag<HardEvent::V_MTE3>(eventId);	802	+ SetFlag<HardEvent::V_MTE3>(eventId);
731	- CopyDataToGm(tileTopkValueGm_[oneRowGmOffset], reuseBuffer2Copy, tileId, 1);	803	+ WaitFlag<HardEvent::V_MTE3>(eventId);
732	- if (topkValueInput_ > 0) {	804	+ CopyDataToGm(tileTopkValueGm_[oneRowGmOffset], reuseBuffer2Copy, tileId, 1);
733	- event_t eventMet3VecId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));	805	+ if (topkValueInput_ > 0) {
734	- SetFlag<HardEvent::MTE3_V>( eventMet3VecId);	806	+ event_t eventMet3VecId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
735	- WaitFlag<HardEvent::MTE3_V>( eventMet3VecId);	807	+ SetFlag<HardEvent::MTE3_V>( eventMet3VecId);
736	- remainTileTopkValue_(tileTopkValueIndex) =	808	+ WaitFlag<HardEvent::MTE3_V>( eventMet3VecId);
737	- tileCusumBuffer[tileTopkOffsetInUb](boundaryBin) - tilePrevCusumValue;	809	+ remainTileTopkValue_(tileTopkValueIndex) =
738	- CopyUb2Ub(reuseBuffer2Copy, remainTileTopkValue_, tileTopkValueIndex);	810	+ tileCusumBuffer[tileTopkOffsetInUb](boundaryBin) - tilePrevCusumValue;
739	- event_t eventIdWaitV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));	811	+ CopyUb2Ub(reuseBuffer2Copy, remainTileTopkValue_, tileTopkValueIndex);
740	- SetFlag<HardEvent::V_MTE3>(eventIdWaitV);	812	+ event_t eventIdWaitV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
741	- WaitFlag<HardEvent::V_MTE3>(eventIdWaitV);	813	+ SetFlag<HardEvent::V_MTE3>(eventIdWaitV);
742	- CopyDataToGm(tileTopkRemainValueGm_[oneRowGmOffset], reuseBuffer2Copy, tileId, 1);	814	+ WaitFlag<HardEvent::V_MTE3>(eventIdWaitV);
		815	+ CopyDataToGm(tileTopkRemainValueGm_[oneRowGmOffset], reuseBuffer2Copy, tileId, 1);
		816	+ }
		817	+ event_t eventIdWaitMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
		818	+ SetFlag<HardEvent::MTE3_V>(eventIdWaitMte3);
		819	+ WaitFlag<HardEvent::MTE3_V>(eventIdWaitMte3);
743	}	820	}
744	- event_t eventIdWaitMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
745	- SetFlag<HardEvent::MTE3_V>(eventIdWaitMte3);
746	- WaitFlag<HardEvent::MTE3_V>(eventIdWaitMte3);
747	}	821	}
748	PipeBarrier<PIPE_ALL>();	822	PipeBarrier<PIPE_ALL>();
749	SyncAll();	823	SyncAll();
750	- for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {	824	+ if (inBlockDimRange_) {
751	- uint32_t tileTopkValueIndex = (tileId / platformCoreNum_);	825	+ for(uint32_t tileId = startTileId; tileId < tileCount; tileId += lastDimRealCore_) {
752	- uint64_t tileTopkOffsetInUb = tileTopkValueIndex * RADIX_SORT_BIN_NUM;	826	+ uint32_t tileTopkValueIndex = (tileId / platformCoreNum_);
753	- T_INDEX tilePrevCusumValue = ((boundaryBin >= 1) ? tileCusumBuffer[tileTopkOffsetInUb](boundaryBin - 1) : 0);	827	+ uint64_t tileTopkOffsetInUb = tileTopkValueIndex * topkV2::RADIX_SORT_BIN_NUM;
754	- if (topkValueInput_ > 0) {	828	+ T_INDEX tilePrevCusumValue = ((boundaryBin >= 1) ? tileCusumBuffer[tileTopkOffsetInUb](boundaryBin - 1) : 0);
755	- // get tile offset	829	+ if (topkValueInput_ > 0) {
756	- // get cusum and store to tileTopkRemainValueGm_	830	+ // get tile offset
757	- T_INDEX cumSumValue = GetTileTopkValueOffset(tileTopkRemainValueGm_, tileId, tileCount, oneRowGmOffset);	831	+ // get cusum and store to tileTopkRemainValueGm_
758	- // update tileTopkValue	832	+ T_INDEX cumSumValue = GetTileTopkValueOffset(tileTopkRemainValueGm_, tileId, tileCount, oneRowGmOffset);
759	- UpdateTileTopkValue(tileCusumBuffer[tileTopkOffsetInUb], cumSumValue, boundaryBin,	833	+ // update tileTopkValue
760	- tilePrevCusumValue, tileId, oneRowGmOffset);	834	+ UpdateTileTopkValue(tileCusumBuffer[tileTopkOffsetInUb], cumSumValue, boundaryBin,
		835	+ tilePrevCusumValue, tileId, oneRowGmOffset);
		836	+ }
761	}	837	}
762	}	838	}
763	PipeBarrier<PIPE_ALL>();	839	PipeBarrier<PIPE_ALL>();
	@@ -895,7 +971,11 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
895	static_cast<uint16_t>(1),	971	static_cast<uint16_t>(1),
896	static_cast<uint32_t>(topKValue * sizeof(T)),	972	static_cast<uint32_t>(topKValue * sizeof(T)),
897	0, 0, 0};	973	0, 0, 0};
898	- AscendC::DataCopyPad(topkValueGm_[gmOffset + tileOffset], topkValueOutLocal, dataCopyParamValue);	974	+ if (needSortWithIndex_) {
		975	+ AscendC::DataCopyPad(topkValuesGm_[gmOffset + tileOffset], topkValueOutLocal, dataCopyParamValue);
		976	+ } else {
		977	+ AscendC::DataCopyPad(topkValueGm_[gmOffset + tileOffset], topkValueOutLocal, dataCopyParamValue);
		978	+ }
899	topkValueQueue_.FreeTensor(topkValueOutLocal);	979	topkValueQueue_.FreeTensor(topkValueOutLocal);
900	// copy sorted value index	980	// copy sorted value index
901	AscendC::LocalTensor<T_INDEX_TO> topkValueOutIndexLocal = topkValueIndexQueue_.DeQue<T_INDEX_TO>();	981	AscendC::LocalTensor<T_INDEX_TO> topkValueOutIndexLocal = topkValueIndexQueue_.DeQue<T_INDEX_TO>();
	@@ -903,7 +983,11 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
903	static_cast<uint16_t>(1),	983	static_cast<uint16_t>(1),
904	static_cast<uint32_t>(topKValue * sizeof(T_INDEX_TO)),	984	static_cast<uint32_t>(topKValue * sizeof(T_INDEX_TO)),
905	0, 0, 0};	985	0, 0, 0};
906	- AscendC::DataCopyPad(topkValueIndexGm_[gmOffset + tileOffset], topkValueOutIndexLocal, dataCopyParamIndex);	986	+ if (needSortWithIndex_) {
		987	+ AscendC::DataCopyPad(topkIndicesGm_[gmOffset + tileOffset], topkValueOutIndexLocal, dataCopyParamIndex);
		988	+ } else {
		989	+ AscendC::DataCopyPad(topkValueIndexGm_[gmOffset + tileOffset], topkValueOutIndexLocal, dataCopyParamIndex);
		990	+ }
907	topkValueIndexQueue_.FreeTensor(topkValueOutIndexLocal);	991	topkValueIndexQueue_.FreeTensor(topkValueOutIndexLocal);
908	}	992	}
909		993
	@@ -993,21 +1077,21 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
993	int32_t round,	1077	int32_t round,
994	uint32_t numTileData)	1078	uint32_t numTileData)
995	{	1079	{
996	- if constexpr (is_same<int64_t, T>::value \|\| is_same<uint64_t, T>::value) {	1080	+ if constexpr (topkV2::is_same<int64_t, T>::value \|\| topkV2::is_same<uint64_t, T>::value) {
997	RadixSortTopKB64<T, uint64_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	1081	RadixSortTopKB64<T, uint64_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
998	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	1082	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
999	involveDataMask, round, numTileData);	1083	involveDataMask, round, numTileData);
1000	- } else if constexpr (is_same<int32_t, T>::value \|\| is_same<uint32_t, T>::value \|\| is_same<float, T>::value) {	1084	+ } else if constexpr (topkV2::is_same<int32_t, T>::value \|\| topkV2::is_same<uint32_t, T>::value \|\| topkV2::is_same<float, T>::value) {
1001	RadixSortTopKB32<T, uint32_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	1085	RadixSortTopKB32<T, uint32_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
1002	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	1086	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
1003	involveDataMask, round, numTileData);	1087	involveDataMask, round, numTileData);
1004	} else if constexpr (	1088	} else if constexpr (
1005	- is_same<half, T>::value \|\| is_same<uint16_t, T>::value \|\| is_same<int16_t, T>::value \|\|	1089	+ topkV2::is_same<half, T>::value \|\| topkV2::is_same<uint16_t, T>::value \|\| topkV2::is_same<int16_t, T>::value \|\|
1006	- is_same<bfloat16_t, T>::value) {	1090	+ topkV2::is_same<bfloat16_t, T>::value) {
1007	RadixSortTopKB16<T, uint16_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	1091	RadixSortTopKB16<T, uint16_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
1008	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	1092	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
1009	involveDataMask, round, numTileData);	1093	involveDataMask, round, numTileData);
1010	- } else if constexpr (is_same<int8_t, T>::value \|\| is_same<uint8_t, T>::value) {	1094	+ } else if constexpr (topkV2::is_same<int8_t, T>::value \|\| topkV2::is_same<uint8_t, T>::value) {
1011	RadixSortTopKB8<T, uint8_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	1095	RadixSortTopKB8<T, uint8_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
1012	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	1096	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
1013	involveDataMask, round, numTileData);	1097	involveDataMask, round, numTileData);
	@@ -1019,27 +1103,27 @@ __aicore__ inline LocalTensor<UNSIGNED_TYPE> RadixSortTopK<T, UNSIGNED_TYPE, NUM
1019	LocalTensor<T> inputX,	1103	LocalTensor<T> inputX,
1020	uint32_t numTileData)	1104	uint32_t numTileData)
1021	{	1105	{
1022	- if constexpr (is_same<int64_t, T>::value) {	1106	+ if constexpr (topkV2::is_same<int64_t, T>::value) {
1023	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	1107	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
1024	radixSortTopK.TwiddleInB64(inputX, inputXCopy_, numTileData);	1108	radixSortTopK.TwiddleInB64(inputX, inputXCopy_, numTileData);
1025	return inputXCopy_;	1109	return inputXCopy_;
1026	- } else if constexpr (is_same<int32_t, T>::value) {	1110	+ } else if constexpr (topkV2::is_same<int32_t, T>::value) {
1027	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	1111	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
1028	radixSortTopK.TwiddleInB32(inputX, inputXCopy_, numTileData);	1112	radixSortTopK.TwiddleInB32(inputX, inputXCopy_, numTileData);
1029	return inputXCopy_;	1113	return inputXCopy_;
1030	- } else if constexpr (is_same<half, T>::value \|\| is_same<bfloat16_t, T>::value) {	1114	+ } else if constexpr (topkV2::is_same<half, T>::value \|\| topkV2::is_same<bfloat16_t, T>::value) {
1031	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	1115	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
1032	radixSortTopK.TwiddleInFp16(inputX, inputXCopy_, numTileData);	1116	radixSortTopK.TwiddleInFp16(inputX, inputXCopy_, numTileData);
1033	return inputXCopy_;	1117	return inputXCopy_;
1034	- } else if constexpr (is_same<float, T>::value) {	1118	+ } else if constexpr (topkV2::is_same<float, T>::value) {
1035	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	1119	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
1036	radixSortTopK.TwiddleInFp32(inputX, inputXCopy_, numTileData);	1120	radixSortTopK.TwiddleInFp32(inputX, inputXCopy_, numTileData);
1037	return inputXCopy_;	1121	return inputXCopy_;
1038	- } else if constexpr (is_same<int16_t, T>::value) {	1122	+ } else if constexpr (topkV2::is_same<int16_t, T>::value) {
1039	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	1123	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
1040	radixSortTopK.TwiddleInB16(inputX, inputXCopy_, numTileData);	1124	radixSortTopK.TwiddleInB16(inputX, inputXCopy_, numTileData);
1041	return inputXCopy_;	1125	return inputXCopy_;
1042	- } else if constexpr (is_same<int8_t, T>::value) {	1126	+ } else if constexpr (topkV2::is_same<int8_t, T>::value) {
1043	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	1127	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
1044	radixSortTopK.TwiddleInB8(inputX, inputXCopy_, numTileData);	1128	radixSortTopK.TwiddleInB8(inputX, inputXCopy_, numTileData);
1045	return inputXCopy_;	1129	return inputXCopy_;
	@@ -1058,16 +1142,16 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
1058	LocalTensor<UNSIGNED_TYPE> reverseInputX,	1142	LocalTensor<UNSIGNED_TYPE> reverseInputX,
1059	uint32_t numTileData)	1143	uint32_t numTileData)
1060	{	1144	{
1061	- if constexpr (is_same<uint64_t, T>::value) {	1145	+ if constexpr (topkV2::is_same<uint64_t, T>::value) {
1062	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	1146	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
1063	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	1147	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
1064	- } else if constexpr (is_same<uint32_t, T>::value) {	1148	+ } else if constexpr (topkV2::is_same<uint32_t, T>::value) {
1065	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	1149	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
1066	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	1150	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
1067	- } else if constexpr (is_same<uint16_t, T>::value) {	1151	+ } else if constexpr (topkV2::is_same<uint16_t, T>::value) {
1068	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	1152	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
1069	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	1153	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
1070	- } else if constexpr (is_same<uint8_t, T>::value) {	1154	+ } else if constexpr (topkV2::is_same<uint8_t, T>::value) {
1071	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	1155	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
1072	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	1156	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
1073	}	1157	}
	@@ -1087,7 +1171,7 @@ __aicore__ inline void RadixSortTopK<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_
1087	MicroAPI::DataCopy<uint32_t, MicroAPI::LoadDist::DIST_BRC_B32>(inputVectorOne, inputUbPtr + offset);	1171	MicroAPI::DataCopy<uint32_t, MicroAPI::LoadDist::DIST_BRC_B32>(inputVectorOne, inputUbPtr + offset);
1088	MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputUbPtr,	1172	MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputUbPtr,
1089	inputVectorOne,	1173	inputVectorOne,
1090	- ONE_TIMES_B32_NUM,	1174	+ topkV2::ONE_TIMES_B32_NUM,
1091	predicateDefaultB32);	1175	predicateDefaultB32);
1092	}	1176	}
1093	}	1177	}

Mmath/top_k_v2/op_kernel/arch35/radix_sort_top_k_inter_core_template_optimization.h+8-2

	@@ -21,6 +21,7 @@
21	#include <algorithm>	21	#include <algorithm>
22		22
23	using namespace AscendC;	23	using namespace AscendC;
		24	+using namespace topkV2;
24		25
25	// 类比当前文件	26	// 类比当前文件
26	template <typename T, bool IS_LARGEST, bool IS_SORT, typename T_INDEX_TO>	27	template <typename T, bool IS_LARGEST, bool IS_SORT, typename T_INDEX_TO>
	@@ -88,8 +89,8 @@ __aicore__ inline void RadixSortTopKMultiCoreOptimization<T, IS_LARGEST, IS_SORT
88	pipe.InitBuffer(topkOutValueQueue_, 1, ROUND_UP_AGLIN(topkValueInput_ * sizeof(T)));	89	pipe.InitBuffer(topkOutValueQueue_, 1, ROUND_UP_AGLIN(topkValueInput_ * sizeof(T)));
89	pipe.InitBuffer(tempIndexConversionQueue_, 1, ROUND_UP_AGLIN(topkValueInput_ * sizeof(T_INDEX_TO)));	90	pipe.InitBuffer(tempIndexConversionQueue_, 1, ROUND_UP_AGLIN(topkValueInput_ * sizeof(T_INDEX_TO)));
90	pipe.InitBuffer(topKApiTmpTBuf_, ROUND_UP_AGLIN(topKApiTmpSize_));	91	pipe.InitBuffer(topKApiTmpTBuf_, ROUND_UP_AGLIN(topKApiTmpSize_));
91	- uint32_t oneBlockNum = UB_AGLIN_VALUE / static_cast<uint32_t>(sizeof(T));	92	+ uint32_t oneBlockNum = topkV2::UB_AGLIN_VALUE / static_cast<uint32_t>(sizeof(T));
92	- uint32_t oneBlockNumB32 = UB_AGLIN_VALUE / static_cast<uint32_t>(sizeof(int32_t));	93	+ uint32_t oneBlockNumB32 = topkV2::UB_AGLIN_VALUE / static_cast<uint32_t>(sizeof(int32_t));
93	uint32_t sortResultOffset = CeilDivMul(unsortedDimParallel_ * topkValueInput_ * lastDimTileNum_, oneBlockNum);	94	uint32_t sortResultOffset = CeilDivMul(unsortedDimParallel_ * topkValueInput_ * lastDimTileNum_, oneBlockNum);
94	tempSortResultDataGm_.SetGlobalBuffer((__gm__ T*)workspace_, sortResultOffset);	95	tempSortResultDataGm_.SetGlobalBuffer((__gm__ T*)workspace_, sortResultOffset);
95	workSpaceOffset += sortResultOffset * static_cast<uint32_t>(sizeof(T)) / static_cast<uint32_t>(sizeof(int32_t));	96	workSpaceOffset += sortResultOffset * static_cast<uint32_t>(sizeof(T)) / static_cast<uint32_t>(sizeof(int32_t));
	@@ -135,7 +136,12 @@ __aicore__ inline void RadixSortTopKMultiCoreOptimization<T, IS_LARGEST, IS_SORT
135	int tileCount = (totalDataNum_ + numTileData_ - 1) / numTileData_;	136	int tileCount = (totalDataNum_ + numTileData_ - 1) / numTileData_;
136	uint32_t unsortedAxisId = GetBlockIdx() / lastDimRealCore_;	137	uint32_t unsortedAxisId = GetBlockIdx() / lastDimRealCore_;
137	uint32_t unsortedDimIndex = unsortedAxisId + sortLoopRound * unsortedDimParallel_;	138	uint32_t unsortedDimIndex = unsortedAxisId + sortLoopRound * unsortedDimParallel_;
		139	+ // inUnsortedDimRange决定是否执行本次计算流程，超出Batch轴数量范围不执行
138	bool inUnsortedDimRange = unsortedDimIndex >= unsortedDimNum_ ? false : true;	140	bool inUnsortedDimRange = unsortedDimIndex >= unsortedDimNum_ ? false : true;
		141	+ // 超出需要的核心数不执行
		142	+ if (GetBlockIdx() >= lastDimTileNum_ * unsortedDimParallel_) {
		143	+ inUnsortedDimRange = false;
		144	+ }
139	uint32_t startTileId = GetBlockIdx() % lastDimRealCore_;	145	uint32_t startTileId = GetBlockIdx() % lastDimRealCore_;
140	uint32_t inputXUnsortedAxisOffset = unsortedAxisId * totalDataNum_;	146	uint32_t inputXUnsortedAxisOffset = unsortedAxisId * totalDataNum_;
141	LocalTensor<bool> emptyFinishLocal;	147	LocalTensor<bool> emptyFinishLocal;

Mmath/top_k_v2/op_kernel/arch35/radix_sort_top_k_single_block.h+2-2

Mmath/top_k_v2/op_kernel/arch35/radix_sort_top_k_single_core.h+115-51

	@@ -26,9 +26,10 @@
26	#include "top_k_radix_block_sort_b32.h"	26	#include "top_k_radix_block_sort_b32.h"
27	#include "top_k_radix_block_sort_b16.h"	27	#include "top_k_radix_block_sort_b16.h"
28	#include "top_k_radix_block_sort_b8.h"	28	#include "top_k_radix_block_sort_b8.h"
		29	+#include "sort_with_index_entry.h"
29		30
30	using namespace AscendC;	31	using namespace AscendC;
31	-	32	+using namespace topkV2;
32	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	33	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
33	struct RadixSortTopKSingleCore : public RadixSortTopKBase<T, T_INDEX, T_INDEX_TO> {	34	struct RadixSortTopKSingleCore : public RadixSortTopKBase<T, T_INDEX, T_INDEX_TO> {
34	__aicore__ inline RadixSortTopKSingleCore() {};	35	__aicore__ inline RadixSortTopKSingleCore() {};
	@@ -75,6 +76,19 @@ private:
75	private:	76	private:
76	// 所有块的统计直方图结果	77	// 所有块的统计直方图结果
77	GlobalTensor<T_INDEX> tilesCusumGm_;	78	GlobalTensor<T_INDEX> tilesCusumGm_;
		79	+ // TopKV2 output value
		80	+ GlobalTensor<T> topkValuesGm_;
		81	+ // TopKV2 output index
		82	+ GlobalTensor<T_INDEX_TO> topkIndicesGm_;
		83	+
		84	+ // sortWithindex输入数据
		85	+ GM_ADDR sortWithIndexWorkspace_;
		86	+ GM_ADDR topkValuesGmAddr_;
		87	+ GM_ADDR topkIndicesGmAddr_;
		88	+ GM_ADDR valueAddr_;
		89	+ GM_ADDR indicesAddr_;
		90	+ const TopKV2TilingDataSimd* tilingDataPtr_;
		91	+
78	// 无符号输入	92	// 无符号输入
79	LocalTensor<UNSIGNED_TYPE> unsignedInputXLocal_;	93	LocalTensor<UNSIGNED_TYPE> unsignedInputXLocal_;
80	// tileK每块K的贡献度	94	// tileK每块K的贡献度
	@@ -93,6 +107,7 @@ private:
93	uint32_t tailTileNum_ = 0;	107	uint32_t tailTileNum_ = 0;
94	// B轴尾块大小	108	// B轴尾块大小
95	uint32_t tailBatchNum_ = 0;	109	uint32_t tailBatchNum_ = 0;
		110	+ bool needSortWithIndex_ = false;
96		111
97	UNSIGNED_TYPE histDataMask = 0;	112	UNSIGNED_TYPE histDataMask = 0;
98	UNSIGNED_TYPE highBitMask = 0;	113	UNSIGNED_TYPE highBitMask = 0;
	@@ -108,6 +123,8 @@ private:
108	// 排序时输入索引TBuf	123	// 排序时输入索引TBuf
109	TBuf<TPosition::VECCALC> sortSrcIndexTBuf_;	124	TBuf<TPosition::VECCALC> sortSrcIndexTBuf_;
110	TBuf<TPosition::VECCALC> tileCusumInt64TBuf_;	125	TBuf<TPosition::VECCALC> tileCusumInt64TBuf_;
		126	+
		127	+ uint64_t oneBlock_ = Ops::Base::GetUbBlockSize();
111	};	128	};
112		129
113	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	130	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
	@@ -139,16 +156,50 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
139	this->tPipe_->InitBuffer(this->topKApiTmpTBuf_, ROUND_UP_AGLIN(this->topKApiTmpSize_));	156	this->tPipe_->InitBuffer(this->topKApiTmpTBuf_, ROUND_UP_AGLIN(this->topKApiTmpSize_));
140	this->tPipe_->InitBuffer(this->indicesOutTbuf_, ROUND_UP_AGLIN(outQueueNum * sizeof(int32_t)));	157	this->tPipe_->InitBuffer(this->indicesOutTbuf_, ROUND_UP_AGLIN(outQueueNum * sizeof(int32_t)));
141		158
142	-
143	// 存放所有块统计直方图的结果	159	// 存放所有块统计直方图的结果
144	- tilesCusumGm_.SetGlobalBuffer((__gm__ T_INDEX)workSpace, unsortedDimParallel_ RADIX_SORT_BIN_NUM * tileCount_);	160	+ uint32_t workSpaceOffset = unsortedDimParallel_ * topkV2::RADIX_SORT_BIN_NUM * tileCount_;
145	-	161	+ tilesCusumGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T_INDEX*>(workSpace), int64_t(workSpaceOffset));
		162	+ workSpaceOffset = workSpaceOffset * sizeof(T_INDEX);
		163	+
		164	+ if (IS_SORT) {
		165	+ if (this->k_ <= SUPPORT_SORT_MAX_SIZE && this->k_ * sizeof(T) <= SUPPORT_SORT_MAX_BYTE_SIZE) {
		166	+ needSortWithIndex_ = false;
		167	+ } else {
		168	+ needSortWithIndex_ = true;
		169	+ }
		170	+ }
		171	+ if (needSortWithIndex_) {
		172	+ // sort 尾轴的大小
		173	+ uint32_t lastAxisNumForSort = tilingData->lastAxisNumForSort;
		174	+ // sort 外轴的大小
		175	+ uint32_t unsortedDimNumForSort = tilingData->unsortedDimNumForSort;
		176	+ uint64_t topkIndicesGmOffset = lastAxisNumForSort * unsortedDimNumForSort;
		177	+ topkIndicesGmOffset =
		178	+ CeilAlignDivMul<uint64_t>(int64_t(topkIndicesGmOffset * sizeof(T_INDEX_TO)), int64_t(oneBlock_)) /
		179	+ sizeof(T_INDEX_TO);
		180	+ topkIndicesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T_INDEX_TO*>(workSpace + workSpaceOffset), topkIndicesGmOffset);
		181	+ topkIndicesGmAddr_ = workSpace + workSpaceOffset;
		182	+ workSpaceOffset += topkIndicesGmOffset * sizeof(T_INDEX_TO);
		183	+
		184	+ uint64_t topkValuesGmOffset = lastAxisNumForSort * unsortedDimNumForSort;
		185	+ topkValuesGmOffset = CeilAlignDivMul<uint64_t>(int64_t(topkValuesGmOffset * sizeof(T)), int64_t(oneBlock_)) /
		186	+ sizeof(T);
		187	+ topkValuesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(workSpace + workSpaceOffset), topkValuesGmOffset);
		188	+ topkValuesGmAddr_ = workSpace + workSpaceOffset;
		189	+ workSpaceOffset += topkValuesGmOffset * sizeof(T);
		190	+
		191	+ sortWithIndexWorkspace_ = workSpace + workSpaceOffset;
		192	+ valueAddr_ = value;
		193	+ indicesAddr_ = indices;
		194	+ tilingDataPtr_ = tilingData;
		195	+ }
		196	+
146	// 存放块统计直方图累加和的结果，累加之前先搬运到tileCusumGm_上，然后累加到cusumTBuf	197	// 存放块统计直方图累加和的结果，累加之前先搬运到tileCusumGm_上，然后累加到cusumTBuf
147	- this->tPipe_->InitBuffer(tileCusumTBuf_, RADIX_SORT_BIN_NUM * sizeof(int32_t));	198	+ this->tPipe_->InitBuffer(tileCusumTBuf_, topkV2::RADIX_SORT_BIN_NUM * sizeof(int32_t));
148	- this->tPipe_->InitBuffer(cusumTBuf_, RADIX_SORT_BIN_NUM * sizeof(T_INDEX));	199	+ this->tPipe_->InitBuffer(cusumTBuf_, topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX));
149	this->tPipe_->InitBuffer(tileKTBuf_, ROUND_UP_AGLIN(tileCount_ * sizeof(int32_t)));	200	this->tPipe_->InitBuffer(tileKTBuf_, ROUND_UP_AGLIN(tileCount_ * sizeof(int32_t)));
150	this->tPipe_->InitBuffer(unsignedInputXTBuf_, ROUND_UP_AGLIN(tileNum * sizeof(UNSIGNED_TYPE)));	201	this->tPipe_->InitBuffer(unsignedInputXTBuf_, ROUND_UP_AGLIN(tileNum * sizeof(UNSIGNED_TYPE)));
151	- this->tPipe_->InitBuffer(tileCusumInt64TBuf_, RADIX_SORT_BIN_NUM * sizeof(T_INDEX));	202	+ this->tPipe_->InitBuffer(tileCusumInt64TBuf_, topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX));
152	// 排序时输入索引TBuf	203	// 排序时输入索引TBuf
153	if (IS_SORT && this->k_ * sizeof(T) <= SUPPORT_SORT_MAX_BYTE_SIZE) {	204	if (IS_SORT && this->k_ * sizeof(T) <= SUPPORT_SORT_MAX_BYTE_SIZE) {
154	this->tPipe_->InitBuffer(sortSrcIndexTBuf_, ROUND_UP_AGLIN(this->k_ * sizeof(T_INDEX_TO)));	205	this->tPipe_->InitBuffer(sortSrcIndexTBuf_, ROUND_UP_AGLIN(this->k_ * sizeof(T_INDEX_TO)));
	@@ -170,6 +221,11 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
170	for(int32_t i = 0; i < loopCount; i++) {	221	for(int32_t i = 0; i < loopCount; i++) {
171	ProcessSingleTopK(i);	222	ProcessSingleTopK(i);
172	}	223	}
		224	+ if (needSortWithIndex_) {
		225	+ this->tPipe_->Reset();
		226	+ sortwithindexForTopK<T_INDEX_TO>(topkValuesGmAddr_, topkIndicesGmAddr_, valueAddr_, indicesAddr_,
		227	+ sortWithIndexWorkspace_, tilingDataPtr_, this->tPipe_);
		228	+ }
173	}	229	}
174		230
175	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	231	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
	@@ -218,17 +274,17 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
218	LocalTensor<T_INDEX> tileCusumInt64Tmp;	274	LocalTensor<T_INDEX> tileCusumInt64Tmp;
219		275
220	// 每块K的贡献度清零	276	// 每块K的贡献度清零
221	- Duplicate(tileKLocal_, static_cast<int32_t>(CLEAR_UB_VALUE), tileCount_);	277	+ Duplicate(tileKLocal_, static_cast<int32_t>(topkV2::CLEAR_UB_VALUE), tileCount_);
222	// 该核输入的offset	278	// 该核输入的offset
223	uint64_t inputGmOffset = unsortedDimParallel_ * loopTime * this->lastAxisNum_ + this->blockIndex_ * this->lastAxisNum_;	279	uint64_t inputGmOffset = unsortedDimParallel_ * loopTime * this->lastAxisNum_ + this->blockIndex_ * this->lastAxisNum_;
224	// 每个核统计直方图的offset	280	// 每个核统计直方图的offset
225	- uint32_t cusumGmOffset = this->blockIndex_ * RADIX_SORT_BIN_NUM * tileCount_;	281	+ uint32_t cusumGmOffset = this->blockIndex_ * topkV2::RADIX_SORT_BIN_NUM * tileCount_;
226	for(int32_t round = (NUM_PASS - 1); round >= 0; round--) {	282	for(int32_t round = (NUM_PASS - 1); round >= 0; round--) {
227	if (updatedK_ > 0) {	283	if (updatedK_ > 0) {
228	// 每块统计直方图累加结果，需清零	284	// 每块统计直方图累加结果，需清零
229	- Duplicate(cusumLocal_, static_cast<T_INDEX>(CLEAR_UB_VALUE), RADIX_SORT_BIN_NUM);	285	+ Duplicate(cusumLocal_, static_cast<T_INDEX>(topkV2::CLEAR_UB_VALUE), topkV2::RADIX_SORT_BIN_NUM);
230	// 每块统计直方图的结果，需清零	286	// 每块统计直方图的结果，需清零
231	- Duplicate(tileCusumLocal, static_cast<int32_t>(CLEAR_UB_VALUE), RADIX_SORT_BIN_NUM);	287	+ Duplicate(tileCusumLocal, static_cast<int32_t>(topkV2::CLEAR_UB_VALUE), topkV2::RADIX_SORT_BIN_NUM);
232		288
233	// 计算该高8位每块统计方图的结果，并累加到cusumLocal_，并搬运到tilesCusumGm_上，方便后面使用	289	// 计算该高8位每块统计方图的结果，并累加到cusumLocal_，并搬运到tilesCusumGm_上，方便后面使用
234	for(uint32_t tileId = 0; tileId < tileCount_; tileId++) {	290	for(uint32_t tileId = 0; tileId < tileCount_; tileId++) {
	@@ -249,20 +305,20 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
249	// tileCusumInt64TBuf_	305	// tileCusumInt64TBuf_
250	if (IsSameType<T_INDEX, int64_t>::value) {	306	if (IsSameType<T_INDEX, int64_t>::value) {
251	tileCusumInt64Tmp = tileCusumInt64TBuf_.AllocTensor<T_INDEX>();	307	tileCusumInt64Tmp = tileCusumInt64TBuf_.AllocTensor<T_INDEX>();
252	- AscendC::Cast<T_INDEX, int32_t>(tileCusumInt64Tmp, tileCusumLocal, RoundMode::CAST_NONE, static_cast<int32_t>(RADIX_SORT_BIN_NUM));	308	+ AscendC::Cast<T_INDEX, int32_t>(tileCusumInt64Tmp, tileCusumLocal, RoundMode::CAST_NONE, static_cast<int32_t>(topkV2::RADIX_SORT_BIN_NUM));
253	- Add(cusumLocal_, cusumLocal_, tileCusumInt64Tmp, RADIX_SORT_BIN_NUM);	309	+ Add(cusumLocal_, cusumLocal_, tileCusumInt64Tmp, topkV2::RADIX_SORT_BIN_NUM);
254	} else {	310	} else {
255	- Add(cusumLocal_.template ReinterpretCast<int32_t>(), cusumLocal_.template ReinterpretCast<int32_t>(), tileCusumLocal, RADIX_SORT_BIN_NUM);	311	+ Add(cusumLocal_.template ReinterpretCast<int32_t>(), cusumLocal_.template ReinterpretCast<int32_t>(), tileCusumLocal, topkV2::RADIX_SORT_BIN_NUM);
256	}	312	}
257	event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));	313	event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
258	SetFlag<HardEvent::V_MTE3>(eventId);	314	SetFlag<HardEvent::V_MTE3>(eventId);
259	WaitFlag<HardEvent::V_MTE3>(eventId);	315	WaitFlag<HardEvent::V_MTE3>(eventId);
260	DataCopyExtParams dataCopyParam{static_cast<uint16_t>(1),	316	DataCopyExtParams dataCopyParam{static_cast<uint16_t>(1),
261	- static_cast<uint32_t>(RADIX_SORT_BIN_NUM * sizeof(T_INDEX)), 0, 0, 0};	317	+ static_cast<uint32_t>(topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX)), 0, 0, 0};
262	if (IsSameType<T_INDEX, int64_t>::value) {	318	if (IsSameType<T_INDEX, int64_t>::value) {
263	- DataCopyPad(tilesCusumGm_[cusumGmOffset + tileId * RADIX_SORT_BIN_NUM], tileCusumInt64Tmp, dataCopyParam);	319	+ DataCopyPad(tilesCusumGm_[cusumGmOffset + tileId * topkV2::RADIX_SORT_BIN_NUM], tileCusumInt64Tmp, dataCopyParam);
264	} else {	320	} else {
265	- DataCopyPad(tilesCusumGm_[cusumGmOffset + tileId * RADIX_SORT_BIN_NUM], tileCusumLocal.template ReinterpretCast<T_INDEX>(), dataCopyParam);	321	+ DataCopyPad(tilesCusumGm_[cusumGmOffset + tileId * topkV2::RADIX_SORT_BIN_NUM], tileCusumLocal.template ReinterpretCast<T_INDEX>(), dataCopyParam);
266	}	322	}
267	event_t eventIdV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));	323	event_t eventIdV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
268	SetFlag<HardEvent::MTE3_V>(eventIdV);	324	SetFlag<HardEvent::MTE3_V>(eventIdV);
	@@ -272,8 +328,8 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
272	// 计算该高8位是否满足TopK的要求，如果满足计算出来TopK的边界值	328	// 计算该高8位是否满足TopK的要求，如果满足计算出来TopK的边界值
273	PipeBarrier<PIPE_ALL>();	329	PipeBarrier<PIPE_ALL>();
274	FindBoundary(boundaryBin, boundaryBinPrev, boundaryBinCuSum, boundaryBinPrevCuSum);	330	FindBoundary(boundaryBin, boundaryBinPrev, boundaryBinCuSum, boundaryBinPrevCuSum);
275	- involvedDataMask += static_cast<UNSIGNED_TYPE>(boundaryBin) << (round * SHIFT_BIT_NUM);	331	+ involvedDataMask += static_cast<UNSIGNED_TYPE>(boundaryBin) << (round * topkV2::SHIFT_BIT_NUM);
276	- andDataMask += static_cast<UNSIGNED_TYPE>(0xFF) << (round * SHIFT_BIT_NUM);	332	+ andDataMask += static_cast<UNSIGNED_TYPE>(0xFF) << (round * topkV2::SHIFT_BIT_NUM);
277	updatedK_ -= boundaryBinPrevCuSum;	333	updatedK_ -= boundaryBinPrevCuSum;
278	PipeBarrier<PIPE_ALL>();	334	PipeBarrier<PIPE_ALL>();
279		335
	@@ -282,11 +338,11 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
282	if (boundaryBinPrev >= 0) {	338	if (boundaryBinPrev >= 0) {
283	DataCopyPadExtParams<T_INDEX> padParams{true, 0, 0, static_cast<T_INDEX>(0)};	339	DataCopyPadExtParams<T_INDEX> padParams{true, 0, 0, static_cast<T_INDEX>(0)};
284	DataCopyExtParams dataCopyParam{static_cast<uint16_t>(1),	340	DataCopyExtParams dataCopyParam{static_cast<uint16_t>(1),
285	- static_cast<uint32_t>(RADIX_SORT_BIN_NUM * sizeof(T_INDEX)), 0, 0, 0};	341	+ static_cast<uint32_t>(topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX)), 0, 0, 0};
286	if (IsSameType<T_INDEX, int64_t>::value) {	342	if (IsSameType<T_INDEX, int64_t>::value) {
287	- DataCopyPad(tileCusumInt64Tmp, tilesCusumGm_[cusumGmOffset + tileId * RADIX_SORT_BIN_NUM], dataCopyParam, padParams);	343	+ DataCopyPad(tileCusumInt64Tmp, tilesCusumGm_[cusumGmOffset + tileId * topkV2::RADIX_SORT_BIN_NUM], dataCopyParam, padParams);
288	} else {	344	} else {
289	- DataCopyPad(tileCusumLocal.template ReinterpretCast<T_INDEX>(), tilesCusumGm_[cusumGmOffset + tileId * RADIX_SORT_BIN_NUM], dataCopyParam, padParams);	345	+ DataCopyPad(tileCusumLocal.template ReinterpretCast<T_INDEX>(), tilesCusumGm_[cusumGmOffset + tileId * topkV2::RADIX_SORT_BIN_NUM], dataCopyParam, padParams);
290	}	346	}
291	event_t eventIdScalar = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));	347	event_t eventIdScalar = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
292	SetFlag<HardEvent::MTE2_S>(eventIdScalar);	348	SetFlag<HardEvent::MTE2_S>(eventIdScalar);
	@@ -310,9 +366,9 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
310	// 调用高阶API计算TopTileK，并搬运到输出Gm上	366	// 调用高阶API计算TopTileK，并搬运到输出Gm上
311	CalTileTopK2CopyOut(loopTime);	367	CalTileTopK2CopyOut(loopTime);
312		368
313	- // 如果需要排序，则在核内进行排序	369	+ // 如果需要，则在核内进行排序
314	- if (IS_SORT && this->k_ <= SUPPORT_SORT_MAX_SIZE && this->k_ * sizeof(T) <= SUPPORT_SORT_MAX_BYTE_SIZE) {	370	+ if (IS_SORT && !needSortWithIndex_) {
315	- SortTopKRes(loopTime);	371	+ SortTopKRes(loopTime);
316	}	372	}
317	}	373	}
318		374
	@@ -324,12 +380,12 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
324	return;	380	return;
325	}	381	}
326		382
327	- uint32_t cusumGmOffset = this->blockIndex_ * RADIX_SORT_BIN_NUM * tileCount_;	383	+ uint32_t cusumGmOffset = this->blockIndex_ * topkV2::RADIX_SORT_BIN_NUM * tileCount_;
328	for(uint32_t tileId = 0; tileId < tileCount_; tileId++) {	384	for(uint32_t tileId = 0; tileId < tileCount_; tileId++) {
329	DataCopyPadExtParams<T_INDEX> padParams{true, 0, 0, static_cast<T_INDEX>(0)};	385	DataCopyPadExtParams<T_INDEX> padParams{true, 0, 0, static_cast<T_INDEX>(0)};
330	DataCopyExtParams dataCopyParam{static_cast<uint16_t>(1),	386	DataCopyExtParams dataCopyParam{static_cast<uint16_t>(1),
331	- static_cast<uint32_t>(RADIX_SORT_BIN_NUM * sizeof(T_INDEX)), 0, 0, 0};	387	+ static_cast<uint32_t>(topkV2::RADIX_SORT_BIN_NUM * sizeof(T_INDEX)), 0, 0, 0};
332	- DataCopyPad(tileCusumLocal, tilesCusumGm_[cusumGmOffset + tileId * RADIX_SORT_BIN_NUM], dataCopyParam, padParams);	388	+ DataCopyPad(tileCusumLocal, tilesCusumGm_[cusumGmOffset + tileId * topkV2::RADIX_SORT_BIN_NUM], dataCopyParam, padParams);
333	event_t eventIdScalar = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));	389	event_t eventIdScalar = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
334	SetFlag<HardEvent::MTE2_S>(eventIdScalar);	390	SetFlag<HardEvent::MTE2_S>(eventIdScalar);
335	WaitFlag<HardEvent::MTE2_S>(eventIdScalar);	391	WaitFlag<HardEvent::MTE2_S>(eventIdScalar);
	@@ -449,14 +505,22 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
449	AscendC::LocalTensor<T> valuesLocal = this->valuesQue_.template DeQue<T>();	505	AscendC::LocalTensor<T> valuesLocal = this->valuesQue_.template DeQue<T>();
450	AscendC::DataCopyExtParams dataCopyParamValue{static_cast<uint16_t>(1),	506	AscendC::DataCopyExtParams dataCopyParamValue{static_cast<uint16_t>(1),
451	static_cast<uint32_t>(k * sizeof(T)), 0, 0, 0};	507	static_cast<uint32_t>(k * sizeof(T)), 0, 0, 0};
452	- AscendC::DataCopyPad(this->valuesGm_[offset], valuesLocal, dataCopyParamValue);	508	+ if (needSortWithIndex_) {
		509	+ AscendC::DataCopyPad(topkValuesGm_[offset], valuesLocal, dataCopyParamValue);
		510	+ } else {
		511	+ AscendC::DataCopyPad(this->valuesGm_[offset], valuesLocal, dataCopyParamValue);
		512	+ }
453	this->valuesQue_.template FreeTensor(valuesLocal);	513	this->valuesQue_.template FreeTensor(valuesLocal);
454		514
455	// copy sorted value index	515	// copy sorted value index
456	AscendC::LocalTensor<T_INDEX_TO> indicesLocal = this->indicesQue_.template DeQue<T_INDEX_TO>();	516	AscendC::LocalTensor<T_INDEX_TO> indicesLocal = this->indicesQue_.template DeQue<T_INDEX_TO>();
457	AscendC::DataCopyExtParams dataCopyParamIndex{static_cast<uint16_t>(1),	517	AscendC::DataCopyExtParams dataCopyParamIndex{static_cast<uint16_t>(1),
458	static_cast<uint32_t>(k * sizeof(T_INDEX_TO)), 0, 0, 0};	518	static_cast<uint32_t>(k * sizeof(T_INDEX_TO)), 0, 0, 0};
459	- AscendC::DataCopyPad(this->indicesGm_[offset], indicesLocal, dataCopyParamIndex);	519	+ if (needSortWithIndex_) {
		520	+ AscendC::DataCopyPad(topkIndicesGm_[offset], indicesLocal, dataCopyParamIndex);
		521	+ } else {
		522	+ AscendC::DataCopyPad(this->indicesGm_[offset], indicesLocal, dataCopyParamIndex);
		523	+ }
460	this->indicesQue_.template FreeTensor(indicesLocal);	524	this->indicesQue_.template FreeTensor(indicesLocal);
461	}	525	}
462		526
	@@ -534,21 +598,21 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
534	int32_t round,	598	int32_t round,
535	uint32_t numTileData)	599	uint32_t numTileData)
536	{	600	{
537	- if constexpr (is_same<int64_t, T>::value \|\| is_same<uint64_t, T>::value) {	601	+ if constexpr (topkV2::is_same<int64_t, T>::value \|\| topkV2::is_same<uint64_t, T>::value) {
538	RadixSortTopKB64<T, uint64_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	602	RadixSortTopKB64<T, uint64_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
539	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	603	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
540	involveDataMask, round, numTileData);	604	involveDataMask, round, numTileData);
541	- } else if constexpr (is_same<int32_t, T>::value \|\| is_same<uint32_t, T>::value	605	+ } else if constexpr (topkV2::is_same<int32_t, T>::value \|\| topkV2::is_same<uint32_t, T>::value
542	- \|\| is_same<float, T>::value) {	606	+ \|\| topkV2::is_same<float, T>::value) {
543	RadixSortTopKB32<T, uint32_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	607	RadixSortTopKB32<T, uint32_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
544	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	608	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
545	involveDataMask, round, numTileData);	609	involveDataMask, round, numTileData);
546	- } else if constexpr (is_same<half, T>::value \|\| is_same<uint16_t, T>::value	610	+ } else if constexpr (topkV2::is_same<half, T>::value \|\| topkV2::is_same<uint16_t, T>::value
547	- \|\| is_same<int16_t, T>::value \|\| is_same<bfloat16_t, T>::value) {	611	+ \|\| topkV2::is_same<int16_t, T>::value \|\| topkV2::is_same<bfloat16_t, T>::value) {
548	RadixSortTopKB16<T, uint16_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	612	RadixSortTopKB16<T, uint16_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
549	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	613	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
550	involveDataMask, round, numTileData);	614	involveDataMask, round, numTileData);
551	- } else if constexpr (is_same<int8_t, T>::value \|\| is_same<uint8_t, T>::value) {	615	+ } else if constexpr (topkV2::is_same<int8_t, T>::value \|\| topkV2::is_same<uint8_t, T>::value) {
552	RadixSortTopKB8<T, uint8_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;	616	RadixSortTopKB8<T, uint8_t, NUM_PASS, IS_LARGEST, int32_t> radixSortTopK;
553	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,	617	radixSortTopK.GetCumSum(inputX, cumSumHist, andDataMask,
554	involveDataMask, round, numTileData);	618	involveDataMask, round, numTileData);
	@@ -560,27 +624,27 @@ __aicore__ inline LocalTensor<UNSIGNED_TYPE> RadixSortTopKSingleCore<T, UNSIGNED
560	LocalTensor<T> inputX,	624	LocalTensor<T> inputX,
561	uint32_t numTileData)	625	uint32_t numTileData)
562	{	626	{
563	- if constexpr (is_same<int64_t, T>::value) {	627	+ if constexpr (topkV2::is_same<int64_t, T>::value) {
564	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	628	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
565	radixSortTopK.TwiddleInB64(inputX, unsignedInputXLocal_, numTileData);	629	radixSortTopK.TwiddleInB64(inputX, unsignedInputXLocal_, numTileData);
566	return unsignedInputXLocal_;	630	return unsignedInputXLocal_;
567	- } else if constexpr (is_same<int32_t, T>::value) {	631	+ } else if constexpr (topkV2::is_same<int32_t, T>::value) {
568	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	632	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
569	radixSortTopK.TwiddleInB32(inputX, unsignedInputXLocal_, numTileData);	633	radixSortTopK.TwiddleInB32(inputX, unsignedInputXLocal_, numTileData);
570	return unsignedInputXLocal_;	634	return unsignedInputXLocal_;
571	- } else if constexpr (is_same<half, T>::value \|\| is_same<bfloat16_t, T>::value) {	635	+ } else if constexpr (topkV2::is_same<half, T>::value \|\| topkV2::is_same<bfloat16_t, T>::value) {
572	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	636	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
573	radixSortTopK.TwiddleInFp16(inputX, unsignedInputXLocal_, numTileData);	637	radixSortTopK.TwiddleInFp16(inputX, unsignedInputXLocal_, numTileData);
574	return unsignedInputXLocal_;	638	return unsignedInputXLocal_;
575	- } else if constexpr (is_same<float, T>::value) {	639	+ } else if constexpr (topkV2::is_same<float, T>::value) {
576	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	640	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
577	radixSortTopK.TwiddleInFp32(inputX, unsignedInputXLocal_, numTileData);	641	radixSortTopK.TwiddleInFp32(inputX, unsignedInputXLocal_, numTileData);
578	return unsignedInputXLocal_;	642	return unsignedInputXLocal_;
579	- } else if constexpr (is_same<int16_t, T>::value) {	643	+ } else if constexpr (topkV2::is_same<int16_t, T>::value) {
580	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	644	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
581	radixSortTopK.TwiddleInB16(inputX, unsignedInputXLocal_, numTileData);	645	radixSortTopK.TwiddleInB16(inputX, unsignedInputXLocal_, numTileData);
582	return unsignedInputXLocal_;	646	return unsignedInputXLocal_;
583	- } else if constexpr (is_same<int8_t, T>::value) {	647	+ } else if constexpr (topkV2::is_same<int8_t, T>::value) {
584	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;	648	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixSortTopK;
585	radixSortTopK.TwiddleInB8(inputX, unsignedInputXLocal_, numTileData);	649	radixSortTopK.TwiddleInB8(inputX, unsignedInputXLocal_, numTileData);
586	return unsignedInputXLocal_;	650	return unsignedInputXLocal_;
	@@ -599,16 +663,16 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
599	LocalTensor<UNSIGNED_TYPE> reverseInputX,	663	LocalTensor<UNSIGNED_TYPE> reverseInputX,
600	uint32_t numTileData)	664	uint32_t numTileData)
601	{	665	{
602	- if constexpr (is_same<uint64_t, T>::value) {	666	+ if constexpr (topkV2::is_same<uint64_t, T>::value) {
603	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	667	RadixBlockSortSimdB64<T, uint64_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
604	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	668	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
605	- } else if constexpr (is_same<uint32_t, T>::value) {	669	+ } else if constexpr (topkV2::is_same<uint32_t, T>::value) {
606	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	670	RadixBlockSortSimdB32<T, uint32_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
607	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	671	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
608	- } else if constexpr (is_same<uint16_t, T>::value) {	672	+ } else if constexpr (topkV2::is_same<uint16_t, T>::value) {
609	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	673	RadixBlockSortSimdB16<T, uint16_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
610	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	674	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
611	- } else if constexpr (is_same<uint8_t, T>::value) {	675	+ } else if constexpr (topkV2::is_same<uint8_t, T>::value) {
612	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;	676	RadixBlockSortSimdB8<T, uint8_t, NUM_PASS, IS_LARGEST, T_INDEX> radixBlockSort;
613	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);	677	radixBlockSort.ReverseInputData(inputX, reverseInputX, numTileData);
614	}	678	}
	@@ -621,9 +685,9 @@ __aicore__ inline void RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LA
621	T_INDEX& boundaryBinCuSum,	685	T_INDEX& boundaryBinCuSum,
622	T_INDEX& boundaryBinPrevCuSum)	686	T_INDEX& boundaryBinPrevCuSum)
623	{	687	{
624	- if (cusumLocal_(RADIX_SORT_BIN_NUM - 1) <= updatedK_) {	688	+ if (cusumLocal_(topkV2::RADIX_SORT_BIN_NUM - 1) <= updatedK_) {
625	boundaryBin = -1;	689	boundaryBin = -1;
626	- boundaryBinPrev = RADIX_SORT_BIN_NUM - 1;	690	+ boundaryBinPrev = topkV2::RADIX_SORT_BIN_NUM - 1;
627	boundaryBinCuSum = -1;	691	boundaryBinCuSum = -1;
628	boundaryBinPrevCuSum = cusumLocal_(boundaryBinPrev);	692	boundaryBinPrevCuSum = cusumLocal_(boundaryBinPrev);
629	return ;	693	return ;
	@@ -646,16 +710,16 @@ template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_LARGEST,
646	__aicore__ inline int32_t RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_SORT, T_INDEX, T_INDEX_TO>::BinarySearch()	710	__aicore__ inline int32_t RadixSortTopKSingleCore<T, UNSIGNED_TYPE, NUM_PASS, IS_LARGEST, IS_SORT, T_INDEX, T_INDEX_TO>::BinarySearch()
647	{	711	{
648	int32_t left = 0;	712	int32_t left = 0;
649	- int32_t right = RADIX_SORT_BIN_NUM - 1;	713	+ int32_t right = topkV2::RADIX_SORT_BIN_NUM - 1;
650	while(left <= right) {	714	while(left <= right) {
651	int mid = (right + left) / 2;	715	int mid = (right + left) / 2;
652	if (cusumLocal_(mid) == updatedK_) {	716	if (cusumLocal_(mid) == updatedK_) {
653	- if ((mid + 1) < RADIX_SORT_BIN_NUM && cusumLocal_(mid + 1) > updatedK_) {	717	+ if ((mid + 1) < topkV2::RADIX_SORT_BIN_NUM && cusumLocal_(mid + 1) > updatedK_) {
654	return mid;	718	return mid;
655	} else {	719	} else {
656	left = mid + 1;	720	left = mid + 1;
657	}	721	}
658	- } else if ((mid + 1) < RADIX_SORT_BIN_NUM && cusumLocal_(mid + 1) > updatedK_ && cusumLocal_(mid) < updatedK_) {	722	+ } else if ((mid + 1) < topkV2::RADIX_SORT_BIN_NUM && cusumLocal_(mid + 1) > updatedK_ && cusumLocal_(mid) < updatedK_) {
659	return mid;	723	return mid;
660	} else if (cusumLocal_(mid) < updatedK_) {	724	} else if (cusumLocal_(mid) < updatedK_) {
661	left = mid + 1;	725	left = mid + 1;

Mmath/top_k_v2/op_kernel/arch35/radix_sort_topk_b16.h+1-0

	@@ -16,6 +16,7 @@
16	#include "kernel_operator.h"	16	#include "kernel_operator.h"
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
18	using namespace AscendC;	18	using namespace AscendC;
		19	+using namespace topkV2;
19	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	20	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
20	struct RadixSortTopKB16{	21	struct RadixSortTopKB16{
21	__aicore__ inline RadixSortTopKB16() {}	22	__aicore__ inline RadixSortTopKB16() {}

Mmath/top_k_v2/op_kernel/arch35/radix_sort_topk_b32.h+2-0

	@@ -15,6 +15,8 @@
15	#define RADIX_SORT_TOPK_B32_H	15	#define RADIX_SORT_TOPK_B32_H
16	#include "kernel_operator.h"	16	#include "kernel_operator.h"
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
		18	+
		19	+using namespace topkV2;
18	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	20	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
19	struct RadixSortTopKB32{	21	struct RadixSortTopKB32{
20	__aicore__ inline RadixSortTopKB32() {}	22	__aicore__ inline RadixSortTopKB32() {}

Mmath/top_k_v2/op_kernel/arch35/radix_sort_topk_b64.h+2-0

	@@ -15,6 +15,8 @@
15	#define RADIX_SORT_TOPK_B64_H	15	#define RADIX_SORT_TOPK_B64_H
16	#include "kernel_operator.h"	16	#include "kernel_operator.h"
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
		18	+
		19	+using namespace topkV2;
18	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	20	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
19	struct RadixSortTopKB64{	21	struct RadixSortTopKB64{
20	__aicore__ inline RadixSortTopKB64() {}	22	__aicore__ inline RadixSortTopKB64() {}

Mmath/top_k_v2/op_kernel/arch35/radix_sort_topk_b8.h+1-0

	@@ -16,6 +16,7 @@
16	#include "kernel_operator.h"	16	#include "kernel_operator.h"
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
18	using namespace AscendC;	18	using namespace AscendC;
		19	+using namespace topkV2;
19	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	20	template <typename T, typename UNSIGNED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
20	struct RadixSortTopKB8{	21	struct RadixSortTopKB8{
21	__aicore__ inline RadixSortTopKB8() {}	22	__aicore__ inline RadixSortTopKB8() {}

Amath/top_k_v2/op_kernel/arch35/sort_and_top_k_more_core.h+321-0

@@ -0,0 +1,321 @@
	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 sort_and_top_k_more_core.h
	13	+* \brief sort_and_top_k morecore mode impl
	14	+*/
	15	+
	16	+#ifndef SORT_AND_TOP_K_MORE_CORE_H
	17	+#define SORT_AND_TOP_K_MORE_CORE_H
	18	+
	19	+#include <cmath>
	20	+#include "kernel_operator.h"
	21	+#include "../../sort/arch35/sort_radix_sort_more_core.h"
	22	+#include "../../sort/arch35/sort_tiling_data.h" // sort_radix_sort_more_core.h 里面引用了 sort_tiling_data.h
	23	+#include "../../sort/arch35/util_type_simd.h" // 使用 ROUND_UP_AGLIN , DoubleBufferSimd
	24	+
	25	+namespace SortAndTopK {
	26	+const uint64_t AGLIN_FACTOR = 32;
	27	+
	28	+using namespace AscendC;
	29	+
	30	+// T1-输入x dtype,T2-输出Idx dtype, UT-无符号的数据类型
	31	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	32	+class SortAndTopKMoreCore : public Sort::SortRadixMoreCore<T1, T2, UT, T3, isDescend> {
	33	+public:
	34	+ __aicore__ inline SortAndTopKMoreCore(){};
	35	+ __aicore__ inline void InitParam(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,
	36	+ const TopKV2TilingDataSimd* tilingData, TPipe *pipe);
	37	+ __aicore__ inline void ProcessTopK();
	38	+private:
	39	+ GlobalTensor<T1> sortOutValueGM_;
	40	+ GlobalTensor<uint32_t> sortOutIdxGM_;
	41	+
	42	+ TQueBind<QuePosition::VECIN, QuePosition::VECOUT, 1> topkResQueue_;
	43	+
	44	+ const TopKV2TilingDataSimd *tilingData_;
	45	+
	46	+ int64_t topKRealValue_ = 0;
	47	+ uint32_t tileDataSize_ = 0;
	48	+ uint32_t blockTileNum_ = 0;
	49	+ uint32_t tailTileNum_ = 0;
	50	+
	51	+ __aicore__ inline void ProcessSortAndTopK(GlobalTensor<T1> inputXGm, int64_t gmOffset, uint32_t sortLoopRound);
	52	+ __aicore__ inline void ParserTilingData();
	53	+ __aicore__ inline void GetValueResData(uint32_t dataCopyLoopTimes, bool hasLastTile, int64_t outGmOffset,
	54	+ int64_t sortOutGmOffset);
	55	+ __aicore__ inline void GetIndexResData(uint32_t dataCopyLoopTimes, bool hasLastTile, int64_t outGmOffset,
	56	+ int64_t sortOutGmOffset);
	57	+ __aicore__ inline void GetTopKRes(uint32_t sortLoopRound);
	58	+};
	59	+
	60	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	61	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::InitParam(GM_ADDR x, GM_ADDR value,
	62	+ GM_ADDR sortIndex, GM_ADDR workspace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe)
	63	+{
	64	+ this->blockIdx_ = GetBlockIdx();
	65	+ this->pipe_ = pipe;
	66	+ tilingData_ = tilingData;
	67	+ ParserTilingData();
	68	+ this->realCoreNum_ = GetBlockNum();
	69	+ if constexpr (sizeof(T3) == sizeof(int64_t)) {
	70	+ this->factor_ = Sort::CONST_2;
	71	+ }
	72	+
	73	+ this->inputXGm_.SetGlobalBuffer((__gm__ T1 *)x);
	74	+ this->outValueGm_.SetGlobalBuffer((__gm__ T1 *)value);
	75	+ this->outIdxGm_.SetGlobalBuffer((__gm__ uint32_t *)sortIndex);
	76	+ uint64_t wkOffset = this->clearCoreSize0_ * this->clearCore0_;
	77	+ uint64_t oneBlockNumB32 = this->oneBlock_ / sizeof(int32_t);
	78	+ if constexpr (sizeof(T3) == sizeof(int64_t)) {
	79	+ wkOffset = wkOffset * Sort::CONST_2;
	80	+ }
	81	+ wkOffset = this->CeilDivMul(wkOffset, oneBlockNumB32);
	82	+ this->excusiveBinsGmWk_.SetGlobalBuffer((__gm__ uint32_t *)workspace, wkOffset);
	83	+ wkOffset = wkOffset * sizeof(uint32_t);
	84	+
	85	+ uint64_t histOffset = this->clearCout_ * this->clearSize_ * this->clearCore1_;
	86	+ if constexpr (sizeof(T3) == sizeof(int64_t)) {
	87	+ histOffset = histOffset * Sort::CONST_2;
	88	+ }
	89	+ histOffset = this->CeilDivMul(histOffset, oneBlockNumB32);
	90	+ this->globalHistGmWk_.SetGlobalBuffer((__gm__ uint32_t *)(workspace + wkOffset), histOffset);
	91	+ wkOffset = wkOffset + histOffset * sizeof(uint32_t);
	92	+
	93	+ uint64_t dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
	94	+ if constexpr (sizeof(T3) == sizeof(int64_t)) {
	95	+ dbOffset = dbOffset * Sort::CONST_2;
	96	+ }
	97	+ dbOffset = this->CeilDivMul(dbOffset, oneBlockNumB32);
	98	+ this->outIdxDbWK_.SetGlobalBuffer((__gm__ uint32_t *)(workspace + wkOffset), dbOffset);
	99	+ wkOffset = wkOffset + dbOffset * sizeof(uint32_t);
	100	+
	101	+ dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
	102	+ if constexpr (sizeof(T2) == sizeof(int64_t)) {
	103	+ dbOffset = dbOffset * Sort::CONST_2;
	104	+ }
	105	+ dbOffset = this->CeilDivMul(dbOffset, oneBlockNumB32);
	106	+ sortOutIdxGM_.SetGlobalBuffer((__gm__ uint32_t *)(workspace + wkOffset), dbOffset);
	107	+ wkOffset = wkOffset + dbOffset * sizeof(uint32_t);
	108	+
	109	+ uint64_t histTileOffset = this->lastDimTileNum_ * Sort::RADIX_SORT_NUM * this->unsortedDimParallel_;
	110	+ this->histTileGmWk_.SetGlobalBuffer((__gm__ uint16_t *)(workspace + wkOffset), histTileOffset);
	111	+ wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);
	112	+ this->histCumsumTileGmWk_.SetGlobalBuffer((__gm__ uint16_t *)(workspace + wkOffset), histTileOffset);
	113	+ wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);
	114	+
	115	+ uint64_t xB8Offset = this->lastDimTileNum_ * this->numTileData_ * this->unsortedDimParallel_;
	116	+ xB8Offset = this->CeilDivMul(xB8Offset, this->oneBlock_);
	117	+ this->xB8GmWk_.SetGlobalBuffer((__gm__ uint8_t *)(workspace + wkOffset), xB8Offset);
	118	+ wkOffset = wkOffset + xB8Offset * sizeof(uint8_t);
	119	+
	120	+ dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
	121	+ dbOffset = this->CeilDivMul(dbOffset * sizeof(T1), this->oneBlock_) / sizeof(T1);
	122	+ this->outValueDbWK_.SetGlobalBuffer((__gm__ T1 *)(workspace + wkOffset), dbOffset);
	123	+ wkOffset = wkOffset + dbOffset * sizeof(T1);
	124	+ sortOutValueGM_.SetGlobalBuffer((__gm__ T1 *)(workspace + wkOffset), dbOffset);
	125	+
	126	+ this->pipe_->InitBuffer(this->inQueueX_, 1, this->numTileData_ * sizeof(T1));
	127	+ this->pipe_->InitBuffer(this->inQueueIndex_, 1, this->numTileData_ * sizeof(T3));
	128	+ this->pipe_->InitBuffer(this->inQueueGlobalHist_, 1, Sort::RADIX_SORT_NUM * sizeof(T3));
	129	+ this->pipe_->InitBuffer(this->outValueQueue_, 1, this->numTileData_);
	130	+ this->pipe_->InitBuffer(this->blockExcusiveInQue_, 1, Sort::RADIX_SORT_NUM * sizeof(uint16_t));
	131	+ this->pipe_->InitBuffer(this->blockHistInQue_, 1, Sort::RADIX_SORT_NUM * sizeof(uint16_t));
	132	+ this->pipe_->InitBuffer(this->blockUbFlagQue_, 1, Sort::RADIX_SORT_NUM * sizeof(T3));
	133	+ this->pipe_->InitBuffer(this->inputB8Que_, 1, this->numTileData_);
	134	+ this->pipe_->InitBuffer(this->outIdxQueue_, 1, this->numTileData_ * sizeof(uint32_t));
	135	+ this->pipe_->InitBuffer(this->tmpUb_, this->tmpUbSize_);
	136	+ this->pipe_->InitBuffer(this->blockHistFlagUbQue_, 1, Sort::RADIX_SORT_NUM * sizeof(T3));
	137	+
	138	+ this->globalHistGmWkTmp_ = this->globalHistGmWk_.template ReinterpretCast<T3>();
	139	+ this->excusiveBinsGmWkTmp_ = this->excusiveBinsGmWk_.template ReinterpretCast<T3>();
	140	+}
	141	+
	142	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	143	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::ParserTilingData()
	144	+{
	145	+ this->totalDataNum_ = tilingData_->lastAxisNum; // h轴大小
	146	+ this->numTileData_ = tilingData_->numTileDataSize; // ub循环块大小
	147	+ this->unsortedDimNum_ = tilingData_->unsortedDimNum; // b轴大小
	148	+ this->unsortedDimParallel_ = tilingData_->unsortedDimParallel; // b轴使用的核数
	149	+ this->lastDimTileNum_ = tilingData_->lastDimTileNum; // h轴循环次数
	150	+ this->sortLoopTimes_ = tilingData_->sortLoopTimes; // b轴循环次数
	151	+ this->lastDimRealCore_ = tilingData_->lastDimNeedCore; // h轴需要的核数
	152	+ this->tmpUbSize_ = tilingData_->tmpUbSize; // 高级api需要用的ub大小
	153	+ topKRealValue_ = tilingData_->topKRealValue; // k值
	154	+ tileDataSize_ = tilingData_->sortAndTopkTileDataSize; // 获取topk时tile大小
	155	+ blockTileNum_ = tilingData_->sortAndTopkBlockTileNum; // 获取topk时单核分配tile数量
	156	+ tailTileNum_ = tilingData_->sortAndTopkTailTileNum; // 获取topk时尾部tile数量
	157	+
	158	+ this->clearCore1_ = tilingData_->keyParams0; // 用于清零的globalHistGmWk_的核
	159	+ this->clearCore0_ = tilingData_->keyParams1; // 用于清零excusiveBinsGmWk_的核
	160	+ this->clearSize_ = tilingData_->keyParams2; // 每次清零的ub大小，按照大的globalHistGmWk_所需ub算
	161	+ this->clearCout_ = tilingData_->keyParams3; // 清零globalHistGmWk_ ub循环次数
	162	+ this->clearCoreSize0_ = tilingData_->keyParams4; // 清零excusiveBinsGmWk_,每个核处理多少个数
	163	+ this->clearCoreSize1_ = tilingData_->keyParams5; // 清零globalHistGmWk_，每个核处理多少
	164	+}
	165	+
	166	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	167	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::GetValueResData(
	168	+ uint32_t dataCopyLoopTimes, bool hasLastTile, int64_t outGmOffset, int64_t sortOutGmOffset)
	169	+{
	170	+ if (dataCopyLoopTimes == 0)
	171	+ {
	172	+ return;
	173	+ }
	174	+ int64_t sortOutValueGMOffset;
	175	+ int64_t outValueGmOffset;
	176	+ uint32_t tileDataSize = tileDataSize_;
	177	+ LocalTensor<T1> topkResTensor = topkResQueue_.AllocTensor<T1>();
	178	+ for (uint32_t i = 0; i < dataCopyLoopTimes; i++) {
	179	+ if (hasLastTile && i == dataCopyLoopTimes - 1) {
	180	+ tileDataSize = topKRealValue_ % tileDataSize_;
	181	+ }
	182	+ sortOutValueGMOffset = sortOutGmOffset + i * tileDataSize_;
	183	+ outValueGmOffset = outGmOffset + i * tileDataSize_;
	184	+ DataCopyPadExtParams<T1> padParams{ false, 0, 0, 0 };
	185	+ DataCopyExtParams dataCopyParam;
	186	+ dataCopyParam.blockCount = 1;
	187	+ dataCopyParam.blockLen = tileDataSize * sizeof(T1);
	188	+ dataCopyParam.srcStride = 0;
	189	+ dataCopyParam.dstStride = 0;
	190	+ DataCopyPad(topkResTensor, sortOutValueGM_[sortOutValueGMOffset], dataCopyParam, padParams);
	191	+ topkResQueue_.EnQue(topkResTensor);
	192	+ topkResTensor = topkResQueue_.DeQue<T1>();
	193	+ DataCopyPad(this->outValueGm_[outValueGmOffset], topkResTensor, dataCopyParam);
	194	+ }
	195	+ topkResQueue_.FreeTensor(topkResTensor);
	196	+}
	197	+
	198	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	199	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::GetIndexResData(
	200	+ uint32_t dataCopyLoopTimes, bool hasLastTile, int64_t outGmOffset, int64_t sortOutGmOffset)
	201	+{
	202	+ if (dataCopyLoopTimes == 0)
	203	+ {
	204	+ return;
	205	+ }
	206	+ GlobalTensor<T2> sortOutIdxTmpGM_ = sortOutIdxGM_.template ReinterpretCast<T2>();
	207	+ GlobalTensor<T2> outIdxTmpGm_ = this->outIdxGm_.template ReinterpretCast<T2>();
	208	+ int64_t sortOutIdxGMOffset;
	209	+ int64_t outIdxTmpGmOffset;
	210	+ uint32_t tileDataSize = tileDataSize_;
	211	+ LocalTensor<T2> topkResTensor = topkResQueue_.AllocTensor<T2>();
	212	+ for (uint32_t i = 0; i < dataCopyLoopTimes; i++) {
	213	+ if (hasLastTile && i == dataCopyLoopTimes - 1) {
	214	+ tileDataSize = topKRealValue_ % tileDataSize_;
	215	+ }
	216	+ sortOutIdxGMOffset = sortOutGmOffset + i * tileDataSize_;
	217	+ outIdxTmpGmOffset = outGmOffset + i * tileDataSize_;
	218	+ DataCopyPadExtParams<T2> padParams{ false, 0, 0, 0 };
	219	+ DataCopyExtParams dataCopyParam;
	220	+ dataCopyParam.blockCount = 1;
	221	+ dataCopyParam.blockLen = tileDataSize * sizeof(T2);
	222	+ dataCopyParam.srcStride = 0;
	223	+ dataCopyParam.dstStride = 0;
	224	+ DataCopyPad(topkResTensor, sortOutIdxTmpGM_[sortOutIdxGMOffset], dataCopyParam, padParams);
	225	+ topkResQueue_.EnQue(topkResTensor);
	226	+ topkResTensor = topkResQueue_.DeQue<T2>();
	227	+ DataCopyPad(outIdxTmpGm_[outIdxTmpGmOffset], topkResTensor, dataCopyParam);
	228	+ }
	229	+ topkResQueue_.FreeTensor(topkResTensor);
	230	+}
	231	+
	232	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	233	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::GetTopKRes(uint32_t sortLoopRound)
	234	+{
	235	+ // 重置TPipe,申请新的UB空间
	236	+ this->pipe_->Reset();
	237	+ uint32_t maxDTypeSize = sizeof(T1) >= sizeof(T2) ? sizeof(T1) : sizeof(T2);
	238	+ this->pipe_->InitBuffer(topkResQueue_, 1, maxDTypeSize * tileDataSize_);
	239	+
	240	+ // 当前block处理的tile数
	241	+ uint32_t dataCopyLoopTimes =
	242	+ (tailTileNum_ > 0 && this->blockIdx_ < tailTileNum_) ? blockTileNum_ + 1 : blockTileNum_;
	243	+ // 判断当前block是否有尾块需要处理
	244	+ bool hasLastTile = false;
	245	+ if (blockTileNum_ > 0 && this->blockIdx_ == this->realCoreNum_ - 1) {
	246	+ hasLastTile = true;
	247	+ } else if (blockTileNum_ == 0 && this->blockIdx_ == tailTileNum_ - 1) {
	248	+ hasLastTile = true;
	249	+ }
	250	+
	251	+ // 当前block处理的tile的起始索引
	252	+ int64_t blockStartIdx = 0;
	253	+ if (this->blockIdx_ < tailTileNum_) {
	254	+ blockStartIdx = (blockTileNum_ + 1) * this->blockIdx_;
	255	+ } else {
	256	+ blockStartIdx = (blockTileNum_ + 1) * tailTileNum_ + (this->blockIdx_ - tailTileNum_) * blockTileNum_;
	257	+ }
	258	+
	259	+ // 输出GM地址偏移量
	260	+ int64_t outGmOffset = sortLoopRound * this->unsortedDimParallel_ * topKRealValue_ + blockStartIdx * tileDataSize_;
	261	+ // sortOutGM地址偏移量
	262	+ int64_t sortOutGmOffset = blockStartIdx * tileDataSize_;
	263	+ // 超大排序轴场景下，unsortedDimParallel_一般等于1
	264	+ for (uint32_t i = 0; i < this->unsortedDimParallel_; i++) {
	265	+ outGmOffset += i * topKRealValue_;
	266	+ sortOutGmOffset += i * this->totalDataNum_;
	267	+ GetValueResData(dataCopyLoopTimes, hasLastTile, outGmOffset, sortOutGmOffset);
	268	+ SyncAll();
	269	+ GetIndexResData(dataCopyLoopTimes, hasLastTile, outGmOffset, sortOutGmOffset);
	270	+ SyncAll();
	271	+ }
	272	+}
	273	+
	274	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	275	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::ProcessSortAndTopK(GlobalTensor<T1> inputXGm,
	276	+ int64_t gmOffset, uint32_t sortLoopRound)
	277	+{
	278	+ this->ClearWorkSapce();
	279	+ SyncAll();
	280	+
	281	+ if constexpr (sizeof(T2) == sizeof(uint32_t)) {
	282	+ if constexpr (sizeof(T1) == sizeof(int8_t)) {
	283	+ this->inputXDbGm_.SetDoubleBuffer(this->outValueDbWK_, sortOutValueGM_);
	284	+ this->idxDbGm_.SetDoubleBuffer(this->outIdxDbWK_, sortOutIdxGM_);
	285	+ } else {
	286	+ this->inputXDbGm_.SetDoubleBuffer(sortOutValueGM_, this->outValueDbWK_);
	287	+ this->idxDbGm_.SetDoubleBuffer(sortOutIdxGM_, this->outIdxDbWK_);
	288	+ }
	289	+ } else {
	290	+ if constexpr (sizeof(T1) == sizeof(int8_t)) {
	291	+ this->inputXDbGm_.SetDoubleBuffer(this->outValueDbWK_, sortOutValueGM_);
	292	+ this->idxDbGm_.SetDoubleBuffer(this->outIdxDbWK_, sortOutIdxGM_);
	293	+ } else {
	294	+ this->inputXDbGm_.SetDoubleBuffer(sortOutValueGM_, this->outValueDbWK_);
	295	+ this->idxDbGm_.SetDoubleBuffer(sortOutIdxGM_, this->outIdxDbWK_);
	296	+ }
	297	+ }
	298	+ for (uint32_t round = 0; round < static_cast<uint32_t>(sizeof(T1)); round++) {
	299	+ this->GetGlobalExcusiveSum(round, sortLoopRound, inputXGm);
	300	+ SyncAll();
	301	+ this->ComputeOnePass(round, sortLoopRound, inputXGm);
	302	+ SyncAll();
	303	+ }
	304	+ GetTopKRes(sortLoopRound);
	305	+}
	306	+
	307	+template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
	308	+__aicore__ inline void SortAndTopKMoreCore<T1, T2, UT, T3, isDescend>::ProcessTopK()
	309	+{
	310	+ if (this->blockIdx_ > this->realCoreNum_) {
	311	+ return;
	312	+ }
	313	+ for (uint32_t i = 0; i < this->sortLoopTimes_; i++) {
	314	+ int64_t loopOffset = i * this->unsortedDimParallel_ * this->totalDataNum_;
	315	+ ProcessSortAndTopK(this->inputXGm_[loopOffset], loopOffset, i);
	316	+ }
	317	+}
	318	+
	319	+} // namespace SortAndTopK
	320	+
	321	+#endif // SORT_AND_TOP_K_MORE_CORE_H

Amath/top_k_v2/op_kernel/arch35/sort_with_index_entry.h+136-0

@@ -0,0 +1,136 @@
	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 sort_with_index_entry.h
	13	+* \brief sort_with_index kernel entry impl
	14	+*/
	15	+
	16	+#ifndef SORT_WITH_INDEX_ENTRY_H
	17	+#define SORT_WITH_INDEX_ENTRY_H
	18	+
	19	+#include "sort_with_index_merge_sort.h"
	20	+#include "sort_with_index_multi_block.h"
	21	+#include "sort_with_index_single_block.h"
	22	+
	23	+using namespace AscendC;
	24	+
	25	+const uint32_t SORT_WITH_INDEX_SMALL_SIZE_MODE = 1;
	26	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_INT64 = 1004;
	27	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_INT32 = 1003;
	28	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_INT16 = 1002;
	29	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_INT8 = 1001;
	30	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_UINT64 = 2004;
	31	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_UINT32 = 2003;
	32	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_UINT16 = 2002;
	33	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_UINT8 = 2001;
	34	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_FLOAT = 3003;
	35	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_FLOAT16 = 3002;
	36	+const uint32_t SORT_WITH_INDEX_COMMON_TILING_KEY_BF16 = 4002;
	37	+const uint32_t SORT_WITH_INDEX_MERGE_SORT_TILING_KEY_FLOAT = 13003;
	38	+const uint32_t SORT_WITH_INDEX_MERGE_SORT_TILING_KEY_FLOAT16 = 13002;
	39	+const uint32_t SORT_WITH_INDEX_MERGE_SORT_TILING_KEY_BF16 = 14002;
	40	+
	41	+template <typename XType, typename ConvertType, typename UnsignedType, typename IndexType>
	42	+__aicore__ inline void generateOpForTopK(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,
	43	+ GM_ADDR globalWorkGm, const TopKV2TilingDataSimd* tilingData, TPipe* tPipe)
	44	+{
	45	+ bool isDescend = tilingData->isLargest > 0 ? true : false;
	46	+ bool isInt32Range = tilingData->isInInt32RangeForSort == 1 ? true : false;
	47	+ bool isSingleBlock = tilingData->lastDimNeedCoreForSort == 1 ? true : false;
	48	+ uint32_t modeType = tilingData->modeTypeForSort;
	49	+ if (isSingleBlock && SORT_WITH_INDEX_SMALL_SIZE_MODE == modeType) {
	50	+ if (isDescend) {
	51	+ SortWithIndexSingleBlock<XType, true, IndexType> radixSort;
	52	+ radixSort.Init(x, index, y, sortedIndex, globalWorkGm, tilingData, tPipe);
	53	+ radixSort.Process();
	54	+ } else {
	55	+ SortWithIndexSingleBlock<XType, false, IndexType> radixSort;
	56	+ radixSort.Init(x, index, y, sortedIndex, globalWorkGm, tilingData, tPipe);
	57	+ radixSort.Process();
	58	+ }
	59	+ return;
	60	+ }
	61	+
	62	+ if (isInt32Range) {
	63	+ if (isDescend) {
	64	+ SortWithIndexMultiBlock<XType, UnsignedType, true, uint32_t, IndexType> radixSort;
	65	+ radixSort.Init(x, index, y, sortedIndex, globalWorkGm, tilingData, tPipe);
	66	+ radixSort.Process();
	67	+ } else {
	68	+ SortWithIndexMultiBlock<XType, UnsignedType, false, uint32_t, IndexType> radixSort;
	69	+ radixSort.Init(x, index, y, sortedIndex, globalWorkGm, tilingData, tPipe);
	70	+ radixSort.Process();
	71	+ }
	72	+ } else {
	73	+ if (isDescend) {
	74	+ SortWithIndexMultiBlock<XType, UnsignedType, true, int64_t, IndexType> radixSort;
	75	+ radixSort.Init(x, index, y, sortedIndex, globalWorkGm, tilingData, tPipe);
	76	+ radixSort.Process();
	77	+ } else {
	78	+ SortWithIndexMultiBlock<XType, UnsignedType, false, int64_t, IndexType> radixSort;
	79	+ radixSort.Init(x, index, y, sortedIndex, globalWorkGm, tilingData, tPipe);
	80	+ radixSort.Process();
	81	+ }
	82	+ }
	83	+}
	84	+
	85	+// todo MergeSort适配Int64
	86	+template <typename XType, typename ConvertType, typename IndexType>
	87	+__aicore__ inline void generateMergeSortOpForTopK(GM_ADDR x, GM_ADDR index, GM_ADDR values, GM_ADDR indices,
	88	+ GM_ADDR globalWorkGm, const TopKV2TilingDataSimd* tilingData, TPipe* tPipe)
	89	+{
	90	+ bool isDescend = (tilingData->isLargest > 0) ? true : false;
	91	+ if (isDescend) {
	92	+ SortWithIndexMergeSort<XType, ConvertType, TopKV2TilingDataSimd, true, IndexType> mergeSort;
	93	+ mergeSort.Init(x, index, values, indices, globalWorkGm, tilingData, tPipe);
	94	+ mergeSort.ProcessSort();
	95	+ } else {
	96	+ SortWithIndexMergeSort<XType, ConvertType, TopKV2TilingDataSimd, false, IndexType> mergeSort;
	97	+ mergeSort.Init(x, index, values, indices, globalWorkGm, tilingData, tPipe);
	98	+ mergeSort.ProcessSort();
	99	+ }
	100	+}
	101	+
	102	+template <typename T_INDEX_TO>
	103	+__aicore__ void sortwithindexForTopK(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex, GM_ADDR workspace,
	104	+ const TopKV2TilingDataSimd* tilingData, TPipe* tPipe)
	105	+{
	106	+ if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_INT8) {
	107	+ generateOpForTopK<int8_t, uint8_t, uint8_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	108	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_INT16) {
	109	+ generateOpForTopK<int16_t, uint16_t, uint16_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	110	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_INT32) {
	111	+ generateOpForTopK<int32_t, uint32_t, uint32_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	112	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_INT64) {
	113	+ generateOpForTopK<int64_t, uint64_t, uint64_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	114	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_UINT8) {
	115	+ generateOpForTopK<uint8_t, uint8_t, uint8_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	116	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_UINT64) {
	117	+ generateOpForTopK<uint64_t, uint64_t, uint64_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	118	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_UINT32) {
	119	+ generateOpForTopK<uint32_t, uint32_t, uint32_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	120	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_UINT16) {
	121	+ generateOpForTopK<uint16_t, uint16_t, uint16_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	122	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_FLOAT) {
	123	+ generateOpForTopK<float, float, uint32_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	124	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_FLOAT16) {
	125	+ generateOpForTopK<half, half, uint16_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	126	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_COMMON_TILING_KEY_BF16) {
	127	+ generateOpForTopK<bfloat16_t, float, uint16_t, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	128	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_MERGE_SORT_TILING_KEY_FLOAT) {
	129	+ generateMergeSortOpForTopK<float, float, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	130	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_MERGE_SORT_TILING_KEY_FLOAT16) {
	131	+ generateMergeSortOpForTopK<half, half, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	132	+ } else if (tilingData->tilingKeyForSort == SORT_WITH_INDEX_MERGE_SORT_TILING_KEY_BF16) {
	133	+ generateMergeSortOpForTopK<bfloat16_t, float, T_INDEX_TO>(x, index, y, sortedIndex, workspace, tilingData, tPipe);
	134	+ }
	135	+}
	136	+#endif

Amath/top_k_v2/op_kernel/arch35/sort_with_index_merge_sort.h+80-0

@@ -0,0 +1,80 @@
	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 sort_with_index_merge_sort.h
	13	+* \brief sort_with_index mergesort mode impl
	14	+*/
	15	+
	16	+#ifndef SORT_WITH_INDEX_MERGE_SORT_TOP_K_H
	17	+#define SORT_WITH_INDEX_MERGE_SORT_TOP_K_H
	18	+
	19	+#include "../../sort_with_index/arch35/merge_sort_with_index.h"
	20	+
	21	+using namespace AscendC;
	22	+using namespace SortWithIndex;
	23	+
	24	+template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
	25	+class SortWithIndexMergeSort: public SortWithIndex::MergeSortWithIndex<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST, INDEX_TYPE> {
	26	+public:
	27	+ __aicore__ inline SortWithIndexMergeSort(){}
	28	+ __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR presetIndex, GM_ADDR value, GM_ADDR indices,
	29	+ GM_ADDR workSpace, const TILING_DATA_TYPE* tilingData, TPipe* pipe);
	30	+ __aicore__ inline void InitTilingData(const TILING_DATA_TYPE* tilingData);
	31	+ __aicore__ inline void ProcessSort();
	32	+};
	33	+
	34	+template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
	35	+__aicore__ inline void SortWithIndexMergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST, INDEX_TYPE>::Init(
	36	+ GM_ADDR inputValue, GM_ADDR presetIndex, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace,
	37	+ const TILING_DATA_TYPE* tilingData, TPipe* pipe)
	38	+{
	39	+ InitTilingData(tilingData);
	40	+ this->InitBuffers(inputValue, value, indices, workSpace, pipe);
	41	+ this->presetIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(presetIndex));
	42	+ // vbs init
	43	+ this->vbsSortMe.SetPipe(pipe);
	44	+ this->vbsSortMe.MergeSortInitBuffer(
	45	+ this->numTileData_,
	46	+ this->oneCoreRowNum_,
	47	+ this->mergSortAcApiNeedBufferSize_);
	48	+}
	49	+
	50	+template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
	51	+__aicore__ inline void SortWithIndexMergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST, INDEX_TYPE>::
	52	+ InitTilingData(const TILING_DATA_TYPE* tilingData)
	53	+{
	54	+ // 尾轴size 512
	55	+ this->outputLastDimValue_ = tilingData->outputLastDimValueForSort;
	56	+ this->numTileData_ = tilingData->numTileDataSizeForSort;
	57	+ this->unsortedDimNum_ = tilingData->unsortedDimNumForSort;
	58	+ this->sortLoopTimes_ = tilingData->sortLoopTimesForSort;
	59	+ this->unsortedDimParallel_ = tilingData->unsortedDimParallelForSort;
	60	+ this->oneCoreRowNum_ = tilingData->oneCoreRowNumForSort;
	61	+ // 高阶API需要的临时空间大小
	62	+ this->mergSortAcApiNeedBufferSize_ = tilingData->mergSortAcApiNeedBufferSizeForSort;
	63	+ this->radixSortApiNeedSpace_ = tilingData->sortAcApiNeedBufferSizeForSort;
	64	+}
	65	+
	66	+template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
	67	+__aicore__ inline void SortWithIndexMergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST, INDEX_TYPE>::ProcessSort()
	68	+{
	69	+ if (GetBlockIdx() >= this->unsortedDimParallel_) {
	70	+ return;
	71	+ }
	72	+
	73	+ for (int32_t i = 0; i < this->sortLoopTimes_; i++) {
	74	+ this->sortLoopRound_ = i;
	75	+ uint64_t loopOffset = i * this->unsortedDimParallel_ * this->oneCoreRowNum_ * this->numTileData_;
	76	+ this->ProcessSingleBlockSort(this->inputValueGm_[loopOffset], this->presetIndexGm_[loopOffset]);
	77	+ }
	78	+}
	79	+
	80	+#endif // SORT_WITH_INDEX_MERGE_SORT_TOP_K_H

Amath/top_k_v2/op_kernel/arch35/sort_with_index_multi_block.h+178-0

@@ -0,0 +1,178 @@
	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 sort_with_index_multi_block.h
	13	+* \brief sort_with_index multiblock mode impl
	14	+*/
	15	+
	16	+#ifndef SORT_WITH_INDEX_MULTI_BLOCK_H
	17	+#define SORT_WITH_INDEX_MULTI_BLOCK_H
	18	+
	19	+#include "../../sort_with_index/arch35/radix_sort_with_index_multi_block.h"
	20	+
	21	+using namespace AscendC;
	22	+using namespace SortWithIndex;
	23	+
	24	+template <typename XType, typename UnsignedType, bool IsDescend, typename XRangeType, typename IndexType>
	25	+class SortWithIndexMultiBlock : public SortWithIndex::RadixSortWithIndexMultiBlock<XType, UnsignedType, IsDescend, XRangeType, IndexType> {
	26	+public:
	27	+ __aicore__ inline SortWithIndexMultiBlock(){}
	28	+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,
	29	+ GM_ADDR workspace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe);
	30	+ __aicore__ inline void Process();
	31	+
	32	+private:
	33	+ __aicore__ inline void ParserTilingData();
	34	+ __aicore__ inline void ProcessMultiBlock(GlobalTensor<XType> xGm, GlobalTensor<IndexType> indexGm,
	35	+ uint64_t gmOffset, uint64_t loopRound);
	36	+private:
	37	+ const TopKV2TilingDataSimd* tilingData_;
	38	+};
	39	+
	40	+template <typename XType, typename UnsignedType, bool IsDescend, typename XRangeType, typename IndexType>
	41	+__aicore__ inline void SortWithIndexMultiBlock<XType, UnsignedType, IsDescend, XRangeType, IndexType>::Init(
	42	+ GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex, GM_ADDR workspace,
	43	+ const TopKV2TilingDataSimd* tilingData, TPipe* pipe)
	44	+{
	45	+ this->blockIdx_ = GetBlockIdx();
	46	+ this->pipe_ = pipe;
	47	+ tilingData_ = tilingData;
	48	+ ParserTilingData();
	49	+ // 需要使用的核心数
	50	+ this->realCoreNum_ = this->unsortedDimParallel_ * this->lastDimRealCore_;
	51	+ if constexpr (sizeof(XRangeType) == sizeof(int64_t)) {
	52	+ this->factor_ = Sort::CONST_2;
	53	+ }
	54	+ // 输入输出GlobalTensor初始化
	55	+ this->inputXGm_.SetGlobalBuffer((__gm__ XType *)x);
	56	+ this->indexGm_.SetGlobalBuffer((__gm__ IndexType *)index);
	57	+ this->outValueGm_.SetGlobalBuffer((__gm__ XType *)y);
	58	+ this->outIdxGm_.SetGlobalBuffer((__gm__ IndexType *)sortedIndex);
	59	+
	60	+ uint64_t wkOffset = this->clearCoreSize0_ * this->clearCore0_;
	61	+ uint64_t oneBlockNumB32 = this->oneBlock_ / sizeof(int32_t); // oneBlock_ = 32
	62	+ if constexpr (sizeof(XRangeType) == sizeof(int64_t)) {
	63	+ wkOffset = wkOffset * Sort::CONST_2;
	64	+ }
	65	+ wkOffset = this->CeilDivMul(wkOffset, oneBlockNumB32);
	66	+ this->excusiveBinsGmWk_.SetGlobalBuffer((__gm__ uint32_t *)workspace, wkOffset);
	67	+ wkOffset = wkOffset * sizeof(uint32_t);
	68	+
	69	+ uint64_t histOffset = this->clearCout_ * this->clearSize_ * this->clearCore1_;
	70	+ if constexpr (sizeof(XRangeType) == sizeof(int64_t)) {
	71	+ histOffset = histOffset * Sort::CONST_2;
	72	+ }
	73	+ histOffset = this->CeilDivMul(histOffset, oneBlockNumB32);
	74	+ this->globalHistGmWk_.SetGlobalBuffer((__gm__ uint32_t *)(workspace + wkOffset), histOffset);
	75	+ wkOffset = wkOffset + histOffset * sizeof(uint32_t);
	76	+
	77	+ uint64_t indexDbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
	78	+ indexDbOffset = this->CeilDivMul(indexDbOffset, oneBlockNumB32);
	79	+ this->outIdxDbWK_.SetGlobalBuffer((__gm__ IndexType *)(workspace + wkOffset), indexDbOffset);
	80	+ wkOffset = wkOffset + indexDbOffset * sizeof(IndexType);
	81	+
	82	+ uint64_t histTileOffset = this->lastDimTileNum_ * Sort::RADIX_SORT_NUM * this->unsortedDimParallel_;
	83	+ this->histTileGmWk_.SetGlobalBuffer((__gm__ uint16_t *)(workspace + wkOffset), histTileOffset);
	84	+ wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);
	85	+ this->histCumsumTileGmWk_.SetGlobalBuffer((__gm__ uint16_t *)(workspace + wkOffset), histTileOffset);
	86	+ wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);
	87	+
	88	+ uint64_t xB8Offset = this->lastDimTileNum_ * this->numTileData_ * this->unsortedDimParallel_;
	89	+ xB8Offset = this->CeilDivMul(xB8Offset, this->oneBlock_);
	90	+ this->xB8GmWk_.SetGlobalBuffer((__gm__ uint8_t *)(workspace + wkOffset), xB8Offset);
	91	+ wkOffset = wkOffset + xB8Offset * sizeof(uint8_t);
	92	+
	93	+ uint64_t dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
	94	+ dbOffset = this->CeilDivMul(dbOffset * sizeof(XType), this->oneBlock_) / sizeof(XType);
	95	+ this->outValueDbWK_.SetGlobalBuffer((__gm__ XType *)(workspace + wkOffset), dbOffset);
	96	+
	97	+ this->pipe_->InitBuffer(this->inQueueX_, 1, this->numTileData_ * sizeof(XType));
	98	+ this->pipe_->InitBuffer(this->inQueueIndex_, 1, this->numTileData_ * sizeof(IndexType));
	99	+ this->pipe_->InitBuffer(this->inQueueGlobalHist_, 1, Sort::RADIX_SORT_NUM * sizeof(XRangeType));
	100	+ this->pipe_->InitBuffer(this->outValueQueue_, 1, this->numTileData_);
	101	+ this->pipe_->InitBuffer(this->blockExcusiveInQue_, 1, Sort::RADIX_SORT_NUM * sizeof(uint16_t));
	102	+ this->pipe_->InitBuffer(this->blockHistInQue_, 1, Sort::RADIX_SORT_NUM * sizeof(uint16_t));
	103	+ this->pipe_->InitBuffer(this->blockUbFlagQue_, 1, Sort::RADIX_SORT_NUM * sizeof(XRangeType));
	104	+ this->pipe_->InitBuffer(this->inputB8Que_, 1, this->numTileData_);
	105	+ this->pipe_->InitBuffer(this->outIdxQueue_, 1, this->numTileData_ * sizeof(uint32_t));
	106	+ this->pipe_->InitBuffer(this->tmpUb_, this->tmpUbSize_);
	107	+ this->pipe_->InitBuffer(this->blockHistFlagUbQue_, 1, Sort::RADIX_SORT_NUM * sizeof(XRangeType));
	108	+
	109	+ this->globalHistGmWkTmp_ = this->globalHistGmWk_.template ReinterpretCast<XRangeType>();
	110	+ this->excusiveBinsGmWkTmp_ = this->excusiveBinsGmWk_.template ReinterpretCast<XRangeType>();
	111	+}
	112	+
	113	+template <typename XType, typename UnsignedType, bool IsDescend, typename XRangeType, typename IndexType>
	114	+__aicore__ inline void SortWithIndexMultiBlock<XType, UnsignedType, IsDescend, XRangeType, IndexType>::ParserTilingData()
	115	+{
	116	+ this->totalDataNum_ = tilingData_->lastAxisNumForSort; // h轴大小
	117	+ this->numTileData_ = tilingData_->numTileDataSizeForSort; // ub循环块大小
	118	+ this->unsortedDimNum_ = tilingData_->unsortedDimNumForSort; // b轴大小
	119	+ this->unsortedDimParallel_ = tilingData_->unsortedDimParallelForSort; // b轴使用的核数
	120	+ this->lastDimTileNum_ = tilingData_->lastDimTileNumForSort; // h轴循环次数
	121	+ this->sortLoopTimes_ = tilingData_->sortLoopTimesForSort; // b轴循环次数
	122	+ this->lastDimRealCore_ = tilingData_->lastDimNeedCoreForSort; // h轴需要的核数
	123	+ this->tmpUbSize_ = tilingData_->tmpUbSize; // 高级api需要用的ub大小
	124	+
	125	+ this->clearCore1_ = tilingData_->keyParams0; // 用于清零的globalHistGmWk_的核
	126	+ this->clearCore0_ = tilingData_->keyParams1; // 用于清零excusiveBinsGmWk_的核
	127	+ this->clearSize_ = tilingData_->keyParams2; // 每次清零的ub大小，按照大的globalHistGmWk_所需ub算
	128	+ this->clearCout_ = tilingData_->keyParams3; // 清零globalHistGmWk_ ub循环次数
	129	+ this->clearCoreSize0_ = tilingData_->keyParams4; // 清零excusiveBinsGmWk_,每个核处理多少个数
	130	+ this->clearCoreSize1_ = tilingData_->keyParams5; // 清零globalHistGmWk_，每个核处理多少
	131	+}
	132	+
	133	+template <typename XType, typename UnsignedType, bool IsDescend, typename XRangeType, typename IndexType>
	134	+__aicore__ inline void SortWithIndexMultiBlock<XType, UnsignedType, IsDescend, XRangeType, IndexType>::Process()
	135	+{
	136	+ for (uint64_t i = 0; i < this->sortLoopTimes_; i++) {
	137	+ uint64_t loopOffset = i * this->unsortedDimParallel_ * this->totalDataNum_;
	138	+ ProcessMultiBlock(this->inputXGm_[loopOffset], this->indexGm_[loopOffset], loopOffset, i);
	139	+ }
	140	+}
	141	+
	142	+template <typename XType, typename UnsignedType, bool IsDescend, typename XRangeType, typename IndexType>
	143	+__aicore__ inline void SortWithIndexMultiBlock<XType, UnsignedType, IsDescend, XRangeType, IndexType>::ProcessMultiBlock(
	144	+ GlobalTensor<XType> xGm, GlobalTensor<IndexType> indexGm, uint64_t gmOffset, uint64_t loopRound)
	145	+{
	146	+ if (this->blockIdx_ < this->realCoreNum_) {
	147	+ this->ClearWorkSapce();
	148	+ }
	149	+ PipeBarrier<PIPE_ALL>();
	150	+ SyncAll();
	151	+
	152	+ if (this->blockIdx_ < this->realCoreNum_) {
	153	+ uint64_t indexGmOffset = gmOffset;
	154	+ if constexpr (sizeof(XType) == sizeof(int8_t)) {
	155	+ this->inputXDbGm_.SetDoubleBuffer(this->outValueDbWK_, this->outValueGm_[gmOffset]);
	156	+ this->idxDbGm_.SetDoubleBuffer(indexGm, this->outIdxGm_[indexGmOffset]);
	157	+ } else {
	158	+ this->inputXDbGm_.SetDoubleBuffer(this->outValueGm_[gmOffset], this->outValueDbWK_);
	159	+ this->idxDbGm_.SetDoubleBuffer(this->outIdxGm_[indexGmOffset], this->outIdxDbWK_);
	160	+ }
	161	+ }
	162	+
	163	+ for (uint32_t sortRound = 0; sortRound < static_cast<uint32_t>(sizeof(XType)); sortRound++) {
	164	+ if (this->blockIdx_ < this->realCoreNum_) {
	165	+ // 确定 histTileGmWk_(直方图), histCumsumTileGmWk_, excusiveBinsGmWk_
	166	+ this->GetGlobalExcusiveSum(sortRound, loopRound, xGm);
	167	+ }
	168	+ PipeBarrier<PIPE_ALL>();
	169	+ SyncAll();
	170	+ if (this->blockIdx_ < this->realCoreNum_) {
	171	+ this->ComputeOnePass(sortRound, loopRound, xGm, indexGm);
	172	+ }
	173	+ PipeBarrier<PIPE_ALL>();
	174	+ SyncAll();
	175	+ }
	176	+}
	177	+
	178	+#endif // SORT_WITH_INDEX_MULTI_BLOCK_H

Amath/top_k_v2/op_kernel/arch35/sort_with_index_single_block.h+74-0

@@ -0,0 +1,74 @@
	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 sort_with_index_single_block.h
	13	+* \brief sort_with_index singleblock mode impl
	14	+*/
	15	+
	16	+#ifndef SORT_WITH_INDEX_SINGLE_BLOCK_H
	17	+#define SORT_WITH_INDEX_SINGLE_BLOCK_H
	18	+
	19	+#include "../../sort_with_index/arch35/radix_sort_with_index_single_block.h"
	20	+
	21	+using namespace AscendC;
	22	+
	23	+template <typename XType, bool IsDescend, typename IndexType>
	24	+struct SortWithIndexSingleBlock : public RadixSortWithIndexSingleBlock<XType, IsDescend, IndexType> {
	25	+public:
	26	+ __aicore__ inline SortWithIndexSingleBlock() {}
	27	+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,
	28	+ GM_ADDR workspace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe);
	29	+ __aicore__ inline void Process();
	30	+};
	31	+
	32	+template <typename XType, bool IsDescend, typename IndexType>
	33	+__aicore__ inline void SortWithIndexSingleBlock<XType, IsDescend, IndexType>::Init(
	34	+ GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex, GM_ADDR workspace, const TopKV2TilingDataSimd* tilingData,
	35	+ TPipe* pipe)
	36	+{
	37	+ this->xGm_.SetGlobalBuffer((__gm__ XType*)(x));
	38	+ this->indexGm_.SetGlobalBuffer((__gm__ IndexType*)(index));
	39	+ this->yGm_.SetGlobalBuffer((__gm__ XType*)(y));
	40	+ this->sortedIndexGm_.SetGlobalBuffer((__gm__ IndexType*)(sortedIndex));
	41	+
	42	+ this->numTileData_ = tilingData->numTileDataSizeForSort;
	43	+ this->lastDimRealCore_ = tilingData->lastDimNeedCoreForSort;
	44	+ this->totalDataNum_ = tilingData->lastAxisNumForSort;
	45	+ this->unsortedDimNum_ = tilingData->unsortedDimNumForSort;
	46	+ this->sortLoopTimes_ = tilingData->sortLoopTimesForSort;
	47	+ this->lastDimTileNum_ = tilingData->lastDimTileNumForSort;
	48	+ this->unsortedDimParallel_ = tilingData->unsortedDimParallelForSort;
	49	+ this->oneCoreRowNum_ = tilingData->oneCoreRowNumForSort;
	50	+ this->sortAcApiNeedTmpBufferSize_ = tilingData->sortAcApiNeedBufferSizeForSort;
	51	+
	52	+ this->blockIdx_ = GetBlockIdx();
	53	+
	54	+ this->pipe_ = pipe;
	55	+ this->pipe_->InitBuffer(this->inQueueX_, 1, ROUND_UP_AGLIN(this->numTileData_ * sizeof(XType)));
	56	+ this->pipe_->InitBuffer(this->inQueueIndex_, 1, ROUND_UP_AGLIN(this->numTileData_ * sizeof(IndexType)));
	57	+ this->pipe_->InitBuffer(this->yQueue_, 1, ROUND_UP_AGLIN(this->numTileData_ * sizeof(XType)));
	58	+ this->pipe_->InitBuffer(this->sortedIndexQueue_, 1, ROUND_UP_AGLIN(this->numTileData_ * sizeof(IndexType)));
	59	+ this->pipe_->InitBuffer(this->sortedShareMemTbuf_, ROUND_UP_AGLIN(this->sortAcApiNeedTmpBufferSize_));
	60	+}
	61	+
	62	+template <typename XType, bool IsDescend, typename IndexType>
	63	+__aicore__ inline void SortWithIndexSingleBlock<XType, IsDescend, IndexType>::Process()
	64	+{
	65	+ if (GetBlockIdx() >= this->unsortedDimParallel_) {
	66	+ return;
	67	+ }
	68	+
	69	+ for (uint64_t i = 0; i < this->sortLoopTimes_; i++) {
	70	+ uint64_t loopOffset = i * this->unsortedDimParallel_ * this->totalDataNum_ * this->oneCoreRowNum_;
	71	+ this->ProcessSingleBlock(this->xGm_[loopOffset], this->indexGm_[loopOffset], i);
	72	+ }
	73	+}
	74	+#endif // RADIX_SORT_WITH_INDEX_SINGLE_BLOCK_H

Mmath/top_k_v2/op_kernel/arch35/top_k_constant_var_simd.h+2-0

	@@ -13,6 +13,7 @@
13	*/	13	*/
14	#ifndef TOP_K_CONSTANT_VAR_SIMD_H	14	#ifndef TOP_K_CONSTANT_VAR_SIMD_H
15	#define TOP_K_CONSTANT_VAR_SIMD_H	15	#define TOP_K_CONSTANT_VAR_SIMD_H
		16	+namespace topkV2 {
16	const uint32_t THREAD_DIM_NUM = 1024;	17	const uint32_t THREAD_DIM_NUM = 1024;
17	const int32_t ONE_TIMES_B64_NUM = 32;	18	const int32_t ONE_TIMES_B64_NUM = 32;
18	const int32_t ONE_TIMES_B32_NUM = 64;	19	const int32_t ONE_TIMES_B32_NUM = 64;
	@@ -65,4 +66,5 @@ const int16_t XOR_OP_VALUE_HALF = 0x8000;
65	const int32_t SMALL_SORT_MAX_DATA_SIZE = 128;	66	const int32_t SMALL_SORT_MAX_DATA_SIZE = 128;
66	const uint32_t CONCAT_AGLIN_VALUE = 16;	67	const uint32_t CONCAT_AGLIN_VALUE = 16;
67	const uint32_t DOUBLE_BUFFER = 2;	68	const uint32_t DOUBLE_BUFFER = 2;
		69	+}
68	#endif	70	#endif

Mmath/top_k_v2/op_kernel/arch35/top_k_merge_sort.h+3-1

	@@ -18,6 +18,7 @@
18	#include "top_k_constant_var_simd.h"	18	#include "top_k_constant_var_simd.h"
19	#include "top_k_merge_sort_simd.h"	19	#include "top_k_merge_sort_simd.h"
20	using namespace AscendC;	20	using namespace AscendC;
		21	+namespace topkV2 {
21	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>	22	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
22	struct MergeSort {	23	struct MergeSort {
23	__aicore__ inline MergeSort() {}	24	__aicore__ inline MergeSort() {}
	@@ -71,7 +72,7 @@ public:
71	uint32_t platformCoreNum_ = 0;	72	uint32_t platformCoreNum_ = 0;
72	uint32_t outputLastDimValue_ = 0;	73	uint32_t outputLastDimValue_ = 0;
73	// merge sort kernel	74	// merge sort kernel
74	- KernelVbsMergeSort<T, CONVERT_TYPE, IS_LARGEST> vbsSort;	75	+ topkV2::KernelVbsMergeSort<T, CONVERT_TYPE, IS_LARGEST> vbsSort;
75	};	76	};
76		77
77	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>	78	template <typename T, typename CONVERT_TYPE, typename TILING_DATA_TYPE, bool IS_LARGEST, typename INDEX_TYPE>
	@@ -259,4 +260,5 @@ __aicore__ inline void MergeSort<T, CONVERT_TYPE, TILING_DATA_TYPE, IS_LARGEST,
259	outIndexQueue_.FreeTensor(outIndexLocal);	260	outIndexQueue_.FreeTensor(outIndexLocal);
260	outValueQueue_.FreeTensor(outValueLocal);	261	outValueQueue_.FreeTensor(outValueLocal);
261	}	262	}
		263	+}
262	#endif	264	#endif

Mmath/top_k_v2/op_kernel/arch35/top_k_merge_sort_simd.h+2-0

	@@ -18,6 +18,7 @@
18	#include "top_k_constant_var_simd.h"	18	#include "top_k_constant_var_simd.h"
19		19
20	using namespace AscendC;	20	using namespace AscendC;
		21	+namespace topkV2 {
21	template <typename T, typename CONVERT_TYPE, bool IS_DESCEND>	22	template <typename T, typename CONVERT_TYPE, bool IS_DESCEND>
22	struct KernelVbsMergeSort {	23	struct KernelVbsMergeSort {
23	public:	24	public:
	@@ -169,4 +170,5 @@ __aicore__ inline void KernelVbsMergeSort<T, CONVERT_TYPE, IS_DESCEND>::flipSign
169	AscendC::Adds(castTensor, castTensor, XOR_OP_VALUE_HALF, aglinTileSize);	170	AscendC::Adds(castTensor, castTensor, XOR_OP_VALUE_HALF, aglinTileSize);
170	}	171	}
171	}	172	}
		173	+}
172	#endif	174	#endif

Mmath/top_k_v2/op_kernel/arch35/top_k_radix_block_sort_b16.h+1-0

	@@ -17,6 +17,7 @@
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
18	#include "top_k_constant_var_simd.h"	18	#include "top_k_constant_var_simd.h"
19	using namespace AscendC;	19	using namespace AscendC;
		20	+using namespace topkV2;
20	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	21	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
21	struct RadixBlockSortSimdB16 {	22	struct RadixBlockSortSimdB16 {
22	public:	23	public:

Mmath/top_k_v2/op_kernel/arch35/top_k_radix_block_sort_b32.h+1-0

	@@ -17,6 +17,7 @@
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
18	#include "top_k_constant_var_simd.h"	18	#include "top_k_constant_var_simd.h"
19	using namespace AscendC;	19	using namespace AscendC;
		20	+using namespace topkV2;
20	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	21	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
21	struct RadixBlockSortSimdB32 {	22	struct RadixBlockSortSimdB32 {
22	public:	23	public:

Mmath/top_k_v2/op_kernel/arch35/top_k_radix_block_sort_b64.h+1-0

	@@ -17,6 +17,7 @@
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
18	#include "top_k_constant_var_simd.h"	18	#include "top_k_constant_var_simd.h"
19	using namespace AscendC;	19	using namespace AscendC;
		20	+using namespace topkV2;
20	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	21	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
21	struct RadixBlockSortSimdB64 {	22	struct RadixBlockSortSimdB64 {
22	public:	23	public:

Mmath/top_k_v2/op_kernel/arch35/top_k_radix_block_sort_b8.h+1-0

	@@ -17,6 +17,7 @@
17	#include "top_k_util_type_simd.h"	17	#include "top_k_util_type_simd.h"
18	#include "top_k_constant_var_simd.h"	18	#include "top_k_constant_var_simd.h"
19	using namespace AscendC;	19	using namespace AscendC;
		20	+using namespace topkV2;
20	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>	21	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, bool IS_DESCEND, typename T_INDEX>
21	struct RadixBlockSortSimdB8 {	22	struct RadixBlockSortSimdB8 {
22	public:	23	public:

Mmath/top_k_v2/op_kernel/arch35/top_k_util_type_simd.h+12-0

	@@ -15,6 +15,8 @@
15	#define UTIL_TYPE_SIMD_H	15	#define UTIL_TYPE_SIMD_H
16	#include "kernel_operator.h"	16	#include "kernel_operator.h"
17	#include "top_k_constant_var_simd.h"	17	#include "top_k_constant_var_simd.h"
		18	+
		19	+namespace topkV2 {
18	template <typename Tp, Tp v>	20	template <typename Tp, Tp v>
19	struct integral_constant {	21	struct integral_constant {
20	static constexpr Tp value = v;	22	static constexpr Tp value = v;
	@@ -38,6 +40,15 @@ __aicore__ inline uint32_t CeilDivMul(uint32_t a, uint32_t b)
38	return ((a + b - 1) / b) * b;	40	return ((a + b - 1) / b) * b;
39	}	41	}
40		42
		43	+template <typename T>
		44	+__aicore__ auto CeilAlignDivMul(uint64_t a, uint64_t b) ->T const
		45	+{
		46	+ if (b == 0) {
		47	+ return static_cast<T>(a);
		48	+ }
		49	+ return static_cast<T>(((a + b - 1) / b) * b);
		50	+}
		51	+
41	template <typename T>	52	template <typename T>
42	__aicore__ inline T SortGetMin(T left, T right)	53	__aicore__ inline T SortGetMin(T left, T right)
43	{	54	{
	@@ -64,4 +75,5 @@ struct DoubleBufferSimd
64	selector_ = selector_ ^ 1;	75	selector_ = selector_ ^ 1;
65	}	76	}
66	};	77	};
		78	+}
67	#endif	79	#endif

Mmath/top_k_v2/op_kernel/top_k_v2_apt.cpp+52-13

	@@ -13,14 +13,19 @@
13	* \brief top k v2 impl	13	* \brief top k v2 impl
14	*/	14	*/
15		15
		16	+#ifndef TOP_K_V2_APT_H
		17	+#define TOP_K_V2_APT_H
		18	+
16	#include "arch35/radix_sort_top_k.h"	19	#include "arch35/radix_sort_top_k.h"
17	#include "arch35/radix_topk_constant.h"	20	#include "arch35/radix_topk_constant.h"
18	#include "arch35/top_k_merge_sort.h"	21	#include "arch35/top_k_merge_sort.h"
19	#include "arch35/radix_sort_top_k_single_block.h"	22	#include "arch35/radix_sort_top_k_single_block.h"
20	#include "arch35/radix_sort_top_k_single_core.h"	23	#include "arch35/radix_sort_top_k_single_core.h"
21	#include "arch35/radix_sort_top_k_inter_core_template_optimization.h"	24	#include "arch35/radix_sort_top_k_inter_core_template_optimization.h"
		25	+#include "arch35/sort_and_top_k_more_core.h"
22		26
23	using namespace AscendC;	27	using namespace AscendC;
		28	+using namespace SortAndTopK;
24		29
25	#define TOPK_COMMON_TILING_KEY_INT64 1004	30	#define TOPK_COMMON_TILING_KEY_INT64 1004
26	#define TOPK_COMMON_TILING_KEY_INT32 1003	31	#define TOPK_COMMON_TILING_KEY_INT32 1003
	@@ -39,6 +44,7 @@ using namespace AscendC;
39		44
40	const uint32_t SINGLE_CORE_MODE = 1;	45	const uint32_t SINGLE_CORE_MODE = 1;
41	const uint32_t MULT_CORE_OPTIM_MODE = 4;	46	const uint32_t MULT_CORE_OPTIM_MODE = 4;
		47	+const uint32_t SORT_AND_TOP_K_MODE = 5;
42		48
43	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, typename T_INDEX, typename T_INDEX_TO>	49	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, typename T_INDEX, typename T_INDEX_TO>
44	__aicore__ inline void RadixSortTopKOpObject(	50	__aicore__ inline void RadixSortTopKOpObject(
	@@ -72,6 +78,25 @@ __aicore__ inline void RadixSortTopKOpObject(
72	}	78	}
73	}	79	}
74		80
		81	+template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, typename T_INDEX, typename T_INDEX_TO>
		82	+__aicore__ inline void SortAndTopKOpObject(GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR globalWorkGm,
		83	+ GM_ADDR tiling)
		84	+{
		85	+ GET_TILING_DATA(tilingData, tiling);
		86	+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
		87	+ bool isLargest = (tilingData.isLargest > 0) ? true : false;
		88	+ TPipe tPipe;
		89	+ if (isLargest) {
		90	+ SortAndTopK::SortAndTopKMoreCore<T, T_INDEX_TO, UNSINGED_TYPE, T_INDEX, 1> sortAndTopKMoreCore;
		91	+ sortAndTopKMoreCore.InitParam(x, values, indices, globalWorkGm, &tilingData, &tPipe);
		92	+ sortAndTopKMoreCore.ProcessTopK();
		93	+ } else {
		94	+ SortAndTopK::SortAndTopKMoreCore<T, T_INDEX_TO, UNSINGED_TYPE, T_INDEX, 0> sortAndTopKMoreCore;
		95	+ sortAndTopKMoreCore.InitParam(x, values, indices, globalWorkGm, &tilingData, &tPipe);
		96	+ sortAndTopKMoreCore.ProcessTopK();
		97	+ }
		98	+}
		99	+
75	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, typename T_INDEX, typename T_INDEX_TO>	100	template <typename T, typename UNSINGED_TYPE, int32_t NUM_PASS, typename T_INDEX, typename T_INDEX_TO>
76	__aicore__ inline void RadixSortTopKSingleCoreOpObject(	101	__aicore__ inline void RadixSortTopKSingleCoreOpObject(
77	GM_ADDR x, GM_ADDR k, GM_ADDR values, GM_ADDR indices, GM_ADDR globalWorkGm, GM_ADDR tiling)	102	GM_ADDR x, GM_ADDR k, GM_ADDR values, GM_ADDR indices, GM_ADDR globalWorkGm, GM_ADDR tiling)
	@@ -179,6 +204,7 @@ __aicore__ inline void generateOpObject(
179	bool isSort = (tilingData.isSort > 0) ? true : false;	204	bool isSort = (tilingData.isSort > 0) ? true : false;
180	bool isSingleBlock = (tilingData.lastDimNeedCore == 1) ? true : false;	205	bool isSingleBlock = (tilingData.lastDimNeedCore == 1) ? true : false;
181	bool isSingleCore = (tilingData.modeType == SINGLE_CORE_MODE) ? true : false;	206	bool isSingleCore = (tilingData.modeType == SINGLE_CORE_MODE) ? true : false;
		207	+ bool isSortAndTopK = (tilingData.modeType == SORT_AND_TOP_K_MODE) ? true : false;
		J J JimmyLam23月27日是否简化为bool isSortAndTopK = (tilingData.modeType == SORT_AND_TOP_K_MODE); BruceLP 3月31日评论： 3月31日评论：
182	bool isInInt32Range = (tilingData.isInInt32Range > 0) ? true : false;	208	bool isInInt32Range = (tilingData.isInInt32Range > 0) ? true : false;
183	bool isMultiCoreOptimMode = (tilingData.modeType == MULT_CORE_OPTIM_MODE) ? true : false;	209	bool isMultiCoreOptimMode = (tilingData.modeType == MULT_CORE_OPTIM_MODE) ? true : false;
184		210
	@@ -194,6 +220,18 @@ __aicore__ inline void generateOpObject(
194	return;	220	return;
195	}	221	}
196		222
		223	+ // SortAndTopK模板
		224	+ if (isSortAndTopK) {
		225	+ if (isInInt32Range) {
		226	+ SortAndTopKOpObject<T, UNSINGED_TYPE, NUM_PASS, uint32_t, T_INDEX_TO>(
		227	+ x, values, indices, globalWorkGm, tiling);
		228	+ } else {
		229	+ SortAndTopKOpObject<T, UNSINGED_TYPE, NUM_PASS, int64_t, T_INDEX_TO>(
		230	+ x, values, indices, globalWorkGm, tiling);
		231	+ }
		232	+ return;
		233	+ }
		234	+
197	// 单核多次处理模板（930新增模板用于性能优化）	235	// 单核多次处理模板（930新增模板用于性能优化）
198	if (isSingleCore) {	236	if (isSingleCore) {
199	if (isInInt32Range) {	237	if (isInInt32Range) {
	@@ -232,11 +270,11 @@ __aicore__ inline void generateMergeTopKObject(
232	bool isLargest = (tilingData.isLargest > 0) ? true : false;	270	bool isLargest = (tilingData.isLargest > 0) ? true : false;
233	TPipe pipe;	271	TPipe pipe;
234	if (isLargest) {	272	if (isLargest) {
235	- MergeSort<T, CONVERT_TYPE, TopKV2TilingDataSimd, true, INDEX_DTYPE> mergeSort;	273	+ topkV2::MergeSort<T, CONVERT_TYPE, TopKV2TilingDataSimd, true, INDEX_DTYPE> mergeSort;
236	mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);	274	mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);
237	mergeSort.ProcessSort();	275	mergeSort.ProcessSort();
238	} else {	276	} else {
239	- MergeSort<T, CONVERT_TYPE, TopKV2TilingDataSimd, false, INDEX_DTYPE> mergeSort;	277	+ topkV2::MergeSort<T, CONVERT_TYPE, TopKV2TilingDataSimd, false, INDEX_DTYPE> mergeSort;
240	mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);	278	mergeSort.Init(x, values, indices, globalWorkGm, &tilingData, &pipe);
241	mergeSort.ProcessSort();	279	mergeSort.ProcessSort();
242	}	280	}
	@@ -256,56 +294,56 @@ extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR val
256	#if ORIG_DTYPE_X == DT_INT64	294	#if ORIG_DTYPE_X == DT_INT64
257	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT64);	295	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT64);
258	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT64	296	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT64
259	- generateOpObject<int64_t, uint64_t, B64_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	297	+ generateOpObject<int64_t, uint64_t, topkV2::B64_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
260	#endif	298	#endif
261	#endif	299	#endif
262		300
263	#if ORIG_DTYPE_X == DT_INT32	301	#if ORIG_DTYPE_X == DT_INT32
264	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT32);	302	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT32);
265	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT32	303	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT32
266	- generateOpObject<int32_t, uint32_t, B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	304	+ generateOpObject<int32_t, uint32_t, topkV2::B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
267	#endif	305	#endif
268	#endif	306	#endif
269		307
270	#if ORIG_DTYPE_X == DT_INT16	308	#if ORIG_DTYPE_X == DT_INT16
271	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT16);	309	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT16);
272	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT16	310	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT16
273	- generateOpObject<int16_t, uint16_t, B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	311	+ generateOpObject<int16_t, uint16_t, topkV2::B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
274	#endif	312	#endif
275	#endif	313	#endif
276		314
277	#if ORIG_DTYPE_X == DT_INT8	315	#if ORIG_DTYPE_X == DT_INT8
278	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT8);	316	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_INT8);
279	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT8	317	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_INT8
280	- generateOpObject<int8_t, uint8_t, B8_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	318	+ generateOpObject<int8_t, uint8_t, topkV2::B8_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
281	#endif	319	#endif
282	#endif	320	#endif
283		321
284	#if ORIG_DTYPE_X == DT_UINT64	322	#if ORIG_DTYPE_X == DT_UINT64
285	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT64);	323	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT64);
286	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT64	324	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT64
287	- generateOpObject<uint64_t, uint64_t, B64_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	325	+ generateOpObject<uint64_t, uint64_t, topkV2::B64_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
288	#endif	326	#endif
289	#endif	327	#endif
290		328
291	#if ORIG_DTYPE_X == DT_UINT32	329	#if ORIG_DTYPE_X == DT_UINT32
292	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT32);	330	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT32);
293	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT32	331	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT32
294	- generateOpObject<uint32_t, uint32_t, B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	332	+ generateOpObject<uint32_t, uint32_t, topkV2::B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
295	#endif	333	#endif
296	#endif	334	#endif
297		335
298	#if ORIG_DTYPE_X == DT_UINT16	336	#if ORIG_DTYPE_X == DT_UINT16
299	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT16);	337	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT16);
300	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT16	338	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT16
301	- generateOpObject<uint16_t, uint16_t, B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	339	+ generateOpObject<uint16_t, uint16_t, topkV2::B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
302	#endif	340	#endif
303	#endif	341	#endif
304		342
305	#if ORIG_DTYPE_X == DT_UINT8	343	#if ORIG_DTYPE_X == DT_UINT8
306	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT8);	344	TILING_KEY_IS(TOPK_COMMON_TILING_KEY_UINT8);
307	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT8	345	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_UINT8
308	- generateOpObject<uint8_t, uint8_t, B8_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	346	+ generateOpObject<uint8_t, uint8_t, topkV2::B8_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
309	#endif	347	#endif
310	#endif	348	#endif
311		349
	@@ -314,7 +352,7 @@ extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR val
314	TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_FLOAT);	352	TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_FLOAT);
315		353
316	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_FLOAT	354	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_FLOAT
317	- generateOpObject<float, uint32_t, B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	355	+ generateOpObject<float, uint32_t, topkV2::B32_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
318	#elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_FLOAT	356	#elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_FLOAT
319	generateMergeTopKObject<float, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);	357	generateMergeTopKObject<float, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);
320	#endif	358	#endif
	@@ -325,7 +363,7 @@ extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR val
325	TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_FLOAT16);	363	TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_FLOAT16);
326		364
327	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_FLOAT16	365	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_FLOAT16
328	- generateOpObject<half, uint16_t, B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	366	+ generateOpObject<half, uint16_t, topkV2::B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
329	#elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_FLOAT16	367	#elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_FLOAT16
330	generateMergeTopKObject<half, half, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);	368	generateMergeTopKObject<half, half, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);
331	#endif	369	#endif
	@@ -336,9 +374,10 @@ extern "C" __global__ __aicore__ void top_k_v2(GM_ADDR x, GM_ADDR k, GM_ADDR val
336	TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_BF16);	374	TILING_KEY_IS(TOPK_MERGE_SORT_TILING_KEY_BF16);
337		375
338	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_BF16	376	#if TILING_KEY_VAR == TOPK_COMMON_TILING_KEY_BF16
339	- generateOpObject<bfloat16_t, uint16_t, B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);	377	+ generateOpObject<bfloat16_t, uint16_t, topkV2::B16_BITE_SIZE, DTYPE_INDICES>(x, k, values, indices, globalWorkGm, tiling);
340	#elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_BF16	378	#elif TILING_KEY_VAR == TOPK_MERGE_SORT_TILING_KEY_BF16
341	generateMergeTopKObject<bfloat16_t, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);	379	generateMergeTopKObject<bfloat16_t, float, DTYPE_INDICES>(x, values, indices, globalWorkGm, tiling);
342	#endif	380	#endif
343	#endif	381	#endif
344	}	382	}
		383	+#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 ";	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 ";
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);

Mmath/topk/op_host/op_api/topk.cpp+1-21

	@@ -118,12 +118,6 @@ static bool IsAscendCSupport(const aclTensor* self, int64_t k)
118	return false;	118	return false;
119	}	119	}
120		120
121	-// 根据芯片类型、k 和 sotrd 判断是否需要额外增加 SortWithIndex
122	-static bool IsSortWithIndex(int64_t k, bool sorted)
123	-{
124	- return (IsRegBase()) && (k > TWO_THOUSAND) && (sorted == true);
125	-}
126	-
127	// AICORE算子kernel	121	// AICORE算子kernel
128	std::tuple<aclTensor, aclTensor> TopkV2AiCore(	122	std::tuple<aclTensor, aclTensor> TopkV2AiCore(
129	const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values,	123	const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values,
	@@ -145,18 +139,6 @@ std::tuple<aclTensor, aclTensor> TopkV2AiCoreForDavid(
145	return std::tuple<aclTensor, aclTensor>(values, indices);	139	return std::tuple<aclTensor, aclTensor>(values, indices);
146	}	140	}
147		141
148	-// 950 TopK + SortWithIndex
149	-std::tuple<aclTensor, aclTensor> TopKAndSort(
150	- const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values,
151	- aclTensor* indices, aclOpExecutor* executor)
152	-{
153	- L0_DFX(TopKAndSort, self, k, dim, largest, sorted, values, indices);
154	- // 使用框架宏ADD_TO_LAUNCHER_LIST_AICORE，将AiCore TopKV2算子加入任务队列
155	- ADD_TO_LAUNCHER_LIST_AICORE(
156	- TopKV2, OP_INPUT(self, k), OP_OUTPUT(values, indices), OP_ATTR(sorted, dim, largest, op::DataType::DT_INT32));
157	- return SortWithIndex(values, indices, dim, largest, true, executor);
158	-}
159	-
160	std::tuple<aclTensor, aclTensor> TopkV3(	142	std::tuple<aclTensor, aclTensor> TopkV3(
161	const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values,	143	const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values,
162	aclTensor* indices, aclOpExecutor* executor)	144	aclTensor* indices, aclOpExecutor* executor)
	@@ -198,7 +180,7 @@ std::tuple<aclTensor, aclTensor> Topk(
198	const aclTensor* kTensor = executor->ConvertToTensor(kScalar, op::ToOpDataType(ACL_INT32));	180	const aclTensor* kTensor = executor->ConvertToTensor(kScalar, op::ToOpDataType(ACL_INT32));
199	auto valuesOut = executor->AllocTensor(outShape, self->GetDataType(), self->GetStorageFormat());	181	auto valuesOut = executor->AllocTensor(outShape, self->GetDataType(), self->GetStorageFormat());
200	aclTensor* indicesOut = nullptr;	182	aclTensor* indicesOut = nullptr;
201	- if ((IsRegBase()) && !IsSortWithIndex(k, sorted)) {	183	+ if (IsRegBase()) {
202	indicesOut = executor->AllocTensor(outShape, indicesDType, self->GetStorageFormat());	184	indicesOut = executor->AllocTensor(outShape, indicesDType, self->GetStorageFormat());
203	} else {	185	} else {
204	indicesOut = executor->AllocTensor(outShape, op::DataType::DT_INT32, self->GetStorageFormat());	186	indicesOut = executor->AllocTensor(outShape, op::DataType::DT_INT32, self->GetStorageFormat());
	@@ -207,8 +189,6 @@ std::tuple<aclTensor, aclTensor> Topk(
207	if (IsAiCoreSupport(self, k)) {	189	if (IsAiCoreSupport(self, k)) {
208	if (IsAscendCSupport(self, k)) {	190	if (IsAscendCSupport(self, k)) {
209	return TopkV3(self, kTensor, dim, largest, sorted, valuesOut, indicesOut, executor);	191	return TopkV3(self, kTensor, dim, largest, sorted, valuesOut, indicesOut, executor);
210	- } else if (IsSortWithIndex(k, sorted)) {
211	- return TopKAndSort(self, kTensor, dim, largest, sorted, valuesOut, indicesOut, executor);
212	} else {	192	} else {
213	if (IsRegBase()) {	193	if (IsRegBase()) {
214	return TopkV2AiCoreForDavid(	194	return TopkV2AiCoreForDavid(

	@@ -34,6 +34,8 @@ public:
34	__aicore__ inline void Init(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,	34	__aicore__ inline void Init(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR sortedIndex,
35	GM_ADDR workspace, const SortWithIndexTilingDataSimt* tilingData, TPipe* pipe);	35	GM_ADDR workspace, const SortWithIndexTilingDataSimt* tilingData, TPipe* pipe);
36	__aicore__ inline void Process();	36	__aicore__ inline void Process();
		37	+ __aicore__ inline void ComputeOnePass(uint32_t round, uint64_t sortLoopRound, GlobalTensor<XType> inputXGm,
		38	+ GlobalTensor<IndexType> indexGm);
37		39
38	private:	40	private:
39	__aicore__ inline void ParserTilingData();	41	__aicore__ inline void ParserTilingData();
	@@ -41,8 +43,6 @@ private:
41	uint64_t gmOffset, uint64_t loopRound);	43	uint64_t gmOffset, uint64_t loopRound);
42	__aicore__ inline void CopyInputIndexDataIn(GlobalTensor<IndexType> inputIndex, LocalTensor<IndexType> &xLocal,	44	__aicore__ inline void CopyInputIndexDataIn(GlobalTensor<IndexType> inputIndex, LocalTensor<IndexType> &xLocal,
43	uint64_t tileOffset, uint32_t currTileSize);	45	uint64_t tileOffset, uint32_t currTileSize);
44	- __aicore__ inline void ComputeOnePass(uint32_t round, uint64_t sortLoopRound, GlobalTensor<XType> inputXGm,
45	- GlobalTensor<IndexType> indexGm);
46	__aicore__ inline void ScatterKeysGlobal(LocalTensor<XType> xInputValueLocal,	46	__aicore__ inline void ScatterKeysGlobal(LocalTensor<XType> xInputValueLocal,
47	LocalTensor<uint32_t> sortedIndexLocal,	47	LocalTensor<uint32_t> sortedIndexLocal,
48	LocalTensor<IndexType> xInputIndexLocal,	48	LocalTensor<IndexType> xInputIndexLocal,
	@@ -51,7 +51,7 @@ private:
51	LocalTensor<XRangeType> blockDataInGlobalPos,	51	LocalTensor<XRangeType> blockDataInGlobalPos,
52	LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,	52	LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
53	uint32_t sortRound, XRangeType tileDataStart, uint32_t cureTileSize);	53	uint32_t sortRound, XRangeType tileDataStart, uint32_t cureTileSize);
54	-private:	54	+protected:
55	const SortWithIndexTilingDataSimt* tilingData_;	55	const SortWithIndexTilingDataSimt* tilingData_;
56	GlobalTensor<IndexType> indexGm_;	56	GlobalTensor<IndexType> indexGm_;
57	GlobalTensor<IndexType> outIdxGm_;	57	GlobalTensor<IndexType> outIdxGm_;

	@@ -21,7 +21,7 @@
21	#include <algorithm>	21	#include <algorithm>
22		22
23	using namespace AscendC;	23	using namespace AscendC;
24	-	24	+using namespace topkV2;
25	template <typename T, typename UNSIGNED_TYPE, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>	25	template <typename T, typename UNSIGNED_TYPE, bool IS_LARGEST, bool IS_SORT, typename T_INDEX, typename T_INDEX_TO>
26	struct RadixSortTopKSingleBlock : public RadixSortTopKBase<T, T_INDEX, T_INDEX_TO> {	26	struct RadixSortTopKSingleBlock : public RadixSortTopKBase<T, T_INDEX, T_INDEX_TO> {
27	__aicore__ inline RadixSortTopKSingleBlock() {};	27	__aicore__ inline RadixSortTopKSingleBlock() {};
	@@ -195,7 +195,7 @@ __aicore__ inline void RadixSortTopKSingleBlock<T, UNSIGNED_TYPE, IS_LARGEST, IS
195	uint32_t aglinIndicesOffset = ROUND_UP_AGLIN(this->k_ * sizeof(int32_t)) / sizeof(int32_t);	195	uint32_t aglinIndicesOffset = ROUND_UP_AGLIN(this->k_ * sizeof(int32_t)) / sizeof(int32_t);
196	uint32_t blockCastIntervalBytes = (aglinIndicesOffset - this->k_) * sizeof(T_INDEX_TO);	196	uint32_t blockCastIntervalBytes = (aglinIndicesOffset - this->k_) * sizeof(T_INDEX_TO);
197	bool needsCast = IsSameType<T_INDEX_TO, int64_t>::value;	197	bool needsCast = IsSameType<T_INDEX_TO, int64_t>::value;
198	- uint32_t srcStride = needsCast && blockCastIntervalBytes >= UB_AGLIN_VALUE && parallelBatchNum >= 2 ? 1 : 0;	198	+ uint32_t srcStride = needsCast && blockCastIntervalBytes >= topkV2::UB_AGLIN_VALUE && parallelBatchNum >= 2 ? 1 : 0;
199	AscendC::LocalTensor<T_INDEX_TO> indicesLocal = this->indicesQue_.template DeQue<T_INDEX_TO>();	199	AscendC::LocalTensor<T_INDEX_TO> indicesLocal = this->indicesQue_.template DeQue<T_INDEX_TO>();
200	AscendC::DataCopyExtParams dataCopyParamIndex{static_cast<uint16_t>(parallelBatchNum),	200	AscendC::DataCopyExtParams dataCopyParamIndex{static_cast<uint16_t>(parallelBatchNum),
201	static_cast<uint32_t>(this->k_ * sizeof(T_INDEX_TO)), srcStride, 0, 0};	201	static_cast<uint32_t>(this->k_ * sizeof(T_INDEX_TO)), srcStride, 0, 0};