Star1430
1502
代码
代码
Issues73
Pull Requests285
Actions
项目讨论
Wiki
分析
Star1430
1502
已合并
kth_value: 新增radix_select和small_axis_short_rank_select调度策略 #3804
黄晓彬创建于 7月5日
kth_value: 新增radix_select和small_axis_short_rank_select调度策略 #3804
已合并
黄晓彬创建于 7月5日
从ConanHuang/ops-mathmaster合入到cann/ops-mathmaster
共 9 个文件变更+2235-507
Mmath/kth_value/op_host/arch35/kth_value_tiling_arch35.cpp+369-199
@@ -30,27 +30,24 @@ constexpr uint32_t DEFAULT_OUTPUT_ROWS = 1024;
30static ge::graphStatus CheckKthValueDtypes(gert::TilingContext* context, ge::DataType dataType, uint32_t& dtypeSize)30static ge::graphStatus CheckKthValueDtypes(gert::TilingContext* context, ge::DataType dataType, uint32_t& dtypeSize)
31{31{
32 if (!ge::TypeUtils::GetDataTypeLength(dataType, dtypeSize)) {32 if (!ge::TypeUtils::GetDataTypeLength(dataType, dtypeSize)) {
33- OP_LOGE_FOR_INVALID_DTYPE(33+ OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "x", Ops::Base::ToString(dataType).c_str(),
34- context->GetNodeName(), "x", Ops::Base::ToString(dataType).c_str(),34+ "INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, FLOAT, FLOAT16 or BF16");
35- "INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, FLOAT, FLOAT16 or BF16");
36 return ge::GRAPH_FAILED;35 return ge::GRAPH_FAILED;
37 }36 }
38 auto valuesDesc = context->GetOutputDesc(0);37 auto valuesDesc = context->GetOutputDesc(0);
39 OP_CHECK_NULL_WITH_CONTEXT(context, valuesDesc);38 OP_CHECK_NULL_WITH_CONTEXT(context, valuesDesc);
40 auto indicesDesc = context->GetOutputDesc(1);39 auto indicesDesc = context->GetOutputDesc(1);
41 OP_CHECK_NULL_WITH_CONTEXT(context, indicesDesc);40 OP_CHECK_NULL_WITH_CONTEXT(context, indicesDesc);
42- OP_CHECK_IF(41+ OP_CHECK_IF(valuesDesc->GetDataType() != dataType,
43- valuesDesc->GetDataType() != dataType,42+ OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(
44- OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(43+ context->GetNodeName(), "x, values",
45- context->GetNodeName(), "x, values",44+ (Ops::Base::ToString(dataType) + ", " + Ops::Base::ToString(valuesDesc->GetDataType())).c_str(),
46- (Ops::Base::ToString(dataType) + ", " + Ops::Base::ToString(valuesDesc->GetDataType())).c_str(),45+ "The dtype of input x should be the same as output values"),
47- "The dtype of input x should be the same as output values"),46+ return ge::GRAPH_FAILED);
48- return ge::GRAPH_FAILED);47+ OP_CHECK_IF(indicesDesc->GetDataType() != ge::DT_INT64,
49- OP_CHECK_IF(48+ OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "indices",
50- indicesDesc->GetDataType() != ge::DT_INT64,49+ Ops::Base::ToString(indicesDesc->GetDataType()).c_str(), "INT64"),
51- OP_LOGE_FOR_INVALID_DTYPE(50+ return ge::GRAPH_FAILED);
52- context->GetNodeName(), "indices", Ops::Base::ToString(indicesDesc->GetDataType()).c_str(), "INT64"),
53- return ge::GRAPH_FAILED);
54 return ge::GRAPH_SUCCESS;51 return ge::GRAPH_SUCCESS;
55}52}
56 53 
@@ -62,19 +59,19 @@ static ge::graphStatus ValidateKthValueShapes(gert::TilingContext* context, cons
62 OP_CHECK_NULL_WITH_CONTEXT(context, valuesShapePtr);59 OP_CHECK_NULL_WITH_CONTEXT(context, valuesShapePtr);
63 auto indicesShapePtr = context->GetOutputShape(1);60 auto indicesShapePtr = context->GetOutputShape(1);
64 OP_CHECK_NULL_WITH_CONTEXT(context, indicesShapePtr);61 OP_CHECK_NULL_WITH_CONTEXT(context, indicesShapePtr);
65- OP_CHECK_IF(62+ OP_CHECK_IF(xShape->GetStorageShape().GetShapeSize() == 0 ||
66- xShape->GetStorageShape().GetShapeSize() == 0 || valuesShapePtr->GetStorageShape().GetShapeSize() == 0 ||63+ valuesShapePtr->GetStorageShape().GetShapeSize() == 0 ||
67- indicesShapePtr->GetStorageShape().GetShapeSize() == 0,64+ indicesShapePtr->GetStorageShape().GetShapeSize() == 0,
68- OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(65+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
69- context->GetNodeName(), "x, values, indices", "0",66+ context->GetNodeName(), "x, values, indices", "0",
70- "The shape size of input x, output values and output indices should be positive"),67+ "The shape size of input x, output values and output indices should be positive"),
71- return ge::GRAPH_FAILED);68+ return ge::GRAPH_FAILED);
72 xStorageShape = &xShape->GetStorageShape();69 xStorageShape = &xShape->GetStorageShape();
73 return ge::GRAPH_SUCCESS;70 return ge::GRAPH_SUCCESS;
74}71}
75 72 
76-static ge::graphStatus ParseKthValueShapeInfo(73+static ge::graphStatus ParseKthValueShapeInfo(gert::TilingContext* context, const int64_t* kAttr,
77- gert::TilingContext* context, const int64_t* kAttr, const int64_t* dimAttr, SortKthTileInfo& info)74+ const int64_t* dimAttr, SortKthTileInfo& info)
78{75{
79 const gert::Shape* xStorageShape = nullptr;76 const gert::Shape* xStorageShape = nullptr;
80 if (ValidateKthValueShapes(context, xStorageShape) != ge::GRAPH_SUCCESS) {77 if (ValidateKthValueShapes(context, xStorageShape) != ge::GRAPH_SUCCESS) {
@@ -82,9 +79,8 @@ static ge::graphStatus ParseKthValueShapeInfo(
82 }79 }
83 info.rank = xStorageShape->GetDimNum();80 info.rank = xStorageShape->GetDimNum();
84 if (info.rank <= 0) {81 if (info.rank <= 0) {
85- OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(82+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context->GetNodeName(), "x", (std::to_string(info.rank) + "D").c_str(),
86- context->GetNodeName(), "x", (std::to_string(info.rank) + "D").c_str(),83+ "The shape dim of input x should be greater than 0");
87- "The shape dim of input x should be greater than 0");
88 return ge::GRAPH_FAILED;84 return ge::GRAPH_FAILED;
89 }85 }
90 int64_t originSortAxis = (dimAttr == nullptr) ? -1 : *dimAttr;86 int64_t originSortAxis = (dimAttr == nullptr) ? -1 : *dimAttr;
@@ -97,9 +93,8 @@ static ge::graphStatus ParseKthValueShapeInfo(
97 }93 }
98 info.lastAxis = xStorageShape->GetDim(info.sortAxis);94 info.lastAxis = xStorageShape->GetDim(info.sortAxis);
99 if (info.lastAxis <= 0) {95 if (info.lastAxis <= 0) {
100- OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(96+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "x", std::to_string(info.lastAxis).c_str(),
101- context->GetNodeName(), "x", std::to_string(info.lastAxis).c_str(),97+ "The sort axis of input x should be greater than 0");
102- "The sort axis of input x should be greater than 0");
103 return ge::GRAPH_FAILED;98 return ge::GRAPH_FAILED;
104 }99 }
105 if (*kAttr < 1 || *kAttr > info.lastAxis) {100 if (*kAttr < 1 || *kAttr > info.lastAxis) {
@@ -115,27 +110,24 @@ static ge::graphStatus ParseKthValueShapeInfo(
115// =============================================================================110// =============================================================================
116// UB computation and base tiling init111// UB computation and base tiling init
117// =============================================================================112// =============================================================================
118-static ge::graphStatus ComputeKthValueUbInfo(113+static ge::graphStatus ComputeKthValueUbInfo(gert::TilingContext* context,
119- gert::TilingContext* context, const platform_ascendc::PlatformAscendC& ascendcPlatform, SortKthTileInfo& info,114+ const platform_ascendc::PlatformAscendC& ascendcPlatform,
120- bool& oneCoreUbValid)115+ SortKthTileInfo& info, bool& oneCoreUbValid)
121{116{
122 uint64_t ubSize64 = 0;117 uint64_t ubSize64 = 0;
123 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize64);118 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize64);
124- OP_CHECK_IF(119+ OP_CHECK_IF((ubSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())),
125- (ubSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())),120+ OP_LOGE(context->GetNodeName(), "kth_value UB size exceeds uint32 limit."), return ge::GRAPH_FAILED);
126- OP_LOGE(context->GetNodeName(), "kth_value UB size exceeds uint32 limit."), return ge::GRAPH_FAILED);
127 info.ubSize = static_cast<uint32_t>(ubSize64);121 info.ubSize = static_cast<uint32_t>(ubSize64);
128 int64_t int32Max = static_cast<int64_t>(std::numeric_limits<int32_t>::max());122 int64_t int32Max = static_cast<int64_t>(std::numeric_limits<int32_t>::max());
129 info.isInt32 = static_cast<uint32_t>(info.lastAxis <= int32Max);123 info.isInt32 = static_cast<uint32_t>(info.lastAxis <= int32Max);
130- oneCoreUbValid = ComputeRadixOneCoreUbSizes(124+ oneCoreUbValid = ComputeRadixOneCoreUbSizes(info.lastAxis, info.dtypeSize, static_cast<uint32_t>(sizeof(uint32_t)),
131- info.lastAxis, info.dtypeSize, static_cast<uint32_t>(sizeof(uint32_t)), info.blockUbSize,125+ info.blockUbSize, info.xUbSize, info.idxUbSize) &&
132- info.xUbSize, info.idxUbSize) &&126+ (info.isInt32 != 0U);
133- (info.isInt32 != 0U);
134 info.outputRowsPerLoop = static_cast<uint32_t>(std::min<int64_t>(DEFAULT_OUTPUT_ROWS, info.unsortedDim));127 info.outputRowsPerLoop = static_cast<uint32_t>(std::min<int64_t>(DEFAULT_OUTPUT_ROWS, info.unsortedDim));
135- uint32_t compactValueSize =128+ uint32_t compactValueSize = Ops::Base::CeilAlign(info.outputRowsPerLoop * info.dtypeSize, info.blockUbSize);
136- Ops::Base::CeilAlign(info.outputRowsPerLoop * info.dtypeSize, info.blockUbSize);129+ uint32_t compactIndexSize = Ops::Base::CeilAlign(info.outputRowsPerLoop * static_cast<uint32_t>(sizeof(int64_t)),
137- uint32_t compactIndexSize =130+ info.blockUbSize);
138- Ops::Base::CeilAlign(info.outputRowsPerLoop * static_cast<uint32_t>(sizeof(int64_t)), info.blockUbSize);
139 if (oneCoreUbValid) {131 if (oneCoreUbValid) {
140 // One pipeline slot owns input, sorted value/index and compact value/index buffers.132 // One pipeline slot owns input, sorted value/index and compact value/index buffers.
141 // SetRadixOneCoreTiling doubles this complete footprint when bufferNum is 2.133 // SetRadixOneCoreTiling doubles this complete footprint when bufferNum is 2.
@@ -148,8 +140,8 @@ static ge::graphStatus ComputeKthValueUbInfo(
148 return ge::GRAPH_SUCCESS;140 return ge::GRAPH_SUCCESS;
149}141}
150 142 
151-static void InitKthValueBaseTiling(143+static void InitKthValueBaseTiling(KthValueTilingData* tilingData, const SortKthTileInfo& info, bool oneCoreUbValid,
152- KthValueTilingData* tilingData, const SortKthTileInfo& info, bool oneCoreUbValid, int64_t kthIndex)144+ int64_t kthIndex)
153{145{
154 PlanToTilingData(info, tilingData);146 PlanToTilingData(info, tilingData);
155 tilingData->numTileDataSize = oneCoreUbValid ? static_cast<uint32_t>(info.lastAxis) : 0U;147 tilingData->numTileDataSize = oneCoreUbValid ? static_cast<uint32_t>(info.lastAxis) : 0U;
@@ -174,8 +166,7 @@ static void InitKthValueBaseTiling(
174// =============================================================================166// =============================================================================
175// Non-last small axis helpers167// Non-last small axis helpers
176// =============================================================================168// =============================================================================
177-static bool CheckNonLastSmallAxisInput(169+static bool CheckNonLastSmallAxisInput(int64_t axisLen, int64_t outerSize, int64_t innerSize, uint32_t& axisLen32)
178- int64_t axisLen, int64_t outerSize, int64_t innerSize, uint32_t& axisLen32)
179{170{
180 if (axisLen < 2 || axisLen > NON_LAST_SMALL_AXIS_THRESHOLD || outerSize <= 0 || innerSize <= 0) {171 if (axisLen < 2 || axisLen > NON_LAST_SMALL_AXIS_THRESHOLD || outerSize <= 0 || innerSize <= 0) {
181 return false;172 return false;
@@ -184,9 +175,9 @@ static bool CheckNonLastSmallAxisInput(
184 return true;175 return true;
185}176}
186 177 
187-static bool TryComputeKthNonLastSmallAxisLayout(178+static bool TryComputeKthNonLastSmallAxisLayout(const SortKthTileInfo& info, uint32_t innerChunk, uint32_t sortCount,
188- const SortKthTileInfo& info, uint32_t innerChunk, uint32_t sortCount, bool useMergeSort, uint32_t& inputRowBytes,179+ bool useMergeSort, uint32_t& inputRowBytes, uint32_t& valueAxisBytes,
189- uint32_t& valueAxisBytes, uint32_t& indexAxisBytes)180+ uint32_t& indexAxisBytes)
190{181{
191 uint32_t sortDtypeSize = GetNonLastSortDtypeSize(info.dtypeSize, useMergeSort, info.dataType);182 uint32_t sortDtypeSize = GetNonLastSortDtypeSize(info.dtypeSize, useMergeSort, info.dataType);
192 if (!CeilAlignUint32(static_cast<uint64_t>(innerChunk) * info.dtypeSize, info.blockUbSize, inputRowBytes) ||183 if (!CeilAlignUint32(static_cast<uint64_t>(innerChunk) * info.dtypeSize, info.blockUbSize, inputRowBytes) ||
@@ -204,9 +195,8 @@ static bool TryComputeKthNonLastSmallAxisLayout(
204 return true;195 return true;
205}196}
206 197 
207-static bool ComputeKthNonLastSmallAxisPeakUb(198+static bool ComputeKthNonLastSmallAxisPeakUb(const SortKthTileInfo& info, uint32_t innerChunk, uint32_t sortCount,
208- const SortKthTileInfo& info, uint32_t innerChunk, uint32_t sortCount, bool useMergeSort, uint64_t& peakUb,199+ bool useMergeSort, uint64_t& peakUb, NonLastSmallAxisCandidate& plan)
209- NonLastSmallAxisCandidate& plan)
210{200{
211 if (innerChunk == 0U) {201 if (innerChunk == 0U) {
212 return false;202 return false;
@@ -215,8 +205,8 @@ static bool ComputeKthNonLastSmallAxisPeakUb(
215 if (sortDtypeSize == 0U || info.dtypeSize == 0U) {205 if (sortDtypeSize == 0U || info.dtypeSize == 0U) {
216 return false;206 return false;
217 }207 }
218- if (!TryComputeKthNonLastSmallAxisLayout(208+ if (!TryComputeKthNonLastSmallAxisLayout(info, innerChunk, sortCount, useMergeSort, plan.inputRowBytes,
219- info, innerChunk, sortCount, useMergeSort, plan.inputRowBytes, plan.valueAxisBytes, plan.indexAxisBytes)) {209+ plan.valueAxisBytes, plan.indexAxisBytes)) {
220 return false;210 return false;
221 }211 }
222 uint64_t inputRowElems = static_cast<uint64_t>(plan.inputRowBytes) / info.dtypeSize;212 uint64_t inputRowElems = static_cast<uint64_t>(plan.inputRowBytes) / info.dtypeSize;
@@ -232,21 +222,21 @@ static bool ComputeKthNonLastSmallAxisPeakUb(
232 uint64_t compactCastBytes = 0;222 uint64_t compactCastBytes = 0;
233 if (useMergeSort && info.dataType == ge::DT_BF16) {223 if (useMergeSort && info.dataType == ge::DT_BF16) {
234 uint32_t inputValueAxisBytes = 0;224 uint32_t inputValueAxisBytes = 0;
235- if (!CeilAlignUint32(225+ if (!CeilAlignUint32(static_cast<uint64_t>(sortCount) * info.dtypeSize, info.blockUbSize,
236- static_cast<uint64_t>(sortCount) * info.dtypeSize, info.blockUbSize, inputValueAxisBytes)) {226+ inputValueAxisBytes)) {
237 return false;227 return false;
238 }228 }
239 inputCastBytes = static_cast<uint64_t>(innerChunk) * inputValueAxisBytes;229 inputCastBytes = static_cast<uint64_t>(innerChunk) * inputValueAxisBytes;
240- compactCastBytes =230+ compactCastBytes = Ops::Base::CeilAlign<uint64_t>(static_cast<uint64_t>(innerChunk) * sortDtypeSize,
241- Ops::Base::CeilAlign<uint64_t>(static_cast<uint64_t>(innerChunk) * sortDtypeSize, info.blockUbSize);231+ info.blockUbSize);
242 if (compactCastBytes == 0U) {232 if (compactCastBytes == 0U) {
243 return false;233 return false;
244 }234 }
245 }235 }
246- uint64_t compactValueBytes =236+ uint64_t compactValueBytes = Ops::Base::CeilAlign<uint64_t>(static_cast<uint64_t>(innerChunk) * info.dtypeSize,
247- Ops::Base::CeilAlign<uint64_t>(static_cast<uint64_t>(innerChunk) * info.dtypeSize, info.blockUbSize);237+ info.blockUbSize);
248- uint64_t compactIndexBytes =238+ uint64_t compactIndexBytes = Ops::Base::CeilAlign<uint64_t>(static_cast<uint64_t>(innerChunk) * sizeof(int64_t),
249- Ops::Base::CeilAlign<uint64_t>(static_cast<uint64_t>(innerChunk) * sizeof(int64_t), info.blockUbSize);239+ info.blockUbSize);
250 if (compactValueBytes == 0U || compactIndexBytes == 0U) {240 if (compactValueBytes == 0U || compactIndexBytes == 0U) {
251 return false;241 return false;
252 }242 }
@@ -260,12 +250,12 @@ static bool ComputeKthNonLastSmallAxisPeakUb(
260// =============================================================================250// =============================================================================
261// Individual strategy Set functions251// Individual strategy Set functions
262// =============================================================================252// =============================================================================
263-static ge::graphStatus SetRadixOneCoreTiling(253+static ge::graphStatus SetRadixOneCoreTiling(gert::TilingContext* context, const SortKthTileInfo& info,
264- gert::TilingContext* context, const SortKthTileInfo& info, KthValueTilingData* tilingData)254+ KthValueTilingData* tilingData)
265{255{
266- OP_CHECK_IF(256+ OP_CHECK_IF((info.oneBufferQueSize >= info.ubSize),
267- (info.oneBufferQueSize >= info.ubSize),257+ OP_LOGE(context->GetNodeName(), "kth_value radix one-core UB is insufficient."),
268- OP_LOGE(context->GetNodeName(), "kth_value radix one-core UB is insufficient."), return ge::GRAPH_FAILED);258+ return ge::GRAPH_FAILED);
269 tilingData->numTileDataSize = static_cast<uint32_t>(info.lastAxis);259 tilingData->numTileDataSize = static_cast<uint32_t>(info.lastAxis);
270 tilingData->lastDimTileNum = 1;260 tilingData->lastDimTileNum = 1;
271 tilingData->lastDimNeedCore = 1;261 tilingData->lastDimNeedCore = 1;
@@ -275,35 +265,34 @@ static ge::graphStatus SetRadixOneCoreTiling(
275 tilingData->keyParams3 = 1;265 tilingData->keyParams3 = 1;
276 tilingData->keyParams4 = info.outputRowsPerLoop;266 tilingData->keyParams4 = info.outputRowsPerLoop;
277 tilingData->keyParams5 = 0;267 tilingData->keyParams5 = 0;
278- OP_CHECK_IF(268+ OP_CHECK_IF(!QuerySortTmpSizeRadix(info.dataType, static_cast<uint32_t>(info.lastAxis), tilingData->tmpUbSize),
279- !QuerySortTmpSizeRadix(info.dataType, static_cast<uint32_t>(info.lastAxis), tilingData->tmpUbSize),269+ OP_LOGE(context->GetNodeName(), "kth_value get radix sort tmp size failed."), return ge::GRAPH_FAILED);
280- OP_LOGE(context->GetNodeName(), "kth_value get radix sort tmp size failed."), return ge::GRAPH_FAILED);
281 uint64_t remainUb = (info.ubSize - info.oneBufferQueSize) / info.blockUbSize * info.blockUbSize;270 uint64_t remainUb = (info.ubSize - info.oneBufferQueSize) / info.blockUbSize * info.blockUbSize;
282- OP_CHECK_IF(271+ OP_CHECK_IF((static_cast<uint64_t>(tilingData->tmpUbSize) > remainUb),
283- (static_cast<uint64_t>(tilingData->tmpUbSize) > remainUb),272+ OP_LOGE(context->GetNodeName(), "kth_value radix one-core tmp UB is insufficient."),
284- OP_LOGE(context->GetNodeName(), "kth_value radix one-core tmp UB is insufficient."), return ge::GRAPH_FAILED);273+ return ge::GRAPH_FAILED);
285- uint64_t doubleBufferRemainUb =274+ uint64_t doubleBufferRemainUb = info.ubSize > info.oneBufferQueSize * 2 ?
286- info.ubSize > info.oneBufferQueSize * 2 ?275+ (info.ubSize - info.oneBufferQueSize * 2) / info.blockUbSize *
287- (info.ubSize - info.oneBufferQueSize * 2) / info.blockUbSize * info.blockUbSize :276+ info.blockUbSize :
288- 0;277+ 0;
289 if (static_cast<uint64_t>(tilingData->tmpUbSize) <= doubleBufferRemainUb) {278 if (static_cast<uint64_t>(tilingData->tmpUbSize) <= doubleBufferRemainUb) {
290 tilingData->keyParams3 = 2;279 tilingData->keyParams3 = 2;
291 }280 }
292 return ge::GRAPH_SUCCESS;281 return ge::GRAPH_SUCCESS;
293}282}
294 283 
295-static bool ComputeKthValueRadixMoreCoreWorkspace(284+static bool ComputeKthValueRadixMoreCoreWorkspace(int64_t axisLen, uint32_t dtypeSize, uint32_t indexSize,
296- int64_t axisLen, uint32_t dtypeSize, uint32_t indexSize, uint32_t unsortedDimParallel, uint32_t blockUbSize,285+ uint32_t unsortedDimParallel, uint32_t blockUbSize,
297- uint64_t sortWorkspaceSize, uint64_t& workspaceSize)286+ uint64_t sortWorkspaceSize, uint64_t& workspaceSize)
298{287{
299 uint64_t axisLen64 = static_cast<uint64_t>(axisLen);288 uint64_t axisLen64 = static_cast<uint64_t>(axisLen);
300 uint64_t unsortedDimParallel64 = static_cast<uint64_t>(unsortedDimParallel);289 uint64_t unsortedDimParallel64 = static_cast<uint64_t>(unsortedDimParallel);
301 uint64_t blockUbSize64 = static_cast<uint64_t>(blockUbSize);290 uint64_t blockUbSize64 = static_cast<uint64_t>(blockUbSize);
302 291 
303- uint64_t valueWorkspace =292+ uint64_t valueWorkspace = Ops::Base::CeilAlign(axisLen64 * unsortedDimParallel64 * static_cast<uint64_t>(dtypeSize),
304- Ops::Base::CeilAlign(axisLen64 * unsortedDimParallel64 * static_cast<uint64_t>(dtypeSize), blockUbSize64);293+ blockUbSize64);
305- uint64_t indexWorkspace =294+ uint64_t indexWorkspace = Ops::Base::CeilAlign(axisLen64 * unsortedDimParallel64 * static_cast<uint64_t>(indexSize),
306- Ops::Base::CeilAlign(axisLen64 * unsortedDimParallel64 * static_cast<uint64_t>(indexSize), blockUbSize64);295+ blockUbSize64);
307 296 
308 workspaceSize = valueWorkspace + indexWorkspace + sortWorkspaceSize;297 workspaceSize = valueWorkspace + indexWorkspace + sortWorkspaceSize;
309 return true;298 return true;
@@ -311,18 +300,17 @@ static bool ComputeKthValueRadixMoreCoreWorkspace(
311 300 
312static ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t& blockDim)301static ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t& blockDim)
313{302{
314- OP_CHECK_IF(303+ OP_CHECK_IF(!FillRadixMoreCoreInfo(info), OP_LOGE(context->GetNodeName(), "kth_value radix more-core plan failed."),
315- !FillRadixMoreCoreInfo(info), OP_LOGE(context->GetNodeName(), "kth_value radix more-core plan failed."),304+ return ge::GRAPH_FAILED);
316- return ge::GRAPH_FAILED);
317 blockDim = info.coreNumNeed;305 blockDim = info.coreNumNeed;
318- uint32_t indexSize =306+ uint32_t indexSize = info.isInt32 != 0 ? static_cast<uint32_t>(sizeof(int32_t)) :
319- info.isInt32 != 0 ? static_cast<uint32_t>(sizeof(int32_t)) : static_cast<uint32_t>(sizeof(int64_t));307+ static_cast<uint32_t>(sizeof(int64_t));
320 uint64_t totalWorkspace = 0;308 uint64_t totalWorkspace = 0;
321- OP_CHECK_IF(309+ OP_CHECK_IF(!ComputeKthValueRadixMoreCoreWorkspace(info.lastAxis, info.dtypeSize, indexSize,
322- !ComputeKthValueRadixMoreCoreWorkspace(310+ info.unsortedDimParallel, info.blockUbSize,
323- info.lastAxis, info.dtypeSize, indexSize, info.unsortedDimParallel, info.blockUbSize,311+ static_cast<uint64_t>(info.workspaceSize), totalWorkspace),
324- static_cast<uint64_t>(info.workspaceSize), totalWorkspace),312+ OP_LOGE(context->GetNodeName(), "kth_value radix more-core workspace overflow."),
325- OP_LOGE(context->GetNodeName(), "kth_value radix more-core workspace overflow."), return ge::GRAPH_FAILED);313+ return ge::GRAPH_FAILED);
326 size_t* userWorkspaceSize = context->GetWorkspaceSizes(1);314 size_t* userWorkspaceSize = context->GetWorkspaceSizes(1);
327 OP_CHECK_NULL_WITH_CONTEXT(context, userWorkspaceSize);315 OP_CHECK_NULL_WITH_CONTEXT(context, userWorkspaceSize);
328 userWorkspaceSize[0] = static_cast<size_t>(totalWorkspace);316 userWorkspaceSize[0] = static_cast<size_t>(totalWorkspace);
@@ -330,17 +318,16 @@ static ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext* context, Sort
330 return ge::GRAPH_SUCCESS;318 return ge::GRAPH_SUCCESS;
331}319}
332 320 
333-static ge::graphStatus SetKthValueMergeSortTiling(321+static ge::graphStatus SetKthValueMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info,
334- gert::TilingContext* context, SortKthTileInfo& info, uint32_t& blockDim, uint64_t& schId)322+ uint32_t& blockDim, uint64_t& schId)
335{323{
336- OP_CHECK_IF(324+ OP_CHECK_IF(!ComputeMergeSortTiling(context, info, static_cast<uint32_t>(sizeof(uint32_t))),
337- !ComputeMergeSortTiling(context, info, static_cast<uint32_t>(sizeof(uint32_t))),325+ OP_LOGE(context->GetNodeName(), "kth_value merge sort tiling failed."), return ge::GRAPH_FAILED);
338- OP_LOGE(context->GetNodeName(), "kth_value merge sort tiling failed."), return ge::GRAPH_FAILED);
339 blockDim = info.coreNumNeed;326 blockDim = info.coreNumNeed;
340- schId =327+ schId = info.lastAxis <= SORT32_SMALL_AXIS_THRESHOLD ? KTH_VALUE_SCHID_SORT32_SMALL_AXIS :
341- info.lastAxis <= SORT32_SMALL_AXIS_THRESHOLD ? KTH_VALUE_SCHID_SORT32_SMALL_AXIS : KTH_VALUE_SCHID_MERGE_SORT;328+ KTH_VALUE_SCHID_MERGE_SORT;
342- OP_LOGI("KthValueMergeSortTiling", "axis=%ld, unsortedDim=%ld, coreNumNeed=%u, schId=%lu",329+ OP_LOGI("KthValueMergeSortTiling", "axis=%ld, unsortedDim=%ld, coreNumNeed=%u, schId=%lu", info.lastAxis,
343- info.lastAxis, info.unsortedDim, info.coreNumNeed, schId);330+ info.unsortedDim, info.coreNumNeed, schId);
344 return ge::GRAPH_SUCCESS;331 return ge::GRAPH_SUCCESS;
345}332}
346 333 
@@ -349,9 +336,8 @@ static ge::graphStatus SetMergeMoreCoreTiling(gert::TilingContext* context, Sort
349 constexpr uint32_t mergeBytesPerElem = MERGE_SORT_LIST_NUM * MERGE_SORT_DATA_BYTES * 2 +336 constexpr uint32_t mergeBytesPerElem = MERGE_SORT_LIST_NUM * MERGE_SORT_DATA_BYTES * 2 +
350 MERGE_SORT_LIST_NUM * sizeof(uint32_t) +337 MERGE_SORT_LIST_NUM * sizeof(uint32_t) +
351 MERGE_SORT_LIST_NUM * sizeof(int64_t) + MERGE_SORT_LIST_NUM * sizeof(float);338 MERGE_SORT_LIST_NUM * sizeof(int64_t) + MERGE_SORT_LIST_NUM * sizeof(float);
352- OP_CHECK_IF(339+ OP_CHECK_IF(!ComputeMergeMoreCoreTiling(context, info, mergeBytesPerElem),
353- !ComputeMergeMoreCoreTiling(context, info, mergeBytesPerElem),340+ OP_LOGE(context->GetNodeName(), "kth_value merge more-core plan failed."), return ge::GRAPH_FAILED);
354- OP_LOGE(context->GetNodeName(), "kth_value merge more-core plan failed."), return ge::GRAPH_FAILED);
355 blockDim = info.coreNumNeed;341 blockDim = info.coreNumNeed;
356 OP_LOGI("KthValueMergeMoreCoreTiling", "maxDealingNum: %u", info.keyParams0);342 OP_LOGI("KthValueMergeMoreCoreTiling", "maxDealingNum: %u", info.keyParams0);
357 return ge::GRAPH_SUCCESS;343 return ge::GRAPH_SUCCESS;
@@ -359,24 +345,21 @@ static ge::graphStatus SetMergeMoreCoreTiling(gert::TilingContext* context, Sort
359 345 
360static ge::graphStatus SetMergeIntraCoreTiling(gert::TilingContext* context, SortKthTileInfo& info)346static ge::graphStatus SetMergeIntraCoreTiling(gert::TilingContext* context, SortKthTileInfo& info)
361{347{
362- OP_CHECK_IF(348+ OP_CHECK_IF(!ComputeMergeIntraCoreTiling(context, info),
363- !ComputeMergeIntraCoreTiling(context, info),349+ OP_LOGE(context->GetNodeName(), "kth_value merge intra-core plan failed."), return ge::GRAPH_FAILED);
364- OP_LOGE(context->GetNodeName(), "kth_value merge intra-core plan failed."), return ge::GRAPH_FAILED);
365 OP_LOGI("KthValueMergeIntraCoreTiling",350 OP_LOGI("KthValueMergeIntraCoreTiling",
366- "B %ld, N %ld, batchPerCore %u, actualCoreNum %u, blockSortSize %u, extractChunkSize %u, "351+ "B %ld, N %ld, batchPerCore %u, actualCoreNum %u, blockSortSize %u, extractChunkSize %u, "
367- "blocksPerRow %u, alignNum %u, ubSize %u",352+ "blocksPerRow %u, alignNum %u, ubSize %u",
368- info.unsortedDim, info.lastAxis, info.keyParams0, info.coreNumNeed,353+ info.unsortedDim, info.lastAxis, info.keyParams0, info.coreNumNeed, info.numTileDataSize, info.keyParams4,
369- info.numTileDataSize, info.keyParams4, info.lastDimTileNum, info.keyParams3, info.ubSize);354+ info.lastDimTileNum, info.keyParams3, info.ubSize);
370 return ge::GRAPH_SUCCESS;355 return ge::GRAPH_SUCCESS;
371}356}
372 357 
373-static bool SetNonLastSmallAxisTiling(358+static bool SetNonLastSmallAxisTiling(gert::TilingContext* context, const SortKthTileInfo& constInfo,
374- gert::TilingContext* context, const SortKthTileInfo& constInfo, KthValueTilingData* tilingData, uint32_t& blockDim,359+ KthValueTilingData* tilingData, uint32_t& blockDim, uint64_t& schId)
375- uint64_t& schId)
376{360{
377 uint32_t axisLen32 = 0;361 uint32_t axisLen32 = 0;
378- if (!CheckNonLastSmallAxisInput(362+ if (!CheckNonLastSmallAxisInput(constInfo.lastAxis, constInfo.outerSize, constInfo.innerSize, axisLen32)) {
379- constInfo.lastAxis, constInfo.outerSize, constInfo.innerSize, axisLen32)) {
380 return false;363 return false;
381 }364 }
382 uint64_t usableUb = ComputeUbAfterSimtReserve(constInfo.ubSize);365 uint64_t usableUb = ComputeUbAfterSimtReserve(constInfo.ubSize);
@@ -387,9 +370,8 @@ static bool SetNonLastSmallAxisTiling(
387 SortKthTileInfo info = constInfo;370 SortKthTileInfo info = constInfo;
388 info.tmpUbSize = tmpUbSize;371 info.tmpUbSize = tmpUbSize;
389 NonLastSmallAxisCandidate best;372 NonLastSmallAxisCandidate best;
390- auto estimateUb = [sortCount, useMergeSort](373+ auto estimateUb = [sortCount, useMergeSort](SortKthTileInfo& candidateInfo, uint32_t innerChunk, uint64_t& peakUb,
391- SortKthTileInfo& candidateInfo, uint32_t innerChunk, uint64_t& peakUb,374+ NonLastSmallAxisCandidate& candidate) -> bool {
392- NonLastSmallAxisCandidate& candidate) -> bool {
393 return ComputeKthNonLastSmallAxisPeakUb(candidateInfo, innerChunk, sortCount, useMergeSort, peakUb, candidate);375 return ComputeKthNonLastSmallAxisPeakUb(candidateInfo, innerChunk, sortCount, useMergeSort, peakUb, candidate);
394 };376 };
395 if (!SearchNonLastSmallAxisPlan(info, usableUb, estimateUb, best)) {377 if (!SearchNonLastSmallAxisPlan(info, usableUb, estimateUb, best)) {
@@ -399,8 +381,8 @@ static bool SetNonLastSmallAxisTiling(
399 int64_t unsortedDim = constInfo.outerSize * constInfo.innerSize;381 int64_t unsortedDim = constInfo.outerSize * constInfo.innerSize;
400 uint32_t inputValueAxisBytes = 0;382 uint32_t inputValueAxisBytes = 0;
401 if (useMergeSort && constInfo.dataType == ge::DT_BF16 &&383 if (useMergeSort && constInfo.dataType == ge::DT_BF16 &&
402- !CeilAlignUint32(384+ !CeilAlignUint32(static_cast<uint64_t>(sortCount) * constInfo.dtypeSize, constInfo.blockUbSize,
403- static_cast<uint64_t>(sortCount) * constInfo.dtypeSize, constInfo.blockUbSize, inputValueAxisBytes)) {385+ inputValueAxisBytes)) {
404 return false;386 return false;
405 }387 }
406 tilingData->lastAxisNum = constInfo.lastAxis;388 tilingData->lastAxisNum = constInfo.lastAxis;
@@ -427,8 +409,8 @@ static bool SetNonLastSmallAxisTiling(
427// =============================================================================409// =============================================================================
428// Fill functions410// Fill functions
429// =============================================================================411// =============================================================================
430-static void FillSmallAxisTiling(412+static void FillSmallAxisTiling(KthValueTilingData* tilingData, const SmallAxisRoutePlan& plan, uint32_t axisLen,
431- KthValueTilingData* tilingData, const SmallAxisRoutePlan& plan, uint32_t axisLen, uint32_t& blockDim)413+ uint32_t& blockDim)
432{414{
433 tilingData->numTileDataSize = axisLen;415 tilingData->numTileDataSize = axisLen;
434 tilingData->keyParams0 = plan.batchSize;416 tilingData->keyParams0 = plan.batchSize;
@@ -445,8 +427,8 @@ static void FillSmallAxisTiling(
445 blockDim = plan.blockDim;427 blockDim = plan.blockDim;
446}428}
447 429 
448-static bool FillNonLastSmallAxisTiling(430+static bool FillNonLastSmallAxisTiling(KthValueTilingData* tilingData, const SmallAxisRoutePlan& plan,
449- KthValueTilingData* tilingData, const SmallAxisRoutePlan& plan, const SortKthTileInfo& info, uint32_t& blockDim)431+ const SortKthTileInfo& info, uint32_t& blockDim)
450{432{
451 uint32_t innerChunk = static_cast<uint32_t>(std::min<int64_t>(plan.batchSize, info.innerSize));433 uint32_t innerChunk = static_cast<uint32_t>(std::min<int64_t>(plan.batchSize, info.innerSize));
452 if (innerChunk == 0U) {434 if (innerChunk == 0U) {
@@ -483,10 +465,10 @@ static bool FillNonLastSmallAxisTiling(
483// =============================================================================465// =============================================================================
484// Axis-one-copy tiling466// Axis-one-copy tiling
485// =============================================================================467// =============================================================================
486-static ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext *context, SortKthTileInfo &info)468+static ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext* context, SortKthTileInfo& info)
487{469{
488 uint64_t bytesPerElem = static_cast<uint64_t>(2) *470 uint64_t bytesPerElem = static_cast<uint64_t>(2) *
489- (static_cast<uint64_t>(info.dtypeSize) + static_cast<uint64_t>(sizeof(int64_t)));471+ (static_cast<uint64_t>(info.dtypeSize) + static_cast<uint64_t>(sizeof(int64_t)));
490 if (bytesPerElem == 0) {472 if (bytesPerElem == 0) {
491 OP_LOGE(context->GetNodeName(), "bytesPerElem is 0, invalid dtype configuration");473 OP_LOGE(context->GetNodeName(), "bytesPerElem is 0, invalid dtype configuration");
492 return ge::GRAPH_FAILED;474 return ge::GRAPH_FAILED;
@@ -501,8 +483,7 @@ static ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext *context, SortKt
501 return ge::GRAPH_FAILED;483 return ge::GRAPH_FAILED;
502 }484 }
503 uint32_t copyElemsPerLoop = static_cast<uint32_t>(copyElemsPerLoop64);485 uint32_t copyElemsPerLoop = static_cast<uint32_t>(copyElemsPerLoop64);
504- uint64_t totalElems = static_cast<uint64_t>(info.unsortedDim) *486+ uint64_t totalElems = static_cast<uint64_t>(info.unsortedDim) * static_cast<uint64_t>(info.lastAxis);
505- static_cast<uint64_t>(info.lastAxis);
506 uint64_t loopTimes64 = (totalElems + copyElemsPerLoop64 - 1) / copyElemsPerLoop64;487 uint64_t loopTimes64 = (totalElems + copyElemsPerLoop64 - 1) / copyElemsPerLoop64;
507 if (loopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {488 if (loopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {
508 OP_LOGE(context->GetNodeName(), "loopTimes exceeds uint32_t limit");489 OP_LOGE(context->GetNodeName(), "loopTimes exceeds uint32_t limit");
@@ -521,18 +502,18 @@ static ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext *context, SortKt
521 info.sortLoopTimes = Ops::Base::CeilDiv(static_cast<int64_t>(loopTimes), static_cast<int64_t>(coreNumNeed));502 info.sortLoopTimes = Ops::Base::CeilDiv(static_cast<int64_t>(loopTimes), static_cast<int64_t>(coreNumNeed));
522 info.tmpUbSize = 0;503 info.tmpUbSize = 0;
523 504 
524- size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);505+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
525 userWorkSpaceSize[0] = WORK_SPACE_SIZE;506 userWorkSpaceSize[0] = WORK_SPACE_SIZE;
526- OP_LOGI("AxisOneCopyTiling", "totalElems %lu, copyElemsPerLoop %u, loopTimes %u, coreNumNeed %u",507+ OP_LOGI("AxisOneCopyTiling", "totalElems %lu, copyElemsPerLoop %u, loopTimes %u, coreNumNeed %u", totalElems,
527- totalElems, info.keyParams0, info.keyParams1, coreNumNeed);508+ info.keyParams0, info.keyParams1, coreNumNeed);
528 return ge::GRAPH_SUCCESS;509 return ge::GRAPH_SUCCESS;
529}510}
530 511 
531// =============================================================================512// =============================================================================
532// Try functions513// Try functions
533// =============================================================================514// =============================================================================
534-static bool TryRadixOneCore(515+static bool TryRadixOneCore(gert::TilingContext* context, const SortKthTileInfo& info, KthValueTilingData* tilingData,
535- gert::TilingContext* context, const SortKthTileInfo& info, KthValueTilingData* tilingData, uint64_t& schId)516+ uint64_t& schId)
536{517{
537 if (info.oneBufferQueSize >= info.ubSize) {518 if (info.oneBufferQueSize >= info.ubSize) {
538 return false;519 return false;
@@ -550,9 +531,207 @@ static bool TryRadixOneCore(
550 return true;531 return true;
551}532}
552 533 
553-static bool TrySmallAxis(534+static bool IsRadixSelectProfitable(const SortKthTileInfo& info, int64_t kthIndex)
554- gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, uint32_t& blockDim,535+{
555- uint64_t& schId)536+ if (kthIndex < 0 || info.lastAxis <= 0 || info.unsortedDim <= 0) {
537+ return false;
538+ }
539+ uint64_t axisLen = static_cast<uint64_t>(info.lastAxis);
540+ uint64_t kth = static_cast<uint64_t>(kthIndex);
541+ if (kth >= axisLen) {
542+ return false;
543+ }
544+ 
545+ constexpr uint64_t radixSelectMinAxis = 65536UL;
546+ constexpr uint64_t radixSelectHugeAxis = 1000000UL;
547+ constexpr uint64_t radixSelectInt64SmallAxis = 4096UL;
548+ constexpr int64_t radixSelectInt32InteriorMinRows = 16;
549+ constexpr int64_t radixSelectInt64InteriorMinRows = 8;
550+ constexpr int64_t radixSelectInt64SmallAxisMinRows = 32;
551+ constexpr int64_t radixSelectFp16InteriorMinRows = 64;
552+ bool isNearHead = kth <= 1U;
553+ bool isTail = (kth + 1U == axisLen);
554+ bool isBeforeTail = (kth + 2U == axisLen);
555+ bool isInterior = !isNearHead && !isTail && !isBeforeTail;
556+ bool isHugeAxis = axisLen >= radixSelectHugeAxis;
557+ bool isLargeAxisInterior = axisLen >= radixSelectMinAxis && isInterior;
558+ bool isInt64SmallAxisInterior = axisLen >= radixSelectInt64SmallAxis && isInterior;
559+ bool hasEnoughRowsForInt32Interior = info.unsortedDim >= radixSelectInt32InteriorMinRows;
560+ bool hasEnoughRowsForInt64Interior = info.unsortedDim >= radixSelectInt64InteriorMinRows;
561+ bool hasEnoughRowsForInt64SmallAxis = info.unsortedDim >= radixSelectInt64SmallAxisMinRows;
562+ bool hasEnoughRowsForFp16Interior = info.unsortedDim >= radixSelectFp16InteriorMinRows;
563+ 
564+ // RadixSelect pays fixed workspace/reduce cost. Huge axes and near-head k amortize it directly; middle-k cases
565+ // need enough independent rows, otherwise the radix-more-core fallback is faster.
566+ switch (info.dataType) {
567+ case ge::DT_INT32:
568+ case ge::DT_UINT32:
569+ return isHugeAxis || (axisLen >= radixSelectMinAxis && isNearHead) ||
570+ (isLargeAxisInterior && hasEnoughRowsForInt32Interior);
571+ case ge::DT_INT64:
572+ case ge::DT_UINT64:
573+ return isHugeAxis || (isLargeAxisInterior && hasEnoughRowsForInt64Interior) ||
574+ (isInt64SmallAxisInterior && hasEnoughRowsForInt64SmallAxis);
575+ case ge::DT_FLOAT:
576+ return isHugeAxis || (axisLen >= radixSelectMinAxis && isNearHead);
577+ case ge::DT_FLOAT16:
578+ case ge::DT_BF16:
579+ return isInterior && (isHugeAxis || (isLargeAxisInterior && hasEnoughRowsForFp16Interior));
580+ default:
581+ return false;
582+ }
583+}
584+ 
585+struct RadixSelectPlan {
586+ uint64_t tileElems = 0;
587+ uint64_t tileCount = 0;
588+ uint32_t rowsParallel = 0;
589+ uint32_t coresPerRow = 1;
590+ uint32_t blockDim = 0;
591+ uint64_t workspace = 0;
592+};
593+ 
594+static uint64_t GetRadixSelectFixedBytes()
595+{
596+ constexpr uint64_t radixBuckets = 256UL;
597+ constexpr uint64_t radixSelectFindThreads = 128UL;
598+ constexpr uint64_t radixSelectResultWords = 8UL;
599+ constexpr uint64_t radixSelectActiveIndexCap = 4096UL;
600+ constexpr uint64_t radixSelectReserveAlign = 1024UL;
601+ constexpr uint64_t histogramBytes = radixBuckets * sizeof(uint64_t);
602+ constexpr uint64_t reservedRawBytes = radixBuckets * sizeof(uint16_t) + radixSelectFindThreads * sizeof(uint32_t) +
603+ radixSelectResultWords * sizeof(uint64_t) + 2UL * 32UL +
604+ radixBuckets * sizeof(uint64_t) +
605+ radixSelectActiveIndexCap * sizeof(uint32_t);
606+ return histogramBytes +
607+ ((reservedRawBytes + radixSelectReserveAlign - 1UL) / radixSelectReserveAlign) * radixSelectReserveAlign;
608+}
609+ 
610+static bool ComputeRadixSelectTileElems(const SortKthTileInfo& info, uint64_t& tileElems)
611+{
612+ constexpr uint64_t radixSelectMaxTileElems = 32768UL;
613+ constexpr uint64_t radixSelectMinAlignElems = 256UL;
614+ uint64_t usableUb = ComputeUbAfterSimtReserve(info.ubSize);
615+ uint64_t bytesPerElem = static_cast<uint64_t>(info.dtypeSize) * 2UL;
616+ uint64_t fixedBytes = GetRadixSelectFixedBytes();
617+ if (bytesPerElem == 0U || usableUb <= fixedBytes) {
618+ return false;
619+ }
620+ tileElems = std::min<uint64_t>((usableUb - fixedBytes) / bytesPerElem, radixSelectMaxTileElems);
621+ uint64_t alignElems = std::max<uint64_t>(radixSelectMinAlignElems, info.blockUbSize / info.dtypeSize);
622+ tileElems = tileElems / alignElems * alignElems;
623+ tileElems = std::min<uint64_t>(tileElems, static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()));
624+ return tileElems != 0U;
625+}
626+ 
627+static bool ComputeRadixSelectPlan(const SortKthTileInfo& info, RadixSelectPlan& plan)
628+{
629+ if (!ComputeRadixSelectTileElems(info, plan.tileElems)) {
630+ return false;
631+ }
632+ plan.rowsParallel = static_cast<uint32_t>(
633+ std::min<int64_t>(static_cast<int64_t>(info.maxCoreNum), info.unsortedDim));
634+ if (plan.rowsParallel == 0U) {
635+ return false;
636+ }
637+ plan.tileCount = Ops::Base::CeilDiv(static_cast<uint64_t>(info.lastAxis), plan.tileElems);
638+ uint32_t maxCoresPerRow = std::max<uint32_t>(1U, info.maxCoreNum / plan.rowsParallel);
639+ if (plan.tileCount > 1UL) {
640+ plan.coresPerRow = static_cast<uint32_t>(std::min<uint64_t>(maxCoresPerRow, plan.tileCount));
641+ }
642+ uint64_t blockDim = static_cast<uint64_t>(plan.rowsParallel) * plan.coresPerRow;
643+ if (blockDim > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {
644+ return false;
645+ }
646+ plan.blockDim = static_cast<uint32_t>(blockDim);
647+ return true;
648+}
649+ 
650+static bool ComputeRadixSelectWorkspace(const SortKthTileInfo& info, RadixSelectPlan& plan)
651+{
652+ constexpr uint64_t radixBuckets = 256UL;
653+ constexpr uint64_t radixSelectResultWords = 8UL;
654+ uint64_t histogramWorkspace = static_cast<uint64_t>(plan.blockDim) * radixBuckets * sizeof(uint64_t);
655+ uint64_t groupStateWorkspace = static_cast<uint64_t>(plan.rowsParallel) * radixSelectResultWords * sizeof(uint64_t);
656+ if (histogramWorkspace > static_cast<uint64_t>(std::numeric_limits<size_t>::max()) - groupStateWorkspace) {
657+ return false;
658+ }
659+ uint64_t workspaceRaw = histogramWorkspace + groupStateWorkspace;
660+ plan.workspace = Ops::Base::CeilAlign(workspaceRaw, static_cast<uint64_t>(info.blockUbSize));
661+ return plan.workspace >= workspaceRaw &&
662+ plan.workspace <= static_cast<uint64_t>(std::numeric_limits<size_t>::max()) - WORK_SPACE_SIZE;
663+}
664+ 
665+static bool TryRadixSelect(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData,
666+ uint32_t& blockDim, uint64_t& schId)
667+{
668+ if (!IsRadixSelectProfitable(info, tilingData->kthIndex) || info.isNonLastAxis ||
669+ info.lastAxis <= static_cast<int64_t>(SMALL_AXIS_THRESHOLD) || info.maxCoreNum == 0U) {
670+ return false;
671+ }
672+ RadixSelectPlan plan;
673+ if (!ComputeRadixSelectPlan(info, plan) || !ComputeRadixSelectWorkspace(info, plan)) {
674+ return false;
675+ }
676+ size_t* userWorkspaceSize = context->GetWorkspaceSizes(1);
677+ if (userWorkspaceSize == nullptr) {
678+ return false;
679+ }
680+ 
681+ SortKthTileInfo candidate = info;
682+ candidate.numTileDataSize = static_cast<uint32_t>(plan.tileElems);
683+ candidate.unsortedDimParallel = plan.rowsParallel;
684+ candidate.lastDimTileNum = static_cast<uint32_t>(
685+ std::min<uint64_t>(plan.tileCount, static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())));
686+ candidate.lastDimNeedCore = plan.coresPerRow;
687+ uint64_t rowLoops = Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim),
688+ static_cast<uint64_t>(plan.rowsParallel));
689+ candidate.sortLoopTimes = static_cast<uint32_t>(
690+ std::min<uint64_t>(rowLoops, static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())));
691+ candidate.tmpUbSize = 0U;
692+ info = candidate;
693+ blockDim = plan.blockDim;
694+ PlanToTilingData(info, tilingData);
695+ schId = KTH_VALUE_SCHID_RADIX_SELECT;
696+ userWorkspaceSize[0] = static_cast<size_t>(WORK_SPACE_SIZE + plan.workspace);
697+ context->SetScheduleMode(1);
698+ return true;
699+}
700+ 
701+static bool TrySmallAxisShortRankSelect(const SortKthTileInfo& info, KthValueTilingData* tilingData, uint32_t& blockDim,
702+ uint64_t& schId)
703+{
704+ constexpr int64_t shortRankSelectMaxAxis = 32;
705+ constexpr uint64_t shortRankSelectMaxRank = 8;
706+ if (info.isNonLastAxis || info.lastAxis <= 1 || info.lastAxis > shortRankSelectMaxAxis || info.unsortedDim <= 0 ||
707+ info.maxCoreNum == 0U || tilingData->kthIndex < 0 || tilingData->kthIndex >= info.lastAxis) {
708+ return false;
709+ }
710+ if (info.dataType != ge::DT_INT64 && info.dataType != ge::DT_UINT64) {
711+ return false;
712+ }
713+ 
714+ uint64_t kthIndex = static_cast<uint64_t>(tilingData->kthIndex);
715+ uint64_t axisLen = static_cast<uint64_t>(info.lastAxis);
716+ uint64_t shortRank = std::min(kthIndex + 1U, axisLen - kthIndex);
717+ if (shortRank == 0U || shortRank > shortRankSelectMaxRank ||
718+ static_cast<uint64_t>(info.unsortedDim) < static_cast<uint64_t>(info.maxCoreNum)) {
719+ return false;
720+ }
721+ 
722+ SmallAxisRoutePlan plan;
723+ if (!SelectSmallAxisRoute(info, plan) || plan.kind != SmallAxisRouteKind::TWO_STAGE) {
724+ return false;
725+ }
726+ FillSmallAxisTiling(tilingData, plan, static_cast<uint32_t>(info.lastAxis), blockDim);
727+ tilingData->keyParams2 = static_cast<uint32_t>(shortRank);
728+ tilingData->tmpUbSize = 0;
729+ schId = KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT;
730+ return true;
731+}
732+ 
733+static bool TrySmallAxis(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData,
734+ uint32_t& blockDim, uint64_t& schId)
556{735{
557 if (info.lastAxis > static_cast<int64_t>(SMALL_AXIS_THRESHOLD)) {736 if (info.lastAxis > static_cast<int64_t>(SMALL_AXIS_THRESHOLD)) {
558 return false;737 return false;
@@ -567,9 +746,7 @@ static bool TrySmallAxis(
567 return true;746 return true;
568 }747 }
569 SmallAxisRoutePlan plan;748 SmallAxisRoutePlan plan;
570- bool selected = info.isNonLastAxis ?749+ bool selected = info.isNonLastAxis ? SelectNonLastSmallAxisRoute(info, plan) : SelectSmallAxisRoute(info, plan);
571- SelectNonLastSmallAxisRoute(info, plan) :
572- SelectSmallAxisRoute(info, plan);
573 if (!selected) {750 if (!selected) {
574 return false;751 return false;
575 }752 }
@@ -581,13 +758,12 @@ static bool TrySmallAxis(
581 FillSmallAxisTiling(tilingData, plan, static_cast<uint32_t>(info.lastAxis), blockDim);758 FillSmallAxisTiling(tilingData, plan, static_cast<uint32_t>(info.lastAxis), blockDim);
582 }759 }
583 schId = plan.kind == SmallAxisRouteKind::TWO_STAGE ? KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE :760 schId = plan.kind == SmallAxisRouteKind::TWO_STAGE ? KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE :
584- KTH_VALUE_SCHID_SMALL_AXIS_INSERTION;761+ KTH_VALUE_SCHID_SMALL_AXIS_INSERTION;
585 return true;762 return true;
586}763}
587 764 
588-static bool TryMerge(765+static bool TryMerge(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData,
589- gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, uint32_t& blockDim,766+ uint32_t& blockDim, uint64_t& schId)
590- uint64_t& schId)
591{767{
592 if (IsMergeSortSupported(info.dataType, info.lastAxis)) {768 if (IsMergeSortSupported(info.dataType, info.lastAxis)) {
593 SortKthTileInfo candidate = info;769 SortKthTileInfo candidate = info;
@@ -609,9 +785,8 @@ static bool TryMerge(
609 return false;785 return false;
610}786}
611 787 
612-static bool TryMergeIntraCore(788+static bool TryMergeIntraCore(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData,
613- gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, uint32_t& blockDim,789+ uint32_t& blockDim, uint64_t& schId)
614- uint64_t& schId)
615{790{
616 if (!IsMergeIntraCoreSupported(info.dataType, info.lastAxis, info.unsortedDim, info.maxCoreNum, info.ubSize)) {791 if (!IsMergeIntraCoreSupported(info.dataType, info.lastAxis, info.unsortedDim, info.maxCoreNum, info.ubSize)) {
617 return false;792 return false;
@@ -627,9 +802,8 @@ static bool TryMergeIntraCore(
627 return true;802 return true;
628}803}
629 804 
630-static bool TryNonLastSmallAxis(805+static bool TryNonLastSmallAxis(gert::TilingContext* context, const SortKthTileInfo& info,
631- gert::TilingContext* context, const SortKthTileInfo& info, KthValueTilingData* tilingData, uint32_t& blockDim,806+ KthValueTilingData* tilingData, uint32_t& blockDim, uint64_t& schId)
632- uint64_t& schId)
633{807{
634 if (!info.isNonLastAxis) {808 if (!info.isNonLastAxis) {
635 return false;809 return false;
@@ -640,10 +814,12 @@ static bool TryNonLastSmallAxis(
640// =============================================================================814// =============================================================================
641// Route selection and finalization815// Route selection and finalization
642// =============================================================================816// =============================================================================
643-static ge::graphStatus SelectKthValueRoute(817+static ge::graphStatus SelectKthValueRoute(gert::TilingContext* context, SortKthTileInfo& info,
644- gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, uint32_t& blockDim,818+ KthValueTilingData* tilingData, uint32_t& blockDim, uint64_t& schId)
645- uint64_t& schId)
646{819{
820+ if (TrySmallAxisShortRankSelect(info, tilingData, blockDim, schId)) {
821+ return ge::GRAPH_SUCCESS;
822+ }
647 if (TrySmallAxis(context, info, tilingData, blockDim, schId)) {823 if (TrySmallAxis(context, info, tilingData, blockDim, schId)) {
648 return ge::GRAPH_SUCCESS;824 return ge::GRAPH_SUCCESS;
649 }825 }
@@ -654,22 +830,22 @@ static ge::graphStatus SelectKthValueRoute(
654 OP_LOGE(context->GetNodeName(), "non-last kth_value axis does not meet no-transpose schedule constraints");830 OP_LOGE(context->GetNodeName(), "non-last kth_value axis does not meet no-transpose schedule constraints");
655 return ge::GRAPH_FAILED;831 return ge::GRAPH_FAILED;
656 }832 }
657- if (TryMerge(context, info, tilingData, blockDim, schId) ||833+ if (TryMerge(context, info, tilingData, blockDim, schId) || TryRadixOneCore(context, info, tilingData, schId) ||
658- TryRadixOneCore(context, info, tilingData, schId) ||834+ TryRadixSelect(context, info, tilingData, blockDim, schId) ||
659 TryMergeIntraCore(context, info, tilingData, blockDim, schId)) {835 TryMergeIntraCore(context, info, tilingData, blockDim, schId)) {
660 return ge::GRAPH_SUCCESS;836 return ge::GRAPH_SUCCESS;
661 }837 }
662- OP_CHECK_IF(838+ OP_CHECK_IF((SetRadixMoreCoreTiling(context, info, blockDim) != ge::GRAPH_SUCCESS),
663- (SetRadixMoreCoreTiling(context, info, blockDim) != ge::GRAPH_SUCCESS),839+ OP_LOGE(context->GetNodeName(), "kth_value radix more-core tiling failed."), return ge::GRAPH_FAILED);
664- OP_LOGE(context->GetNodeName(), "kth_value radix more-core tiling failed."), return ge::GRAPH_FAILED);
665 PlanToTilingData(info, tilingData);840 PlanToTilingData(info, tilingData);
666 schId = KTH_VALUE_SCHID_RADIX_MORE_CORE;841 schId = KTH_VALUE_SCHID_RADIX_MORE_CORE;
667 return ge::GRAPH_SUCCESS;842 return ge::GRAPH_SUCCESS;
668}843}
669 844 
670-static ge::graphStatus FinalizeKthValueRoute(845+static ge::graphStatus FinalizeKthValueRoute(gert::TilingContext* context,
671- gert::TilingContext* context, const platform_ascendc::PlatformAscendC& ascendcPlatform, const SortKthTileInfo& info,846+ const platform_ascendc::PlatformAscendC& ascendcPlatform,
672- KthValueTilingData* tilingData, uint64_t schId, uint32_t& blockDim)847+ const SortKthTileInfo& info, KthValueTilingData* tilingData,
848+ uint64_t schId, uint32_t& blockDim)
673{849{
674 if (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) {850 if (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) {
675 blockDim = static_cast<uint32_t>(std::min<int64_t>(ascendcPlatform.GetCoreNumAiv(), info.unsortedDim));851 blockDim = static_cast<uint32_t>(std::min<int64_t>(ascendcPlatform.GetCoreNumAiv(), info.unsortedDim));
@@ -677,12 +853,12 @@ static ge::graphStatus FinalizeKthValueRoute(
677 OP_LOGE(context->GetNodeName(), "kth_value blockDim is zero.");853 OP_LOGE(context->GetNodeName(), "kth_value blockDim is zero.");
678 return ge::GRAPH_FAILED;854 return ge::GRAPH_FAILED;
679 }855 }
680- uint64_t maxRowsPerCore =856+ uint64_t maxRowsPerCore = Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim),
681- Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), static_cast<uint64_t>(blockDim));857+ static_cast<uint64_t>(blockDim));
682 uint64_t sortLoopTimes = Ops::Base::CeilDiv(maxRowsPerCore, static_cast<uint64_t>(info.outputRowsPerLoop));858 uint64_t sortLoopTimes = Ops::Base::CeilDiv(maxRowsPerCore, static_cast<uint64_t>(info.outputRowsPerLoop));
683- OP_CHECK_IF(859+ OP_CHECK_IF((sortLoopTimes > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())),
684- (sortLoopTimes > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())),860+ OP_LOGE(context->GetNodeName(), "kth_value sortLoopTimes exceeds uint32 limit."),
685- OP_LOGE(context->GetNodeName(), "kth_value sortLoopTimes exceeds uint32 limit."), return ge::GRAPH_FAILED);861+ return ge::GRAPH_FAILED);
686 tilingData->unsortedDimParallel = blockDim;862 tilingData->unsortedDimParallel = blockDim;
687 tilingData->sortLoopTimes = static_cast<uint32_t>(sortLoopTimes);863 tilingData->sortLoopTimes = static_cast<uint32_t>(sortLoopTimes);
688 }864 }
@@ -690,20 +866,21 @@ static ge::graphStatus FinalizeKthValueRoute(
690 OP_CHECK_NULL_WITH_CONTEXT(context, userWorkspaceSize);866 OP_CHECK_NULL_WITH_CONTEXT(context, userWorkspaceSize);
691 if (schId != KTH_VALUE_SCHID_MERGE_MORE_CORE && schId != KTH_VALUE_SCHID_MERGE_INTRA_CORE &&867 if (schId != KTH_VALUE_SCHID_MERGE_MORE_CORE && schId != KTH_VALUE_SCHID_MERGE_INTRA_CORE &&
692 schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS && schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX &&868 schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS && schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX &&
693- schId != KTH_VALUE_SCHID_RADIX_MORE_CORE) {869+ schId != KTH_VALUE_SCHID_RADIX_MORE_CORE && schId != KTH_VALUE_SCHID_RADIX_SELECT) {
694 userWorkspaceSize[0] = WORK_SPACE_SIZE;870 userWorkspaceSize[0] = WORK_SPACE_SIZE;
695 }871 }
696 return ge::GRAPH_SUCCESS;872 return ge::GRAPH_SUCCESS;
697}873}
698 874 
699-static void SetKthValueTilingContext(875+static void SetKthValueTilingContext(gert::TilingContext* context, uint64_t schId, const SortKthTileInfo& info,
700- gert::TilingContext* context, uint64_t schId, const SortKthTileInfo& info, uint32_t blockDim)876+ uint32_t blockDim)
701{877{
702 uint64_t tilingKeyIsInt32 = schId == KTH_VALUE_SCHID_RADIX_MORE_CORE ? info.isInt32 : 1U;878 uint64_t tilingKeyIsInt32 = schId == KTH_VALUE_SCHID_RADIX_MORE_CORE ? info.isInt32 : 1U;
703 context->SetTilingKey(GET_TPL_TILING_KEY(schId, tilingKeyIsInt32));879 context->SetTilingKey(GET_TPL_TILING_KEY(schId, tilingKeyIsInt32));
704 context->SetBlockDim(blockDim);880 context->SetBlockDim(blockDim);
705 if (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION || schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE ||881 if (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION || schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE ||
706- schId == KTH_VALUE_SCHID_RADIX_MORE_CORE || schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS ||882+ schId == KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT || schId == KTH_VALUE_SCHID_RADIX_MORE_CORE ||
883+ schId == KTH_VALUE_SCHID_RADIX_SELECT || schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS ||
707 schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX) {884 schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX) {
708 context->SetLocalMemorySize(info.ubSize - SIMT_UB);885 context->SetLocalMemorySize(info.ubSize - SIMT_UB);
709 } else {886 } else {
@@ -721,9 +898,8 @@ static ge::graphStatus Tiling4KthValue(gert::TilingContext* context)
721 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);898 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
722 auto tilingData = context->GetTilingData<KthValueTilingData>();899 auto tilingData = context->GetTilingData<KthValueTilingData>();
723 OP_CHECK_NULL_WITH_CONTEXT(context, tilingData);900 OP_CHECK_NULL_WITH_CONTEXT(context, tilingData);
724- OP_CHECK_IF(901+ OP_CHECK_IF((memset_s(tilingData, sizeof(KthValueTilingData), 0, sizeof(KthValueTilingData)) != EOK),
725- (memset_s(tilingData, sizeof(KthValueTilingData), 0, sizeof(KthValueTilingData)) != EOK),902+ OP_LOGE(context->GetNodeName(), "memset tilingdata failed"), return ge::GRAPH_FAILED);
726- OP_LOGE(context->GetNodeName(), "memset tilingdata failed"), return ge::GRAPH_FAILED);
727 auto attrs = context->GetAttrs();903 auto attrs = context->GetAttrs();
728 OP_CHECK_NULL_WITH_CONTEXT(context, attrs);904 OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
729 const int64_t* kAttr = attrs->GetAttrPointer<int64_t>(0);905 const int64_t* kAttr = attrs->GetAttrPointer<int64_t>(0);
@@ -733,39 +909,34 @@ static ge::graphStatus Tiling4KthValue(gert::TilingContext* context)
733 OP_CHECK_NULL_WITH_CONTEXT(context, inputDesc);909 OP_CHECK_NULL_WITH_CONTEXT(context, inputDesc);
734 ge::DataType dataType = inputDesc->GetDataType();910 ge::DataType dataType = inputDesc->GetDataType();
735 uint32_t dtypeSize = 0;911 uint32_t dtypeSize = 0;
736- OP_CHECK_IF(912+ OP_CHECK_IF((CheckKthValueDtypes(context, dataType, dtypeSize) != ge::GRAPH_SUCCESS),
737- (CheckKthValueDtypes(context, dataType, dtypeSize) != ge::GRAPH_SUCCESS),913+ OP_LOGE(context->GetNodeName(), "kth_value dtype check failed."), return ge::GRAPH_FAILED);
738- OP_LOGE(context->GetNodeName(), "kth_value dtype check failed."), return ge::GRAPH_FAILED);
739 SortKthTileInfo info;914 SortKthTileInfo info;
740 info.dataType = dataType;915 info.dataType = dataType;
741 info.dtypeSize = dtypeSize;916 info.dtypeSize = dtypeSize;
742 info.y2DtypeSize = static_cast<uint32_t>(sizeof(uint32_t));917 info.y2DtypeSize = static_cast<uint32_t>(sizeof(uint32_t));
743 info.blockUbSize = Ops::Base::GetUbBlockSize(context);918 info.blockUbSize = Ops::Base::GetUbBlockSize(context);
744 info.maxCoreNum = ascendcPlatform.GetCoreNumAiv();919 info.maxCoreNum = ascendcPlatform.GetCoreNumAiv();
745- OP_CHECK_IF(920+ OP_CHECK_IF((ParseKthValueShapeInfo(context, kAttr, dimAttr, info) != ge::GRAPH_SUCCESS),
746- (ParseKthValueShapeInfo(context, kAttr, dimAttr, info) != ge::GRAPH_SUCCESS),921+ OP_LOGE(context->GetNodeName(), "kth_value shape parse failed."), return ge::GRAPH_FAILED);
747- OP_LOGE(context->GetNodeName(), "kth_value shape parse failed."), return ge::GRAPH_FAILED);
748 info.isNonLastAxis = (info.sortAxis != info.rank - 1);922 info.isNonLastAxis = (info.sortAxis != info.rank - 1);
749 bool oneCoreUbValid = false;923 bool oneCoreUbValid = false;
750- OP_CHECK_IF(924+ OP_CHECK_IF((ComputeKthValueUbInfo(context, ascendcPlatform, info, oneCoreUbValid) != ge::GRAPH_SUCCESS),
751- (ComputeKthValueUbInfo(context, ascendcPlatform, info, oneCoreUbValid) != ge::GRAPH_SUCCESS),925+ OP_LOGE(context->GetNodeName(), "kth_value UB info compute failed."), return ge::GRAPH_FAILED);
752- OP_LOGE(context->GetNodeName(), "kth_value UB info compute failed."), return ge::GRAPH_FAILED);
753 InitKthValueBaseTiling(tilingData, info, oneCoreUbValid, *kAttr - 1);926 InitKthValueBaseTiling(tilingData, info, oneCoreUbValid, *kAttr - 1);
754 KthValueTilingData candidateTilingData = *tilingData;927 KthValueTilingData candidateTilingData = *tilingData;
755 uint64_t schId = KTH_VALUE_SCHID_RADIX_MORE_CORE;928 uint64_t schId = KTH_VALUE_SCHID_RADIX_MORE_CORE;
756 uint32_t blockDim = 1;929 uint32_t blockDim = 1;
757- OP_CHECK_IF(930+ OP_CHECK_IF((SelectKthValueRoute(context, info, &candidateTilingData, blockDim, schId) != ge::GRAPH_SUCCESS),
758- (SelectKthValueRoute(context, info, &candidateTilingData, blockDim, schId) != ge::GRAPH_SUCCESS),931+ OP_LOGE(context->GetNodeName(), "kth_value route selection failed."), return ge::GRAPH_FAILED);
759- OP_LOGE(context->GetNodeName(), "kth_value route selection failed."), return ge::GRAPH_FAILED);932+ OP_CHECK_IF((FinalizeKthValueRoute(context, ascendcPlatform, info, &candidateTilingData, schId, blockDim) !=
760- OP_CHECK_IF(933+ ge::GRAPH_SUCCESS),
761- (FinalizeKthValueRoute(context, ascendcPlatform, info, &candidateTilingData, schId, blockDim) !=934+ OP_LOGE(context->GetNodeName(), "kth_value route finalize failed."), return ge::GRAPH_FAILED);
762- ge::GRAPH_SUCCESS),
763- OP_LOGE(context->GetNodeName(), "kth_value route finalize failed."), return ge::GRAPH_FAILED);
764 *tilingData = candidateTilingData;935 *tilingData = candidateTilingData;
765 OP_LOGI(context->GetNodeName(),936 OP_LOGI(context->GetNodeName(),
766- "KthValueTiling: schId=%lu, blockDim=%u, lastAxis=%ld, unsortedDim=%ld, "937+ "KthValueTiling: schId=%lu, blockDim=%u, lastAxis=%ld, unsortedDim=%ld, "
767- "isNonLastAxis=%d, dtypeSize=%u",938+ "isNonLastAxis=%d, dtypeSize=%u",
768- schId, blockDim, info.lastAxis, info.unsortedDim, static_cast<int>(info.isNonLastAxis), info.dtypeSize);939+ schId, blockDim, info.lastAxis, info.unsortedDim, static_cast<int>(info.isNonLastAxis), info.dtypeSize);
769 SetKthValueTilingContext(context, schId, info, blockDim);940 SetKthValueTilingContext(context, schId, info, blockDim);
770 return ge::GRAPH_SUCCESS;941 return ge::GRAPH_SUCCESS;
771}942}
@@ -775,9 +946,8 @@ static ge::graphStatus TilingPrepare4KthValue(gert::TilingParseContext* context)
775 auto platformInfo = context->GetPlatformInfo();946 auto platformInfo = context->GetPlatformInfo();
776 OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);947 OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
777 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);948 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
778- OP_CHECK_IF(949+ OP_CHECK_IF((ascendcPlatform.GetCoreNumAiv() <= 0), OP_LOGE(context->GetNodeName(), "The core num is invalid."),
779- (ascendcPlatform.GetCoreNumAiv() <= 0), OP_LOGE(context->GetNodeName(), "The core num is invalid."),950+ return ge::GRAPH_FAILED);
780- return ge::GRAPH_FAILED);
781 return ge::GRAPH_SUCCESS;951 return ge::GRAPH_SUCCESS;
782}952}
783 953 
Amath/kth_value/op_kernel/arch35/kth_value_radix_select.h+1117-0
@@ -0,0 +1,1117 @@
1+/**
C
Ccpy_1234567月6日

在issue中记录下当前详细的优化点和核心代码块

likedislike
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef KTH_VALUE_RADIX_SELECT_H
12+#define KTH_VALUE_RADIX_SELECT_H
13+ 
14+#include "kernel_operator.h"
15+#include "op_kernel/platform_util.h"
16+#include "kth_value_tiling_data.h"
17+#include "../../sort/arch35/common/radix_sort_simd_utils.h"
18+ 
19+// Radix select: narrows the kth element by histogramming one byte at a time (MSB to LSB),
20+// selecting one bucket per round instead of fully sorting. Supports multi-core per row.
21+namespace KthValue {
22+using namespace AscendC;
23+using namespace RadixSortCommon;
24+ 
25+constexpr uint32_t RADIX_SELECT_FIND_THREADS = 128U;
26+constexpr uint32_t RADIX_SELECT_RESULT_WORDS = 8U;
27+constexpr uint32_t RADIX_SELECT_ACTIVE_INDEX_CAP = 4096U;
28+constexpr uint32_t RADIX_SELECT_ACTIVE_MODE_THRESHOLD = 512U;
29+constexpr uint32_t RADIX_SELECT_STATE_LEFT_K_IDX = 2U;
30+constexpr uint32_t RADIX_SELECT_STATE_SELECTED_IDX = 3U;
31+constexpr uint32_t RADIX_SELECT_STATE_SELECTED_COUNT_IDX = 4U;
32+constexpr uint32_t RADIX_SELECT_STATE_TARGET_CORE_IDX = 5U;
33+constexpr uint32_t RADIX_SELECT_STATE_TARGET_LEFT_K_IDX = 6U;
34+ 
35+template <typename UT>
36+__simt_callee__ __aicore__ inline void CountPrefixMatchesInThread(__ubuf__ UT* keys, uint32_t begin, uint32_t end,
37+ UT prefixMask, UT prefixKey,
38+ __ubuf__ uint32_t* threadCounts)
39+{
40+ uint32_t localCount = 0;
41+ for (uint32_t i = begin; i < end; ++i) {
42+ localCount += ((keys[i] & prefixMask) == prefixKey) ? 1U : 0U;
43+ }
44+ threadCounts[threadIdx.x] = localCount;
45+}
46+ 
47+template <typename UT>
48+__simt_callee__ __aicore__ inline int64_t LocatePrefixTarget(__ubuf__ UT* keys, uint32_t begin, uint32_t end,
49+ UT prefixMask, UT prefixKey, uint64_t localTarget)
50+{
51+ uint64_t seen = 0;
52+ for (uint32_t i = begin; i < end; ++i) {
53+ if ((keys[i] & prefixMask) != prefixKey) {
54+ continue;
55+ }
56+ if (seen == localTarget) {
57+ return static_cast<int64_t>(i);
58+ }
59+ ++seen;
60+ }
61+ return -1;
62+}
63+ 
64+__aicore__ inline void StoreRadixByteHistogram(__local_mem__ uint16_t* histogramPtr,
65+ MicroAPI::RegTensor<uint16_t>& hist0,
66+ MicroAPI::RegTensor<uint16_t>& hist1, MicroAPI::MaskReg maskB16)
67+{
68+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramPtr, hist0, VF_LEN_B16, maskB16);
69+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramPtr, hist1, VF_LEN_B16, maskB16);
70+}
71+ 
72+__aicore__ inline void AddSelectedByteHistogram(MicroAPI::RegTensor<uint16_t>& hist0,
73+ MicroAPI::RegTensor<uint16_t>& hist1,
74+ MicroAPI::RegTensor<uint8_t>& bytes,
75+ MicroAPI::RegTensor<uint8_t>& flagBytes, MicroAPI::MaskReg histMask)
76+{
77+ MicroAPI::MaskReg selectedMask;
78+ MicroAPI::CompareScalar<uint8_t, CMPMODE::EQ>(selectedMask, flagBytes, 1, histMask);
79+ MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0, MicroAPI::HistogramsType::FREQUENCY>(
80+ hist0, bytes, selectedMask);
81+ MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1, MicroAPI::HistogramsType::FREQUENCY>(
82+ hist1, bytes, selectedMask);
83+}
84+ 
85+template <typename VT>
86+__aicore__ inline void SelectPrefixFlag(MicroAPI::RegTensor<VT>& flag, MicroAPI::RegTensor<VT>& input,
87+ MicroAPI::RegTensor<VT>& masked, MicroAPI::RegTensor<VT>& prefixMaskReg,
88+ VT prefixKey, MicroAPI::RegTensor<VT>& one, MicroAPI::RegTensor<VT>& zero,
89+ MicroAPI::MaskReg mask)
90+{
91+ MicroAPI::MaskReg selected;
92+ MicroAPI::And(masked, input, prefixMaskReg, mask);
93+ MicroAPI::CompareScalar<VT, CMPMODE::EQ>(selected, masked, prefixKey, mask);
94+ MicroAPI::Select(flag, one, zero, selected);
95+}
96+ 
97+template <typename DstT, typename SrcT>
98+__aicore__ inline void DeInterleaveRadixBytes(MicroAPI::RegTensor<DstT>& dst0, MicroAPI::RegTensor<DstT>& dst1,
99+ MicroAPI::RegTensor<SrcT>& src0, MicroAPI::RegTensor<SrcT>& src1)
100+{
101+ MicroAPI::DeInterleave(dst0, dst1, (MicroAPI::RegTensor<DstT>&)src0, (MicroAPI::RegTensor<DstT>&)src1);
102+}
103+ 
104+__aicore__ inline void PackB32RadixBytes(MicroAPI::RegTensor<uint8_t>& bytes, MicroAPI::RegTensor<uint8_t>& highBytes,
105+ MicroAPI::RegTensor<uint32_t>& shifted0,
106+ MicroAPI::RegTensor<uint32_t>& shifted1,
107+ MicroAPI::RegTensor<uint32_t>& shifted2,
108+ MicroAPI::RegTensor<uint32_t>& shifted3)
109+{
110+ MicroAPI::RegTensor<uint16_t> data16_0, data16_1, data16_2, data16_3;
111+ DeInterleaveRadixBytes(data16_0, data16_1, shifted0, shifted1);
112+ DeInterleaveRadixBytes(data16_2, data16_3, shifted2, shifted3);
113+ DeInterleaveRadixBytes(bytes, highBytes, data16_0, data16_2);
114+}
115+ 
116+__aicore__ inline void PackB32FlagBytes(MicroAPI::RegTensor<uint8_t>& flagBytes,
117+ MicroAPI::RegTensor<uint8_t>& highFlags, MicroAPI::RegTensor<uint32_t>& flag0,
118+ MicroAPI::RegTensor<uint32_t>& flag1, MicroAPI::RegTensor<uint32_t>& flag2,
119+ MicroAPI::RegTensor<uint32_t>& flag3)
120+{
121+ MicroAPI::RegTensor<uint16_t> flag16_0, flag16_1, flag16_2, flag16_3;
122+ DeInterleaveRadixBytes(flag16_0, flag16_1, flag0, flag1);
123+ DeInterleaveRadixBytes(flag16_2, flag16_3, flag2, flag3);
124+ DeInterleaveRadixBytes(flagBytes, highFlags, flag16_0, flag16_2);
125+}
126+ 
127+__aicore__ inline void PackB64RadixBytes(
128+ MicroAPI::RegTensor<uint8_t>& bytes, MicroAPI::RegTensor<uint8_t>& highBytes,
129+ MicroAPI::RegTensor<uint64_t>& shifted0, MicroAPI::RegTensor<uint64_t>& shifted1,
130+ MicroAPI::RegTensor<uint64_t>& shifted2, MicroAPI::RegTensor<uint64_t>& shifted3,
131+ MicroAPI::RegTensor<uint64_t>& shifted4, MicroAPI::RegTensor<uint64_t>& shifted5,
132+ MicroAPI::RegTensor<uint64_t>& shifted6, MicroAPI::RegTensor<uint64_t>& shifted7)
133+{
134+ MicroAPI::RegTensor<uint32_t> data32_0, data32_1, data32_2, data32_3;
135+ MicroAPI::RegTensor<uint32_t> data32_4, data32_5, data32_6, data32_7;
136+ DeInterleaveRadixBytes(data32_0, data32_1, shifted0, shifted1);
137+ DeInterleaveRadixBytes(data32_2, data32_3, shifted2, shifted3);
138+ DeInterleaveRadixBytes(data32_4, data32_5, shifted4, shifted5);
139+ DeInterleaveRadixBytes(data32_6, data32_7, shifted6, shifted7);
140+ PackB32RadixBytes(bytes, highBytes, data32_0, data32_2, data32_4, data32_6);
141+}
142+ 
143+__aicore__ inline void PackB64FlagBytes(MicroAPI::RegTensor<uint8_t>& flagBytes,
144+ MicroAPI::RegTensor<uint8_t>& highFlags, MicroAPI::RegTensor<uint64_t>& flag0,
145+ MicroAPI::RegTensor<uint64_t>& flag1, MicroAPI::RegTensor<uint64_t>& flag2,
146+ MicroAPI::RegTensor<uint64_t>& flag3, MicroAPI::RegTensor<uint64_t>& flag4,
147+ MicroAPI::RegTensor<uint64_t>& flag5, MicroAPI::RegTensor<uint64_t>& flag6,
148+ MicroAPI::RegTensor<uint64_t>& flag7)
149+{
150+ MicroAPI::RegTensor<uint32_t> flag32_0, flag32_1, flag32_2, flag32_3;
151+ MicroAPI::RegTensor<uint32_t> flag32_4, flag32_5, flag32_6, flag32_7;
152+ DeInterleaveRadixBytes(flag32_0, flag32_1, flag0, flag1);
153+ DeInterleaveRadixBytes(flag32_2, flag32_3, flag2, flag3);
154+ DeInterleaveRadixBytes(flag32_4, flag32_5, flag4, flag5);
155+ DeInterleaveRadixBytes(flag32_6, flag32_7, flag6, flag7);
156+ PackB32FlagBytes(flagBytes, highFlags, flag32_0, flag32_2, flag32_4, flag32_6);
157+}
158+ 
159+// SIMT: each of 128 threads counts prefix matches in its slice, then thread 0 does
160+// prefix-sum to locate which thread owns the target, and searches within that slice.
161+template <typename UT>
162+__simt_vf__ LAUNCH_BOUND(RADIX_SELECT_FIND_THREADS) __aicore__
163+ void FindKthMatchInTile(__ubuf__ UT* keys, uint32_t count, UT prefixMask, UT prefixKey, uint64_t target,
164+ __ubuf__ uint32_t* threadCounts, __ubuf__ int64_t* result)
165+{
166+ uint32_t tid = threadIdx.x;
167+ uint32_t elemsPerThread = (count + RADIX_SELECT_FIND_THREADS - 1U) / RADIX_SELECT_FIND_THREADS;
168+ uint32_t begin = tid * elemsPerThread;
169+ uint32_t end = begin + elemsPerThread;
170+ end = end < count ? end : count;
171+ CountPrefixMatchesInThread(keys, begin, end, prefixMask, prefixKey, threadCounts);
172+ asc_syncthreads();
173+ if (tid != 0U) {
174+ return;
175+ }
176+ uint64_t accumulated = 0;
177+ int64_t found = -1;
178+ for (uint32_t thread = 0; thread < RADIX_SELECT_FIND_THREADS; ++thread) {
179+ uint64_t next = accumulated + threadCounts[thread];
180+ if (target < next) {
181+ uint64_t localTarget = target - accumulated;
182+ uint32_t targetBegin = thread * elemsPerThread;
183+ uint32_t targetEnd = targetBegin + elemsPerThread;
184+ targetEnd = targetEnd < count ? targetEnd : count;
185+ found = LocatePrefixTarget(keys, targetBegin, targetEnd, prefixMask, prefixKey, localTarget);
186+ break;
187+ }
188+ accumulated = next;
189+ }
190+ uint64_t total = 0;
191+ for (uint32_t thread = 0; thread < RADIX_SELECT_FIND_THREADS; ++thread) {
192+ total += threadCounts[thread];
193+ }
194+ result[0] = found;
195+ result[1] = static_cast<int64_t>(total);
196+}
197+ 
198+// SIMT: collect all prefix-matching indices into activeIndices via exclusive prefix-sum offsets.
199+// Aborts if total exceeds ACTIVE_INDEX_CAP (4096); result[1] signals success.
200+template <typename UT>
201+__simt_vf__ LAUNCH_BOUND(RADIX_SELECT_FIND_THREADS) __aicore__
202+ void CollectActiveIndices(__ubuf__ UT* keys, uint32_t count, UT prefixMask, UT prefixKey,
203+ __ubuf__ uint32_t* threadCounts, __ubuf__ uint32_t* activeIndices,
204+ __ubuf__ int64_t* result)
205+{
206+ uint32_t tid = threadIdx.x;
207+ uint32_t elemsPerThread = (count + RADIX_SELECT_FIND_THREADS - 1U) / RADIX_SELECT_FIND_THREADS;
208+ uint32_t begin = tid * elemsPerThread;
209+ uint32_t end = begin + elemsPerThread;
210+ end = end < count ? end : count;
211+ CountPrefixMatchesInThread(keys, begin, end, prefixMask, prefixKey, threadCounts);
212+ asc_syncthreads();
213+ if (tid == 0U) {
214+ uint32_t accumulated = 0;
215+ for (uint32_t thread = 0; thread < RADIX_SELECT_FIND_THREADS; ++thread) {
216+ uint32_t countInThread = threadCounts[thread];
217+ threadCounts[thread] = accumulated;
218+ accumulated += countInThread;
219+ }
220+ result[0] = static_cast<int64_t>(accumulated);
221+ result[1] = accumulated <= RADIX_SELECT_ACTIVE_INDEX_CAP ? 1 : 0;
222+ }
223+ asc_syncthreads();
224+ if (result[1] == 0) {
225+ return;
226+ }
227+ uint32_t writePos = threadCounts[tid];
228+ for (uint32_t i = begin; i < end; ++i) {
229+ if ((keys[i] & prefixMask) == prefixKey) {
230+ activeIndices[writePos] = i;
231+ ++writePos;
232+ }
233+ }
234+}
235+ 
236+// Active-mode histogram: single-threaded scan over activeIndices only (count <= 512).
237+template <typename UT>
238+__simt_vf__ LAUNCH_BOUND(RADIX_SELECT_FIND_THREADS) __aicore__
239+ void BuildActiveHistogram(__ubuf__ UT* keys, __ubuf__ uint32_t* activeIndices, uint32_t activeCount, uint32_t shift,
240+ __ubuf__ uint64_t* histogram)
241+{
242+ if (threadIdx.x != 0U) {
243+ return;
244+ }
245+ for (uint32_t bucket = 0; bucket < 256U; ++bucket) {
246+ histogram[bucket] = 0UL;
247+ }
248+ for (uint32_t i = 0; i < activeCount; ++i) {
249+ UT key = keys[activeIndices[i]];
250+ uint32_t bucket = static_cast<uint32_t>((key >> shift) & static_cast<UT>(0xFFU));
251+ histogram[bucket] += 1UL;
252+ }
253+}
254+ 
255+// Active-mode compaction: filter activeIndices in-place by the expanded prefix.
256+template <typename UT>
257+__simt_vf__ LAUNCH_BOUND(RADIX_SELECT_FIND_THREADS) __aicore__
258+ void CompactActiveIndices(__ubuf__ UT* keys, __ubuf__ uint32_t* activeIndices, uint32_t activeCount, UT prefixMask,
259+ UT prefixKey, __ubuf__ int64_t* result)
260+{
261+ if (threadIdx.x != 0U) {
262+ return;
263+ }
264+ uint32_t writePos = 0;
265+ for (uint32_t i = 0; i < activeCount; ++i) {
266+ uint32_t index = activeIndices[i];
267+ if ((keys[index] & prefixMask) == prefixKey) {
268+ activeIndices[writePos] = index;
269+ ++writePos;
270+ }
271+ }
272+ result[0] = static_cast<int64_t>(writePos);
273+}
274+ 
275+// Selects one radix bucket per byte instead of scattering a fully sorted row.
276+// Cores are grouped by row. Each core builds the histogram for a contiguous axis
277+// slice; the leader reduces the group histogram and broadcasts the selected prefix.
278+// Optimizations: retained tile (avoid re-loading when slice fits in one tile),
279+// active mode (track matching indices when count <= 512 to skip full scans).
280+template <typename T, typename UT>
281+class KthValueRadixSelect {
282+public:
283+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace,
284+ const KthValueTilingData* tiling, TPipe* pipe);
285+ __aicore__ inline void Process();
286+ 
287+private:
288+ __aicore__ inline LocalTensor<UT> MakeSortKeys(LocalTensor<T>& input, uint32_t count);
289+ __aicore__ inline void LoadTile(int64_t offset, uint32_t count, LocalTensor<T>& input);
290+ __aicore__ inline void CountByte(int64_t rowOffset, int64_t sliceStart, int64_t sliceCount, uint32_t shift,
291+ UT prefixMask, UT prefixKey, LocalTensor<uint64_t>& histogram);
292+ __aicore__ inline void ClearHistogram(LocalTensor<uint64_t>& histogram);
293+ __aicore__ inline void AccumulateHistogram(LocalTensor<uint64_t>& histogram, LocalTensor<uint64_t>& src);
294+ __aicore__ inline void AccumulateTileHistogram(LocalTensor<uint64_t>& histogram,
295+ LocalTensor<uint16_t>& tileHistogram);
296+ __aicore__ inline void BuildTileHistogram(LocalTensor<UT>& keys, uint32_t count, uint32_t shift, UT prefixMask,
297+ UT prefixKey, LocalTensor<uint16_t>& tileHistogram);
298+ __aicore__ inline void BuildTileHistogramB8(LocalTensor<UT>& keys, uint32_t count, UT prefixMask, UT prefixKey,
299+ LocalTensor<uint16_t>& tileHistogram);
300+ __aicore__ inline void BuildTileHistogramB16(LocalTensor<UT>& keys, uint32_t count, uint32_t shift, UT prefixMask,
301+ UT prefixKey, LocalTensor<uint16_t>& tileHistogram);
302+ __aicore__ inline void BuildTileHistogramB32(LocalTensor<UT>& keys, uint32_t count, uint32_t shift, UT prefixMask,
303+ UT prefixKey, LocalTensor<uint16_t>& tileHistogram);
304+ __aicore__ inline void BuildTileHistogramB64(LocalTensor<UT>& keys, uint32_t count, uint32_t shift, UT prefixMask,
305+ UT prefixKey, LocalTensor<uint16_t>& tileHistogram);
306+ __aicore__ inline void WriteOutput(int64_t row, T value, int64_t index);
307+ __aicore__ inline void SelectOneRow(int64_t row, uint32_t group, uint32_t coreInGroup, bool writeOutput);
308+ __aicore__ inline void UpdateActiveIndices(LocalTensor<UT>& retainedKeys, LocalTensor<uint32_t>& activeIndices,
309+ uint32_t retainedCount, UT prefixMask, UT prefixKey, bool& activeMode,
310+ uint32_t& activeCount);
311+ __aicore__ inline void StoreCoreHistogram(LocalTensor<uint64_t>& histogram);
312+ __aicore__ inline uint64_t LoadCoreHistogramBucket(uint32_t group, uint32_t coreInGroup, uint32_t bucket,
313+ LocalTensor<uint64_t>& scratch);
314+ __aicore__ inline void ReduceGroupHistogram(uint32_t group, LocalTensor<uint64_t>& histogram,
315+ LocalTensor<uint64_t>& scratch);
316+ __aicore__ inline void StoreGroupState(uint32_t group, LocalTensor<uint64_t>& state);
317+ __aicore__ inline void LoadGroupState(uint32_t group, LocalTensor<uint64_t>& state);
318+ 
319+ static constexpr uint32_t RADIX = 256U;
320+ static constexpr uint32_t BYTE_BITS = 8U;
321+ 
322+ GlobalTensor<T> xGm_;
323+ GlobalTensor<T> valuesGm_;
324+ GlobalTensor<int64_t> indicesGm_;
325+ GlobalTensor<uint64_t> coreHistogramGm_;
326+ GlobalTensor<uint64_t> groupStateGm_;
327+ TPipe* pipe_{nullptr};
328+ TQue<QuePosition::VECIN, 1> inputQueue_;
329+ TBuf<TPosition::VECCALC> keyBuf_;
330+ TBuf<TPosition::VECCALC> histogramBuf_;
331+ TBuf<TPosition::VECCALC> tileHistogramBuf_;
332+ TBuf<TPosition::VECCALC> outputValueBuf_;
333+ TBuf<TPosition::VECCALC> outputIndexBuf_;
334+ TBuf<TPosition::VECCALC> findCountBuf_;
335+ TBuf<TPosition::VECCALC> findResultBuf_;
336+ TBuf<TPosition::VECCALC> groupReduceBuf_;
337+ TBuf<TPosition::VECCALC> activeIndexBuf_;
338+ 
339+ uint32_t blockIdx_{0};
340+ uint32_t blockNum_{0};
341+ uint32_t tileElems_{0};
342+ uint32_t rowsParallel_{0};
343+ uint32_t coresPerRow_{0};
344+ uint32_t sortLoopTimes_{0};
345+ int64_t kthIndex_{0};
346+ int64_t axisLen_{0};
347+ int64_t rowCount_{0};
348+};
349+ 
350+template <typename T, typename UT>
351+__aicore__ inline void KthValueRadixSelect<T, UT>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace,
352+ const KthValueTilingData* tiling, TPipe* pipe)
353+{
354+ if (tiling == nullptr || pipe == nullptr) {
355+ return;
356+ }
357+ blockIdx_ = GetBlockIdx();
358+ blockNum_ = GetBlockNum();
359+ pipe_ = pipe;
360+ tileElems_ = tiling->numTileDataSize;
361+ rowsParallel_ = tiling->unsortedDimParallel;
362+ coresPerRow_ = tiling->lastDimNeedCore;
363+ sortLoopTimes_ = tiling->sortLoopTimes;
364+ kthIndex_ = tiling->kthIndex;
365+ axisLen_ = tiling->lastAxisNum;
366+ rowCount_ = tiling->unsortedDimNum;
367+ xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
368+ valuesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(values));
369+ indicesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(indices));
370+ uint64_t histogramWords = static_cast<uint64_t>(blockNum_) * RADIX;
371+ coreHistogramGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t*>(workspace), histogramWords);
372+ groupStateGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t*>(workspace) + histogramWords,
373+ static_cast<uint64_t>(rowsParallel_) * RADIX_SELECT_RESULT_WORDS);
374+ pipe_->InitBuffer(inputQueue_, 1, tileElems_ * sizeof(T));
375+ pipe_->InitBuffer(keyBuf_, tileElems_ * sizeof(UT));
376+ pipe_->InitBuffer(histogramBuf_, RADIX * sizeof(uint64_t));
377+ pipe_->InitBuffer(tileHistogramBuf_, RADIX * sizeof(uint16_t));
378+ pipe_->InitBuffer(outputValueBuf_, Ops::Base::GetUbBlockSize());
379+ pipe_->InitBuffer(outputIndexBuf_, Ops::Base::GetUbBlockSize());
380+ pipe_->InitBuffer(findCountBuf_, RADIX_SELECT_FIND_THREADS * sizeof(uint32_t));
381+ pipe_->InitBuffer(findResultBuf_, RADIX_SELECT_RESULT_WORDS * sizeof(uint64_t));
382+ pipe_->InitBuffer(groupReduceBuf_, RADIX * sizeof(uint64_t));
383+ pipe_->InitBuffer(activeIndexBuf_, RADIX_SELECT_ACTIVE_INDEX_CAP * sizeof(uint32_t));
384+}
385+ 
386+template <typename T, typename UT>
387+__aicore__ inline LocalTensor<UT> KthValueRadixSelect<T, UT>::MakeSortKeys(LocalTensor<T>& input, uint32_t count)
388+{
389+ LocalTensor<UT> keys = keyBuf_.Get<UT>();
390+ if constexpr (IsSameType<int8_t, T>::value) {
391+ TwiddleInB8<T, UT, 0>(input, keys, count);
392+ } else if constexpr (IsSameType<int16_t, T>::value) {
393+ TwiddleInB16<T, UT, 0>(input, keys, count);
394+ } else if constexpr (IsSameType<int32_t, T>::value) {
395+ TwiddleInB32<T, UT, 0>(input, keys, count);
396+ } else if constexpr (IsSameType<int64_t, T>::value) {
397+ TwiddleInB64<T, UT, 0>(input, keys, count);
398+ } else if constexpr (IsSameType<half, T>::value || IsSameType<bfloat16_t, T>::value) {
399+ TwiddleInFp16<T, UT, 0>(input, keys, count);
400+ } else if constexpr (IsSameType<float, T>::value) {
401+ TwiddleInFp32<T, UT, 0>(input, keys, count);
402+ } else {
403+ // Ascending unsigned integers already have lexicographically sortable bit patterns.
404+ keys = input.template ReinterpretCast<UT>();
405+ }
406+ return keys;
407+}
408+ 
409+template <typename T, typename UT>
410+__aicore__ inline void KthValueRadixSelect<T, UT>::LoadTile(int64_t offset, uint32_t count, LocalTensor<T>& input)
411+{
412+ DataCopyExtParams copyParams{1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
413+ DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
414+ DataCopyPad(input, xGm_[offset], copyParams, padParams);
415+}
416+ 
417+template <typename T, typename UT>
418+__aicore__ inline void KthValueRadixSelect<T, UT>::BuildTileHistogram(LocalTensor<UT>& keys, uint32_t count,
419+ uint32_t shift, UT prefixMask, UT prefixKey,
420+ LocalTensor<uint16_t>& tileHistogram)
421+{
422+ if constexpr (sizeof(UT) == sizeof(uint8_t)) {
423+ BuildTileHistogramB8(keys, count, prefixMask, prefixKey, tileHistogram);
424+ } else if constexpr (sizeof(UT) == sizeof(uint16_t)) {
425+ BuildTileHistogramB16(keys, count, shift, prefixMask, prefixKey, tileHistogram);
426+ } else if constexpr (sizeof(UT) == sizeof(uint32_t)) {
427+ BuildTileHistogramB32(keys, count, shift, prefixMask, prefixKey, tileHistogram);
428+ } else {
429+ BuildTileHistogramB64(keys, count, shift, prefixMask, prefixKey, tileHistogram);
430+ }
431+}
432+ 
433+template <typename T, typename UT>
434+__aicore__ inline void KthValueRadixSelect<T, UT>::BuildTileHistogramB8(LocalTensor<UT>& keys, uint32_t count,
435+ UT prefixMask, UT prefixKey,
436+ LocalTensor<uint16_t>& tileHistogram)
437+{
438+ __local_mem__ UT* keyPtr = (__ubuf__ UT*)keys.GetPhyAddr();
439+ __local_mem__ uint16_t* histogramPtr = (__ubuf__ uint16_t*)tileHistogram.GetPhyAddr();
440+ uint32_t remain = count;
441+ uint16_t repeats = CeilDivision(count, VF_LEN_B8);
442+ __VEC_SCOPE__
443+ {
444+ MicroAPI::RegTensor<uint16_t> hist0, hist1;
445+ MicroAPI::RegTensor<uint8_t> input, prefixMaskReg, masked;
446+ MicroAPI::MaskReg maskB8 = MicroAPI::CreateMask<uint8_t>();
447+ MicroAPI::MaskReg maskB16 = MicroAPI::CreateMask<uint16_t>();
448+ MicroAPI::Duplicate(hist0, 0, maskB16);
449+ MicroAPI::Duplicate(hist1, 0, maskB16);
450+ MicroAPI::Duplicate(prefixMaskReg, prefixMask, maskB8);
451+ for (uint16_t i = 0; i < repeats; ++i) {
452+ MicroAPI::MaskReg validMask = MicroAPI::UpdateMask<uint8_t>(remain);
453+ MicroAPI::DataCopy<uint8_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input, keyPtr, VF_LEN_B8);
454+ MicroAPI::And(masked, input, prefixMaskReg, validMask);
455+ MicroAPI::MaskReg selectedMask;
456+ MicroAPI::CompareScalar<uint8_t, CMPMODE::EQ>(selectedMask, masked, prefixKey, validMask);
457+ MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
458+ MicroAPI::HistogramsType::FREQUENCY>(hist0, input, selectedMask);
459+ MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
460+ MicroAPI::HistogramsType::FREQUENCY>(hist1, input, selectedMask);
461+ }
462+ StoreRadixByteHistogram(histogramPtr, hist0, hist1, maskB16);
463+ }
464+}
465+ 
466+template <typename T, typename UT>
467+__aicore__ inline void KthValueRadixSelect<T, UT>::BuildTileHistogramB16(LocalTensor<UT>& keys, uint32_t count,
468+ uint32_t shift, UT prefixMask, UT prefixKey,
469+ LocalTensor<uint16_t>& tileHistogram)
470+{
471+ __local_mem__ UT* keyPtr = (__ubuf__ UT*)keys.GetPhyAddr();
472+ __local_mem__ uint16_t* histogramPtr = (__ubuf__ uint16_t*)tileHistogram.GetPhyAddr();
473+ uint32_t remain = count;
474+ uint16_t repeats = CeilDivision(count, VF_LEN_B8);
475+ __VEC_SCOPE__
476+ {
477+ MicroAPI::RegTensor<uint16_t> hist0, hist1, input0, input1, shifted0, shifted1;
478+ MicroAPI::RegTensor<uint16_t> prefixMaskReg, masked0, masked1, flag0, flag1, one, zero;
479+ MicroAPI::MaskReg maskB16 = MicroAPI::CreateMask<uint16_t>();
480+ MicroAPI::Duplicate(hist0, 0, maskB16);
481+ MicroAPI::Duplicate(hist1, 0, maskB16);
482+ MicroAPI::Duplicate(prefixMaskReg, prefixMask, maskB16);
483+ MicroAPI::Duplicate(one, 1, maskB16);
484+ MicroAPI::Duplicate(zero, 0, maskB16);
485+ for (uint16_t i = 0; i < repeats; ++i) {
486+ MicroAPI::MaskReg histMask = MicroAPI::UpdateMask<uint8_t>(remain);
487+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input0, keyPtr, VF_LEN_B16);
488+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input1, keyPtr, VF_LEN_B16);
489+ MicroAPI::And(masked0, input0, prefixMaskReg, maskB16);
490+ MicroAPI::And(masked1, input1, prefixMaskReg, maskB16);
491+ MicroAPI::MaskReg selected0, selected1;
492+ MicroAPI::CompareScalar<uint16_t, CMPMODE::EQ>(selected0, masked0, prefixKey, maskB16);
493+ MicroAPI::CompareScalar<uint16_t, CMPMODE::EQ>(selected1, masked1, prefixKey, maskB16);
494+ MicroAPI::Select(flag0, one, zero, selected0);
495+ MicroAPI::Select(flag1, one, zero, selected1);
496+ MicroAPI::ShiftRights<uint16_t, int16_t>(shifted0, input0, shift, maskB16);
497+ MicroAPI::ShiftRights<uint16_t, int16_t>(shifted1, input1, shift, maskB16);
498+ MicroAPI::RegTensor<uint8_t> bytes, highBytes, flagBytes, highFlags;
499+ MicroAPI::DeInterleave(bytes, highBytes, (MicroAPI::RegTensor<uint8_t>&)shifted0,
500+ (MicroAPI::RegTensor<uint8_t>&)shifted1);
501+ MicroAPI::DeInterleave(flagBytes, highFlags, (MicroAPI::RegTensor<uint8_t>&)flag0,
502+ (MicroAPI::RegTensor<uint8_t>&)flag1);
503+ AddSelectedByteHistogram(hist0, hist1, bytes, flagBytes, histMask);
504+ }
505+ StoreRadixByteHistogram(histogramPtr, hist0, hist1, maskB16);
506+ }
507+}
508+ 
509+template <typename T, typename UT>
510+__aicore__ inline void KthValueRadixSelect<T, UT>::BuildTileHistogramB32(LocalTensor<UT>& keys, uint32_t count,
511+ uint32_t shift, UT prefixMask, UT prefixKey,
512+ LocalTensor<uint16_t>& tileHistogram)
513+{
514+ __local_mem__ UT* keyPtr = (__ubuf__ UT*)keys.GetPhyAddr();
515+ __local_mem__ uint16_t* histogramPtr = (__ubuf__ uint16_t*)tileHistogram.GetPhyAddr();
516+ uint32_t remain = count;
517+ uint16_t repeats = CeilDivision(count, VF_LEN_B8);
518+ __VEC_SCOPE__
519+ {
520+ MicroAPI::RegTensor<uint16_t> hist0, hist1;
521+ MicroAPI::RegTensor<uint32_t> input0, input1, input2, input3, shifted0, shifted1, shifted2, shifted3;
522+ MicroAPI::RegTensor<uint32_t> prefixMaskReg, masked, flag0, flag1, flag2, flag3, one, zero;
523+ MicroAPI::MaskReg maskB16 = MicroAPI::CreateMask<uint16_t>();
524+ MicroAPI::MaskReg maskB32 = MicroAPI::CreateMask<uint32_t>();
525+ MicroAPI::Duplicate(hist0, 0, maskB16);
526+ MicroAPI::Duplicate(hist1, 0, maskB16);
527+ MicroAPI::Duplicate(prefixMaskReg, prefixMask, maskB32);
528+ MicroAPI::Duplicate(one, 1, maskB32);
529+ MicroAPI::Duplicate(zero, 0, maskB32);
530+ for (uint16_t i = 0; i < repeats; ++i) {
531+ MicroAPI::MaskReg histMask = MicroAPI::UpdateMask<uint8_t>(remain);
532+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input0, keyPtr, VF_LEN_B32);
533+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input1, keyPtr, VF_LEN_B32);
534+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input2, keyPtr, VF_LEN_B32);
535+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input3, keyPtr, VF_LEN_B32);
536+ SelectPrefixFlag(flag0, input0, masked, prefixMaskReg, prefixKey, one, zero, maskB32);
537+ SelectPrefixFlag(flag1, input1, masked, prefixMaskReg, prefixKey, one, zero, maskB32);
538+ SelectPrefixFlag(flag2, input2, masked, prefixMaskReg, prefixKey, one, zero, maskB32);
539+ SelectPrefixFlag(flag3, input3, masked, prefixMaskReg, prefixKey, one, zero, maskB32);
540+ MicroAPI::ShiftRights<uint32_t, int16_t>(shifted0, input0, shift, maskB32);
541+ MicroAPI::ShiftRights<uint32_t, int16_t>(shifted1, input1, shift, maskB32);
542+ MicroAPI::ShiftRights<uint32_t, int16_t>(shifted2, input2, shift, maskB32);
543+ MicroAPI::ShiftRights<uint32_t, int16_t>(shifted3, input3, shift, maskB32);
544+ MicroAPI::RegTensor<uint8_t> bytes, highBytes;
545+ MicroAPI::RegTensor<uint8_t> flagBytes, highFlags;
546+ PackB32RadixBytes(bytes, highBytes, shifted0, shifted1, shifted2, shifted3);
547+ PackB32FlagBytes(flagBytes, highFlags, flag0, flag1, flag2, flag3);
548+ AddSelectedByteHistogram(hist0, hist1, bytes, flagBytes, histMask);
549+ }
550+ StoreRadixByteHistogram(histogramPtr, hist0, hist1, maskB16);
551+ }
552+}
553+ 
554+template <typename T, typename UT>
555+__aicore__ inline void KthValueRadixSelect<T, UT>::BuildTileHistogramB64(LocalTensor<UT>& keys, uint32_t count,
556+ uint32_t shift, UT prefixMask, UT prefixKey,
557+ LocalTensor<uint16_t>& tileHistogram)
558+{
559+ __local_mem__ UT* keyPtr = (__ubuf__ UT*)keys.GetPhyAddr();
560+ uint16_t repeats = CeilDivision(count, VF_LEN_B8);
561+ uint32_t remain = count;
562+ __local_mem__ uint16_t* histogramPtr = (__ubuf__ uint16_t*)tileHistogram.GetPhyAddr();
563+ __VEC_SCOPE__
564+ {
565+ MicroAPI::RegTensor<uint16_t> hist0, hist1;
566+ MicroAPI::RegTensor<uint64_t> input0, input1, input2, input3, input4, input5, input6, input7;
567+ MicroAPI::RegTensor<uint64_t> shifted0, shifted1, shifted2, shifted3, shifted4, shifted5, shifted6, shifted7;
568+ MicroAPI::RegTensor<uint64_t> flag0, flag1, flag2, flag3, flag4, flag5, flag6, flag7;
569+ MicroAPI::RegTensor<uint64_t> prefixMaskReg, masked, one, zero;
570+ MicroAPI::MaskReg maskB16 = MicroAPI::CreateMask<uint16_t>();
571+ MicroAPI::MaskReg maskB64 = MicroAPI::CreateMask<uint64_t>();
572+ MicroAPI::Duplicate(hist0, 0, maskB16);
573+ MicroAPI::Duplicate(hist1, 0, maskB16);
574+ MicroAPI::Duplicate(prefixMaskReg, prefixMask, maskB64);
575+ MicroAPI::Duplicate(one, 1, maskB64);
576+ MicroAPI::Duplicate(zero, 0, maskB64);
577+ for (uint16_t i = 0; i < repeats; ++i) {
578+ MicroAPI::MaskReg histMask = MicroAPI::UpdateMask<uint8_t>(remain);
579+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input0, keyPtr, VF_LEN_B64);
580+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input1, keyPtr, VF_LEN_B64);
581+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input2, keyPtr, VF_LEN_B64);
582+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input3, keyPtr, VF_LEN_B64);
583+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input4, keyPtr, VF_LEN_B64);
584+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input5, keyPtr, VF_LEN_B64);
585+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input6, keyPtr, VF_LEN_B64);
586+ MicroAPI::DataCopy<uint64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(input7, keyPtr, VF_LEN_B64);
587+ SelectPrefixFlag(flag0, input0, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
588+ SelectPrefixFlag(flag1, input1, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
589+ SelectPrefixFlag(flag2, input2, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
590+ SelectPrefixFlag(flag3, input3, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
591+ SelectPrefixFlag(flag4, input4, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
592+ SelectPrefixFlag(flag5, input5, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
593+ SelectPrefixFlag(flag6, input6, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
594+ SelectPrefixFlag(flag7, input7, masked, prefixMaskReg, prefixKey, one, zero, maskB64);
595+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted0, input0, shift, maskB64);
596+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted1, input1, shift, maskB64);
597+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted2, input2, shift, maskB64);
598+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted3, input3, shift, maskB64);
599+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted4, input4, shift, maskB64);
600+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted5, input5, shift, maskB64);
601+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted6, input6, shift, maskB64);
602+ MicroAPI::ShiftRights<uint64_t, int16_t>(shifted7, input7, shift, maskB64);
603+ MicroAPI::RegTensor<uint8_t> bytes, highBytes;
604+ MicroAPI::RegTensor<uint8_t> flagBytes, highFlags;
605+ PackB64RadixBytes(bytes, highBytes, shifted0, shifted1, shifted2, shifted3, shifted4, shifted5, shifted6,
606+ shifted7);
607+ PackB64FlagBytes(flagBytes, highFlags, flag0, flag1, flag2, flag3, flag4, flag5, flag6, flag7);
608+ AddSelectedByteHistogram(hist0, hist1, bytes, flagBytes, histMask);
609+ }
610+ StoreRadixByteHistogram(histogramPtr, hist0, hist1, maskB16);
611+ }
612+}
613+ 
614+template <typename T, typename UT>
615+__aicore__ inline void KthValueRadixSelect<T, UT>::ClearHistogram(LocalTensor<uint64_t>& histogram)
616+{
617+ __local_mem__ int64_t* histogramPtr = (__ubuf__ int64_t*)histogram.GetPhyAddr();
618+ __VEC_SCOPE__
619+ {
620+ MicroAPI::MaskReg maskB64 = MicroAPI::CreateMask<int64_t>();
621+ MicroAPI::RegTensor<int64_t> zero;
622+ MicroAPI::Duplicate(zero, 0, maskB64);
623+ for (uint16_t i = 0; i < RADIX / VF_LEN_B64; ++i) {
624+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramPtr, zero, VF_LEN_B64,
625+ maskB64);
626+ }
627+ }
628+}
629+ 
630+template <typename T, typename UT>
631+__aicore__ inline void KthValueRadixSelect<T, UT>::AccumulateHistogram(LocalTensor<uint64_t>& histogram,
632+ LocalTensor<uint64_t>& src)
633+{
634+ __local_mem__ int64_t* histogramReadPtr = (__ubuf__ int64_t*)histogram.GetPhyAddr();
635+ __local_mem__ int64_t* histogramWritePtr = histogramReadPtr;
636+ __local_mem__ int64_t* srcPtr = (__ubuf__ int64_t*)src.GetPhyAddr();
637+ __VEC_SCOPE__
638+ {
639+ MicroAPI::MaskReg maskB64 = MicroAPI::CreateMask<int64_t>();
640+ MicroAPI::RegTensor<int64_t> dstReg, srcReg;
641+ for (uint16_t i = 0; i < RADIX / VF_LEN_B64; ++i) {
642+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(dstReg, histogramReadPtr, VF_LEN_B64);
643+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(srcReg, srcPtr, VF_LEN_B64);
644+ MicroAPI::Add(dstReg, dstReg, srcReg, maskB64);
645+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, dstReg, VF_LEN_B64,
646+ maskB64);
647+ }
648+ }
649+}
650+ 
651+template <typename T, typename UT>
652+__aicore__ inline void KthValueRadixSelect<T, UT>::AccumulateTileHistogram(LocalTensor<uint64_t>& histogram,
653+ LocalTensor<uint16_t>& tileHistogram)
654+{
655+ __local_mem__ uint16_t* tilePtr = (__ubuf__ uint16_t*)tileHistogram.GetPhyAddr();
656+ __local_mem__ int64_t* histogramReadPtr = (__ubuf__ int64_t*)histogram.GetPhyAddr();
657+ __local_mem__ int64_t* histogramWritePtr = histogramReadPtr;
658+ __VEC_SCOPE__
659+ {
660+ MicroAPI::MaskReg maskB16 = MicroAPI::CreateMask<uint16_t>();
661+ MicroAPI::MaskReg maskB64 = MicroAPI::CreateMask<int64_t>();
662+ MicroAPI::RegTensor<uint16_t> hist0, hist1, zero16;
663+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(hist0, tilePtr, VF_LEN_B16);
664+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(hist1, tilePtr, VF_LEN_B16);
665+ MicroAPI::Duplicate(zero16, 0, maskB16);
666+ 
667+ MicroAPI::RegTensor<uint32_t> hist32_0, hist32_1, hist32_2, hist32_3;
668+ MicroAPI::Interleave((MicroAPI::RegTensor<uint16_t>&)hist32_0, (MicroAPI::RegTensor<uint16_t>&)hist32_1, hist0,
669+ zero16);
670+ MicroAPI::Interleave((MicroAPI::RegTensor<uint16_t>&)hist32_2, (MicroAPI::RegTensor<uint16_t>&)hist32_3, hist1,
671+ zero16);
672+ 
673+ MicroAPI::RegTensor<int64_t> hist64_0, hist64_1, hist64_2, hist64_3;
674+ MicroAPI::RegTensor<int64_t> hist64_4, hist64_5, hist64_6, hist64_7;
675+ MicroAPI::Interleave((MicroAPI::RegTensor<uint32_t>&)hist64_0, (MicroAPI::RegTensor<uint32_t>&)hist64_1,
676+ hist32_0, (MicroAPI::RegTensor<uint32_t>&)zero16);
677+ MicroAPI::Interleave((MicroAPI::RegTensor<uint32_t>&)hist64_2, (MicroAPI::RegTensor<uint32_t>&)hist64_3,
678+ hist32_1, (MicroAPI::RegTensor<uint32_t>&)zero16);
679+ MicroAPI::Interleave((MicroAPI::RegTensor<uint32_t>&)hist64_4, (MicroAPI::RegTensor<uint32_t>&)hist64_5,
680+ hist32_2, (MicroAPI::RegTensor<uint32_t>&)zero16);
681+ MicroAPI::Interleave((MicroAPI::RegTensor<uint32_t>&)hist64_6, (MicroAPI::RegTensor<uint32_t>&)hist64_7,
682+ hist32_3, (MicroAPI::RegTensor<uint32_t>&)zero16);
683+ 
684+ MicroAPI::RegTensor<int64_t> old0, old1, old2, old3, old4, old5, old6, old7;
C
Ccpy_1234567月6日

命名不规范

likedislike
黄
黄晓彬
7月6日 评论：
685+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old0, histogramReadPtr, VF_LEN_B64);
686+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old1, histogramReadPtr, VF_LEN_B64);
687+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old2, histogramReadPtr, VF_LEN_B64);
688+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old3, histogramReadPtr, VF_LEN_B64);
689+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old4, histogramReadPtr, VF_LEN_B64);
690+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old5, histogramReadPtr, VF_LEN_B64);
691+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old6, histogramReadPtr, VF_LEN_B64);
692+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(old7, histogramReadPtr, VF_LEN_B64);
693+ MicroAPI::Add(old0, old0, hist64_0, maskB64);
694+ MicroAPI::Add(old1, old1, hist64_1, maskB64);
695+ MicroAPI::Add(old2, old2, hist64_2, maskB64);
696+ MicroAPI::Add(old3, old3, hist64_3, maskB64);
697+ MicroAPI::Add(old4, old4, hist64_4, maskB64);
698+ MicroAPI::Add(old5, old5, hist64_5, maskB64);
699+ MicroAPI::Add(old6, old6, hist64_6, maskB64);
700+ MicroAPI::Add(old7, old7, hist64_7, maskB64);
701+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old0, VF_LEN_B64,
702+ maskB64);
703+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old1, VF_LEN_B64,
704+ maskB64);
705+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old2, VF_LEN_B64,
706+ maskB64);
707+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old3, VF_LEN_B64,
708+ maskB64);
709+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old4, VF_LEN_B64,
710+ maskB64);
711+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old5, VF_LEN_B64,
712+ maskB64);
713+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old6, VF_LEN_B64,
714+ maskB64);
715+ MicroAPI::DataCopy<int64_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(histogramWritePtr, old7, VF_LEN_B64,
716+ maskB64);
717+ }
718+}
719+ 
720+template <typename T, typename UT>
721+__aicore__ inline void KthValueRadixSelect<T, UT>::CountByte(int64_t rowOffset, int64_t sliceStart, int64_t sliceCount,
722+ uint32_t shift, UT prefixMask, UT prefixKey,
723+ LocalTensor<uint64_t>& histogram)
724+{
725+ ClearHistogram(histogram);
726+ for (int64_t start = 0; start < sliceCount; start += static_cast<int64_t>(tileElems_)) {
727+ int64_t remain = sliceCount - start;
728+ uint32_t count = remain < static_cast<int64_t>(tileElems_) ? static_cast<uint32_t>(remain) : tileElems_;
729+ LocalTensor<T> input = inputQueue_.AllocTensor<T>();
730+ LoadTile(rowOffset + sliceStart + start, count, input);
731+ inputQueue_.EnQue<T>(input);
732+ input = inputQueue_.DeQue<T>();
733+ LocalTensor<UT> keys = MakeSortKeys(input, count);
734+ LocalTensor<uint16_t> tileHistogram = tileHistogramBuf_.Get<uint16_t>();
735+ BuildTileHistogram(keys, count, shift, prefixMask, prefixKey, tileHistogram);
736+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
737+ SetFlag<HardEvent::V_S>(eventId);
738+ WaitFlag<HardEvent::V_S>(eventId);
739+ AccumulateTileHistogram(histogram, tileHistogram);
740+ inputQueue_.FreeTensor(input);
741+ }
742+}
743+ 
744+template <typename T, typename UT>
745+__aicore__ inline void KthValueRadixSelect<T, UT>::StoreCoreHistogram(LocalTensor<uint64_t>& histogram)
746+{
747+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));
748+ SetFlag<HardEvent::S_MTE3>(eventId);
749+ WaitFlag<HardEvent::S_MTE3>(eventId);
750+ DataCopyExtParams params{1, RADIX * static_cast<uint32_t>(sizeof(uint64_t)), 0, 0, 0};
751+ DataCopyPad(coreHistogramGm_[static_cast<uint64_t>(blockIdx_) * RADIX], histogram, params);
752+ event_t doneEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE3_S));
753+ SetFlag<HardEvent::MTE3_S>(doneEvent);
754+ WaitFlag<HardEvent::MTE3_S>(doneEvent);
755+}
756+ 
757+template <typename T, typename UT>
758+__aicore__ inline uint64_t KthValueRadixSelect<T, UT>::LoadCoreHistogramBucket(uint32_t group, uint32_t coreInGroup,
759+ uint32_t bucket,
760+ LocalTensor<uint64_t>& scratch)
761+{
762+ DataCopyExtParams params{1, static_cast<uint32_t>(sizeof(uint64_t)), 0, 0, 0};
763+ DataCopyPadExtParams<uint64_t> padParams{false, 0, 0, 0};
764+ uint64_t coreOffset = (static_cast<uint64_t>(group) * coresPerRow_ + coreInGroup) * RADIX;
765+ DataCopyPad(scratch, coreHistogramGm_[coreOffset + bucket], params, padParams);
766+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_S));
767+ SetFlag<HardEvent::MTE2_S>(eventId);
768+ WaitFlag<HardEvent::MTE2_S>(eventId);
769+ return scratch.GetValue(0);
770+}
771+ 
772+template <typename T, typename UT>
773+__aicore__ inline void KthValueRadixSelect<T, UT>::ReduceGroupHistogram(uint32_t group,
774+ LocalTensor<uint64_t>& histogram,
775+ LocalTensor<uint64_t>& scratch)
776+{
777+ ClearHistogram(histogram);
778+ DataCopyExtParams params{1, RADIX * static_cast<uint32_t>(sizeof(uint64_t)), 0, 0, 0};
779+ DataCopyPadExtParams<uint64_t> padParams{false, 0, 0, 0};
780+ for (uint32_t core = 0; core < coresPerRow_; ++core) {
781+ uint64_t coreOffset = (static_cast<uint64_t>(group) * coresPerRow_ + core) * RADIX;
782+ DataCopyPad(scratch, coreHistogramGm_[coreOffset], params, padParams);
783+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_S));
784+ SetFlag<HardEvent::MTE2_S>(eventId);
785+ WaitFlag<HardEvent::MTE2_S>(eventId);
786+ AccumulateHistogram(histogram, scratch);
787+ event_t accSyncEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2));
788+ SetFlag<HardEvent::V_MTE2>(accSyncEvent);
789+ WaitFlag<HardEvent::V_MTE2>(accSyncEvent);
790+ }
791+}
792+ 
793+template <typename T, typename UT>
794+__aicore__ inline void KthValueRadixSelect<T, UT>::StoreGroupState(uint32_t group, LocalTensor<uint64_t>& state)
795+{
796+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));
797+ SetFlag<HardEvent::S_MTE3>(eventId);
798+ WaitFlag<HardEvent::S_MTE3>(eventId);
799+ DataCopyExtParams params{1, RADIX_SELECT_RESULT_WORDS * static_cast<uint32_t>(sizeof(uint64_t)), 0, 0, 0};
800+ DataCopyPad(groupStateGm_[static_cast<uint64_t>(group) * RADIX_SELECT_RESULT_WORDS], state, params);
801+ event_t doneEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE3_S));
802+ SetFlag<HardEvent::MTE3_S>(doneEvent);
803+ WaitFlag<HardEvent::MTE3_S>(doneEvent);
804+}
805+ 
806+template <typename T, typename UT>
807+__aicore__ inline void KthValueRadixSelect<T, UT>::LoadGroupState(uint32_t group, LocalTensor<uint64_t>& state)
808+{
809+ DataCopyExtParams params{1, RADIX_SELECT_RESULT_WORDS * static_cast<uint32_t>(sizeof(uint64_t)), 0, 0, 0};
810+ DataCopyPadExtParams<uint64_t> padParams{false, 0, 0, 0};
811+ DataCopyPad(state, groupStateGm_[static_cast<uint64_t>(group) * RADIX_SELECT_RESULT_WORDS], params, padParams);
812+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_S));
813+ SetFlag<HardEvent::MTE2_S>(eventId);
814+ WaitFlag<HardEvent::MTE2_S>(eventId);
815+}
816+ 
817+template <typename T, typename UT>
818+__aicore__ inline void KthValueRadixSelect<T, UT>::WriteOutput(int64_t row, T value, int64_t index)
819+{
820+ LocalTensor<T> outputValue = outputValueBuf_.Get<T>();
821+ LocalTensor<int64_t> outputIndex = outputIndexBuf_.Get<int64_t>();
822+ outputValue.SetValue(0, value);
823+ outputIndex.SetValue(0, index);
824+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));
825+ SetFlag<HardEvent::S_MTE3>(eventId);
826+ WaitFlag<HardEvent::S_MTE3>(eventId);
827+ DataCopyExtParams valueParams{1, static_cast<uint32_t>(sizeof(T)), 0, 0, 0};
828+ DataCopyExtParams indexParams{1, static_cast<uint32_t>(sizeof(int64_t)), 0, 0, 0};
829+ DataCopyPad(valuesGm_[row], outputValue, valueParams);
830+ DataCopyPad(indicesGm_[row], outputIndex, indexParams);
831+}
832+ 
833+// Enter or update active mode: if already active, compact existing indices;
834+// otherwise collect from full data. Stay active only if count in (0, 512].
835+template <typename T, typename UT>
836+__aicore__ inline void KthValueRadixSelect<T, UT>::UpdateActiveIndices(LocalTensor<UT>& retainedKeys,
837+ LocalTensor<uint32_t>& activeIndices,
838+ uint32_t retainedCount, UT prefixMask,
839+ UT prefixKey, bool& activeMode,
840+ uint32_t& activeCount)
841+{
842+ LocalTensor<int64_t> activeResult = findResultBuf_.Get<int64_t>();
843+ if (activeMode) {
844+ asc_vf_call<CompactActiveIndices<UT>>(
845+ dim3(RADIX_SELECT_FIND_THREADS), reinterpret_cast<__ubuf__ UT*>(retainedKeys.GetPhyAddr()),
846+ reinterpret_cast<__ubuf__ uint32_t*>(activeIndices.GetPhyAddr()), activeCount, prefixMask, prefixKey,
847+ reinterpret_cast<__ubuf__ int64_t*>(activeResult.GetPhyAddr()));
848+ } else {
849+ LocalTensor<uint32_t> threadCounts = findCountBuf_.Get<uint32_t>();
850+ asc_vf_call<CollectActiveIndices<UT>>(
851+ dim3(RADIX_SELECT_FIND_THREADS), reinterpret_cast<__ubuf__ UT*>(retainedKeys.GetPhyAddr()), retainedCount,
852+ prefixMask, prefixKey, reinterpret_cast<__ubuf__ uint32_t*>(threadCounts.GetPhyAddr()),
853+ reinterpret_cast<__ubuf__ uint32_t*>(activeIndices.GetPhyAddr()),
854+ reinterpret_cast<__ubuf__ int64_t*>(activeResult.GetPhyAddr()));
855+ }
856+ event_t activeEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
857+ SetFlag<HardEvent::V_S>(activeEvent);
858+ WaitFlag<HardEvent::V_S>(activeEvent);
859+ activeCount = static_cast<uint32_t>(activeResult.GetValue(0));
860+ activeMode = activeCount > 0U && activeCount <= RADIX_SELECT_ACTIVE_MODE_THRESHOLD;
861+}
862+ 
863+// Core algorithm: process one row by narrowing prefix byte-by-byte (MSB to LSB).
864+// Each round: histogram current byte of prefix-matching elements, find which bucket
865+// contains the kth element, extend prefix. Multi-core: each core handles a slice,
866+// leader reduces histograms and broadcasts the selected bucket via workspace GM.
867+template <typename T, typename UT>
868+__aicore__ inline void KthValueRadixSelect<T, UT>::SelectOneRow(int64_t row, uint32_t group, uint32_t coreInGroup,
869+ bool writeOutput)
870+{
871+ LocalTensor<uint64_t> histogram = histogramBuf_.Get<uint64_t>();
C
Ccpy_1234567月6日

超大函数，注意整改，若涉及性能劣化，合入后及时屏蔽

likedislike
黄
黄晓彬
7月6日 评论：
872+ LocalTensor<uint64_t> reduceScratch = groupReduceBuf_.Get<uint64_t>();
873+ uint64_t leftK = static_cast<uint64_t>(kthIndex_); // remaining elements to skip
874+ UT prefixMask = static_cast<UT>(0); // bits already determined
875+ UT prefixKey = static_cast<UT>(0); // values of those bits
876+ int64_t rowOffset = row * axisLen_;
877+ int64_t sliceSize = Ops::Base::CeilDiv(axisLen_, static_cast<int64_t>(coresPerRow_));
878+ int64_t sliceStart = static_cast<int64_t>(coreInGroup) * sliceSize;
879+ int64_t sliceRemain = axisLen_ - sliceStart;
880+ int64_t sliceCount = sliceStart < axisLen_ ? (sliceSize < sliceRemain ? sliceSize : sliceRemain) : 0;
881+ uint32_t finalTargetCore = coreInGroup;
882+ uint64_t finalTargetLeftK = leftK;
883+ bool keepFinalTile = sliceCount > 0 && sliceCount <= static_cast<int64_t>(tileElems_);
884+ bool hasRetainedTile = false;
885+ bool activeMode = false;
886+ uint32_t activeCount = 0;
887+ uint32_t retainedCount = 0;
888+ LocalTensor<T> retainedInput;
889+ LocalTensor<UT> retainedKeys;
890+ LocalTensor<uint32_t> activeIndices = activeIndexBuf_.Get<uint32_t>();
891+ 
892+ for (int32_t byte = static_cast<int32_t>(sizeof(UT)) - 1; byte >= 0; --byte) {
893+ uint32_t shift = static_cast<uint32_t>(byte) * BYTE_BITS;
894+ if (keepFinalTile) {
895+ if (!hasRetainedTile) {
896+ retainedCount = static_cast<uint32_t>(sliceCount);
897+ retainedInput = inputQueue_.AllocTensor<T>();
898+ LoadTile(rowOffset + sliceStart, retainedCount, retainedInput);
899+ inputQueue_.EnQue<T>(retainedInput);
900+ retainedInput = inputQueue_.DeQue<T>();
901+ retainedKeys = MakeSortKeys(retainedInput, retainedCount);
902+ hasRetainedTile = true;
903+ }
904+ if (activeMode) {
905+ asc_vf_call<BuildActiveHistogram<UT>>(
906+ dim3(RADIX_SELECT_FIND_THREADS), reinterpret_cast<__ubuf__ UT*>(retainedKeys.GetPhyAddr()),
907+ reinterpret_cast<__ubuf__ uint32_t*>(activeIndices.GetPhyAddr()), activeCount, shift,
908+ reinterpret_cast<__ubuf__ uint64_t*>(histogram.GetPhyAddr()));
909+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
910+ SetFlag<HardEvent::V_S>(eventId);
911+ WaitFlag<HardEvent::V_S>(eventId);
912+ } else {
913+ ClearHistogram(histogram);
914+ event_t clearEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
915+ SetFlag<HardEvent::V_S>(clearEvent);
916+ WaitFlag<HardEvent::V_S>(clearEvent);
917+ LocalTensor<uint16_t> tileHistogram = tileHistogramBuf_.Get<uint16_t>();
918+ BuildTileHistogram(retainedKeys, retainedCount, shift, prefixMask, prefixKey, tileHistogram);
919+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
920+ SetFlag<HardEvent::V_S>(eventId);
921+ WaitFlag<HardEvent::V_S>(eventId);
922+ AccumulateTileHistogram(histogram, tileHistogram);
923+ }
924+ } else {
925+ CountByte(rowOffset, sliceStart, sliceCount, shift, prefixMask, prefixKey, histogram);
926+ }
927+ event_t histSyncEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
928+ SetFlag<HardEvent::V_S>(histSyncEvent);
929+ WaitFlag<HardEvent::V_S>(histSyncEvent);
930+ if (coresPerRow_ == 1U) {
931+ uint64_t accumulated = 0;
932+ uint32_t selected = 0;
933+ for (; selected < RADIX; ++selected) {
934+ uint64_t count = histogram.GetValue(selected);
935+ if (leftK < accumulated + count) {
936+ leftK -= accumulated;
937+ break;
938+ }
939+ accumulated += count;
940+ }
941+ if (selected >= RADIX) {
942+ if (hasRetainedTile) {
943+ inputQueue_.FreeTensor(retainedInput);
944+ }
945+ return;
946+ }
947+ uint64_t selectedCount = histogram.GetValue(selected);
948+ UT byteMask = static_cast<UT>(static_cast<UT>(0xFFU) << shift);
949+ prefixMask = static_cast<UT>(prefixMask | byteMask);
950+ prefixKey = static_cast<UT>(prefixKey | (static_cast<UT>(selected) << shift));
951+ if (keepFinalTile && selectedCount <= RADIX_SELECT_ACTIVE_MODE_THRESHOLD) {
952+ UpdateActiveIndices(retainedKeys, activeIndices, retainedCount, prefixMask, prefixKey, activeMode,
953+ activeCount);
954+ }
955+ if (selectedCount == 1UL) {
956+ break;
957+ }
958+ continue;
959+ }
960+ uint32_t selected = 0;
961+ uint64_t selectedCount = 0;
962+ StoreCoreHistogram(histogram);
963+ SyncAll();
964+ LocalTensor<uint64_t> state = findResultBuf_.Get<uint64_t>();
965+ if (coreInGroup == 0U) {
966+ ReduceGroupHistogram(group, histogram, reduceScratch);
967+ event_t reduceSyncEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
968+ SetFlag<HardEvent::V_S>(reduceSyncEvent);
969+ WaitFlag<HardEvent::V_S>(reduceSyncEvent);
970+ uint64_t accumulated = 0;
971+ for (; selected < RADIX; ++selected) {
972+ uint64_t count = histogram.GetValue(selected);
973+ if (leftK < accumulated + count) {
974+ leftK -= accumulated;
975+ break;
976+ }
977+ accumulated += count;
978+ }
979+ if (byte == 0 && selected < RADIX) {
980+ uint64_t coreAccumulated = 0;
981+ uint32_t targetCore = coresPerRow_;
982+ uint64_t targetLeftK = leftK;
983+ for (uint32_t core = 0; core < coresPerRow_; ++core) {
984+ uint64_t count = LoadCoreHistogramBucket(group, core, selected, reduceScratch);
985+ if (leftK < coreAccumulated + count) {
986+ targetCore = core;
987+ targetLeftK = leftK - coreAccumulated;
988+ break;
989+ }
990+ coreAccumulated += count;
991+ }
992+ state.SetValue(RADIX_SELECT_STATE_TARGET_CORE_IDX, targetCore);
993+ state.SetValue(RADIX_SELECT_STATE_TARGET_LEFT_K_IDX, targetLeftK);
994+ }
995+ state.SetValue(RADIX_SELECT_STATE_LEFT_K_IDX, leftK);
996+ state.SetValue(RADIX_SELECT_STATE_SELECTED_IDX, selected);
997+ state.SetValue(RADIX_SELECT_STATE_SELECTED_COUNT_IDX,
998+ selected < RADIX ? histogram.GetValue(selected) : 0UL);
999+ StoreGroupState(group, state);
1000+ }
1001+ SyncAll();
1002+ LoadGroupState(group, state);
1003+ selected = static_cast<uint32_t>(state.GetValue(RADIX_SELECT_STATE_SELECTED_IDX));
1004+ if (selected >= RADIX) {
1005+ if (hasRetainedTile) {
1006+ inputQueue_.FreeTensor(retainedInput);
1007+ }
1008+ return;
1009+ }
1010+ leftK = state.GetValue(RADIX_SELECT_STATE_LEFT_K_IDX);
1011+ selectedCount = state.GetValue(RADIX_SELECT_STATE_SELECTED_COUNT_IDX);
1012+ if (byte == 0) {
1013+ finalTargetCore = static_cast<uint32_t>(state.GetValue(RADIX_SELECT_STATE_TARGET_CORE_IDX));
1014+ finalTargetLeftK = state.GetValue(RADIX_SELECT_STATE_TARGET_LEFT_K_IDX);
1015+ }
1016+ UT byteMask = static_cast<UT>(static_cast<UT>(0xFFU) << shift);
1017+ prefixMask = static_cast<UT>(prefixMask | byteMask);
1018+ prefixKey = static_cast<UT>(prefixKey | (static_cast<UT>(selected) << shift));
1019+ if (keepFinalTile && selectedCount <= RADIX_SELECT_ACTIVE_MODE_THRESHOLD) {
1020+ UpdateActiveIndices(retainedKeys, activeIndices, retainedCount, prefixMask, prefixKey, activeMode,
1021+ activeCount);
1022+ }
1023+ }
1024+ 
1025+ if (coresPerRow_ > 1U) {
1026+ if (coreInGroup != finalTargetCore) {
1027+ if (hasRetainedTile) {
1028+ inputQueue_.FreeTensor(retainedInput);
1029+ }
1030+ return;
1031+ }
1032+ leftK = finalTargetLeftK;
1033+ }
1034+ 
1035+ // Final location: three paths depending on retained tile and active mode
1036+ // Path A: active mode + retained tile → O(1) lookup from activeIndices
1037+ if (hasRetainedTile && activeMode && leftK < static_cast<uint64_t>(activeCount)) {
1038+ uint32_t found = activeIndices.GetValue(static_cast<uint32_t>(leftK));
1039+ if (writeOutput) {
1040+ WriteOutput(row, retainedInput.GetValue(found), sliceStart + found);
1041+ }
1042+ inputQueue_.FreeTensor(retainedInput);
1043+ return;
1044+ }
1045+ 
1046+ if (hasRetainedTile) {
1047+ LocalTensor<uint32_t> threadCounts = findCountBuf_.Get<uint32_t>();
1048+ LocalTensor<int64_t> findResult = findResultBuf_.Get<int64_t>();
1049+ asc_vf_call<FindKthMatchInTile<UT>>(
1050+ dim3(RADIX_SELECT_FIND_THREADS), reinterpret_cast<__ubuf__ UT*>(retainedKeys.GetPhyAddr()), retainedCount,
1051+ prefixMask, prefixKey, leftK, reinterpret_cast<__ubuf__ uint32_t*>(threadCounts.GetPhyAddr()),
1052+ reinterpret_cast<__ubuf__ int64_t*>(findResult.GetPhyAddr()));
1053+ event_t findEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
1054+ SetFlag<HardEvent::V_S>(findEvent);
1055+ WaitFlag<HardEvent::V_S>(findEvent);
1056+ int64_t found = findResult.GetValue(0);
1057+ if (found >= 0 && writeOutput) {
1058+ WriteOutput(row, retainedInput.GetValue(static_cast<uint32_t>(found)), sliceStart + found);
1059+ }
1060+ inputQueue_.FreeTensor(retainedInput);
1061+ return;
1062+ }
1063+ 
1064+ uint64_t matched = 0;
1065+ for (int64_t start = 0; start < sliceCount; start += static_cast<int64_t>(tileElems_)) {
1066+ int64_t remain = sliceCount - start;
1067+ uint32_t count = remain < static_cast<int64_t>(tileElems_) ? static_cast<uint32_t>(remain) : tileElems_;
1068+ LocalTensor<T> input = inputQueue_.AllocTensor<T>();
1069+ LoadTile(rowOffset + sliceStart + start, count, input);
1070+ inputQueue_.EnQue<T>(input);
1071+ input = inputQueue_.DeQue<T>();
1072+ LocalTensor<UT> keys = MakeSortKeys(input, count);
1073+ event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
1074+ SetFlag<HardEvent::V_S>(eventId);
1075+ WaitFlag<HardEvent::V_S>(eventId);
1076+ LocalTensor<uint32_t> threadCounts = findCountBuf_.Get<uint32_t>();
1077+ LocalTensor<int64_t> findResult = findResultBuf_.Get<int64_t>();
1078+ uint64_t tileTarget = leftK >= matched ? leftK - matched : 0UL;
1079+ asc_vf_call<FindKthMatchInTile<UT>>(
1080+ dim3(RADIX_SELECT_FIND_THREADS), reinterpret_cast<__ubuf__ UT*>(keys.GetPhyAddr()), count, prefixMask,
1081+ prefixKey, tileTarget, reinterpret_cast<__ubuf__ uint32_t*>(threadCounts.GetPhyAddr()),
1082+ reinterpret_cast<__ubuf__ int64_t*>(findResult.GetPhyAddr()));
1083+ event_t findEvent = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
1084+ SetFlag<HardEvent::V_S>(findEvent);
1085+ WaitFlag<HardEvent::V_S>(findEvent);
1086+ int64_t found = findResult.GetValue(0);
1087+ uint64_t tileMatches = static_cast<uint64_t>(findResult.GetValue(1));
1088+ if (found >= 0 && leftK >= matched) {
1089+ if (writeOutput) {
1090+ WriteOutput(row, input.GetValue(static_cast<uint32_t>(found)), sliceStart + start + found);
1091+ }
1092+ inputQueue_.FreeTensor(input);
1093+ break;
1094+ }
1095+ matched += tileMatches;
1096+ inputQueue_.FreeTensor(input);
1097+ }
1098+}
1099+ 
1100+template <typename T, typename UT>
1101+__aicore__ inline void KthValueRadixSelect<T, UT>::Process()
1102+{
1103+ if (tileElems_ == 0U || blockNum_ == 0U || rowsParallel_ == 0U || coresPerRow_ == 0U || sortLoopTimes_ == 0U) {
1104+ return;
1105+ }
1106+ uint32_t group = blockIdx_ / coresPerRow_;
1107+ uint32_t coreInGroup = blockIdx_ % coresPerRow_;
1108+ for (uint32_t loop = 0; loop < sortLoopTimes_; ++loop) {
1109+ int64_t logicalRow = static_cast<int64_t>(loop) * rowsParallel_ + group;
1110+ bool writeOutput = logicalRow < rowCount_;
1111+ int64_t row = writeOutput ? logicalRow : 0;
1112+ SelectOneRow(row, group, coreInGroup, writeOutput);
1113+ }
1114+}
1115+} // namespace KthValue
1116+ 
1117+#endif
Amath/kth_value/op_kernel/arch35/kth_value_small_axis_short_rank_select.h+296-0
@@ -0,0 +1,296 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
C
Ccpy_1234567月6日

修改过大，已经不是问题单的修改量，后续这种还是尽量走需求，由测试人员进行测试兜底

likedislike
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef KTH_VALUE_SMALL_AXIS_SHORT_RANK_SELECT_H
12+#define KTH_VALUE_SMALL_AXIS_SHORT_RANK_SELECT_H
13+ 
14+#include "kernel_operator.h"
15+#include "op_kernel/platform_util.h"
16+#include "simt_api/asc_simt.h"
17+#include "kth_value_tiling_data.h"
18+ 
19+// Top-K partial selection for short axis (<=32) kth_value.
20+// Maintains a sorted candidate array of size K, scans once per row via SIMT.
21+namespace KthValue {
22+using namespace AscendC;
23+ 
24+constexpr uint32_t SHORT_RANK_SELECT_THREAD_NUM = 1024;
25+constexpr uint32_t SHORT_RANK_SELECT_MAX_CANDIDATES = 8;
26+ 
27+// Returns true if (value, index) is strictly better than (baseValue, baseIndex).
28+template <typename T, bool SelectLargest>
29+__simt_callee__ __aicore__ inline bool IsShortRankSelectBetter(T value, uint32_t index, T baseValue, uint32_t baseIndex)
30+{
31+ if constexpr (SelectLargest) {
32+ return (value > baseValue) || ((value == baseValue) && (index > baseIndex));
33+ } else {
34+ return (value < baseValue) || ((value == baseValue) && (index < baseIndex));
35+ }
36+}
37+ 
38+// Bubble-up: swap candidate at pos with pos-1 if it is better (insertion-sort step).
39+template <typename T, bool SelectLargest>
40+__simt_callee__ __aicore__ inline void PromoteShortRankSelectCandidate(T* candidateValues, uint32_t* candidateIndices,
41+ uint32_t pos)
42+{
43+ T rightValue = candidateValues[pos];
44+ uint32_t rightIndex = candidateIndices[pos];
45+ if (!IsShortRankSelectBetter<T, SelectLargest>(rightValue, rightIndex, candidateValues[pos - 1U],
46+ candidateIndices[pos - 1U])) {
47+ return;
48+ }
49+ candidateValues[pos] = candidateValues[pos - 1U];
50+ candidateIndices[pos] = candidateIndices[pos - 1U];
51+ candidateValues[pos - 1U] = rightValue;
52+ candidateIndices[pos - 1U] = rightIndex;
53+}
54+ 
55+// SIMT kernel: each thread does top-K selection on one row. CandidateNum is compile-time constant.
56+template <typename T, uint32_t CandidateNum, bool SelectLargest>
57+__simt_vf__ LAUNCH_BOUND(SHORT_RANK_SELECT_THREAD_NUM) __aicore__
58+ void SimtSelectShortRankFromUbFixed(uint32_t validSegs, uint32_t segmentLen, __ubuf__ T* input,
59+ __ubuf__ T* outputValue, __ubuf__ int64_t* outputIndex)
60+{
61+ // Grid-stride loop: each thread processes rows with stride = thread count
62+ for (uint32_t row = static_cast<uint32_t>(threadIdx.x); row < validSegs; row += SHORT_RANK_SELECT_THREAD_NUM) {
63+ T candidateValues[SHORT_RANK_SELECT_MAX_CANDIDATES];
64+ uint32_t candidateIndices[SHORT_RANK_SELECT_MAX_CANDIDATES];
65+ uint32_t rowOffset = row * segmentLen;
66+ 
67+ // Phase 1: Initialize candidate array with first K elements
68+ for (uint32_t col = 0U; col < CandidateNum; ++col) {
69+ candidateValues[col] = input[rowOffset + col];
70+ candidateIndices[col] = col;
71+ }
72+ // Phase 2: Sort initial candidates using insertion sort (candidateValues[0] = best, [K-1] = worst)
73+ for (uint32_t pass = 1U; pass < CandidateNum; ++pass) {
74+ for (uint32_t pos = CandidateNum - 1U; pos >= pass; --pos) {
75+ PromoteShortRankSelectCandidate<T, SelectLargest>(candidateValues, candidateIndices, pos);
76+ }
77+ }
78+ 
79+ // Phase 3: Scan remaining elements, maintain top-K heap
80+ for (uint32_t col = CandidateNum; col < segmentLen; ++col) {
81+ T value = input[rowOffset + col];
82+ // Skip if not better than the worst candidate (threshold)
83+ if (!IsShortRankSelectBetter<T, SelectLargest>(value, col, candidateValues[CandidateNum - 1U],
84+ candidateIndices[CandidateNum - 1U])) {
85+ continue;
86+ }
87+ 
88+ // Replace worst candidate and bubble up to correct position
89+ candidateValues[CandidateNum - 1U] = value;
90+ candidateIndices[CandidateNum - 1U] = col;
91+ for (uint32_t pos = CandidateNum - 1U; pos > 0U; --pos) {
92+ PromoteShortRankSelectCandidate<T, SelectLargest>(candidateValues, candidateIndices, pos);
93+ }
94+ }
95+ 
96+ // Output the K-th element (worst among top-K = the answer)
97+ outputValue[row] = candidateValues[CandidateNum - 1U];
98+ outputIndex[row] = static_cast<int64_t>(candidateIndices[CandidateNum - 1U]);
99+ }
100+}
101+ 
102+// Recursive template dispatch: maps runtime candidateNum to compile-time CandidateNum for loop unrolling.
103+template <typename T, bool SelectLargest, uint32_t CandidateNum = 1U>
104+__aicore__ inline void DispatchShortRankSelectFixed(uint32_t candidateNum, uint32_t validSegs, uint32_t segmentLen,
105+ __ubuf__ T* input, __ubuf__ T* outputValue,
106+ __ubuf__ int64_t* outputIndex)
107+{
108+ if (candidateNum == CandidateNum) {
109+ asc_vf_call<SimtSelectShortRankFromUbFixed<T, CandidateNum, SelectLargest>>(
110+ dim3(SHORT_RANK_SELECT_THREAD_NUM), validSegs, segmentLen, input, outputValue, outputIndex);
111+ return;
112+ }
113+ if constexpr (CandidateNum < SHORT_RANK_SELECT_MAX_CANDIDATES) {
114+ DispatchShortRankSelectFixed<T, SelectLargest, CandidateNum + 1U>(candidateNum, validSegs, segmentLen, input,
115+ outputValue, outputIndex);
116+ }
117+}
118+ 
119+template <typename T>
120+class KthValueSmallAxisShortRankSelect {
121+public:
122+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
123+ TPipe* pipe);
124+ __aicore__ inline void Process();
125+ 
126+private:
127+ __aicore__ inline bool IsProcessInvalid() const;
128+ __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchId) const;
129+ __aicore__ inline int64_t GetInputStart(uint32_t batchId) const;
130+ __aicore__ inline void CopyIn(uint32_t batchId, uint32_t validSegs);
131+ __aicore__ inline void SelectShortRank(uint32_t validSegs, LocalTensor<T>& input, LocalTensor<T>& outputValue,
132+ LocalTensor<int64_t>& outputIndex);
133+ __aicore__ inline void CopyOut(uint32_t batchId, uint32_t validSegs, LocalTensor<T>& outputValue,
134+ LocalTensor<int64_t>& outputIndex);
135+ 
136+ TQue<QuePosition::VECIN, 1> inputQueue_;
137+ TQue<QuePosition::VECOUT, 1> outputValueQueue_;
138+ TQue<QuePosition::VECOUT, 1> outputIndexQueue_;
139+ 
140+ GlobalTensor<T> inputGm_;
141+ GlobalTensor<T> valueGm_;
142+ GlobalTensor<int64_t> indexGm_;
143+ const KthValueTilingData* tiling_ = nullptr;
144+ TPipe* pipe_ = nullptr;
145+ uint32_t blockIdx_ = 0;
146+ uint32_t blockDim_ = 0;
147+ uint32_t batchSize_ = 0;
148+ uint32_t batchNum_ = 0;
149+ uint32_t segmentLen_ = 0;
150+ uint32_t kthIndex_ = 0;
151+ int64_t totalSegs_ = 0;
152+};
153+ 
154+template <typename T>
155+__aicore__ inline void KthValueSmallAxisShortRankSelect<T>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
156+ const KthValueTilingData* tiling, TPipe* pipe)
157+{
158+ if (tiling == nullptr || pipe == nullptr) {
159+ return;
160+ }
161+ tiling_ = tiling;
162+ pipe_ = pipe;
163+ blockIdx_ = GetBlockIdx();
164+ blockDim_ = GetBlockNum();
165+ batchSize_ = tiling_->keyParams0;
166+ batchNum_ = tiling_->keyParams1;
167+ segmentLen_ = tiling_->numTileDataSize;
168+ kthIndex_ = static_cast<uint32_t>(tiling_->kthIndex);
169+ totalSegs_ = tiling_->unsortedDimNum;
170+ 
171+ inputGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
172+ valueGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(values));
173+ indexGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(indices));
174+ 
175+ if (batchSize_ == 0U || segmentLen_ == 0U) {
176+ return;
177+ }
178+ 
179+ pipe_->InitBuffer(inputQueue_, 1, ROUND_UP_AGLIN(batchSize_ * segmentLen_ * sizeof(T)));
180+ pipe_->InitBuffer(outputValueQueue_, 1, ROUND_UP_AGLIN(batchSize_ * sizeof(T)));
181+ pipe_->InitBuffer(outputIndexQueue_, 1, ROUND_UP_AGLIN(batchSize_ * sizeof(int64_t)));
182+}
183+ 
184+template <typename T>
185+__aicore__ inline bool KthValueSmallAxisShortRankSelect<T>::IsProcessInvalid() const
186+{
187+ if (tiling_ == nullptr || pipe_ == nullptr || blockIdx_ >= blockDim_ || blockDim_ == 0U || batchSize_ == 0U ||
188+ batchNum_ == 0U || segmentLen_ == 0U || totalSegs_ <= 0 || segmentLen_ > 32U || kthIndex_ >= segmentLen_) {
189+ return true;
190+ }
191+ uint32_t rankFromHead = kthIndex_ + 1U;
192+ uint32_t rankFromTail = segmentLen_ - kthIndex_;
193+ uint32_t candidateNum = rankFromTail < rankFromHead ? rankFromTail : rankFromHead;
194+ return candidateNum == 0U || candidateNum > SHORT_RANK_SELECT_MAX_CANDIDATES;
195+}
196+ 
197+template <typename T>
198+__aicore__ inline uint32_t KthValueSmallAxisShortRankSelect<T>::ComputeValidSegs(uint32_t batchId) const
199+{
200+ int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);
201+ int64_t remain = totalSegs_ - segStart;
202+ if (remain <= 0) {
203+ return 0;
204+ }
205+ return remain >= static_cast<int64_t>(batchSize_) ? batchSize_ : static_cast<uint32_t>(remain);
206+}
207+ 
208+template <typename T>
209+__aicore__ inline int64_t KthValueSmallAxisShortRankSelect<T>::GetInputStart(uint32_t batchId) const
210+{
211+ return static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_) * static_cast<int64_t>(segmentLen_);
212+}
213+ 
214+template <typename T>
215+__aicore__ inline void KthValueSmallAxisShortRankSelect<T>::CopyIn(uint32_t batchId, uint32_t validSegs)
216+{
217+ LocalTensor<T> input = inputQueue_.AllocTensor<T>();
218+ uint32_t totalElems = validSegs * segmentLen_;
219+ DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
220+ DataCopyExtParams copyParam{1, static_cast<uint32_t>(totalElems * sizeof(T)), 0, 0, 0};
221+ DataCopyPad(input, inputGm_[GetInputStart(batchId)], copyParam, padParams);
222+ inputQueue_.EnQue<T>(input);
223+}
224+ 
225+template <typename T>
226+__aicore__ inline void KthValueSmallAxisShortRankSelect<T>::SelectShortRank(uint32_t validSegs, LocalTensor<T>& input,
227+ LocalTensor<T>& outputValue,
228+ LocalTensor<int64_t>& outputIndex)
229+{
230+ if (segmentLen_ == 0U || segmentLen_ > 32U || kthIndex_ >= segmentLen_) {
231+ return;
232+ }
233+ // Pick the shorter direction: find K-th smallest from head or K-th largest from tail.
234+ uint32_t rankFromHead = kthIndex_ + 1U;
235+ uint32_t rankFromTail = segmentLen_ - kthIndex_;
236+ bool selectLargest = rankFromTail < rankFromHead;
237+ uint32_t candidateNum = selectLargest ? rankFromTail : rankFromHead;
238+ if (candidateNum == 0U || candidateNum > SHORT_RANK_SELECT_MAX_CANDIDATES) {
239+ return;
240+ }
atomgit-bot
atomgit-botatomgit-bot7月6日

🟡 Medium Priority

SelectShortRank（行 208-232）在以下条件不满足时直接 return 而不填充 outputValue / outputIndex：

而 Process()（行 245-270）在调用 SelectShortRank 之后，无条件地将 outputValue / outputIndex 通过 CopyOut 写入 GM。IsProcessInvalid() 仅检查 segmentLen_ == 0U，不覆盖 segmentLen_ > 32U 或 kthIndex_ >= segmentLen_ 或 candidateNum > 8 的情况。

当 tiling 因 bug 将不满足前置条件的 case 路由到本 kernel 时，UB 中的未初始化数据被静默写入 GM 输出，造成难以排查的精度/正确性问题。

建议：将 SelectShortRank 中的全部前置条件检查统一提升到 IsProcessInvalid() 中（或在 SelectShortRank 提前返回前对 output 做清零/哨兵填充），确保绝不会将未初始化的 UB 数据拷贝到 GM 输出。最简单且安全的方案：在 IsProcessInvalid() 中增加 segmentLen_ > 32U || kthIndex_ >= segmentLen_ 的检查，这样 Process() 直接返回，不会进入 CopyOut。

likedislike
不准确？
黄
黄晓彬
7月6日 评论：
241+ 
242+ __ubuf__ T* inputPtr = (__ubuf__ T*)input.GetPhyAddr();
243+ __ubuf__ T* outputValuePtr = (__ubuf__ T*)outputValue.GetPhyAddr();
244+ __ubuf__ int64_t* outputIndexPtr = (__ubuf__ int64_t*)outputIndex.GetPhyAddr();
245+ if (selectLargest) {
246+ DispatchShortRankSelectFixed<T, true>(candidateNum, validSegs, segmentLen_, inputPtr, outputValuePtr,
247+ outputIndexPtr);
248+ } else {
249+ DispatchShortRankSelectFixed<T, false>(candidateNum, validSegs, segmentLen_, inputPtr, outputValuePtr,
250+ outputIndexPtr);
251+ }
252+}
253+ 
254+template <typename T>
255+__aicore__ inline void KthValueSmallAxisShortRankSelect<T>::CopyOut(uint32_t batchId, uint32_t validSegs,
256+ LocalTensor<T>& outputValue,
257+ LocalTensor<int64_t>& outputIndex)
258+{
259+ int64_t outputStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);
260+ DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(validSegs * sizeof(T)), 0, 0, 0};
261+ DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(validSegs * sizeof(int64_t)), 0, 0, 0};
262+ DataCopyPad(valueGm_[outputStart], outputValue, valueCopyParam);
263+ DataCopyPad(indexGm_[outputStart], outputIndex, indexCopyParam);
264+}
265+ 
266+template <typename T>
267+__aicore__ inline void KthValueSmallAxisShortRankSelect<T>::Process()
268+{
269+ if (IsProcessInvalid()) {
270+ return;
271+ }
272+ // Grid-stride loop: each core processes ceil(batchNum/blockDim) batches.
273+ for (uint32_t batchId = blockIdx_; batchId < batchNum_; batchId += blockDim_) {
274+ uint32_t validSegs = ComputeValidSegs(batchId);
275+ if (validSegs == 0U) {
276+ continue;
277+ }
278+ CopyIn(batchId, validSegs);
279+ LocalTensor<T> input = inputQueue_.DeQue<T>();
280+ LocalTensor<T> outputValue = outputValueQueue_.AllocTensor<T>();
281+ LocalTensor<int64_t> outputIndex = outputIndexQueue_.AllocTensor<int64_t>();
282+ SelectShortRank(validSegs, input, outputValue, outputIndex);
283+ inputQueue_.FreeTensor(input);
284+ outputValueQueue_.EnQue<T>(outputValue);
285+ outputIndexQueue_.EnQue<int64_t>(outputIndex);
286+ outputValue = outputValueQueue_.DeQue<T>();
287+ outputIndex = outputIndexQueue_.DeQue<int64_t>();
288+ CopyOut(batchId, validSegs, outputValue, outputIndex);
289+ outputValueQueue_.FreeTensor(outputValue);
290+ outputIndexQueue_.FreeTensor(outputIndex);
291+ }
292+}
293+ 
294+} // namespace KthValue
295+ 
296+#endif
Mmath/kth_value/op_kernel/arch35/kth_value_tiling_key.h+30-18
@@ -24,14 +24,17 @@
24#define KTH_VALUE_SCHID_AXIS_ONE_COPY 724#define KTH_VALUE_SCHID_AXIS_ONE_COPY 7
25#define KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS 925#define KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS 9
26#define KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX 1026#define KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX 10
27+#define KTH_VALUE_SCHID_RADIX_SELECT 11
28+#define KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT 12
27 29 
28-#define KTH_VALUE_TPL_KEY_DECL() \30+#define KTH_VALUE_TPL_KEY_DECL() \
29- ASCENDC_TPL_UINT_DECL( \31+ ASCENDC_TPL_UINT_DECL(schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_MERGE_SORT, \
30- schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_MERGE_SORT, KTH_VALUE_SCHID_RADIX_ONE_CORE, \32+ KTH_VALUE_SCHID_RADIX_ONE_CORE, KTH_VALUE_SCHID_RADIX_MORE_CORE, \
31- KTH_VALUE_SCHID_RADIX_MORE_CORE, KTH_VALUE_SCHID_MERGE_MORE_CORE, KTH_VALUE_SCHID_MERGE_INTRA_CORE, \33+ KTH_VALUE_SCHID_MERGE_MORE_CORE, KTH_VALUE_SCHID_MERGE_INTRA_CORE, \
32- KTH_VALUE_SCHID_SMALL_AXIS_INSERTION, KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE, KTH_VALUE_SCHID_AXIS_ONE_COPY, \34+ KTH_VALUE_SCHID_SMALL_AXIS_INSERTION, KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE, \
33- KTH_VALUE_SCHID_SORT32_SMALL_AXIS, KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS, \35+ KTH_VALUE_SCHID_AXIS_ONE_COPY, KTH_VALUE_SCHID_SORT32_SMALL_AXIS, \
34- KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX), \36+ KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS, KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX, \
37+ KTH_VALUE_SCHID_RADIX_SELECT, KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT), \
35 ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1)38 ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1)
36 39 
37#define KTH_VALUE_TPL_MERGE_SORT_KEY_SEL() \40#define KTH_VALUE_TPL_MERGE_SORT_KEY_SEL() \
@@ -74,18 +77,27 @@
74 ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX), \77 ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX), \
75 ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1)78 ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1)
76 79 
80+#define KTH_VALUE_TPL_RADIX_SELECT_KEY_SEL() \
81+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_RADIX_SELECT), \
82+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1)
83+ 
84+#define KTH_VALUE_TPL_SMALL_AXIS_SHORT_RANK_SELECT_KEY_SEL() \
85+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT), \
86+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1)
87+ 
77ASCENDC_TPL_ARGS_DECL(KthValue, KTH_VALUE_TPL_KEY_DECL());88ASCENDC_TPL_ARGS_DECL(KthValue, KTH_VALUE_TPL_KEY_DECL());
78 89 
79-ASCENDC_TPL_SEL(90+ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_MERGE_SORT_KEY_SEL()),
80- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_MERGE_SORT_KEY_SEL()),91+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_ONE_CORE_KEY_SEL()),
81- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_ONE_CORE_KEY_SEL()),92+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_MORE_CORE_KEY_SEL()),
82- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_MORE_CORE_KEY_SEL()),93+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_MERGE_MORE_CORE_KEY_SEL()),
83- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_MERGE_MORE_CORE_KEY_SEL()),94+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_MERGE_INTRA_CORE_KEY_SEL()),
84- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_MERGE_INTRA_CORE_KEY_SEL()),95+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_INSERTION_KEY_SEL()),
85- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_INSERTION_KEY_SEL()),96+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_TWO_STAGE_KEY_SEL()),
86- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_TWO_STAGE_KEY_SEL()),97+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_AXIS_ONE_COPY_KEY_SEL()),
87- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_AXIS_ONE_COPY_KEY_SEL()),98+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_NON_LAST_SMALL_AXIS_KEY_SEL()),
88- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_NON_LAST_SMALL_AXIS_KEY_SEL()),99+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_NON_LAST_SMALL_AXIS_RADIX_KEY_SEL()),
89- ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_NON_LAST_SMALL_AXIS_RADIX_KEY_SEL()));100+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_SELECT_KEY_SEL()),
101+ ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_SHORT_RANK_SELECT_KEY_SEL()));
90 102 
91#endif103#endif
Mmath/kth_value/op_kernel/kth_value_apt.cpp+74-23
@@ -17,7 +17,9 @@
17#include "arch35/kth_value_non_last_small_axis.h"17#include "arch35/kth_value_non_last_small_axis.h"
18#include "arch35/kth_value_radix_more_core.h"18#include "arch35/kth_value_radix_more_core.h"
19#include "arch35/kth_value_radix_one_core.h"19#include "arch35/kth_value_radix_one_core.h"
20+#include "arch35/kth_value_radix_select.h"
20#include "arch35/kth_value_small_axis_insertion.h"21#include "arch35/kth_value_small_axis_insertion.h"
22+#include "arch35/kth_value_small_axis_short_rank_select.h"
21#include "arch35/kth_value_small_axis_two_stage.h"23#include "arch35/kth_value_small_axis_two_stage.h"
22#include "arch35/kth_value_tiling_data.h"24#include "arch35/kth_value_tiling_data.h"
23#include "arch35/kth_value_tiling_key.h"25#include "arch35/kth_value_tiling_key.h"
@@ -40,8 +42,8 @@ __aicore__ inline void RunMergeSortRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthV
40}42}
41 43 
42template <uint64_t isInt32>44template <uint64_t isInt32>
43-__aicore__ inline void RunRadixMoreCoreRoute(45+__aicore__ inline void RunRadixMoreCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
44- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData, TPipe* pipe)46+ KthValueTilingData* tilingData, TPipe* pipe)
45{47{
46 if constexpr (sizeof(DTYPE_X) == 1) {48 if constexpr (sizeof(DTYPE_X) == 1) {
47 if constexpr (isInt32 == 1) {49 if constexpr (isInt32 == 1) {
@@ -89,8 +91,8 @@ __aicore__ inline void RunRadixMoreCoreRoute(
89 }91 }
90}92}
91 93 
92-__aicore__ inline void RunSmallAxisInsertionRoute(94+__aicore__ inline void RunSmallAxisInsertionRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
93- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData, TPipe* pipe)95+ TPipe* pipe)
94{96{
95 if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) {97 if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) {
96 KthValue::KthValueSmallAxisInsertion<DTYPE_X, float> op;98 KthValue::KthValueSmallAxisInsertion<DTYPE_X, float> op;
@@ -103,8 +105,8 @@ __aicore__ inline void RunSmallAxisInsertionRoute(
103 }105 }
104}106}
105 107 
106-__aicore__ inline void RunMergeMoreCoreRoute(108+__aicore__ inline void RunMergeMoreCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
107- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData, TPipe* pipe)109+ KthValueTilingData* tilingData, TPipe* pipe)
108{110{
109 if constexpr (IsSameType<float, DTYPE_X>::value) {111 if constexpr (IsSameType<float, DTYPE_X>::value) {
110 KthValue::KthValueMergeSortMoreCore<DTYPE_X, DTYPE_X, false, int64_t> op;112 KthValue::KthValueMergeSortMoreCore<DTYPE_X, DTYPE_X, false, int64_t> op;
@@ -113,8 +115,8 @@ __aicore__ inline void RunMergeMoreCoreRoute(
113 }115 }
114}116}
115 117 
116-__aicore__ inline void RunMergeIntraCoreRoute(118+__aicore__ inline void RunMergeIntraCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
117- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData, TPipe* pipe)119+ KthValueTilingData* tilingData, TPipe* pipe)
118{120{
119 if constexpr (IsSameType<float, DTYPE_X>::value) {121 if constexpr (IsSameType<float, DTYPE_X>::value) {
120 KthValue::KthValueMergeIntraCore<DTYPE_X, int64_t, false> op;122 KthValue::KthValueMergeIntraCore<DTYPE_X, int64_t, false> op;
@@ -124,13 +126,12 @@ __aicore__ inline void RunMergeIntraCoreRoute(
124}126}
125 127 
126template <bool useMergeSort>128template <bool useMergeSort>
127-__aicore__ inline void RunNonLastSmallAxisRoute(129+__aicore__ inline void RunNonLastSmallAxisRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
128- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData, TPipe* pipe)130+ KthValueTilingData* tilingData, TPipe* pipe)
129{131{
130 if constexpr (useMergeSort) {132 if constexpr (useMergeSort) {
131- if constexpr (133+ if constexpr (IsSameType<DTYPE_X, float>::value || IsSameType<DTYPE_X, half>::value ||
132- IsSameType<DTYPE_X, float>::value || IsSameType<DTYPE_X, half>::value ||134+ IsSameType<DTYPE_X, bfloat16_t>::value) {
133- IsSameType<DTYPE_X, bfloat16_t>::value) {
134 KthValue::KthValueNonLastSmallAxis<DTYPE_X, false, true> op;135 KthValue::KthValueNonLastSmallAxis<DTYPE_X, false, true> op;
135 op.Init(x, y1, y2, workspace, tilingData, pipe);136 op.Init(x, y1, y2, workspace, tilingData, pipe);
136 op.Process();137 op.Process();
@@ -142,30 +143,81 @@ __aicore__ inline void RunNonLastSmallAxisRoute(
142 }143 }
143}144}
144 145 
145-__aicore__ inline void RunAxisOneCopyRoute(146+__aicore__ inline void RunAxisOneCopyRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
146- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData, TPipe* pipe)147+ KthValueTilingData* tilingData, TPipe* pipe)
147{148{
148 KthValue::KthValueAxisOneCopy<DTYPE_X> op;149 KthValue::KthValueAxisOneCopy<DTYPE_X> op;
149 op.Init(x, y1, y2, workspace, tilingData, pipe);150 op.Init(x, y1, y2, workspace, tilingData, pipe);
150 op.Process();151 op.Process();
151}152}
152 153 
153-__aicore__ inline void RunRadixOneCoreRoute(154+__aicore__ inline void RunRadixOneCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
154- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData, TPipe* pipe)155+ TPipe* pipe)
155{156{
156 KthValue::KthValueRadixOneCore<DTYPE_X> op;157 KthValue::KthValueRadixOneCore<DTYPE_X> op;
157 op.Init(x, y1, y2, tilingData, pipe);158 op.Init(x, y1, y2, tilingData, pipe);
158 op.Process();159 op.Process();
159}160}
160 161 
161-__aicore__ inline void RunSmallAxisTwoStageRoute(162+__aicore__ inline void RunRadixSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
162- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData, TPipe* pipe)163+ KthValueTilingData* tilingData, TPipe* pipe)
164+{
165+ if constexpr (sizeof(DTYPE_X) == 1) {
166+ KthValue::KthValueRadixSelect<DTYPE_X, uint8_t> op;
167+ op.Init(x, y1, y2, workspace, tilingData, pipe);
168+ op.Process();
169+ } else if constexpr (sizeof(DTYPE_X) == 2) {
170+ KthValue::KthValueRadixSelect<DTYPE_X, uint16_t> op;
171+ op.Init(x, y1, y2, workspace, tilingData, pipe);
172+ op.Process();
173+ } else if constexpr (sizeof(DTYPE_X) == 4) {
174+ KthValue::KthValueRadixSelect<DTYPE_X, uint32_t> op;
175+ op.Init(x, y1, y2, workspace, tilingData, pipe);
176+ op.Process();
177+ } else if constexpr (sizeof(DTYPE_X) == 8) {
178+ KthValue::KthValueRadixSelect<DTYPE_X, uint64_t> op;
179+ op.Init(x, y1, y2, workspace, tilingData, pipe);
180+ op.Process();
181+ }
182+}
183+ 
184+__aicore__ inline void RunSmallAxisTwoStageRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
185+ TPipe* pipe)
163{186{
164 KthValue::KthValueSmallAxisTwoStage<DTYPE_X> op;187 KthValue::KthValueSmallAxisTwoStage<DTYPE_X> op;
165 op.Init(x, y1, y2, tilingData, pipe);188 op.Init(x, y1, y2, tilingData, pipe);
166 op.Process();189 op.Process();
167}190}
168 191 
192+__aicore__ inline void RunSmallAxisShortRankSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2,
193+ KthValueTilingData* tilingData, TPipe* pipe)
194+{
195+ if constexpr (sizeof(DTYPE_X) == 8) {
196+ KthValue::KthValueSmallAxisShortRankSelect<DTYPE_X> op;
197+ op.Init(x, y1, y2, tilingData, pipe);
198+ op.Process();
199+ }
200+}
201+ 
202+template <uint64_t schId>
203+__aicore__ inline bool TryRunSmallAxisRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
204+ TPipe* pipe)
205+{
206+ if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION) {
207+ RunSmallAxisInsertionRoute(x, y1, y2, tilingData, pipe);
208+ return true;
209+ }
210+ if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE) {
211+ RunSmallAxisTwoStageRoute(x, y1, y2, tilingData, pipe);
212+ return true;
213+ }
214+ if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT) {
215+ RunSmallAxisShortRankSelectRoute(x, y1, y2, tilingData, pipe);
216+ return true;
217+ }
218+ return false;
219+}
220+ 
169template <uint64_t schId, uint64_t isInt32>221template <uint64_t schId, uint64_t isInt32>
170__global__ __aicore__ void kth_value(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, GM_ADDR tiling)222__global__ __aicore__ void kth_value(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, GM_ADDR tiling)
171{223{
@@ -190,12 +242,11 @@ __global__ __aicore__ void kth_value(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR
190 RunRadixMoreCoreRoute<isInt32>(x, y1, y2, usrWorkspace, &tilingData, &pipe);242 RunRadixMoreCoreRoute<isInt32>(x, y1, y2, usrWorkspace, &tilingData, &pipe);
191 return;243 return;
192 }244 }
193- if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION) {245+ if constexpr (schId == KTH_VALUE_SCHID_RADIX_SELECT) {
194- RunSmallAxisInsertionRoute(x, y1, y2, &tilingData, &pipe);246+ RunRadixSelectRoute(x, y1, y2, usrWorkspace, &tilingData, &pipe);
195 return;247 return;
196 }248 }
197- if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE) {249+ if (TryRunSmallAxisRoute<schId>(x, y1, y2, &tilingData, &pipe)) {
198- RunSmallAxisTwoStageRoute(x, y1, y2, &tilingData, &pipe);
199 return;250 return;
200 }251 }
201 if constexpr (schId == KTH_VALUE_SCHID_MERGE_MORE_CORE) {252 if constexpr (schId == KTH_VALUE_SCHID_MERGE_MORE_CORE) {
Mmath/kth_value/tests/ut/op_host/arch35/test_kth_value_tiling.cpp+112-33
@@ -32,9 +32,10 @@ namespace {
32constexpr size_t WORK_SPACE_SIZE = 16777216;32constexpr size_t WORK_SPACE_SIZE = 16777216;
33KthValueCompileInfo g_compileInfo = {64};33KthValueCompileInfo g_compileInfo = {64};
34 34 
35-gert::TilingContextPara MakeKthValueTilingContext(35+gert::TilingContextPara MakeKthValueTilingContext(const gert::StorageShape& xShape,
36- const gert::StorageShape& xShape, const gert::StorageShape& valuesShape, const gert::StorageShape& indicesShape,36+ const gert::StorageShape& valuesShape,
37- ge::DataType xDtype, int64_t k, int64_t dim = -1)37+ const gert::StorageShape& indicesShape, ge::DataType xDtype,
38+ int64_t k, int64_t dim = -1)
38{39{
39 return gert::TilingContextPara(40 return gert::TilingContextPara(
40 "KthValue",41 "KthValue",
@@ -55,8 +56,8 @@ gert::TilingContextPara MakeKthValueTilingContext(
55 56 
56TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_2x1024)57TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_2x1024)
57{58{
58- auto tilingContextPara =59+ auto tilingContextPara = MakeKthValueTilingContext({{2, 1024}, {2, 1024}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
59- MakeKthValueTilingContext({{2, 1024}, {2, 1024}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, ge::DT_FLOAT, 5);60+ ge::DT_FLOAT, 5);
60 61 
61 TilingInfo tilingInfo;62 TilingInfo tilingInfo;
62 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));63 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -67,8 +68,8 @@ TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_2x1024)
67 68 
68TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_fp32_100x1)69TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_fp32_100x1)
69{70{
70- auto tilingContextPara =71+ auto tilingContextPara = MakeKthValueTilingContext({{100, 1}, {100, 1}}, {{100, 1}, {100, 1}}, {{100, 1}, {100, 1}},
71- MakeKthValueTilingContext({{100, 1}, {100, 1}}, {{100, 1}, {100, 1}}, {{100, 1}, {100, 1}}, ge::DT_FLOAT, 1, 1);72+ ge::DT_FLOAT, 1, 1);
72 73 
73 TilingInfo tilingInfo;74 TilingInfo tilingInfo;
74 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));75 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -79,8 +80,8 @@ TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_fp32_100x1)
79 80 
80TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_bf16_8x512)81TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_bf16_8x512)
81{82{
82- auto tilingContextPara =83+ auto tilingContextPara = MakeKthValueTilingContext({{8, 512}, {8, 512}}, {{8, 1}, {8, 1}}, {{8, 1}, {8, 1}},
83- MakeKthValueTilingContext({{8, 512}, {8, 512}}, {{8, 1}, {8, 1}}, {{8, 1}, {8, 1}}, ge::DT_BF16, 100);84+ ge::DT_BF16, 100);
84 85 
85 TilingInfo tilingInfo;86 TilingInfo tilingInfo;
86 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));87 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -91,8 +92,8 @@ TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_bf16_8x512)
91 92 
92TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_3d_8x16x1024)93TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_3d_8x16x1024)
93{94{
94- auto tilingContextPara = MakeKthValueTilingContext(95+ auto tilingContextPara = MakeKthValueTilingContext({{8, 16, 1024}, {8, 16, 1024}}, {{8, 16, 1}, {8, 16, 1}},
95- {{8, 16, 1024}, {8, 16, 1024}}, {{8, 16, 1}, {8, 16, 1}}, {{8, 16, 1}, {8, 16, 1}}, ge::DT_FLOAT, 500, 2);96+ {{8, 16, 1}, {8, 16, 1}}, ge::DT_FLOAT, 500, 2);
96 97 
97 TilingInfo tilingInfo;98 TilingInfo tilingInfo;
98 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));99 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -103,8 +104,8 @@ TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_3d_8x16x1024)
103 104 
104TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_int32_64x1)105TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_int32_64x1)
105{106{
106- auto tilingContextPara =107+ auto tilingContextPara = MakeKthValueTilingContext({{64, 1}, {64, 1}}, {{64, 1}, {64, 1}}, {{64, 1}, {64, 1}},
107- MakeKthValueTilingContext({{64, 1}, {64, 1}}, {{64, 1}, {64, 1}}, {{64, 1}, {64, 1}}, ge::DT_INT32, 1, 1);108+ ge::DT_INT32, 1, 1);
108 109 
109 TilingInfo tilingInfo;110 TilingInfo tilingInfo;
110 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));111 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -116,9 +117,9 @@ TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_int32_64x1)
116TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_allows_unsorted_dim_over_uint32)117TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_allows_unsorted_dim_over_uint32)
117{118{
118 int64_t largeBatch = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + 1;119 int64_t largeBatch = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + 1;
119- auto tilingContextPara = MakeKthValueTilingContext(120+ auto tilingContextPara = MakeKthValueTilingContext({{largeBatch, 1}, {largeBatch, 1}},
120- {{largeBatch, 1}, {largeBatch, 1}}, {{largeBatch, 1}, {largeBatch, 1}}, {{largeBatch, 1}, {largeBatch, 1}},121+ {{largeBatch, 1}, {largeBatch, 1}},
121- ge::DT_FLOAT, 1, 1);122+ {{largeBatch, 1}, {largeBatch, 1}}, ge::DT_FLOAT, 1, 1);
122 123 
123 TilingInfo tilingInfo;124 TilingInfo tilingInfo;
124 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));125 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -133,9 +134,9 @@ TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_allows_unsorted_dim_over_
133TEST_F(KthValueTilingTest, test_kthvalue_merge_intra_core_allows_unsorted_dim_over_uint32)134TEST_F(KthValueTilingTest, test_kthvalue_merge_intra_core_allows_unsorted_dim_over_uint32)
134{135{
135 int64_t largeBatch = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + 1;136 int64_t largeBatch = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + 1;
136- auto tilingContextPara = MakeKthValueTilingContext(137+ auto tilingContextPara = MakeKthValueTilingContext({{largeBatch, 4097}, {largeBatch, 4097}},
137- {{largeBatch, 4097}, {largeBatch, 4097}}, {{largeBatch, 1}, {largeBatch, 1}},138+ {{largeBatch, 1}, {largeBatch, 1}},
138- {{largeBatch, 1}, {largeBatch, 1}}, ge::DT_FLOAT, 64);139+ {{largeBatch, 1}, {largeBatch, 1}}, ge::DT_FLOAT, 64);
139 140 
140 TilingInfo tilingInfo;141 TilingInfo tilingInfo;
141 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));142 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -157,10 +158,80 @@ TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_1d_512)
157 EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE);158 EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE);
158}159}
159 160 
161+TEST_F(KthValueTilingTest, test_kthvalue_radix_select_int32_large_axis_mid_many_rows)
162+{
163+ auto tilingContextPara = MakeKthValueTilingContext({{16, 65536}, {16, 65536}}, {{16, 1}, {16, 1}},
164+ {{16, 1}, {16, 1}}, ge::DT_INT32, 32768);
165+ 
166+ TilingInfo tilingInfo;
167+ ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
168+ EXPECT_EQ(tilingInfo.tilingKey, 267);
169+ ASSERT_GE(tilingInfo.tilingDataSize, sizeof(KthValueTilingData));
170+ const auto* tilingData = reinterpret_cast<const KthValueTilingData*>(tilingInfo.tilingData.get());
171+ EXPECT_GT(tilingData->numTileDataSize, 0U);
172+ EXPECT_EQ(tilingData->unsortedDimParallel, 16U);
173+ EXPECT_GT(tilingData->lastDimNeedCore, 1U);
174+ EXPECT_EQ(tilingInfo.blockNum, tilingData->unsortedDimParallel * tilingData->lastDimNeedCore);
175+ ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1);
176+ size_t histogramWorkspace = static_cast<size_t>(tilingInfo.blockNum) * 256UL * sizeof(uint64_t);
177+ size_t groupStateWorkspace = static_cast<size_t>(tilingData->unsortedDimParallel) * 8UL * sizeof(uint64_t);
178+ EXPECT_GE(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE + histogramWorkspace + groupStateWorkspace);
179+ EXPECT_LT(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE + 2UL * histogramWorkspace);
180+}
181+ 
182+TEST_F(KthValueTilingTest, test_kthvalue_radix_select_int64_small_axis_many_rows_targets)
183+{
184+ std::vector<gert::StorageShape> xShapes = {
185+ {{33, 1, 1, 1, 1, 1, 1, 7154}, {33, 1, 1, 1, 1, 1, 1, 7154}},
186+ {{45, 8670}, {45, 8670}},
187+ {{59, 6931}, {59, 6931}},
188+ {{42, 1, 1, 1, 7496}, {42, 1, 1, 1, 7496}},
189+ {{56, 1, 9479}, {56, 1, 9479}},
190+ };
191+ std::vector<gert::StorageShape> outputShapes = {
192+ {{33, 1, 1, 1, 1, 1, 1, 1}, {33, 1, 1, 1, 1, 1, 1, 1}},
193+ {{45, 1}, {45, 1}},
194+ {{59, 1}, {59, 1}},
195+ {{42, 1, 1, 1, 1}, {42, 1, 1, 1, 1}},
196+ {{56, 1, 1}, {56, 1, 1}},
197+ };
198+ std::vector<int64_t> kValues = {3264, 5033, 4147, 1663, 2441};
199+ 
200+ for (size_t i = 0; i < xShapes.size(); ++i) {
201+ auto tilingContextPara = MakeKthValueTilingContext(xShapes[i], outputShapes[i], outputShapes[i], ge::DT_INT64,
202+ kValues[i]);
203+ TilingInfo tilingInfo;
204+ ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
205+ EXPECT_EQ(tilingInfo.tilingKey, 267);
206+ }
207+}
208+ 
209+TEST_F(KthValueTilingTest, test_kthvalue_int64_small_axis_short_rank_select_16936x23)
210+{
211+ auto tilingContextPara = MakeKthValueTilingContext({{16936, 23}, {16936, 23}}, {{16936, 1}, {16936, 1}},
212+ {{16936, 1}, {16936, 1}}, ge::DT_INT64, 18);
213+ 
214+ TilingInfo tilingInfo;
215+ ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
216+ EXPECT_EQ(tilingInfo.tilingKey, 268);
217+ EXPECT_EQ(tilingInfo.blockNum, 64);
218+ ASSERT_GE(tilingInfo.tilingDataSize, sizeof(KthValueTilingData));
219+ const auto* tilingData = reinterpret_cast<const KthValueTilingData*>(tilingInfo.tilingData.get());
220+ EXPECT_EQ(tilingData->numTileDataSize, 23U);
221+ EXPECT_EQ(tilingData->unsortedDimParallel, 64U);
222+ EXPECT_GT(tilingData->keyParams0, 0U);
223+ EXPECT_GT(tilingData->keyParams1, 0U);
224+ EXPECT_EQ(tilingData->keyParams2, 6U);
225+ EXPECT_EQ(tilingData->kthIndex, 17);
226+ EXPECT_EQ(tilingData->unsortedDimNum, 16936);
227+ ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1);
228+ EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE);
229+}
230+ 
160TEST_F(KthValueTilingTest, test_kthvalue_bf16_nonlast_axis1840_radix)231TEST_F(KthValueTilingTest, test_kthvalue_bf16_nonlast_axis1840_radix)
161{232{
162- auto tilingContextPara = MakeKthValueTilingContext(233+ auto tilingContextPara = MakeKthValueTilingContext({{2, 1840, 1}, {2, 1840, 1}}, {{2, 1, 1}, {2, 1, 1}},
163- {{2, 1840, 1}, {2, 1840, 1}}, {{2, 1, 1}, {2, 1, 1}}, {{2, 1, 1}, {2, 1, 1}}, ge::DT_BF16, 920, 1);234+ {{2, 1, 1}, {2, 1, 1}}, ge::DT_BF16, 920, 1);
164 235 
165 TilingInfo tilingInfo;236 TilingInfo tilingInfo;
166 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));237 ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo));
@@ -171,48 +242,56 @@ TEST_F(KthValueTilingTest, test_kthvalue_bf16_nonlast_axis1840_radix)
171 242 
172TEST_F(KthValueTilingTest, test_kthvalue_rejects_invalid_dtype)243TEST_F(KthValueTilingTest, test_kthvalue_rejects_invalid_dtype)
173{244{
174- auto tilingContextPara =245+ auto tilingContextPara = MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
175- MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, ge::DT_BOOL, 5);246+ ge::DT_BOOL, 5);
247+ 
248+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
249+}
250+ 
251+TEST_F(KthValueTilingTest, test_kthvalue_radix_select_rejects_invalid_dtype)
252+{
253+ auto tilingContextPara = MakeKthValueTilingContext({{2, 100000}, {2, 100000}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
254+ ge::DT_BOOL, 5);
176 255 
177 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);256 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
178}257}
179 258 
180TEST_F(KthValueTilingTest, test_kthvalue_rejects_empty_shape)259TEST_F(KthValueTilingTest, test_kthvalue_rejects_empty_shape)
181{260{
182- auto tilingContextPara =261+ auto tilingContextPara = MakeKthValueTilingContext({{0, 100}, {0, 100}}, {{0, 1}, {0, 1}}, {{0, 1}, {0, 1}},
183- MakeKthValueTilingContext({{0, 100}, {0, 100}}, {{0, 1}, {0, 1}}, {{0, 1}, {0, 1}}, ge::DT_FLOAT, 5);262+ ge::DT_FLOAT, 5);
184 263 
185 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);264 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
186}265}
187 266 
188TEST_F(KthValueTilingTest, test_kthvalue_rejects_k_zero)267TEST_F(KthValueTilingTest, test_kthvalue_rejects_k_zero)
189{268{
190- auto tilingContextPara =269+ auto tilingContextPara = MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
191- MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, ge::DT_FLOAT, 0);270+ ge::DT_FLOAT, 0);
192 271 
193 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);272 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
194}273}
195 274 
196TEST_F(KthValueTilingTest, test_kthvalue_rejects_k_exceeds_axis)275TEST_F(KthValueTilingTest, test_kthvalue_rejects_k_exceeds_axis)
197{276{
198- auto tilingContextPara =277+ auto tilingContextPara = MakeKthValueTilingContext({{2, 10}, {2, 10}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
199- MakeKthValueTilingContext({{2, 10}, {2, 10}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, ge::DT_FLOAT, 11);278+ ge::DT_FLOAT, 11);
200 279 
201 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);280 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
202}281}
203 282 
204TEST_F(KthValueTilingTest, test_kthvalue_rejects_dim_out_of_range)283TEST_F(KthValueTilingTest, test_kthvalue_rejects_dim_out_of_range)
205{284{
206- auto tilingContextPara =285+ auto tilingContextPara = MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
207- MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, ge::DT_FLOAT, 5, 3);286+ ge::DT_FLOAT, 5, 3);
208 287 
209 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);288 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
210}289}
211 290 
212TEST_F(KthValueTilingTest, test_kthvalue_rejects_negative_dim_out_of_range)291TEST_F(KthValueTilingTest, test_kthvalue_rejects_negative_dim_out_of_range)
213{292{
214- auto tilingContextPara =293+ auto tilingContextPara = MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}},
215- MakeKthValueTilingContext({{2, 100}, {2, 100}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, ge::DT_FLOAT, 5, -3);294+ ge::DT_FLOAT, 5, -3);
216 295 
217 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);296 ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
218}297}
Mmath/sort/op_host/arch35/sort_tiling_arch35.cpp+164-168
@@ -34,7 +34,7 @@ namespace optiling {
34// =============================================================================34// =============================================================================
35// Parameter validation35// Parameter validation
36// =============================================================================36// =============================================================================
37-ge::graphStatus CheckInputAndOutput(gert::TilingContext *context, SortKthTileInfo &sortTileInfo)37+ge::graphStatus CheckInputAndOutput(gert::TilingContext* context, SortKthTileInfo& sortTileInfo)
38{38{
39 auto platformInfo = context->GetPlatformInfo();39 auto platformInfo = context->GetPlatformInfo();
40 OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);40 OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
@@ -42,28 +42,31 @@ ge::graphStatus CheckInputAndOutput(gert::TilingContext *context, SortKthTileInf
42 uint64_t ubSize = 0;42 uint64_t ubSize = 0;
43 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);43 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
44 OP_CHECK_IF(ubSize <= static_cast<uint64_t>(SIMT_UB),44 OP_CHECK_IF(ubSize <= static_cast<uint64_t>(SIMT_UB),
45- OP_LOGE(context->GetNodeName(), "ubSize must be greater than %u, but is %lu", SIMT_UB, ubSize),45+ OP_LOGE(context->GetNodeName(), "ubSize must be greater than %u, but is %lu", SIMT_UB, ubSize),
46- return ge::GRAPH_FAILED);46+ return ge::GRAPH_FAILED);
47 sortTileInfo.blockUbSize = Ops::Base::GetUbBlockSize(context);47 sortTileInfo.blockUbSize = Ops::Base::GetUbBlockSize(context);
48 OP_LOGI(context->GetNodeName(), "ubSize is %ld, blockUbSize %u", ubSize, sortTileInfo.blockUbSize);48 OP_LOGI(context->GetNodeName(), "ubSize is %ld, blockUbSize %u", ubSize, sortTileInfo.blockUbSize);
49 sortTileInfo.ubSize = static_cast<uint32_t>(ubSize);49 sortTileInfo.ubSize = static_cast<uint32_t>(ubSize);
50 auto inputShapePtr = context->GetInputShape(0);50 auto inputShapePtr = context->GetInputShape(0);
51 OP_CHECK_NULL_WITH_CONTEXT(context, inputShapePtr);51 OP_CHECK_NULL_WITH_CONTEXT(context, inputShapePtr);
52- const gert::Shape &inputShape = Ops::Base::EnsureNotScalar(inputShapePtr->GetStorageShape());52+ const gert::Shape& inputShape = Ops::Base::EnsureNotScalar(inputShapePtr->GetStorageShape());
53 auto yStorage = context->GetOutputShape(0);53 auto yStorage = context->GetOutputShape(0);
54 OP_CHECK_NULL_WITH_CONTEXT(context, yStorage);54 OP_CHECK_NULL_WITH_CONTEXT(context, yStorage);
55- const gert::Shape &outShape = Ops::Base::EnsureNotScalar(yStorage->GetStorageShape());55+ const gert::Shape& outShape = Ops::Base::EnsureNotScalar(yStorage->GetStorageShape());
56 auto yStorage1 = context->GetOutputShape(1);56 auto yStorage1 = context->GetOutputShape(1);
57 OP_CHECK_NULL_WITH_CONTEXT(context, yStorage1);57 OP_CHECK_NULL_WITH_CONTEXT(context, yStorage1);
58- const gert::Shape &outShape1 = Ops::Base::EnsureNotScalar(yStorage1->GetStorageShape());58+ const gert::Shape& outShape1 = Ops::Base::EnsureNotScalar(yStorage1->GetStorageShape());
59 OP_CHECK_IF(inputShape.GetShapeSize() == 0 || outShape.GetShapeSize() == 0,59 OP_CHECK_IF(inputShape.GetShapeSize() == 0 || outShape.GetShapeSize() == 0,
60- OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "x, y1",60+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "x, y1", "0",
61- "0", "The shape size of input x and output y1 should be positive"),61+ "The shape size of input x and output y1 should be positive"),
62- return ge::GRAPH_FAILED);62+ return ge::GRAPH_FAILED);
63- OP_CHECK_IF(outShape != outShape1 || outShape != inputShape,63+ OP_CHECK_IF(
64+ outShape != outShape1 || outShape != inputShape,
64 OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "x, y1, y2",65 OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "x, y1, y2",
65- (Ops::Base::ToString(inputShape) + ", " + Ops::Base::ToString(outShape) + ", " + Ops::Base::ToString(outShape1)).c_str(),66+ (Ops::Base::ToString(inputShape) + ", " + Ops::Base::ToString(outShape) +
66- "The shape of input x, output y1 and y2 should be the same"),67+ ", " + Ops::Base::ToString(outShape1))
68+ .c_str(),
69+ "The shape of input x, output y1 and y2 should be the same"),
67 return ge::GRAPH_FAILED);70 return ge::GRAPH_FAILED);
68 int32_t xDimNum = inputShape.GetDimNum();71 int32_t xDimNum = inputShape.GetDimNum();
69 sortTileInfo.rank = xDimNum;72 sortTileInfo.rank = xDimNum;
@@ -72,8 +75,8 @@ ge::graphStatus CheckInputAndOutput(gert::TilingContext *context, SortKthTileInf
72 return ge::GRAPH_SUCCESS;75 return ge::GRAPH_SUCCESS;
73}76}
74 77 
75-ge::graphStatus CheckSortOutputDtypes(78+ge::graphStatus CheckSortOutputDtypes(gert::TilingContext* context, SortKthTileInfo& sortTileInfo,
76- gert::TilingContext *context, SortKthTileInfo &sortTileInfo, ge::DataType dataType)79+ ge::DataType dataType)
77{80{
78 auto outDescPtr = context->GetOutputDesc(1);81 auto outDescPtr = context->GetOutputDesc(1);
79 OP_CHECK_NULL_WITH_CONTEXT(context, outDescPtr);82 OP_CHECK_NULL_WITH_CONTEXT(context, outDescPtr);
@@ -81,58 +84,61 @@ ge::graphStatus CheckSortOutputDtypes(
81 auto outDescPtr0 = context->GetOutputDesc(0);84 auto outDescPtr0 = context->GetOutputDesc(0);
82 OP_CHECK_NULL_WITH_CONTEXT(context, outDescPtr0);85 OP_CHECK_NULL_WITH_CONTEXT(context, outDescPtr0);
83 auto y1DType = outDescPtr0->GetDataType();86 auto y1DType = outDescPtr0->GetDataType();
84- OP_CHECK_IF((y2DType != ge::DT_INT64) && (y2DType != ge::DT_INT32),87+ OP_CHECK_IF(
85- OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "y2",88+ (y2DType != ge::DT_INT64) && (y2DType != ge::DT_INT32),
86- Ops::Base::ToString(y2DType).c_str(), "INT32 or INT64"), return ge::GRAPH_FAILED);89+ OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "y2", Ops::Base::ToString(y2DType).c_str(), "INT32 or INT64"),
87- OP_CHECK_IF(y1DType != dataType,
88- OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(context->GetNodeName(), "x, y1",
89- (Ops::Base::ToString(dataType) + ", " + Ops::Base::ToString(y1DType)).c_str(),
90- "The dtype of input x should be the same as output y1"),
91 return ge::GRAPH_FAILED);90 return ge::GRAPH_FAILED);
91+ OP_CHECK_IF(y1DType != dataType,
92+ OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(
93+ context->GetNodeName(), "x, y1",
94+ (Ops::Base::ToString(dataType) + ", " + Ops::Base::ToString(y1DType)).c_str(),
95+ "The dtype of input x should be the same as output y1"),
96+ return ge::GRAPH_FAILED);
92 ge::TypeUtils::GetDataTypeLength(y2DType, sortTileInfo.y2DtypeSize);97 ge::TypeUtils::GetDataTypeLength(y2DType, sortTileInfo.y2DtypeSize);
93 return ge::GRAPH_SUCCESS;98 return ge::GRAPH_SUCCESS;
94}99}
95 100 
96-ge::graphStatus ComputeNonLastAxisLayout(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, int32_t sortAxis)101+ge::graphStatus ComputeNonLastAxisLayout(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, int32_t sortAxis)
97{102{
98 auto inputShapePtr = context->GetInputShape(0);103 auto inputShapePtr = context->GetInputShape(0);
99 OP_CHECK_NULL_WITH_CONTEXT(context, inputShapePtr);104 OP_CHECK_NULL_WITH_CONTEXT(context, inputShapePtr);
100- const gert::Shape &inputShape = Ops::Base::EnsureNotScalar(inputShapePtr->GetStorageShape());105+ const gert::Shape& inputShape = Ops::Base::EnsureNotScalar(inputShapePtr->GetStorageShape());
101 ComputeAxisDimProducts(inputShape, sortAxis, sortTileInfo);106 ComputeAxisDimProducts(inputShape, sortAxis, sortTileInfo);
102 return ge::GRAPH_SUCCESS;107 return ge::GRAPH_SUCCESS;
103}108}
104 109 
105-ge::graphStatus SortCheckParams(gert::TilingContext *context, SortKthTileInfo &sortTileInfo)110+ge::graphStatus SortCheckParams(gert::TilingContext* context, SortKthTileInfo& sortTileInfo)
106{111{
107 OP_CHECK_IF(CheckInputAndOutput(context, sortTileInfo) != ge::GRAPH_SUCCESS,112 OP_CHECK_IF(CheckInputAndOutput(context, sortTileInfo) != ge::GRAPH_SUCCESS,
108- OP_LOGE(context->GetNodeName(), "CheckInputAndOutput failed"), return ge::GRAPH_FAILED);113+ OP_LOGE(context->GetNodeName(), "CheckInputAndOutput failed"), return ge::GRAPH_FAILED);
109 auto inputDescPtr = context->GetInputDesc(0);114 auto inputDescPtr = context->GetInputDesc(0);
110 OP_CHECK_NULL_WITH_CONTEXT(context, inputDescPtr);115 OP_CHECK_NULL_WITH_CONTEXT(context, inputDescPtr);
111 ge::DataType dataType = inputDescPtr->GetDataType();116 ge::DataType dataType = inputDescPtr->GetDataType();
112 sortTileInfo.dataType = dataType;117 sortTileInfo.dataType = dataType;
113- OP_CHECK_IF(!ge::TypeUtils::GetDataTypeLength(dataType, sortTileInfo.dtypeSize),118+ OP_CHECK_IF(
114- OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "x",119+ !ge::TypeUtils::GetDataTypeLength(dataType, sortTileInfo.dtypeSize),
115- Ops::Base::ToString(dataType).c_str(),120+ OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "x", Ops::Base::ToString(dataType).c_str(),
116- "INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, FLOAT, FLOAT16, BF16"), return ge::GRAPH_FAILED);121+ "INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, FLOAT, FLOAT16, BF16"),
122+ return ge::GRAPH_FAILED);
117 OP_CHECK_IF(CheckSortOutputDtypes(context, sortTileInfo, dataType) != ge::GRAPH_SUCCESS,123 OP_CHECK_IF(CheckSortOutputDtypes(context, sortTileInfo, dataType) != ge::GRAPH_SUCCESS,
118- OP_LOGE(context->GetNodeName(), "CheckSortOutputDtypes failed"), return ge::GRAPH_FAILED);124+ OP_LOGE(context->GetNodeName(), "CheckSortOutputDtypes failed"), return ge::GRAPH_FAILED);
119 auto const attrs = context->GetAttrs();125 auto const attrs = context->GetAttrs();
120 OP_CHECK_NULL_WITH_CONTEXT(context, attrs);126 OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
121- const bool *isDescending = attrs->GetAttrPointer<bool>(1);127+ const bool* isDescending = attrs->GetAttrPointer<bool>(1);
122- const int64_t *sortAxisPtr = attrs->GetAttrPointer<int64_t>(0);128+ const int64_t* sortAxisPtr = attrs->GetAttrPointer<int64_t>(0);
123 OP_CHECK_NULL_WITH_CONTEXT(context, isDescending);129 OP_CHECK_NULL_WITH_CONTEXT(context, isDescending);
124 OP_CHECK_NULL_WITH_CONTEXT(context, sortAxisPtr);130 OP_CHECK_NULL_WITH_CONTEXT(context, sortAxisPtr);
125 int32_t sortAxis = static_cast<int32_t>(*sortAxisPtr);131 int32_t sortAxis = static_cast<int32_t>(*sortAxisPtr);
126 sortAxis = sortAxis < 0 ? (sortAxis + static_cast<int32_t>(sortTileInfo.rank)) : sortAxis;132 sortAxis = sortAxis < 0 ? (sortAxis + static_cast<int32_t>(sortTileInfo.rank)) : sortAxis;
127 OP_CHECK_IF(sortAxis < 0 || sortAxis >= static_cast<int32_t>(sortTileInfo.rank),133 OP_CHECK_IF(sortAxis < 0 || sortAxis >= static_cast<int32_t>(sortTileInfo.rank),
128- OP_LOGE_WITH_INVALID_ATTR(context->GetNodeName(), "axis",134+ OP_LOGE_WITH_INVALID_ATTR(context->GetNodeName(), "axis", std::to_string(sortAxis).c_str(),
129- std::to_string(sortAxis).c_str(), "range [-dimNum, dimNum - 1)"),135+ "range [-dimNum, dimNum - 1)"),
130- return ge::GRAPH_FAILED);136+ return ge::GRAPH_FAILED);
131 sortTileInfo.sortAxis = sortAxis;137 sortTileInfo.sortAxis = sortAxis;
132 sortTileInfo.isNonLastAxis = (sortAxis != (static_cast<int32_t>(sortTileInfo.rank) - 1));138 sortTileInfo.isNonLastAxis = (sortAxis != (static_cast<int32_t>(sortTileInfo.rank) - 1));
133 if (sortTileInfo.isNonLastAxis) {139 if (sortTileInfo.isNonLastAxis) {
134 OP_CHECK_IF(ComputeNonLastAxisLayout(context, sortTileInfo, sortAxis) != ge::GRAPH_SUCCESS,140 OP_CHECK_IF(ComputeNonLastAxisLayout(context, sortTileInfo, sortAxis) != ge::GRAPH_SUCCESS,
135- OP_LOGE(context->GetNodeName(), "ComputeNonLastAxisLayout failed"), return ge::GRAPH_FAILED);141+ OP_LOGE(context->GetNodeName(), "ComputeNonLastAxisLayout failed"), return ge::GRAPH_FAILED);
136 }142 }
137 return ge::GRAPH_SUCCESS;143 return ge::GRAPH_SUCCESS;
138}144}
@@ -140,10 +146,10 @@ ge::graphStatus SortCheckParams(gert::TilingContext *context, SortKthTileInfo &s
140// =============================================================================146// =============================================================================
141// UB computation and layout helpers147// UB computation and layout helpers
142// =============================================================================148// =============================================================================
143-void SetSortTmpSize(ge::DataType dataType, uint32_t tileData, bool isDescend, SortKthTileInfo &sortTileInfo)149+void SetSortTmpSize(ge::DataType dataType, uint32_t tileData, bool isDescend, SortKthTileInfo& sortTileInfo)
144{150{
145 int64_t realLen = std::min(sortTileInfo.lastAxis, static_cast<int64_t>(tileData));151 int64_t realLen = std::min(sortTileInfo.lastAxis, static_cast<int64_t>(tileData));
146- std::vector<int64_t> shapeVec = { realLen };152+ std::vector<int64_t> shapeVec = {realLen};
147 ge::Shape srcShape(shapeVec);153 ge::Shape srcShape(shapeVec);
148 AscendC::SortConfig config;154 AscendC::SortConfig config;
149 config.type = AscendC::SortType::RADIX_SORT;155 config.type = AscendC::SortType::RADIX_SORT;
@@ -161,14 +167,14 @@ void SetSortTmpSize(ge::DataType dataType, uint32_t tileData, bool isDescend, So
161// =============================================================================167// =============================================================================
162// Non-last axis helpers168// Non-last axis helpers
163// =============================================================================169// =============================================================================
164-static bool TryGetSortNonLastTileCount(const SortKthTileInfo &info, uint32_t innerChunk,170+static bool TryGetSortNonLastTileCount(const SortKthTileInfo& info, uint32_t innerChunk, uint32_t& innerLoopNum,
165- uint32_t &innerLoopNum, uint64_t &tileCount)171+ uint64_t& tileCount)
166{172{
167 if (innerChunk == 0U || info.innerSize <= 0 || info.outerSize <= 0) {173 if (innerChunk == 0U || info.innerSize <= 0 || info.outerSize <= 0) {
168 return false;174 return false;
169 }175 }
170- uint64_t innerLoopNum64 = Ops::Base::CeilDiv(176+ uint64_t innerLoopNum64 = Ops::Base::CeilDiv(static_cast<uint64_t>(info.innerSize),
171- static_cast<uint64_t>(info.innerSize), static_cast<uint64_t>(innerChunk));177+ static_cast<uint64_t>(innerChunk));
172 if (innerLoopNum64 == 0U || innerLoopNum64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {178 if (innerLoopNum64 == 0U || innerLoopNum64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {
173 return false;179 return false;
174 }180 }
@@ -179,11 +185,11 @@ static bool TryGetSortNonLastTileCount(const SortKthTileInfo &info, uint32_t inn
179 185 
180// Clamp batchSize to innerSize, then compute total tile count across all outer batches.186// Clamp batchSize to innerSize, then compute total tile count across all outer batches.
181// Returns false on overflow or zero tile count.187// Returns false on overflow or zero tile count.
182-bool ComputeSortNonLastBatchPlan(const SortKthTileInfo &sortTileInfo, uint32_t &batchSize,188+bool ComputeSortNonLastBatchPlan(const SortKthTileInfo& sortTileInfo, uint32_t& batchSize, uint32_t& innerLoopNum,
183- uint32_t &innerLoopNum, uint32_t &batchNum)189+ uint32_t& batchNum)
184{190{
185- batchSize = static_cast<uint32_t>(std::min<uint64_t>(191+ batchSize = static_cast<uint32_t>(
186- static_cast<uint64_t>(batchSize), static_cast<uint64_t>(sortTileInfo.innerSize)));192+ std::min<uint64_t>(static_cast<uint64_t>(batchSize), static_cast<uint64_t>(sortTileInfo.innerSize)));
187 if (!ComputeNonLastBatchNum(sortTileInfo.outerSize, sortTileInfo.innerSize, batchSize, batchNum)) {193 if (!ComputeNonLastBatchNum(sortTileInfo.outerSize, sortTileInfo.innerSize, batchSize, batchNum)) {
188 return false;194 return false;
189 }195 }
@@ -198,7 +204,7 @@ bool ComputeSortNonLastBatchPlan(const SortKthTileInfo &sortTileInfo, uint32_t &
198 return batchNum > 0;204 return batchNum > 0;
199}205}
200 206 
201-bool IsRadixSortOneCore(SortKthTileInfo &sortTileInfo)207+bool IsRadixSortOneCore(SortKthTileInfo& sortTileInfo)
202{208{
203 if (sortTileInfo.isInt32 == static_cast<uint32_t>(0)) {209 if (sortTileInfo.isInt32 == static_cast<uint32_t>(0)) {
204 return false;210 return false;
@@ -206,7 +212,8 @@ bool IsRadixSortOneCore(SortKthTileInfo &sortTileInfo)
206 uint32_t xUbSize = 0;212 uint32_t xUbSize = 0;
207 uint32_t y2UbSize = 0;213 uint32_t y2UbSize = 0;
208 if (!ComputeRadixOneCoreUbSizes(sortTileInfo.lastAxis, sortTileInfo.dtypeSize,214 if (!ComputeRadixOneCoreUbSizes(sortTileInfo.lastAxis, sortTileInfo.dtypeSize,
209- static_cast<uint32_t>(sizeof(int32_t)), sortTileInfo.blockUbSize, xUbSize, y2UbSize)) {215+ static_cast<uint32_t>(sizeof(int32_t)), sortTileInfo.blockUbSize, xUbSize,
216+ y2UbSize)) {
210 return false;217 return false;
211 }218 }
212 219 
@@ -226,9 +233,9 @@ bool IsRadixSortOneCore(SortKthTileInfo &sortTileInfo)
226 return false;233 return false;
227 }234 }
228 remainUb = (remainUb / static_cast<int64_t>(sortTileInfo.blockUbSize)) *235 remainUb = (remainUb / static_cast<int64_t>(sortTileInfo.blockUbSize)) *
229- static_cast<int64_t>(sortTileInfo.blockUbSize);236+ static_cast<int64_t>(sortTileInfo.blockUbSize);
230- SetSortTmpSize(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis),237+ SetSortTmpSize(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis), sortTileInfo.isDescend,
231- sortTileInfo.isDescend, sortTileInfo);238+ sortTileInfo);
232 int64_t tmpUb = static_cast<int64_t>(sortTileInfo.tmpUbSize);239 int64_t tmpUb = static_cast<int64_t>(sortTileInfo.tmpUbSize);
233 OP_LOGI("RadixSortTiling", "remainUb is %ld, tmpUb is %ld", remainUb, tmpUb);240 OP_LOGI("RadixSortTiling", "remainUb is %ld, tmpUb is %ld", remainUb, tmpUb);
234 if (tmpUb > remainUb) {241 if (tmpUb > remainUb) {
@@ -237,7 +244,7 @@ bool IsRadixSortOneCore(SortKthTileInfo &sortTileInfo)
237 244 
238 int64_t doubleBufferRemainUb = static_cast<int64_t>(sortTileInfo.ubSize) - oneBufferQueSize * 2;245 int64_t doubleBufferRemainUb = static_cast<int64_t>(sortTileInfo.ubSize) - oneBufferQueSize * 2;
239 doubleBufferRemainUb = (doubleBufferRemainUb / static_cast<int64_t>(sortTileInfo.blockUbSize)) *246 doubleBufferRemainUb = (doubleBufferRemainUb / static_cast<int64_t>(sortTileInfo.blockUbSize)) *
240- static_cast<int64_t>(sortTileInfo.blockUbSize);247+ static_cast<int64_t>(sortTileInfo.blockUbSize);
241 if (tmpUb <= doubleBufferRemainUb) {248 if (tmpUb <= doubleBufferRemainUb) {
242 sortTileInfo.keyParams3 = 2;249 sortTileInfo.keyParams3 = 2;
243 }250 }
@@ -248,16 +255,13 @@ bool IsRadixSortOneCore(SortKthTileInfo &sortTileInfo)
248// =============================================================================255// =============================================================================
249// Individual strategy Set functions256// Individual strategy Set functions
250// =============================================================================257// =============================================================================
251-bool IsAxisOneCopy(const SortKthTileInfo &sortTileInfo)258+bool IsAxisOneCopy(const SortKthTileInfo& sortTileInfo) { return sortTileInfo.lastAxis == static_cast<int64_t>(1); }
252-{
253- return sortTileInfo.lastAxis == static_cast<int64_t>(1);
254-}
255 259 
256-ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext *context, SortKthTileInfo &sortTileInfo)260+ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext* context, SortKthTileInfo& sortTileInfo)
257{261{
258 // double buffer262 // double buffer
259- uint64_t bytesPerElem = static_cast<uint64_t>(2) *263+ uint64_t bytesPerElem = static_cast<uint64_t>(2) * (static_cast<uint64_t>(sortTileInfo.dtypeSize) +
260- (static_cast<uint64_t>(sortTileInfo.dtypeSize) + static_cast<uint64_t>(sortTileInfo.y2DtypeSize));264+ static_cast<uint64_t>(sortTileInfo.y2DtypeSize));
261 if (bytesPerElem == 0) {265 if (bytesPerElem == 0) {
262 OP_LOGE(context->GetNodeName(), "bytesPerElem is 0, invalid dtype configuration");266 OP_LOGE(context->GetNodeName(), "bytesPerElem is 0, invalid dtype configuration");
263 return ge::GRAPH_FAILED;267 return ge::GRAPH_FAILED;
@@ -273,7 +277,7 @@ ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext *context, SortKthTileIn
273 }277 }
274 uint32_t copyElemsPerLoop = static_cast<uint32_t>(copyElemsPerLoop64);278 uint32_t copyElemsPerLoop = static_cast<uint32_t>(copyElemsPerLoop64);
275 uint64_t totalElems = static_cast<uint64_t>(sortTileInfo.unsortedDim) *279 uint64_t totalElems = static_cast<uint64_t>(sortTileInfo.unsortedDim) *
276- static_cast<uint64_t>(sortTileInfo.lastAxis);280+ static_cast<uint64_t>(sortTileInfo.lastAxis);
277 uint64_t loopTimes64 = (totalElems + copyElemsPerLoop64 - 1) / copyElemsPerLoop64;281 uint64_t loopTimes64 = (totalElems + copyElemsPerLoop64 - 1) / copyElemsPerLoop64;
278 if (loopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {282 if (loopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {
279 OP_LOGE(context->GetNodeName(), "loopTimes exceeds uint32_t limit");283 OP_LOGE(context->GetNodeName(), "loopTimes exceeds uint32_t limit");
@@ -292,22 +296,21 @@ ge::graphStatus SetAxisOneCopyTiling(gert::TilingContext *context, SortKthTileIn
292 sortTileInfo.sortLoopTimes = Ops::Base::CeilDiv(static_cast<int64_t>(loopTimes), static_cast<int64_t>(coreNumNeed));296 sortTileInfo.sortLoopTimes = Ops::Base::CeilDiv(static_cast<int64_t>(loopTimes), static_cast<int64_t>(coreNumNeed));
293 sortTileInfo.tmpUbSize = 0;297 sortTileInfo.tmpUbSize = 0;
294 298 
295- size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);299+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
296 userWorkSpaceSize[0] = WORK_SPACE_SIZE;300 userWorkSpaceSize[0] = WORK_SPACE_SIZE;
297- OP_LOGI("AxisOneCopyTiling", "totalElems %lu, copyElemsPerLoop %u, loopTimes %u, coreNumNeed %u",301+ OP_LOGI("AxisOneCopyTiling", "totalElems %lu, copyElemsPerLoop %u, loopTimes %u, coreNumNeed %u", totalElems,
298- totalElems, sortTileInfo.keyParams0, sortTileInfo.keyParams1, coreNumNeed);302+ sortTileInfo.keyParams0, sortTileInfo.keyParams1, coreNumNeed);
299 return ge::GRAPH_SUCCESS;303 return ge::GRAPH_SUCCESS;
300}304}
301 305 
302-void FillSmallAxisBatched(gert::TilingContext *context, SortKthTileInfo &sortTileInfo,306+void FillSmallAxisBatched(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, const SmallAxisRoutePlan& plan)
303- const SmallAxisRoutePlan &plan)
304{307{
305- sortTileInfo.ubSize = sortTileInfo.ubSize - SIMT_UB; // reserve 32KB for SIMT kernel scratch308+ sortTileInfo.ubSize = sortTileInfo.ubSize - SIMT_UB; // reserve 32KB for SIMT kernel scratch
306 sortTileInfo.numTileDataSize = static_cast<uint32_t>(sortTileInfo.lastAxis);309 sortTileInfo.numTileDataSize = static_cast<uint32_t>(sortTileInfo.lastAxis);
307- sortTileInfo.keyParams0 = plan.batchSize; // rows per batch310+ sortTileInfo.keyParams0 = plan.batchSize; // rows per batch
308- sortTileInfo.keyParams1 = plan.batchNum; // total batches311+ sortTileInfo.keyParams1 = plan.batchNum; // total batches
309- sortTileInfo.keyParams2 = plan.useRankInverse ? 1U : 0U; // enable rank-inverse second pass312+ sortTileInfo.keyParams2 = plan.useRankInverse ? 1U : 0U; // enable rank-inverse second pass
310- sortTileInfo.keyParams3 = sortTileInfo.isNonLastAxis ? 1U : 0U; // non-last axis flag313+ sortTileInfo.keyParams3 = sortTileInfo.isNonLastAxis ? 1U : 0U; // non-last axis flag
311 if (sortTileInfo.isNonLastAxis) {314 if (sortTileInfo.isNonLastAxis) {
312 // Recompute tile count for non-last axis: outerSize * batchSize tiles across all cores.315 // Recompute tile count for non-last axis: outerSize * batchSize tiles across all cores.
313 uint64_t tileCount64 = 0;316 uint64_t tileCount64 = 0;
@@ -315,33 +318,33 @@ void FillSmallAxisBatched(gert::TilingContext *context, SortKthTileInfo &sortTil
315 }318 }
316 sortTileInfo.coreNumNeed = plan.blockDim;319 sortTileInfo.coreNumNeed = plan.blockDim;
317 sortTileInfo.tmpUbSize = plan.tmpUbSize;320 sortTileInfo.tmpUbSize = plan.tmpUbSize;
318- size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);321+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
319 userWorkSpaceSize[0] = WORK_SPACE_SIZE;322 userWorkSpaceSize[0] = WORK_SPACE_SIZE;
320}323}
321 324 
322// Compute per-row UB footprint for merge-sort multi-core kernel, then derive325// Compute per-row UB footprint for merge-sort multi-core kernel, then derive
323// how many rows fit in available UB via ComputeMergeMoreCoreTiling.326// how many rows fit in available UB via ComputeMergeMoreCoreTiling.
324-ge::graphStatus SetMergeMoreCoreTiling(gert::TilingContext *context, SortKthTileInfo &info)327+ge::graphStatus SetMergeMoreCoreTiling(gert::TilingContext* context, SortKthTileInfo& info)
325{328{
326- uint32_t byteNum = MERGE_SORT_LIST_NUM * MERGE_SORT_DATA_BYTES * 2; // value double buffer329+ uint32_t byteNum = MERGE_SORT_LIST_NUM * MERGE_SORT_DATA_BYTES * 2; // value double buffer
327- byteNum += MERGE_SORT_LIST_NUM * static_cast<uint32_t>(sizeof(uint32_t)); // int32 index330+ byteNum += MERGE_SORT_LIST_NUM * static_cast<uint32_t>(sizeof(uint32_t)); // int32 index
328 if (info.y2DtypeSize == sizeof(int64_t)) {331 if (info.y2DtypeSize == sizeof(int64_t)) {
329- byteNum += MERGE_SORT_LIST_NUM * static_cast<uint32_t>(sizeof(int64_t)); // int64 index extra332+ byteNum += MERGE_SORT_LIST_NUM * static_cast<uint32_t>(sizeof(int64_t)); // int64 index extra
330 }333 }
331 byteNum += MERGE_SORT_LIST_NUM * info.dtypeSize;334 byteNum += MERGE_SORT_LIST_NUM * info.dtypeSize;
332 OP_CHECK_IF(!ComputeMergeMoreCoreTiling(context, info, byteNum),335 OP_CHECK_IF(!ComputeMergeMoreCoreTiling(context, info, byteNum),
333- OP_LOGE(context->GetNodeName(), "merge more-core plan failed"), return ge::GRAPH_FAILED);336+ OP_LOGE(context->GetNodeName(), "merge more-core plan failed"), return ge::GRAPH_FAILED);
334 OP_LOGI("[mergeSort]", "maxDealingNum: %u", info.keyParams0);337 OP_LOGI("[mergeSort]", "maxDealingNum: %u", info.keyParams0);
335 return ge::GRAPH_SUCCESS;338 return ge::GRAPH_SUCCESS;
336}339}
337 340 
338-ge::graphStatus SetRadixOneCoreTiling(gert::TilingContext *context, SortKthTileInfo &sortTileInfo)341+ge::graphStatus SetRadixOneCoreTiling(gert::TilingContext* context, SortKthTileInfo& sortTileInfo)
339{342{
340 sortTileInfo.lastDimNeedCore = static_cast<uint32_t>(1);343 sortTileInfo.lastDimNeedCore = static_cast<uint32_t>(1);
341 sortTileInfo.numTileDataSize = static_cast<uint32_t>(sortTileInfo.lastAxis);344 sortTileInfo.numTileDataSize = static_cast<uint32_t>(sortTileInfo.lastAxis);
342 sortTileInfo.lastDimTileNum = static_cast<uint32_t>(1);345 sortTileInfo.lastDimTileNum = static_cast<uint32_t>(1);
343- uint64_t sortLoopTimes64 =346+ uint64_t sortLoopTimes64 = Ops::Base::CeilDiv(sortTileInfo.unsortedDim,
344- Ops::Base::CeilDiv(sortTileInfo.unsortedDim, static_cast<int64_t>(sortTileInfo.maxCoreNum));347+ static_cast<int64_t>(sortTileInfo.maxCoreNum));
345 if (sortLoopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {348 if (sortLoopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {
346 OP_LOGE(context->GetNodeName(), "sortLoopTimes exceeds uint32_t limit");349 OP_LOGE(context->GetNodeName(), "sortLoopTimes exceeds uint32_t limit");
347 return ge::GRAPH_FAILED;350 return ge::GRAPH_FAILED;
@@ -354,17 +357,17 @@ ge::graphStatus SetRadixOneCoreTiling(gert::TilingContext *context, SortKthTileI
354 sortTileInfo.coreNumNeed = core == uint32_t(0) ? sortTileInfo.maxCoreNum : core;357 sortTileInfo.coreNumNeed = core == uint32_t(0) ? sortTileInfo.maxCoreNum : core;
355 }358 }
356 sortTileInfo.unsortedDimParallel = sortTileInfo.coreNumNeed;359 sortTileInfo.unsortedDimParallel = sortTileInfo.coreNumNeed;
357- size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);360+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
358 userWorkSpaceSize[0] = WORK_SPACE_SIZE;361 userWorkSpaceSize[0] = WORK_SPACE_SIZE;
359 return ge::GRAPH_SUCCESS;362 return ge::GRAPH_SUCCESS;
360}363}
361 364 
362-ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext *context, SortKthTileInfo &info)365+ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext* context, SortKthTileInfo& info)
363{366{
364- OP_CHECK_IF(!FillRadixMoreCoreInfo(info),367+ OP_CHECK_IF(!FillRadixMoreCoreInfo(info), OP_LOGE(context->GetNodeName(), "radix more-core plan failed"),
365- OP_LOGE(context->GetNodeName(), "radix more-core plan failed"), return ge::GRAPH_FAILED);368+ return ge::GRAPH_FAILED);
366 info.ubSize = info.ubSize - SIMT_UB;369 info.ubSize = info.ubSize - SIMT_UB;
367- size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);370+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
368 userWorkSpaceSize[0] = info.workspaceSize;371 userWorkSpaceSize[0] = info.workspaceSize;
369 context->SetScheduleMode(1);372 context->SetScheduleMode(1);
370 return ge::GRAPH_SUCCESS;373 return ge::GRAPH_SUCCESS;
@@ -373,51 +376,51 @@ ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext *context, SortKthTile
373// =============================================================================376// =============================================================================
374// Fill and print functions377// Fill and print functions
375// =============================================================================378// =============================================================================
376-void FillTilingDataSort(SortKthTileInfo &info, SortRegBaseTilingData *sortTilingData)379+void FillTilingDataSort(SortKthTileInfo& info, SortRegBaseTilingData* sortTilingData)
377{380{
378 PlanToTilingData(info, sortTilingData);381 PlanToTilingData(info, sortTilingData);
379 sortTilingData->outputIndexRowBytes = info.outputIndexRowBytes;382 sortTilingData->outputIndexRowBytes = info.outputIndexRowBytes;
380 return;383 return;
381}384}
382 385 
383-void PrintTilingDataSort(gert::TilingContext *context, SortKthTileInfo &sortTileInfo)386+void PrintTilingDataSort(gert::TilingContext* context, SortKthTileInfo& sortTileInfo)
384{387{
385 OP_LOGI(context->GetNodeName(),388 OP_LOGI(context->GetNodeName(),
386- "realCoreNum %u, numTileDataSize %u, unsortedDimParallel %u, "389+ "realCoreNum %u, numTileDataSize %u, unsortedDimParallel %u, "
387- "lastDimTileNum %u, sortLoopTimes %u, lastDimNeedCore %u, keyParams0 %u, keyParams1 %u "390+ "lastDimTileNum %u, sortLoopTimes %u, lastDimNeedCore %u, keyParams0 %u, keyParams1 %u "
388- "keyParams2 %u, keyParams3 %u, keyParams4 %u, keyParams5 %u, tmpUbSize %u, "391+ "keyParams2 %u, keyParams3 %u, keyParams4 %u, keyParams5 %u, tmpUbSize %u, "
389- "lastAxisNum %ld, unsortedDimNum %ld, outerSize %ld, innerSize %ld, innerChunk %u ",392+ "lastAxisNum %ld, unsortedDimNum %ld, outerSize %ld, innerSize %ld, innerChunk %u ",
390- sortTileInfo.coreNumNeed, sortTileInfo.numTileDataSize, sortTileInfo.unsortedDimParallel,393+ sortTileInfo.coreNumNeed, sortTileInfo.numTileDataSize, sortTileInfo.unsortedDimParallel,
391- sortTileInfo.lastDimTileNum, sortTileInfo.sortLoopTimes, sortTileInfo.lastDimNeedCore, sortTileInfo.keyParams0,394+ sortTileInfo.lastDimTileNum, sortTileInfo.sortLoopTimes, sortTileInfo.lastDimNeedCore,
392- sortTileInfo.keyParams1, sortTileInfo.keyParams2, sortTileInfo.keyParams3, sortTileInfo.keyParams4,395+ sortTileInfo.keyParams0, sortTileInfo.keyParams1, sortTileInfo.keyParams2, sortTileInfo.keyParams3,
393- sortTileInfo.keyParams5, sortTileInfo.tmpUbSize, sortTileInfo.lastAxis, sortTileInfo.unsortedDim,396+ sortTileInfo.keyParams4, sortTileInfo.keyParams5, sortTileInfo.tmpUbSize, sortTileInfo.lastAxis,
394- sortTileInfo.outerSize, sortTileInfo.innerSize, sortTileInfo.innerChunk);397+ sortTileInfo.unsortedDim, sortTileInfo.outerSize, sortTileInfo.innerSize, sortTileInfo.innerChunk);
395 return;398 return;
396}399}
397 400 
398-ge::graphStatus SetMergeSortTiling(gert::TilingContext *context, SortKthTileInfo &info)401+ge::graphStatus SetMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info)
399{402{
400 OP_CHECK_IF(!ComputeMergeSortTiling(context, info, info.y2DtypeSize),403 OP_CHECK_IF(!ComputeMergeSortTiling(context, info, info.y2DtypeSize),
401- OP_LOGE(context->GetNodeName(), "merge sort tiling failed"), return ge::GRAPH_FAILED);404+ OP_LOGE(context->GetNodeName(), "merge sort tiling failed"), return ge::GRAPH_FAILED);
402 return ge::GRAPH_SUCCESS;405 return ge::GRAPH_SUCCESS;
403}406}
404 407 
405-ge::graphStatus SetMergeIntraCoreTiling(gert::TilingContext *context, SortKthTileInfo &info)408+ge::graphStatus SetMergeIntraCoreTiling(gert::TilingContext* context, SortKthTileInfo& info)
406{409{
407 OP_CHECK_IF(!ComputeMergeIntraCoreTiling(context, info),410 OP_CHECK_IF(!ComputeMergeIntraCoreTiling(context, info),
408- OP_LOGE(context->GetNodeName(), "merge intra-core plan computation failed"), return ge::GRAPH_FAILED);411+ OP_LOGE(context->GetNodeName(), "merge intra-core plan computation failed"), return ge::GRAPH_FAILED);
409 OP_LOGI("MergeIntraCoreTiling",412 OP_LOGI("MergeIntraCoreTiling",
410- "B %ld, N %ld, batchPerCore %u, actualCoreNum %u, blockSortSize %u, extractChunkSize %u, "413+ "B %ld, N %ld, batchPerCore %u, actualCoreNum %u, blockSortSize %u, extractChunkSize %u, "
411- "blocksPerRow %u, alignNum %u, ubSize %u",414+ "blocksPerRow %u, alignNum %u, ubSize %u",
412- info.unsortedDim, info.lastAxis, info.keyParams0, info.coreNumNeed,415+ info.unsortedDim, info.lastAxis, info.keyParams0, info.coreNumNeed, info.numTileDataSize, info.keyParams4,
413- info.numTileDataSize, info.keyParams4, info.lastDimTileNum, info.keyParams3, info.ubSize);416+ info.lastDimTileNum, info.keyParams3, info.ubSize);
414 return ge::GRAPH_SUCCESS;417 return ge::GRAPH_SUCCESS;
415}418}
416 419 
417// =============================================================================420// =============================================================================
418// Try functions421// Try functions
419// =============================================================================422// =============================================================================
420-bool TrySmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId)423+bool TrySmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId)
421{424{
422 if (sortTileInfo.lastAxis > static_cast<int64_t>(SMALL_AXIS_THRESHOLD)) {425 if (sortTileInfo.lastAxis > static_cast<int64_t>(SMALL_AXIS_THRESHOLD)) {
423 return false;426 return false;
@@ -430,9 +433,8 @@ bool TrySmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, u
430 return true;433 return true;
431 }434 }
432 SmallAxisRoutePlan smallAxisRoutePlan;435 SmallAxisRoutePlan smallAxisRoutePlan;
433- bool selected = sortTileInfo.isNonLastAxis ?436+ bool selected = sortTileInfo.isNonLastAxis ? SelectNonLastSmallAxisRoute(sortTileInfo, smallAxisRoutePlan) :
434- SelectNonLastSmallAxisRoute(sortTileInfo, smallAxisRoutePlan) :437+ SelectSmallAxisRoute(sortTileInfo, smallAxisRoutePlan);
435- SelectSmallAxisRoute(sortTileInfo, smallAxisRoutePlan);
436 if (!selected) {438 if (!selected) {
437 return false;439 return false;
438 }440 }
@@ -449,7 +451,7 @@ bool TrySmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, u
449 return true;451 return true;
450}452}
451 453 
452-static bool TryAlignBytesToUint32(uint64_t bytes, uint32_t alignBytes, uint32_t &alignedBytes)454+static bool TryAlignBytesToUint32(uint64_t bytes, uint32_t alignBytes, uint32_t& alignedBytes)
453{455{
454 if (bytes > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {456 if (bytes > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) {
455 return false;457 return false;
@@ -463,14 +465,14 @@ static bool TryAlignBytesToUint32(uint64_t bytes, uint32_t alignBytes, uint32_t
463}465}
464 466 
465struct SortNonLastUbLayout {467struct SortNonLastUbLayout {
466- uint32_t inputRowBytes = 0; // one row of input along inner axis468+ uint32_t inputRowBytes = 0; // one row of input along inner axis
467- uint32_t valueAxisBytes = 0; // sort-axis value buffer per inner position469+ uint32_t valueAxisBytes = 0; // sort-axis value buffer per inner position
468- uint32_t indexAxisBytes = 0; // sort-axis index buffer per inner position470+ uint32_t indexAxisBytes = 0; // sort-axis index buffer per inner position
469- uint32_t outputIndexRowBytes = 0; // one row of output index along inner axis471+ uint32_t outputIndexRowBytes = 0; // one row of output index along inner axis
470};472};
471 473 
472-static bool ComputeSortNonLastUbLayout(const SortKthTileInfo &sortTileInfo, uint32_t innerChunk,474+static bool ComputeSortNonLastUbLayout(const SortKthTileInfo& sortTileInfo, uint32_t innerChunk, uint32_t sortCount,
473- uint32_t sortCount, bool useMergeSort, SortNonLastUbLayout &layout)475+ bool useMergeSort, SortNonLastUbLayout& layout)
474{476{
475 uint32_t align = sortTileInfo.blockUbSize;477 uint32_t align = sortTileInfo.blockUbSize;
476 uint32_t sortDtypeSize = GetNonLastSortDtypeSize(sortTileInfo.dtypeSize, useMergeSort, sortTileInfo.dataType);478 uint32_t sortDtypeSize = GetNonLastSortDtypeSize(sortTileInfo.dtypeSize, useMergeSort, sortTileInfo.dataType);
@@ -481,18 +483,18 @@ static bool ComputeSortNonLastUbLayout(const SortKthTileInfo &sortTileInfo, uint
481 valueAxisRawBytes = std::max(valueAxisRawBytes, static_cast<uint64_t>(sortCount) * SORT_STRUCT_BYTES);483 valueAxisRawBytes = std::max(valueAxisRawBytes, static_cast<uint64_t>(sortCount) * SORT_STRUCT_BYTES);
482 }484 }
483 return TryAlignBytesToUint32(static_cast<uint64_t>(innerChunk) * sortTileInfo.dtypeSize, align,485 return TryAlignBytesToUint32(static_cast<uint64_t>(innerChunk) * sortTileInfo.dtypeSize, align,
484- layout.inputRowBytes) &&486+ layout.inputRowBytes) &&
485- TryAlignBytesToUint32(valueAxisRawBytes, align, layout.valueAxisBytes) &&487+ TryAlignBytesToUint32(valueAxisRawBytes, align, layout.valueAxisBytes) &&
486- TryAlignBytesToUint32(static_cast<uint64_t>(sortCount) * sizeof(uint32_t), align, layout.indexAxisBytes) &&488+ TryAlignBytesToUint32(static_cast<uint64_t>(sortCount) * sizeof(uint32_t), align, layout.indexAxisBytes) &&
487- TryAlignBytesToUint32(static_cast<uint64_t>(innerChunk) * sortTileInfo.y2DtypeSize, align,489+ TryAlignBytesToUint32(static_cast<uint64_t>(innerChunk) * sortTileInfo.y2DtypeSize, align,
488- layout.outputIndexRowBytes);490+ layout.outputIndexRowBytes);
489}491}
490 492 
491// Estimate peak UB consumption for a non-last-axis small-axis sort candidate.493// Estimate peak UB consumption for a non-last-axis small-axis sort candidate.
492// innerChunk: number of adjacent inner positions batched into one sort invocation.494// innerChunk: number of adjacent inner positions batched into one sort invocation.
493// On success, writes peak UB bytes to peakUb and the chosen layout to candidate.495// On success, writes peak UB bytes to peakUb and the chosen layout to candidate.
494-static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo &sortTileInfo, uint32_t innerChunk,496+static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo& sortTileInfo, uint32_t innerChunk, uint64_t& peakUb,
495- uint64_t &peakUb, bool useMergeSort, NonLastSmallAxisCandidate &candidate)497+ bool useMergeSort, NonLastSmallAxisCandidate& candidate)
496{498{
497 if (innerChunk == 0U) {499 if (innerChunk == 0U) {
498 return false;500 return false;
@@ -504,15 +506,14 @@ static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo &sortTileInfo, uint32
504 return false;506 return false;
505 }507 }
506 508 
507- // 8/16-bit gather uses uint16_t offsets; verify row strides stay addressable.509+ // 8/16-bit gather uses uint16_t offsets represented through signed 16-bit RangeType bit patterns.
508 uint64_t inputRowElems = static_cast<uint64_t>(layout.inputRowBytes) / sortTileInfo.dtypeSize;510 uint64_t inputRowElems = static_cast<uint64_t>(layout.inputRowBytes) / sortTileInfo.dtypeSize;
509 uint64_t valueAxisElems = static_cast<uint64_t>(layout.valueAxisBytes) /511 uint64_t valueAxisElems = static_cast<uint64_t>(layout.valueAxisBytes) /
510- GetNonLastSortDtypeSize(sortTileInfo.dtypeSize, useMergeSort, sortTileInfo.dataType);512+ GetNonLastSortDtypeSize(sortTileInfo.dtypeSize, useMergeSort, sortTileInfo.dataType);
511 if (sortTileInfo.dtypeSize <= sizeof(uint16_t) &&513 if (sortTileInfo.dtypeSize <= sizeof(uint16_t) &&
512 ((static_cast<uint64_t>(axisLen) - 1U) * inputRowElems > std::numeric_limits<uint16_t>::max() ||514 ((static_cast<uint64_t>(axisLen) - 1U) * inputRowElems > std::numeric_limits<uint16_t>::max() ||
513 static_cast<uint64_t>(innerChunk - 1U) * valueAxisElems > std::numeric_limits<uint16_t>::max())) {515 static_cast<uint64_t>(innerChunk - 1U) * valueAxisElems > std::numeric_limits<uint16_t>::max())) {
514- // 8/16-bit vector gather uses 16-bit offsets after packing, so both source516+ // Keep 16-bit offset bit patterns addressable after reinterpret-casting to uint16_t.
515- // and destination row strides must stay within uint16_t addressable elements.
516 return false;517 return false;
517 }518 }
518 519 
@@ -522,8 +523,8 @@ static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo &sortTileInfo, uint32
522 // BF16 merge uses an extra bf16 staging row before casting to fp32. It is a523 // BF16 merge uses an extra bf16 staging row before casting to fp32. It is a
523 // real producer buffer because its aligned padding is consumed by Cast.524 // real producer buffer because its aligned padding is consumed by Cast.
524 uint32_t inputValueAxisBytes = 0;525 uint32_t inputValueAxisBytes = 0;
525- if (!TryAlignBytesToUint32(static_cast<uint64_t>(sortCount) * sortTileInfo.dtypeSize,526+ if (!TryAlignBytesToUint32(static_cast<uint64_t>(sortCount) * sortTileInfo.dtypeSize, sortTileInfo.blockUbSize,
526- sortTileInfo.blockUbSize, inputValueAxisBytes)) {527+ inputValueAxisBytes)) {
527 return false;528 return false;
528 }529 }
529 bf16CastBytes = static_cast<uint64_t>(innerChunk) * inputValueAxisBytes;530 bf16CastBytes = static_cast<uint64_t>(innerChunk) * inputValueAxisBytes;
@@ -535,8 +536,8 @@ static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo &sortTileInfo, uint32
535 // innerChunk * indexAxisBytes — index axis (one per inner position)536 // innerChunk * indexAxisBytes — index axis (one per inner position)
536 // bf16CastBytes — BF16 staging rows (0 for non-BF16 or non-merge)537 // bf16CastBytes — BF16 staging rows (0 for non-BF16 or non-merge)
537 peakUb = static_cast<uint64_t>(axisLen) * layout.inputRowBytes +538 peakUb = static_cast<uint64_t>(axisLen) * layout.inputRowBytes +
538- static_cast<uint64_t>(innerChunk) * layout.valueAxisBytes * 2U +539+ static_cast<uint64_t>(innerChunk) * layout.valueAxisBytes * 2U +
539- static_cast<uint64_t>(innerChunk) * layout.indexAxisBytes + bf16CastBytes;540+ static_cast<uint64_t>(innerChunk) * layout.indexAxisBytes + bf16CastBytes;
540 if (innerChunk > 1U) {541 if (innerChunk > 1U) {
541 // outputIndex buffer only needed when batching multiple inner positions542 // outputIndex buffer only needed when batching multiple inner positions
542 peakUb += static_cast<uint64_t>(axisLen) * layout.outputIndexRowBytes;543 peakUb += static_cast<uint64_t>(axisLen) * layout.outputIndexRowBytes;
@@ -555,8 +556,8 @@ static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo &sortTileInfo, uint32
555 return true;556 return true;
556}557}
557 558 
558-bool ApplyNonLastSmallAxisResult(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId,559+bool ApplyNonLastSmallAxisResult(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId,
559- bool useMergeSort, const NonLastSmallAxisCandidate &best, uint64_t usableUb)560+ bool useMergeSort, const NonLastSmallAxisCandidate& best, uint64_t usableUb)
560{561{
561 sortTileInfo.innerChunk = best.innerChunk;562 sortTileInfo.innerChunk = best.innerChunk;
562 sortTileInfo.innerLoopNum = best.innerLoopNum;563 sortTileInfo.innerLoopNum = best.innerLoopNum;
@@ -573,21 +574,21 @@ bool ApplyNonLastSmallAxisResult(gert::TilingContext *context, SortKthTileInfo &
573 // Schedule 9 and 10 share the same GM/UB layout. Only the per-row sort primitive574 // Schedule 9 and 10 share the same GM/UB layout. Only the per-row sort primitive
574 // differs, so use the selected sort type to choose the binary.575 // differs, so use the selected sort type to choose the binary.
575 schId = useMergeSort ? SORT_SCHID_9 : SORT_SCHID_10;576 schId = useMergeSort ? SORT_SCHID_9 : SORT_SCHID_10;
576- size_t *userWorkSpaceSize = context->GetWorkspaceSizes(1);577+ size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1);
577 if (userWorkSpaceSize == nullptr) {578 if (userWorkSpaceSize == nullptr) {
578 OP_LOGE(context->GetNodeName(), "get workspace size pointer failed");579 OP_LOGE(context->GetNodeName(), "get workspace size pointer failed");
579 return false;580 return false;
580 }581 }
581 userWorkSpaceSize[0] = WORK_SPACE_SIZE;582 userWorkSpaceSize[0] = WORK_SPACE_SIZE;
582 OP_LOGI(context->GetNodeName(),583 OP_LOGI(context->GetNodeName(),
583- "non-last small-axis no-transpose selected axis=%ld inner=%ld outer=%ld innerChunk=%u tileCount=%lu "584+ "non-last small-axis no-transpose selected axis=%ld inner=%ld outer=%ld innerChunk=%u tileCount=%lu "
584- "activeCore=%u peakUb=%lu usableUb=%lu",585+ "activeCore=%u peakUb=%lu usableUb=%lu",
585- sortTileInfo.lastAxis, sortTileInfo.innerSize, sortTileInfo.outerSize, best.innerChunk, best.tileCount,586+ sortTileInfo.lastAxis, sortTileInfo.innerSize, sortTileInfo.outerSize, best.innerChunk, best.tileCount,
586- best.activeCore, best.peakUb, usableUb);587+ best.activeCore, best.peakUb, usableUb);
587 return true;588 return true;
588}589}
589 590 
590-bool TryNonLastSmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId)591+bool TryNonLastSmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId)
591{592{
592 if (!sortTileInfo.isNonLastAxis) {593 if (!sortTileInfo.isNonLastAxis) {
593 return false;594 return false;
@@ -603,10 +604,8 @@ bool TryNonLastSmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTile
603 return false;604 return false;
604 }605 }
605 606 
606- uint32_t sortCount = GetNonLastSortCount(sortTileInfo.dataType,607+ uint32_t sortCount = GetNonLastSortCount(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis));
607- static_cast<uint32_t>(sortTileInfo.lastAxis));608+ bool useMergeSort = UseNonLastMergeSort(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis));
608- bool useMergeSort = UseNonLastMergeSort(sortTileInfo.dataType,
609- static_cast<uint32_t>(sortTileInfo.lastAxis));
610 // Query Sort tmp with the same aligned sortCount and effective dtype that the609 // Query Sort tmp with the same aligned sortCount and effective dtype that the
611 // kernel will use; BF16 merge sorts as fp32 after the UB cast.610 // kernel will use; BF16 merge sorts as fp32 after the UB cast.
612 uint32_t tmpUb = 0;611 uint32_t tmpUb = 0;
@@ -617,8 +616,8 @@ bool TryNonLastSmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTile
617 SortKthTileInfo selectedInfo = sortTileInfo;616 SortKthTileInfo selectedInfo = sortTileInfo;
618 // Bind useMergeSort into a callback so SearchNonLastSmallAxisPlan can evaluate617 // Bind useMergeSort into a callback so SearchNonLastSmallAxisPlan can evaluate
619 // each innerChunk candidate without knowing the sort-type decision.618 // each innerChunk candidate without knowing the sort-type decision.
620- auto estimateUb = [useMergeSort](SortKthTileInfo &candidateInfo, uint32_t innerChunk,619+ auto estimateUb = [useMergeSort](SortKthTileInfo& candidateInfo, uint32_t innerChunk, uint64_t& peakUb,
621- uint64_t &peakUb, NonLastSmallAxisCandidate &candidate) -> bool {620+ NonLastSmallAxisCandidate& candidate) -> bool {
622 return EstimateSortNonLastSmallAxisUb(candidateInfo, innerChunk, peakUb, useMergeSort, candidate);621 return EstimateSortNonLastSmallAxisUb(candidateInfo, innerChunk, peakUb, useMergeSort, candidate);
623 };622 };
624 if (!SearchNonLastSmallAxisPlan(sortTileInfo, usableUb, estimateUb, best, &selectedInfo)) {623 if (!SearchNonLastSmallAxisPlan(sortTileInfo, usableUb, estimateUb, best, &selectedInfo)) {
@@ -629,20 +628,19 @@ bool TryNonLastSmallAxis(gert::TilingContext *context, SortKthTileInfo &sortTile
629 return ApplyNonLastSmallAxisResult(context, sortTileInfo, schId, useMergeSort, best, usableUb);628 return ApplyNonLastSmallAxisResult(context, sortTileInfo, schId, useMergeSort, best, usableUb);
630}629}
631 630 
632-bool TryMerge(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId)631+bool TryMerge(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId)
633{632{
634 if (IsMergeSortSupported(sortTileInfo.dataType, sortTileInfo.lastAxis)) {633 if (IsMergeSortSupported(sortTileInfo.dataType, sortTileInfo.lastAxis)) {
635 SortKthTileInfo candidate = sortTileInfo;634 SortKthTileInfo candidate = sortTileInfo;
636 if (SetMergeSortTiling(context, candidate) == ge::GRAPH_SUCCESS) {635 if (SetMergeSortTiling(context, candidate) == ge::GRAPH_SUCCESS) {
637 sortTileInfo = candidate;636 sortTileInfo = candidate;
638- schId = (sortTileInfo.lastAxis <= SORT32_SMALL_AXIS_THRESHOLD) ?637+ schId = (sortTileInfo.lastAxis <= SORT32_SMALL_AXIS_THRESHOLD) ? SORT_SCHID_8 : static_cast<uint64_t>(0);
639- SORT_SCHID_8 : static_cast<uint64_t>(0);
640 return true;638 return true;
641 }639 }
642 }640 }
643 641 
644- if (IsMergeMoreCoreSupported(sortTileInfo.dataType, sortTileInfo.lastAxis,642+ if (IsMergeMoreCoreSupported(sortTileInfo.dataType, sortTileInfo.lastAxis, sortTileInfo.unsortedDim,
645- sortTileInfo.unsortedDim, sortTileInfo.maxCoreNum)) {643+ sortTileInfo.maxCoreNum)) {
646 SortKthTileInfo candidate = sortTileInfo;644 SortKthTileInfo candidate = sortTileInfo;
647 if (SetMergeMoreCoreTiling(context, candidate) == ge::GRAPH_SUCCESS) {645 if (SetMergeMoreCoreTiling(context, candidate) == ge::GRAPH_SUCCESS) {
648 sortTileInfo = candidate;646 sortTileInfo = candidate;
@@ -654,11 +652,10 @@ bool TryMerge(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint6
654 return false;652 return false;
655}653}
656 654 
657-bool TryRadixOneCore(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId)655+bool TryRadixOneCore(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId)
658{656{
659 SortKthTileInfo candidate = sortTileInfo;657 SortKthTileInfo candidate = sortTileInfo;
660- if (!IsRadixSortOneCore(candidate) ||658+ if (!IsRadixSortOneCore(candidate) || SetRadixOneCoreTiling(context, candidate) != ge::GRAPH_SUCCESS) {
661- SetRadixOneCoreTiling(context, candidate) != ge::GRAPH_SUCCESS) {
662 return false;659 return false;
663 }660 }
664 sortTileInfo = candidate;661 sortTileInfo = candidate;
@@ -666,10 +663,10 @@ bool TryRadixOneCore(gert::TilingContext *context, SortKthTileInfo &sortTileInfo
666 return true;663 return true;
667}664}
668 665 
669-bool TryMergeIntraCore(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId)666+bool TryMergeIntraCore(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId)
670{667{
671- if (!IsMergeIntraCoreSupported(sortTileInfo.dataType, sortTileInfo.lastAxis,668+ if (!IsMergeIntraCoreSupported(sortTileInfo.dataType, sortTileInfo.lastAxis, sortTileInfo.unsortedDim,
672- sortTileInfo.unsortedDim, sortTileInfo.maxCoreNum, sortTileInfo.ubSize)) {669+ sortTileInfo.maxCoreNum, sortTileInfo.ubSize)) {
673 return false;670 return false;
674 }671 }
675 SortKthTileInfo candidate = sortTileInfo;672 SortKthTileInfo candidate = sortTileInfo;
@@ -684,7 +681,7 @@ bool TryMergeIntraCore(gert::TilingContext *context, SortKthTileInfo &sortTileIn
684// =============================================================================681// =============================================================================
685// Route selection682// Route selection
686// =============================================================================683// =============================================================================
687-ge::graphStatus SelectSortSchedule(gert::TilingContext *context, SortKthTileInfo &sortTileInfo, uint64_t &schId)684+ge::graphStatus SelectSortSchedule(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId)
688{685{
689 if (TrySmallAxis(context, sortTileInfo, schId)) {686 if (TrySmallAxis(context, sortTileInfo, schId)) {
690 return ge::GRAPH_SUCCESS;687 return ge::GRAPH_SUCCESS;
@@ -698,43 +695,42 @@ ge::graphStatus SelectSortSchedule(gert::TilingContext *context, SortKthTileInfo
698 OP_LOGE(context->GetNodeName(), "non-last sort axis does not meet no-transpose schedule constraints");695 OP_LOGE(context->GetNodeName(), "non-last sort axis does not meet no-transpose schedule constraints");
699 return ge::GRAPH_FAILED;696 return ge::GRAPH_FAILED;
700 }697 }
701- if (TryMerge(context, sortTileInfo, schId) ||698+ if (TryMerge(context, sortTileInfo, schId) || TryRadixOneCore(context, sortTileInfo, schId) ||
702- TryRadixOneCore(context, sortTileInfo, schId) ||
703 TryMergeIntraCore(context, sortTileInfo, schId)) {699 TryMergeIntraCore(context, sortTileInfo, schId)) {
704 return ge::GRAPH_SUCCESS;700 return ge::GRAPH_SUCCESS;
705 }701 }
706 702 
707 schId = SORT_SCHID_2;703 schId = SORT_SCHID_2;
708 OP_CHECK_IF(SetRadixMoreCoreTiling(context, sortTileInfo) != ge::GRAPH_SUCCESS,704 OP_CHECK_IF(SetRadixMoreCoreTiling(context, sortTileInfo) != ge::GRAPH_SUCCESS,
709- OP_LOGE(context->GetNodeName(), "radix more-core tiling failed"), return ge::GRAPH_FAILED);705+ OP_LOGE(context->GetNodeName(), "radix more-core tiling failed"), return ge::GRAPH_FAILED);
710 return ge::GRAPH_SUCCESS;706 return ge::GRAPH_SUCCESS;
711}707}
712 708 
713// =============================================================================709// =============================================================================
714// Main entry710// Main entry
715// =============================================================================711// =============================================================================
716-ge::graphStatus RadixSortTiling(gert::TilingContext *context, int32_t maxCoreNum)712+ge::graphStatus RadixSortTiling(gert::TilingContext* context, int32_t maxCoreNum)
717{713{
718- SortRegBaseTilingData *sortTilingData{ nullptr };714+ SortRegBaseTilingData* sortTilingData{nullptr};
719 sortTilingData = context->GetTilingData<SortRegBaseTilingData>();715 sortTilingData = context->GetTilingData<SortRegBaseTilingData>();
720- OP_CHECK_IF(sortTilingData == nullptr,716+ OP_CHECK_IF(sortTilingData == nullptr, OP_LOGE(context->GetNodeName(), "get tilingdata ptr failed"),
721- OP_LOGE(context->GetNodeName(), "get tilingdata ptr failed"), return ge::GRAPH_FAILED);717+ return ge::GRAPH_FAILED);
722 OP_CHECK_IF((memset_s(sortTilingData, sizeof(SortRegBaseTilingData), 0, sizeof(SortRegBaseTilingData)) != EOK),718 OP_CHECK_IF((memset_s(sortTilingData, sizeof(SortRegBaseTilingData), 0, sizeof(SortRegBaseTilingData)) != EOK),
723- OP_LOGE(context->GetNodeName(), "memset tilingdata failed"), return ge::GRAPH_FAILED);719+ OP_LOGE(context->GetNodeName(), "memset tilingdata failed"), return ge::GRAPH_FAILED);
724 SortKthTileInfo sortTileInfo;720 SortKthTileInfo sortTileInfo;
725 OP_CHECK_IF(SortCheckParams(context, sortTileInfo) != ge::GRAPH_SUCCESS,721 OP_CHECK_IF(SortCheckParams(context, sortTileInfo) != ge::GRAPH_SUCCESS,
726- OP_LOGE(context->GetNodeName(), "check params failed"), return ge::GRAPH_FAILED);722+ OP_LOGE(context->GetNodeName(), "check params failed"), return ge::GRAPH_FAILED);
727 sortTileInfo.maxCoreNum = static_cast<uint32_t>(maxCoreNum);723 sortTileInfo.maxCoreNum = static_cast<uint32_t>(maxCoreNum);
728 int64_t int32Max = static_cast<int64_t>(std::numeric_limits<int32_t>::max());724 int64_t int32Max = static_cast<int64_t>(std::numeric_limits<int32_t>::max());
729 uint64_t isInt32 = static_cast<uint64_t>((sortTileInfo.lastAxis <= int32Max));725 uint64_t isInt32 = static_cast<uint64_t>((sortTileInfo.lastAxis <= int32Max));
730- const bool *isDescending = context->GetAttrs()->GetAttrPointer<bool>(1);726+ const bool* isDescending = context->GetAttrs()->GetAttrPointer<bool>(1);
731 uint64_t isDescend = *isDescending;727 uint64_t isDescend = *isDescending;
732 sortTileInfo.isDescend = static_cast<bool>(isDescend);728 sortTileInfo.isDescend = static_cast<bool>(isDescend);
733 sortTileInfo.isInt32 = static_cast<uint32_t>(isInt32);729 sortTileInfo.isInt32 = static_cast<uint32_t>(isInt32);
734 OP_LOGI(context->GetNodeName(), "isInt32 is %lu, isDescend is %lu", isInt32, isDescend);730 OP_LOGI(context->GetNodeName(), "isInt32 is %lu, isDescend is %lu", isInt32, isDescend);
735 uint64_t schId = static_cast<uint64_t>(0);731 uint64_t schId = static_cast<uint64_t>(0);
736 OP_CHECK_IF(SelectSortSchedule(context, sortTileInfo, schId) != ge::GRAPH_SUCCESS,732 OP_CHECK_IF(SelectSortSchedule(context, sortTileInfo, schId) != ge::GRAPH_SUCCESS,
737- OP_LOGE(context->GetNodeName(), "select sort schedule failed"), return ge::GRAPH_FAILED);733+ OP_LOGE(context->GetNodeName(), "select sort schedule failed"), return ge::GRAPH_FAILED);
738 const uint64_t tilingKey = GET_TPL_TILING_KEY(schId, isInt32, isDescend);734 const uint64_t tilingKey = GET_TPL_TILING_KEY(schId, isInt32, isDescend);
739 OP_LOGI(context->GetNodeName(), "tilingKey is %lu, maxCoreNum %d, schId %lu", tilingKey, maxCoreNum, schId);735 OP_LOGI(context->GetNodeName(), "tilingKey is %lu, maxCoreNum %d, schId %lu", tilingKey, maxCoreNum, schId);
740 context->SetTilingKey(tilingKey);736 context->SetTilingKey(tilingKey);
@@ -746,8 +742,8 @@ ge::graphStatus RadixSortTiling(gert::TilingContext *context, int32_t maxCoreNum
746 return ge::GRAPH_SUCCESS;742 return ge::GRAPH_SUCCESS;
747}743}
748 744 
749-ge::graphStatus SortTilingSimt(gert::TilingContext *context, int32_t maxCoreNum)745+ge::graphStatus SortTilingSimt(gert::TilingContext* context, int32_t maxCoreNum)
750{746{
751 return RadixSortTiling(context, maxCoreNum);747 return RadixSortTiling(context, maxCoreNum);
752}748}
753-}749+} // namespace optiling
Mmath/sort/op_host/arch35/sort_tiling_common.h+40-44
@@ -92,8 +92,8 @@ constexpr SmallAxisRule kSmallAxisRules[] = {
92 // dtype, insertionMaxN, twoStageMaxN, insertion tiers, two-stage tiers92 // dtype, insertionMaxN, twoStageMaxN, insertion tiers, two-stage tiers
93 {ge::DT_INT64, 16, 512, {{8, 1}, {16, 4}, {0, 0}}, {{15, 8}, {128, 4}, {512, 8}, {0, 0}}},93 {ge::DT_INT64, 16, 512, {{8, 1}, {16, 4}, {0, 0}}, {{15, 8}, {128, 4}, {512, 8}, {0, 0}}},
94 {ge::DT_UINT64, 16, 512, {{8, 1}, {16, 4}, {0, 0}}, {{15, 8}, {128, 4}, {512, 8}, {0, 0}}},94 {ge::DT_UINT64, 16, 512, {{8, 1}, {16, 4}, {0, 0}}, {{15, 8}, {128, 4}, {512, 8}, {0, 0}}},
95- {ge::DT_INT32, 11, 384, {{8, 2}, {11, 4}, {0, 0}}, {{11, 8}, {64, 4}, {384, 12}, {0, 0}}},95+ {ge::DT_INT32, 11, 384, {{8, 2}, {11, 4}, {0, 0}}, {{11, 8}, {64, 4}, {384, 8}, {0, 0}}},
96- {ge::DT_UINT32, 11, 384, {{8, 2}, {11, 4}, {0, 0}}, {{11, 8}, {64, 4}, {384, 12}, {0, 0}}},96+ {ge::DT_UINT32, 11, 384, {{8, 2}, {11, 4}, {0, 0}}, {{11, 8}, {64, 4}, {384, 8}, {0, 0}}},
97 {ge::DT_INT16, 8, 192, {{4, 2}, {8, 4}, {0, 0}}, {{7, 8}, {64, 4}, {192, 12}, {0, 0}}},97 {ge::DT_INT16, 8, 192, {{4, 2}, {8, 4}, {0, 0}}, {{7, 8}, {64, 4}, {192, 12}, {0, 0}}},
98 {ge::DT_UINT16, 8, 192, {{4, 2}, {8, 4}, {0, 0}}, {{7, 8}, {64, 4}, {192, 12}, {0, 0}}},98 {ge::DT_UINT16, 8, 192, {{4, 2}, {8, 4}, {0, 0}}, {{7, 8}, {64, 4}, {192, 12}, {0, 0}}},
99 {ge::DT_INT8, 8, 128, {{4, 2}, {8, 7}, {0, 0}}, {{3, 8}, {64, 7}, {128, 16}, {0, 0}}},99 {ge::DT_INT8, 8, 128, {{4, 2}, {8, 7}, {0, 0}}, {{3, 8}, {64, 7}, {128, 16}, {0, 0}}},
@@ -107,8 +107,8 @@ constexpr size_t kNumSmallAxisRules = sizeof(kSmallAxisRules) / sizeof(kSmallAxi
107uint32_t LookupMinSegs(const uint32_t (*tiers)[2], uint32_t axisNum);107uint32_t LookupMinSegs(const uint32_t (*tiers)[2], uint32_t axisNum);
108const SmallAxisRule* FindSmallAxisRule(ge::DataType dataType);108const SmallAxisRule* FindSmallAxisRule(ge::DataType dataType);
109bool UseTwoStageRankInverse(uint32_t axisLen);109bool UseTwoStageRankInverse(uint32_t axisLen);
110-uint32_t ComputeInsertionBytesPerSeg(110+uint32_t ComputeInsertionBytesPerSeg(ge::DataType dataType, uint32_t axisLen, uint32_t dtypeSize,
111- ge::DataType dataType, uint32_t axisLen, uint32_t dtypeSize, uint32_t indexDtypeSize, uint32_t blockUbSize);111+ uint32_t indexDtypeSize, uint32_t blockUbSize);
112 112 
113// =============================================================================113// =============================================================================
114// SortKthTileInfo — central tiling information struct114// SortKthTileInfo — central tiling information struct
@@ -169,8 +169,8 @@ ge::DataType GetNonLastSortDtype(ge::DataType dataType, bool useMergeSort);
169uint32_t GetNonLastSortDtypeSize(uint32_t dtypeSize, bool useMergeSort, ge::DataType dataType);169uint32_t GetNonLastSortDtypeSize(uint32_t dtypeSize, bool useMergeSort, ge::DataType dataType);
170 170 
171uint32_t GetNonLastSortCount(ge::DataType dataType, uint32_t axisLen);171uint32_t GetNonLastSortCount(ge::DataType dataType, uint32_t axisLen);
172-bool GetNonLastSortTmpSize(172+bool GetNonLastSortTmpSize(ge::DataType dataType, uint32_t sortCount, bool useMergeSort, bool isDescend,
173- ge::DataType dataType, uint32_t sortCount, bool useMergeSort, bool isDescend, uint32_t& tmpUbSize);173+ uint32_t& tmpUbSize);
174bool ComputeNonLastBatchNum(int64_t outerSize, int64_t innerSize, uint32_t innerChunk, uint32_t& batchNum);174bool ComputeNonLastBatchNum(int64_t outerSize, int64_t innerSize, uint32_t innerChunk, uint32_t& batchNum);
175 175 
176struct NonLastSmallAxisCandidate {176struct NonLastSmallAxisCandidate {
@@ -222,28 +222,26 @@ bool QuerySortTmpSizeRadix(ge::DataType dataType, uint32_t sortAxisNum, uint32_t
222 222 
223uint32_t ComputeRadixRemainUb(uint32_t usableUb, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor);223uint32_t ComputeRadixRemainUb(uint32_t usableUb, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor);
224 224 
225-void AdjustRadixTmpUb(225+void AdjustRadixTmpUb(uint32_t usableUb, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor, uint32_t blockUbSize,
226- uint32_t usableUb, uint32_t tileData, uint32_t ubExtra, uint32_t tileFactor, uint32_t blockUbSize,226+ uint32_t& tmpUbSize);
227- uint32_t& tmpUbSize);227+bool ComputeRadixTileDataForAllCore(int64_t axisLen, uint32_t maxCoreNum, uint32_t usableUb, uint32_t ubExtra,
228-bool ComputeRadixTileDataForAllCore(228+ uint32_t tileFactor, uint32_t blockUbSize, uint32_t lastDimTileNum,
229- int64_t axisLen, uint32_t maxCoreNum, uint32_t usableUb, uint32_t ubExtra, uint32_t tileFactor,229+ uint32_t& tileData, uint32_t& tmpUbSize);
230- uint32_t blockUbSize, uint32_t lastDimTileNum, uint32_t& tileData, uint32_t& tmpUbSize);230+bool NeedAdjustRadixTileData(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t usableUb,
231-bool NeedAdjustRadixTileData(231+ uint32_t ubExtra, uint32_t tileFactor, uint32_t blockUbSize, uint32_t& tileData,
232- int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t usableUb, uint32_t ubExtra, uint32_t tileFactor,232+ uint32_t lastDimTileNum, uint32_t& tmpUbSize, bool& adjusted);
233- uint32_t blockUbSize, uint32_t& tileData, uint32_t lastDimTileNum, uint32_t& tmpUbSize, bool& adjusted);233+bool ComputeRadixTileData(int64_t axisLen, int64_t unsortedDim, uint32_t dtypeSize, uint32_t indexSize,
234-bool ComputeRadixTileData(234+ uint32_t maxCoreNum, uint32_t usableUb, uint32_t blockUbSize, uint32_t& tileData,
235- int64_t axisLen, int64_t unsortedDim, uint32_t dtypeSize, uint32_t indexSize, uint32_t maxCoreNum,235+ uint32_t& tmpUbSize);
236- uint32_t usableUb, uint32_t blockUbSize, uint32_t& tileData, uint32_t& tmpUbSize);236+bool FillRadixKernelParams(uint32_t dtypeSize, uint32_t indexSize, uint32_t coreNumNeed, uint32_t lastDimTileNum,
237-bool FillRadixKernelParams(237+ uint32_t unsortedDimParallel, uint32_t blockUbSize, uint32_t tmpUbSize,
238- uint32_t dtypeSize, uint32_t indexSize, uint32_t coreNumNeed, uint32_t lastDimTileNum, uint32_t unsortedDimParallel,238+ RadixClearParams& out);
239- uint32_t blockUbSize, uint32_t tmpUbSize, RadixClearParams& out);239+bool ComputeRadixSortWorkspace(int64_t axisLen, uint32_t dtypeSize, uint32_t indexSize, uint32_t lastDimTileNum,
240-bool ComputeRadixSortWorkspace(240+ uint32_t numTileDataSize, uint32_t unsortedDimParallel, uint32_t keyParams0,
241- int64_t axisLen, uint32_t dtypeSize, uint32_t indexSize, uint32_t lastDimTileNum, uint32_t numTileDataSize,241+ uint32_t keyParams1, uint32_t keyParams2, uint32_t keyParams3, uint32_t keyParams4,
242- uint32_t unsortedDimParallel, uint32_t keyParams0, uint32_t keyParams1, uint32_t keyParams2, uint32_t keyParams3,242+ uint32_t blockUbSize, uint64_t& workspaceSize);
243- uint32_t keyParams4, uint32_t blockUbSize, uint64_t& workspaceSize);243+bool ComputeRadixOneCoreUbSizes(int64_t lastAxis, uint32_t dtypeSize, uint32_t indexElemSize, uint32_t blockUbSize,
244-bool ComputeRadixOneCoreUbSizes(244+ uint32_t& xUbSize, uint32_t& idxUbSize);
245- int64_t lastAxis, uint32_t dtypeSize, uint32_t indexElemSize, uint32_t blockUbSize, uint32_t& xUbSize,
246- uint32_t& idxUbSize);
247bool FillRadixMoreCoreInfo(SortKthTileInfo& info);245bool FillRadixMoreCoreInfo(SortKthTileInfo& info);
248 246 
249// =============================================================================247// =============================================================================
@@ -266,9 +264,8 @@ struct MergeSortPlan {
266 uint32_t coreNumNeed = 0;264 uint32_t coreNumNeed = 0;
267};265};
268 266 
269-bool ComputeMergeSortPlan(267+bool ComputeMergeSortPlan(int64_t axisLen, int64_t unsortedDim, uint32_t blockUbSize, uint32_t tileDataNum,
270- int64_t axisLen, int64_t unsortedDim, uint32_t blockUbSize, uint32_t tileDataNum, uint32_t maxCoreNum,268+ uint32_t maxCoreNum, MergeSortPlan& plan);
271- MergeSortPlan& plan);
272bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t concatTmpSize);269bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t concatTmpSize);
273 270 
274// =============================================================================271// =============================================================================
@@ -285,8 +282,8 @@ struct MergeMoreCorePlan {
285};282};
286 283 
287bool IsMergeMoreCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum);284bool IsMergeMoreCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum);
288-bool ComputeMergeMoreCorePlan(285+bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t mergeBytesPerElem,
289- int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t mergeBytesPerElem, MergeMoreCorePlan& plan);286+ MergeMoreCorePlan& plan);
290bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem);287bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem);
291 288 
292// =============================================================================289// =============================================================================
@@ -304,22 +301,21 @@ struct MergeIntraCorePlan {
304 uint32_t alignNum = 0;301 uint32_t alignNum = 0;
305};302};
306 303 
307-bool IsMergeIntraCoreSupported(304+bool IsMergeIntraCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum,
308- ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t ubSize);305+ uint32_t ubSize);
309-bool ComputeMergeIntraCorePlan(306+bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum,
310- int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, MergeIntraCorePlan& plan);307+ MergeIntraCorePlan& plan);
311bool FillMergeIntraCoreInfo(SortKthTileInfo& info);308bool FillMergeIntraCoreInfo(SortKthTileInfo& info);
312 309 
313// =============================================================================310// =============================================================================
314// Two-stage sort311// Two-stage sort
315// =============================================================================312// =============================================================================
316uint32_t MaxTwoStageU16SafeBatch(uint32_t axisLen);313uint32_t MaxTwoStageU16SafeBatch(uint32_t axisLen);
317-bool ComputeTwoStageSortTmpUb(314+bool ComputeTwoStageSortTmpUb(ge::DataType dataType, uint32_t axisLen, uint32_t totalElems, uint32_t blockUbSize,
318- ge::DataType dataType, uint32_t axisLen, uint32_t totalElems, uint32_t blockUbSize, uint32_t& tmpUbSize);315+ uint32_t& tmpUbSize);
319uint64_t EstimateTwoStageUbBytes(const SortKthTileInfo& info, uint32_t totalElems, uint32_t sortTmpUb);316uint64_t EstimateTwoStageUbBytes(const SortKthTileInfo& info, uint32_t totalElems, uint32_t sortTmpUb);
320-bool PrepareTwoStageBatchCandidate(317+bool PrepareTwoStageBatchCandidate(const SortKthTileInfo& info, uint32_t candidate, uint32_t& totalElems,
321- const SortKthTileInfo& info, uint32_t candidate, uint32_t& totalElems, uint32_t& tmpUbSize, bool& useRankInverse,318+ uint32_t& tmpUbSize, bool& useRankInverse, uint64_t& totalBytes);
322- uint64_t& totalBytes);
323 319 
324struct TwoStageBatchPlan {320struct TwoStageBatchPlan {
325 uint32_t batchSize = 0;321 uint32_t batchSize = 0;
@@ -328,8 +324,8 @@ struct TwoStageBatchPlan {
328 uint32_t tmpUbSize = 0;324 uint32_t tmpUbSize = 0;
329};325};
330 326 
331-bool SearchTwoStageBatchPlan(327+bool SearchTwoStageBatchPlan(uint32_t maxBatch, std::function<bool(uint32_t, TwoStageBatchPlan&)> tryCandidate,
332- uint32_t maxBatch, std::function<bool(uint32_t, TwoStageBatchPlan&)> tryCandidate, TwoStageBatchPlan& result);328+ TwoStageBatchPlan& result);
333 329 
334// =============================================================================330// =============================================================================
335// Tiling data conversion & top-level tiling computation331// Tiling data conversion & top-level tiling computation
Mmath/sort/op_kernel/arch35/common/non_last_small_axis_base.h+33-22
@@ -44,9 +44,8 @@ constexpr uint32_t NON_LAST_MERGE_SORT_ALIGN = 32;
44 * @tparam UseMergeSort Whether to use MERGE_SORT instead of RADIX_SORT for row sorting44 * @tparam UseMergeSort Whether to use MERGE_SORT instead of RADIX_SORT for row sorting
45 * @tparam IsBf16Merge Whether bf16 input needs an intermediate cast buffer for merge-sort path45 * @tparam IsBf16Merge Whether bf16 input needs an intermediate cast buffer for merge-sort path
46 */46 */
47-template <47+template <typename Derived, typename T, typename SortT, typename RangeType, typename IdxType, typename CastType,
48- typename Derived, typename T, typename SortT, typename RangeType, typename IdxType, typename CastType,48+ bool IsDescend, bool UseMergeSort, bool IsBf16Merge>
49- bool IsDescend, bool UseMergeSort, bool IsBf16Merge>
50class NonLastSmallAxisBase {49class NonLastSmallAxisBase {
51public:50public:
52 __aicore__ inline void Process()51 __aicore__ inline void Process()
@@ -70,8 +69,9 @@ public:
70 continue;69 continue;
71 }70 }
72 int64_t innerStart = static_cast<int64_t>(innerTileId) * static_cast<int64_t>(this->innerChunk_);71 int64_t innerStart = static_cast<int64_t>(innerTileId) * static_cast<int64_t>(this->innerChunk_);
73- int64_t inputOffset =72+ int64_t inputOffset = static_cast<int64_t>(outerId) * static_cast<int64_t>(this->axisLen_) *
74- static_cast<int64_t>(outerId) * static_cast<int64_t>(this->axisLen_) * this->innerSize_ + innerStart;73+ this->innerSize_ +
74+ innerStart;
75 int64_t outputOffset = static_cast<int64_t>(outerId) * this->innerSize_ + innerStart;75 int64_t outputOffset = static_cast<int64_t>(outerId) * this->innerSize_ + innerStart;
76 this->LoadTile(inputOffset, curInnerChunk);76 this->LoadTile(inputOffset, curInnerChunk);
77 this->TransposeToSortMajor(curInnerChunk);77 this->TransposeToSortMajor(curInnerChunk);
@@ -119,6 +119,18 @@ protected:
119 uint32_t sortCount_ = 0;119 uint32_t sortCount_ = 0;
120 uint32_t tmpUbSize_ = 0;120 uint32_t tmpUbSize_ = 0;
121 121 
122+ __aicore__ inline RangeType ToRangeScalar(uint32_t value) const
123+ {
124+ if constexpr (std::is_same_v<RangeType, int16_t>) {
125+ uint32_t valueU16 = value & 0xFFFFU;
126+ int32_t signedValue = (valueU16 <= 32767U) ? static_cast<int32_t>(valueU16) :
127+ static_cast<int32_t>(valueU16) - 65536;
128+ return static_cast<int16_t>(signedValue);
129+ } else {
130+ return static_cast<RangeType>(value);
131+ }
132+ }
133+ 
122 __aicore__ inline uint32_t GetCurrentInnerChunk(uint32_t innerTileId) const134 __aicore__ inline uint32_t GetCurrentInnerChunk(uint32_t innerTileId) const
123 {135 {
124 int64_t start = static_cast<int64_t>(this->innerChunk_) * static_cast<int64_t>(innerTileId);136 int64_t start = static_cast<int64_t>(this->innerChunk_) * static_cast<int64_t>(innerTileId);
@@ -136,8 +148,9 @@ protected:
136 {148 {
137 uint32_t curBytes = curInnerChunk * sizeof(T);149 uint32_t curBytes = curInnerChunk * sizeof(T);
138 uint32_t curAlignedBytes = ROUND_UP_AGLIN(curBytes);150 uint32_t curAlignedBytes = ROUND_UP_AGLIN(curBytes);
139- uint32_t dstStride =151+ uint32_t dstStride = (this->inputRowBytes_ > curAlignedBytes) ?
140- (this->inputRowBytes_ > curAlignedBytes) ? (this->inputRowBytes_ - curAlignedBytes) / UB_BLOCK_SIZE : 0;152+ (this->inputRowBytes_ - curAlignedBytes) / UB_BLOCK_SIZE :
153+ 0;
141 uint32_t rightPadding = this->inputRowElems_ > curInnerChunk ? this->inputRowElems_ - curInnerChunk : 0;154 uint32_t rightPadding = this->inputRowElems_ > curInnerChunk ? this->inputRowElems_ - curInnerChunk : 0;
142 int64_t gmStride = (this->innerSize_ - static_cast<int64_t>(curInnerChunk)) * static_cast<int64_t>(sizeof(T));155 int64_t gmStride = (this->innerSize_ - static_cast<int64_t>(curInnerChunk)) * static_cast<int64_t>(sizeof(T));
143 DataCopyExtParams copyParam{static_cast<uint16_t>(this->axisLen_), curBytes, gmStride, dstStride, 0};156 DataCopyExtParams copyParam{static_cast<uint16_t>(this->axisLen_), curBytes, gmStride, dstStride, 0};
@@ -160,10 +173,9 @@ protected:
160 TransposeTileByGather(curInnerChunk);173 TransposeTileByGather(curInnerChunk);
161 if constexpr (IsBf16Merge) {174 if constexpr (IsBf16Merge) {
162 for (uint32_t inner = 0; inner < curInnerChunk; ++inner) {175 for (uint32_t inner = 0; inner < curInnerChunk; ++inner) {
163- Cast(176+ Cast(this->sortInput_[inner * this->valueAxisElems_],
164- this->sortInput_[inner * this->valueAxisElems_],177+ this->inputCast_[inner * this->inputValueAxisElems_], RoundMode::CAST_NONE,
165- this->inputCast_[inner * this->inputValueAxisElems_], RoundMode::CAST_NONE,178+ this->inputValueAxisElems_);
166- this->inputValueAxisElems_);
167 }179 }
168 }180 }
169 }181 }
@@ -181,19 +193,19 @@ protected:
181 {193 {
182 constexpr uint16_t vlSize = static_cast<uint16_t>(Ops::Base::GetVRegSize() / sizeof(CastType));194 constexpr uint16_t vlSize = static_cast<uint16_t>(Ops::Base::GetVRegSize() / sizeof(CastType));
183 __ubuf__ T* inputAddr = (__ubuf__ T*)this->inputTile_.GetPhyAddr();195 __ubuf__ T* inputAddr = (__ubuf__ T*)this->inputTile_.GetPhyAddr();
184- __ubuf__ T* outputAddr =196+ __ubuf__ T* outputAddr = (__ubuf__ T*)(IsBf16Merge ? this->inputCast_.GetPhyAddr() :
185- (__ubuf__ T*)(IsBf16Merge ? this->inputCast_.GetPhyAddr() : this->sortInput_.GetPhyAddr());197+ this->sortInput_.GetPhyAddr());
186 uint32_t outputValueAxisElems = IsBf16Merge ? this->inputValueAxisElems_ : this->valueAxisElems_;198 uint32_t outputValueAxisElems = IsBf16Merge ? this->inputValueAxisElems_ : this->valueAxisElems_;
187 __VEC_SCOPE__199 __VEC_SCOPE__
188 {200 {
189 AscendC::MicroAPI::RegTensor<CastType> dataReg;201 AscendC::MicroAPI::RegTensor<CastType> dataReg;
190 AscendC::MicroAPI::RegTensor<RangeType> baseIdxReg;202 AscendC::MicroAPI::RegTensor<RangeType> baseIdxReg;
191 AscendC::MicroAPI::RegTensor<RangeType> idxReg;203 AscendC::MicroAPI::RegTensor<RangeType> idxReg;
192- AscendC::MicroAPI::MaskReg idxMask =204+ AscendC::MicroAPI::MaskReg
193- AscendC::MicroAPI::CreateMask<RangeType, AscendC::MicroAPI::MaskPattern::ALL>();205+ idxMask = AscendC::MicroAPI::CreateMask<RangeType, AscendC::MicroAPI::MaskPattern::ALL>();
194 206 
195 AscendC::MicroAPI::Arange(baseIdxReg, 0);207 AscendC::MicroAPI::Arange(baseIdxReg, 0);
196- AscendC::MicroAPI::Muls(baseIdxReg, baseIdxReg, static_cast<RangeType>(this->inputRowElems_), idxMask);208+ AscendC::MicroAPI::Muls(baseIdxReg, baseIdxReg, ToRangeScalar(this->inputRowElems_), idxMask);
197 for (uint16_t axisBase = 0; axisBase < this->axisLen_;209 for (uint16_t axisBase = 0; axisBase < this->axisLen_;
198 axisBase = static_cast<uint16_t>(axisBase + vlSize)) {210 axisBase = static_cast<uint16_t>(axisBase + vlSize)) {
199 uint32_t curCount = this->axisLen_ - axisBase;211 uint32_t curCount = this->axisLen_ - axisBase;
@@ -201,16 +213,15 @@ protected:
201 curCount = vlSize;213 curCount = vlSize;
202 }214 }
203 AscendC::MicroAPI::MaskReg dataMask = AscendC::MicroAPI::UpdateMask<CastType>(curCount);215 AscendC::MicroAPI::MaskReg dataMask = AscendC::MicroAPI::UpdateMask<CastType>(curCount);
204- AscendC::MicroAPI::Adds(216+ AscendC::MicroAPI::Adds(idxReg, baseIdxReg, ToRangeScalar(axisBase * this->inputRowElems_), idxMask);
205- idxReg, baseIdxReg, static_cast<RangeType>(axisBase * this->inputRowElems_), idxMask);
206 for (uint16_t inner = 0; inner < curInnerChunk; ++inner) {217 for (uint16_t inner = 0; inner < curInnerChunk; ++inner) {
207 // Gather: read data from inputAddr using transpose indices in idxReg218 // Gather: read data from inputAddr using transpose indices in idxReg
208- AscendC::MicroAPI::DataCopyGather(219+ AscendC::MicroAPI::DataCopyGather(dataReg, inputAddr + inner,
209- dataReg, inputAddr + inner, (AscendC::MicroAPI::RegTensor<IdxType>&)idxReg, dataMask);220+ (AscendC::MicroAPI::RegTensor<IdxType>&)idxReg, dataMask);
210 if constexpr (sizeof(T) != 1) {221 if constexpr (sizeof(T) != 1) {
211 // Non-int8: write directly to UB output address222 // Non-int8: write directly to UB output address
212- AscendC::MicroAPI::DataCopy(223+ AscendC::MicroAPI::DataCopy(outputAddr + inner * outputValueAxisElems + axisBase, dataReg,
213- outputAddr + inner * outputValueAxisElems + axisBase, dataReg, dataMask);224+ dataMask);
214 } else {225 } else {
215 // int8: pack into b16 for compact UB write226 // int8: pack into b16 for compact UB write
216 __local_mem__ CastType* outputAddrB16 = reinterpret_cast<__local_mem__ CastType*>(227 __local_mem__ CastType* outputAddrB16 = reinterpret_cast<__local_mem__ CastType*>(