Star1727
1645
Code
Code
Issues482
Pull Requests544
Actions
Discussion
Wiki
Analysis
Star1727
1645
Merged
[9.2.0] refactor/perf: 回合算子迁移清理与 Sort 性能优化 #5809
黄晓彬Add on Sep 18
[9.2.0] refactor/perf: 回合算子迁移清理与 Sort 性能优化 #5809
Merged
黄晓彬Add on Sep 18
Merged fromConanHuang/ops-mathbackport/9.2.0-median-sort-20260918intocann/ops-math9.2.0
Total: 81 files changed+1864-9919
Mclassify_rule.yaml+0-5
@@ -633,9 +633,6 @@ sort@ops-math:
633 - ops/ops-math/math/sort/633 - ops/ops-math/math/sort/
634 - ops/ops-math/math/sort_with_index/634 - ops/ops-math/math/sort_with_index/
635 - ops/ops-math/math/top_k_v2/635 - ops/ops-math/math/top_k_v2/
636- - ops/ops-math/math/kth_value/
637- - ops/ops-math/math/median/
638- - ops/ops-math/math/nan_median/
639 opensource_style: null636 opensource_style: null
640 kernel_style: null637 kernel_style: null
641 unrelease:638 unrelease:
@@ -646,8 +643,6 @@ sort@ops-math:
646 - ops/ops-math/math/sort_with_index/tests/643 - ops/ops-math/math/sort_with_index/tests/
647 - ops/ops-math/math/top_k_v2/examples/644 - ops/ops-math/math/top_k_v2/examples/
648 - ops/ops-math/math/top_k_v2/tests/645 - ops/ops-math/math/top_k_v2/tests/
649- - ops/ops-math/math/median/tests/
650- - ops/ops-math/math/nan_median/tests/
651 non_delivery: null646 non_delivery: null
652 open_source: null647 open_source: null
653 648 
Dmath/kth_value/CMakeLists.txt+0-23
@@ -1,23 +0,0 @@
1-# ----------------------------------------------------------------------------
2-# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# CANN Open Software License Agreement Version 2.0 (the "License").
5-# Please refer to the License for details. Please see LICENSE for the full license text.
6-# This software is provided on an "AS IS" basis, without warranties.
7-# ----------------------------------------------------------------------------
8- 
9-# 设置算子定义时支持的芯片类型
10-set(SUPPORT_COMPUTE_UNIT "ascend950")
11-# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
12-set(SUPPORT_TILING_DIR "arch35")
13-add_all_modules_sources(OPTYPE kth_value ACLNNTYPE aclnn_exclude
14- DEPENDENCIES sort
15- COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT}
16- TILING_DIR ${SUPPORT_TILING_DIR}
17- DISABLE_IN_OPP TRUE)
18- 
19-# Keep graph_base loaded for the Host UT shared library under --as-needed linkers.
20-if(OP_HOST_UT AND TARGET ${OP_TILING_MODULE_NAME}_cases_obj)
21- target_link_libraries(${OP_TILING_MODULE_NAME}_cases_obj INTERFACE
22- "-Wl,--no-as-needed,$<TARGET_FILE:graph_base>,--as-needed")
23-endif()
Dmath/kth_value/op_api/kth_value.cpp+0-108
@@ -1,108 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#include "kth_value.h"
12- 
13-#include "op_api/aclnn_check.h"
14-#include "opdev/make_op_executor.h"
15-#include "opdev/op_dfx.h"
16-#include "opdev/op_executor.h"
17-#include "opdev/op_log.h"
18- 
19-#include <initializer_list>
20- 
21-using namespace op;
22- 
23-namespace l0op {
24-OP_TYPE_REGISTER(KthValue);
25- 
26-static const std::initializer_list<op::DataType> KTH_VALUE_DTYPE_SUPPORT_LIST = {
27- op::DataType::DT_FLOAT16, op::DataType::DT_FLOAT, op::DataType::DT_INT16, op::DataType::DT_INT8,
28- op::DataType::DT_UINT8, op::DataType::DT_INT32, op::DataType::DT_INT64, op::DataType::DT_BF16,
29- op::DataType::DT_UINT32, op::DataType::DT_UINT16, op::DataType::DT_UINT64};
30- 
31-static constexpr int64_t NON_LAST_SMALL_AXIS_MIN = 2;
32-static constexpr int64_t NON_LAST_SMALL_AXIS_MAX = 2048;
33- 
34-static bool IsLastAxisOrSupportedNonLastAxis(int64_t rank, int64_t normDim, int64_t axisLen)
35-{
36- if (normDim == rank - 1) {
37- return true;
38- }
39- return axisLen >= NON_LAST_SMALL_AXIS_MIN && axisLen <= NON_LAST_SMALL_AXIS_MAX;
40-}
41- 
42-static bool CheckParams(const aclTensor* self, int64_t k, int64_t dim)
43-{
44- OP_CHECK_NULL(self, return false);
45- OP_CHECK_DTYPE_NOT_SUPPORT(self, KTH_VALUE_DTYPE_SUPPORT_LIST, return false);
46- auto rank = static_cast<int64_t>(self->GetViewShape().GetDimNum());
47- if (rank <= 0) {
48- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "kth_value only supports tensor with rank > 0.");
49- return false;
50- }
51- if (dim < -rank || dim >= rank) {
52- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "dim should be in range [%ld, %ld].", -rank, rank - 1);
53- return false;
54- }
55- auto normDim = dim < 0 ? dim + rank : dim;
56- auto axisLen = static_cast<int64_t>(self->GetViewShape().GetDim(static_cast<size_t>(normDim)));
57- if (axisLen <= 0) {
58- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "kth_value axis length must be greater than 0.");
59- return false;
60- }
61- if (!IsLastAxisOrSupportedNonLastAxis(rank, normDim, axisLen)) {
62- OP_LOGE(
63- ACLNN_ERR_PARAM_INVALID,
64- "Dim must be the last axis or a supported non-last axis with axis size in [2, 2048]. Current dim is %ld.",
65- dim);
66- return false;
67- }
68- if (k < 1 || k > axisLen) {
69- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "k should be in range [1, %ld].", axisLen);
70- return false;
71- }
72- return true;
73-}
74- 
75-static op::Shape MakeKthShape(const op::Shape& shape, int64_t dim)
76-{
77- auto outShape = shape;
78- outShape.SetDim(static_cast<size_t>(dim), 1);
79- return outShape;
80-}
81- 
82-static std::tuple<aclTensor*, aclTensor*> KthValueAiCore(const aclTensor* self, int64_t k, int64_t dim,
83- aclOpExecutor* executor)
84-{
85- auto rank = static_cast<int64_t>(self->GetViewShape().GetDimNum());
86- auto normDim = dim < 0 ? dim + rank : dim;
87- auto outShape = MakeKthShape(self->GetViewShape(), normDim);
88- auto values = executor->AllocTensor(outShape, self->GetDataType(), op::Format::FORMAT_ND);
89- auto indices = executor->AllocTensor(outShape, op::DataType::DT_INT64, op::Format::FORMAT_ND);
90- OP_CHECK_NULL(values, return {});
91- OP_CHECK_NULL(indices, return {});
92- 
93- auto ret = ADD_TO_LAUNCHER_LIST_AICORE(KthValue, OP_INPUT(self), OP_OUTPUT(values, indices), OP_ATTR(k, dim));
94- OP_CHECK(ret == ACLNN_SUCCESS, OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "KthValue ADD_TO_LAUNCHER_LIST_AICORE failed."),
95- return {});
96- return std::tie(values, indices);
97-}
98- 
99-const std::tuple<aclTensor*, aclTensor*> KthValue(const aclTensor* self, int64_t k, int64_t dim,
100- aclOpExecutor* executor)
101-{
102- L0_DFX(KthValue, self, k, dim);
103- if (!CheckParams(self, k, dim)) {
104- return {};
105- }
106- return KthValueAiCore(self, k, dim, executor);
107-}
108-} // namespace l0op
Dmath/kth_value/op_api/kth_value.h+0-21
@@ -1,21 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_KTH_VALUE_OP_H_
12-#define PTA_NPU_OP_API_INC_LEVEL0_OP_KTH_VALUE_OP_H_
13- 
14-#include "opdev/op_executor.h"
15- 
16-namespace l0op {
17-const std::tuple<aclTensor*, aclTensor*> KthValue(
18- const aclTensor* self, int64_t k, int64_t dim, aclOpExecutor* executor);
19-}
20- 
21-#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_KTH_VALUE_OP_H_
Dmath/kth_value/op_host/arch35/kth_value_tiling_arch35.h+0-33
@@ -1,33 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file kth_value_tiling_arch35.h
13- * \brief kth_value tiling
14- */
15-#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_KTH_VALUE_TILING_ARCH35_H
16-#define AIR_CXX_RUNTIME_V2_OP_IMPL_KTH_VALUE_TILING_ARCH35_H
17- 
18-#include "register/op_def_registry.h"
19-#include "register/op_impl_registry.h"
20-#include "register/tilingdata_base.h"
21- 
22-namespace optiling {
23-struct KthValueCompileInfo {
24- int32_t coreNum;
25-};
26- 
27-ge::graphStatus Tiling4Median(gert::TilingContext* context);
28-ge::graphStatus TilingPrepare4Median(gert::TilingParseContext* context);
29-ge::graphStatus Tiling4NanMedian(gert::TilingContext* context);
30-ge::graphStatus TilingPrepare4NanMedian(gert::TilingParseContext* context);
31-} // namespace optiling
32- 
33-#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_KTH_VALUE_TILING_ARCH35_H
Dmath/kth_value/op_host/config/ascend950/kth_value_binary.json+0-544
@@ -1,544 +0,0 @@
1-{
2- "op_type": "KthValue",
3- "op_list": [
4- {
5- "bin_filename": "KthValue_x_float16_indices_int64",
6- "inputs": [
7- {
8- "name": "x",
9- "index": 0,
10- "dtype": "float16",
11- "format": "ND",
12- "paramType": "required",
13- "shape": [
14- -2
15- ]
16- }
17- ],
18- "outputs": [
19- {
20- "name": "y",
21- "index": 0,
22- "dtype": "float16",
23- "format": "ND",
24- "paramType": "required",
25- "shape": [
26- -2
27- ]
28- },
29- {
30- "name": "indices",
31- "index": 1,
32- "dtype": "int64",
33- "format": "ND",
34- "paramType": "required",
35- "shape": [
36- -2
37- ]
38- }
39- ],
40- "attrs": [
41- {
42- "name": "k",
43- "dtype": "int",
44- "value": null
45- },
46- {
47- "name": "dim",
48- "dtype": "int",
49- "value": null
50- }
51- ]
52- },
53- {
54- "bin_filename": "KthValue_x_float32_indices_int64",
55- "inputs": [
56- {
57- "name": "x",
58- "index": 0,
59- "dtype": "float32",
60- "format": "ND",
61- "paramType": "required",
62- "shape": [
63- -2
64- ]
65- }
66- ],
67- "outputs": [
68- {
69- "name": "y",
70- "index": 0,
71- "dtype": "float32",
72- "format": "ND",
73- "paramType": "required",
74- "shape": [
75- -2
76- ]
77- },
78- {
79- "name": "indices",
80- "index": 1,
81- "dtype": "int64",
82- "format": "ND",
83- "paramType": "required",
84- "shape": [
85- -2
86- ]
87- }
88- ],
89- "attrs": [
90- {
91- "name": "k",
92- "dtype": "int",
93- "value": null
94- },
95- {
96- "name": "dim",
97- "dtype": "int",
98- "value": null
99- }
100- ]
Dmath/kth_value/op_host/config/ascend950/kth_value_simplified_key.ini+0-3
@@ -1,3 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值。
2-[KthValue]
3-default=0
Dmath/kth_value/op_host/kth_value_def.cpp+0-61
@@ -1,61 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#include "register/op_def_registry.h"
12- 
13-namespace ops {
14-static const std::vector<ge::DataType> DataTypeXAndValues = {ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT16, ge::DT_INT8,
15- ge::DT_UINT8, ge::DT_INT32, ge::DT_INT64, ge::DT_BF16,
16- ge::DT_UINT32, ge::DT_UINT16, ge::DT_UINT64};
17- 
18-static const std::vector<ge::DataType> DataTypeIndices = {ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
19- ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
20- ge::DT_INT64, ge::DT_INT64, ge::DT_INT64};
21- 
22-static const std::vector<ge::Format> FormatNd = {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
23- ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
24- ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND};
25- 
26-class KthValue : public OpDef {
27-public:
28- explicit KthValue(const char* name) : OpDef(name)
29- {
30- this->Input("x")
31- .ParamType(REQUIRED)
32- .DataType(DataTypeXAndValues)
33- .Format(FormatNd)
34- .UnknownShapeFormat(FormatNd);
35- this->Output("y")
36- .ParamType(REQUIRED)
37- .DataType(DataTypeXAndValues)
38- .Format(FormatNd)
39- .UnknownShapeFormat(FormatNd);
40- this->Output("indices")
41- .ParamType(REQUIRED)
42- .DataType(DataTypeIndices)
43- .Format(FormatNd)
44- .UnknownShapeFormat(FormatNd);
45- this->Attr("k").AttrType(REQUIRED).Int();
46- this->Attr("dim").AttrType(OPTIONAL).Int(-1);
47- 
48- OpAICoreConfig aicoreConfig;
49- aicoreConfig.DynamicCompileStaticFlag(true)
50- .DynamicFormatFlag(false)
51- .DynamicRankSupportFlag(true)
52- .DynamicShapeSupportFlag(true)
53- .NeedCheckSupportFlag(false)
54- .ExtendCfgInfo("opFile.value", "kth_value_apt")
55- .ExtendCfgInfo("opInterface.value", "kth_value");
56- this->AICore().AddConfig("ascend950", aicoreConfig);
57- }
58-};
59- 
60-OP_ADD(KthValue);
61-} // namespace ops
Dmath/kth_value/op_host/kth_value_infershape.cpp+0-56
@@ -1,56 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#include <string>
12- 
13-#include "register/op_impl_registry.h"
14-#include "log/log.h"
15-#include "util/shape_util.h"
16- 
17-namespace ops {
18-static ge::graphStatus KthValueInferShapeFunc(gert::InferShapeContext* context)
19-{
20- const gert::Shape* xShape = context->GetInputShape(0);
21- OP_CHECK_NULL_WITH_CONTEXT(context, xShape);
22- auto* valuesShape = context->GetOutputShape(0);
23- auto* indicesShape = context->GetOutputShape(1);
24- OP_CHECK_NULL_WITH_CONTEXT(context, valuesShape);
25- OP_CHECK_NULL_WITH_CONTEXT(context, indicesShape);
26- 
27- if (Ops::Base::IsUnknownRank(*xShape)) {
28- Ops::Base::SetUnknownRank(*valuesShape);
29- Ops::Base::SetUnknownRank(*indicesShape);
30- return ge::GRAPH_SUCCESS;
31- }
32- 
33- const int64_t rank = static_cast<int64_t>(xShape->GetDimNum());
34- if (rank <= 0) {
35- return ge::GRAPH_FAILED;
36- }
37- auto attrs = context->GetAttrs();
38- OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
39- const int64_t* dimAttr = attrs->GetAttrPointer<int64_t>(1);
40- int64_t dim = (dimAttr == nullptr) ? -1 : *dimAttr;
41- int64_t normDim = dim < 0 ? dim + rank : dim;
42- if (normDim < 0 || normDim >= rank) {
43- std::string dimValue = std::to_string(dim);
44- std::string dimRange = "[" + std::to_string(-rank) + ", " + std::to_string(rank - 1) + "]";
45- OP_LOGE_WITH_INVALID_ATTR(context->GetNodeName(), "dim", dimValue.c_str(), dimRange.c_str());
46- return ge::GRAPH_FAILED;
47- }
48- *valuesShape = *xShape;
49- *indicesShape = *xShape;
50- valuesShape->SetDim(normDim, 1);
51- indicesShape->SetDim(normDim, 1);
52- return ge::GRAPH_SUCCESS;
53-}
54- 
55-IMPL_OP_INFERSHAPE(KthValue).InferShape(KthValueInferShapeFunc);
56-} // namespace ops
Dmath/kth_value/op_kernel/arch35/kth_value_axis_one_copy.h+0-127
@@ -1,127 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#ifndef KTH_VALUE_AXIS_ONE_COPY_H
12-#define KTH_VALUE_AXIS_ONE_COPY_H
13- 
14-#include "kernel_operator.h"
15-#include "kth_value_tiling_data.h"
16- 
17-namespace KthValue {
18-using namespace AscendC;
19- 
20-template <typename T>
21-class KthValueAxisOneCopy {
22-public:
23- __aicore__ inline KthValueAxisOneCopy() {}
24- __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,
25- const KthValueTilingData *tilingData, TPipe *pipe);
26- __aicore__ inline void Process();
27- 
28-private:
29- __aicore__ inline void CopyValue(int64_t offset, uint32_t validElems);
30- __aicore__ inline void CopyIndex(int64_t offset, uint32_t validElems);
31- 
32- static constexpr uint32_t BUFFER_NUM = 2;
33- 
34- TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> valueQueue_;
35- TQue<QuePosition::VECOUT, BUFFER_NUM> idxQueue_;
36- 
37- GlobalTensor<T> inputXGm_;
38- GlobalTensor<T> outValueGm_;
39- GlobalTensor<int64_t> outIdxGm_;
40- 
41- uint32_t blockIdx_ = 0;
42- uint32_t blockDim_ = 0;
43- uint32_t copyElemsPerLoop_ = 0;
44- uint32_t loopTimes_ = 0;
45- int64_t totalElems_ = 0;
46-};
47- 
48-template <typename T>
49-__aicore__ inline void KthValueAxisOneCopy<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,
50- const KthValueTilingData *tilingData, TPipe *pipe)
51-{
52- (void)workspace;
53- if (tilingData == nullptr || pipe == nullptr) {
54- return;
55- }
56- 
57- blockIdx_ = GetBlockIdx();
58- blockDim_ = GetBlockNum();
59- copyElemsPerLoop_ = tilingData->keyParams0;
60- loopTimes_ = tilingData->keyParams1;
61- totalElems_ = tilingData->unsortedDimNum * tilingData->lastAxisNum;
62- 
63- inputXGm_.SetGlobalBuffer((__gm__ T *)x);
64- outValueGm_.SetGlobalBuffer((__gm__ T *)y);
65- outIdxGm_.SetGlobalBuffer((__gm__ int64_t *)idx);
66- 
67- if (copyElemsPerLoop_ == 0) {
68- return;
69- }
70- 
71- pipe->InitBuffer(valueQueue_, BUFFER_NUM, copyElemsPerLoop_ * sizeof(T));
72- pipe->InitBuffer(idxQueue_, BUFFER_NUM, copyElemsPerLoop_ * sizeof(int64_t));
73-}
74- 
75-template <typename T>
76-__aicore__ inline void KthValueAxisOneCopy<T>::CopyValue(int64_t offset, uint32_t validElems)
77-{
78- LocalTensor<T> valueLocal = valueQueue_.AllocTensor<T>();
79- DataCopyExtParams copyParam{ 1, static_cast<uint32_t>(validElems * sizeof(T)), 0, 0, 0 };
80- DataCopyPadExtParams<T> padParams{ false, 0, 0, 0 };
81- DataCopyPad(valueLocal, inputXGm_[offset], copyParam, padParams);
82- valueQueue_.EnQue(valueLocal);
83- 
84- valueLocal = valueQueue_.DeQue<T>();
85- DataCopyPad(outValueGm_[offset], valueLocal, copyParam);
86- valueQueue_.FreeTensor(valueLocal);
87-}
88- 
89-template <typename T>
90-__aicore__ inline void KthValueAxisOneCopy<T>::CopyIndex(int64_t offset, uint32_t validElems)
91-{
92- LocalTensor<int64_t> idxLocal = idxQueue_.AllocTensor<int64_t>();
93- Duplicate(idxLocal, static_cast<int64_t>(0), validElems);
94- idxQueue_.EnQue(idxLocal);
95- 
96- idxLocal = idxQueue_.DeQue<int64_t>();
97- DataCopyExtParams copyParam{ 1, static_cast<uint32_t>(validElems * sizeof(int64_t)), 0, 0, 0 };
98- DataCopyPad(outIdxGm_[offset], idxLocal, copyParam);
99- idxQueue_.FreeTensor(idxLocal);
100-}
101- 
102-template <typename T>
103-__aicore__ inline void KthValueAxisOneCopy<T>::Process()
104-{
105- if (blockIdx_ >= blockDim_ || blockDim_ == 0 || copyElemsPerLoop_ == 0 || totalElems_ <= 0) {
106- return;
107- }
108- 
109- uint32_t loopsPerCore = (loopTimes_ + blockDim_ - 1U) / blockDim_;
110- uint32_t startLoop = blockIdx_ * loopsPerCore;
111- uint32_t endLoop = (startLoop + loopsPerCore < loopTimes_) ? (startLoop + loopsPerCore) : loopTimes_;
112- for (uint32_t loopId = startLoop; loopId < endLoop; ++loopId) {
113- int64_t offset = static_cast<int64_t>(loopId) * static_cast<int64_t>(copyElemsPerLoop_);
114- int64_t remain = totalElems_ - offset;
115- if (remain <= 0) {
116- continue;
117- }
118- uint32_t validElems = remain >= static_cast<int64_t>(copyElemsPerLoop_) ?
119- copyElemsPerLoop_ : static_cast<uint32_t>(remain);
120- CopyValue(offset, validElems);
121- CopyIndex(offset, validElems);
122- }
123-}
124- 
125-} // namespace KthValue
126- 
127-#endif
Dmath/kth_value/op_kernel/arch35/kth_value_kernel_dispatch.h+0-267
@@ -1,267 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#ifndef KTH_VALUE_KERNEL_DISPATCH_H
12-#define KTH_VALUE_KERNEL_DISPATCH_H
13- 
14-#include "kth_value_axis_one_copy.h"
15-#include "kth_value_median_utils.h"
16-#include "kth_value_merge_intra_core.h"
17-#include "kth_value_merge_sort_one_core.h"
18-#include "kth_value_merge_sort_more_core.h"
19-#include "kth_value_non_last_small_axis.h"
20-#include "kth_value_radix_more_core.h"
21-#include "kth_value_radix_one_core.h"
22-#include "kth_value_radix_select.h"
23-#include "kth_value_small_axis_insertion.h"
24-#include "kth_value_small_axis_short_rank_select.h"
25-#include "kth_value_small_axis_two_stage.h"
26-#include "kth_value_tiling_data.h"
27-#include "kth_value_tiling_key.h"
28- 
29-namespace KthValue {
30-using namespace AscendC;
31- 
32-template <bool EnableMedian, uint64_t schId>
33-__aicore__ inline void RunMergeSortRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData, TPipe* pipe)
34-{
35- constexpr uint64_t isSort32SmallAxis = (schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS);
36- if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) {
37- KthValueMergeSortOneCore<DTYPE_X, float, isSort32SmallAxis, EnableMedian> op;
38- op.Init(x, y1, y2, tilingData, pipe);
39- op.Process();
40- } else if constexpr (IsSameType<float, DTYPE_X>::value || IsSameType<half, DTYPE_X>::value) {
41- KthValueMergeSortOneCore<DTYPE_X, DTYPE_X, isSort32SmallAxis, EnableMedian> op;
42- op.Init(x, y1, y2, tilingData, pipe);
43- op.Process();
44- }
45-}
46- 
47-template <bool EnableMedian, uint64_t isInt32>
48-__aicore__ inline void RunRadixMoreCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
49- KthValueTilingData* tilingData, TPipe* pipe)
50-{
51- constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>;
52- if constexpr (sizeof(DTYPE_X) == sizeof(uint8_t)) {
53- if constexpr (isInt32 == 1) {
54- KthValueRadixMoreCore<DTYPE_X, uint32_t, uint8_t, enableNanMode> op;
55- op.Init(x, y1, y2, workspace, tilingData, pipe);
56- op.Process();
57- } else {
58- KthValueRadixMoreCore<DTYPE_X, int64_t, uint8_t, enableNanMode> op;
59- op.Init(x, y1, y2, workspace, tilingData, pipe);
60- op.Process();
61- }
62- } else if constexpr (sizeof(DTYPE_X) == sizeof(uint16_t)) {
63- if constexpr (isInt32 == 1) {
64- KthValueRadixMoreCore<DTYPE_X, uint32_t, uint16_t, enableNanMode> op;
65- op.Init(x, y1, y2, workspace, tilingData, pipe);
66- op.Process();
67- } else {
68- KthValueRadixMoreCore<DTYPE_X, int64_t, uint16_t, enableNanMode> op;
69- op.Init(x, y1, y2, workspace, tilingData, pipe);
70- op.Process();
71- }
72- } else if constexpr (sizeof(DTYPE_X) == sizeof(uint32_t)) {
73- if constexpr (isInt32 == 1) {
74- KthValueRadixMoreCore<DTYPE_X, uint32_t, uint32_t, enableNanMode> op;
75- op.Init(x, y1, y2, workspace, tilingData, pipe);
76- op.Process();
77- } else {
78- KthValueRadixMoreCore<DTYPE_X, int64_t, uint32_t, enableNanMode> op;
79- op.Init(x, y1, y2, workspace, tilingData, pipe);
80- op.Process();
81- }
82- } else if constexpr (sizeof(DTYPE_X) == sizeof(uint64_t)) {
83- if constexpr (isInt32 == 1) {
84- KthValueRadixMoreCore<DTYPE_X, uint32_t, uint64_t, enableNanMode> op;
85- op.Init(x, y1, y2, workspace, tilingData, pipe);
86- op.Process();
87- } else {
88- KthValueRadixMoreCore<DTYPE_X, int64_t, uint64_t, enableNanMode> op;
89- op.Init(x, y1, y2, workspace, tilingData, pipe);
90- op.Process();
91- }
92- }
93-}
94- 
95-template <bool EnableMedian>
96-__aicore__ inline void RunSmallAxisInsertionRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
97- TPipe* pipe)
98-{
99- constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>;
100- if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) {
101- KthValueSmallAxisInsertion<DTYPE_X, float, enableNanMode> op;
102- op.Init(x, y1, y2, tilingData, pipe);
103- op.Process();
104- } else {
105- KthValueSmallAxisInsertion<DTYPE_X, DTYPE_X, enableNanMode> op;
106- op.Init(x, y1, y2, tilingData, pipe);
107- op.Process();
108- }
109-}
110- 
111-template <bool EnableMedian>
112-__aicore__ inline void RunMergeMoreCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
113- KthValueTilingData* tilingData, TPipe* pipe)
114-{
115- if constexpr (IsSameType<float, DTYPE_X>::value) {
116- KthValueMergeSortMoreCore<DTYPE_X, DTYPE_X, false, int64_t, EnableMedian> op;
117- op.Init(x, y1, y2, workspace, tilingData, pipe);
118- op.Process();
119- }
120-}
121- 
122-template <bool EnableMedian>
123-__aicore__ inline void RunMergeIntraCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
124- KthValueTilingData* tilingData, TPipe* pipe)
125-{
126- if constexpr (IsSameType<float, DTYPE_X>::value) {
127- KthValueMergeIntraCore<DTYPE_X, int64_t, false, EnableMedian> op;
128- op.Init(x, y1, y2, workspace, tilingData, pipe);
129- op.Process();
130- }
131-}
132- 
133-template <bool EnableMedian, bool useMergeSort>
134-__aicore__ inline void RunNonLastSmallAxisRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
135- KthValueTilingData* tilingData, TPipe* pipe)
136-{
137- constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>;
138- if constexpr (useMergeSort) {
139- if constexpr (IsSameType<DTYPE_X, float>::value || IsSameType<DTYPE_X, half>::value ||
140- IsSameType<DTYPE_X, bfloat16_t>::value) {
141- KthValueNonLastSmallAxis<DTYPE_X, false, true, enableNanMode> op;
142- op.Init(x, y1, y2, workspace, tilingData, pipe);
143- op.Process();
144- }
145- } else {
146- KthValueNonLastSmallAxis<DTYPE_X, false, false, enableNanMode> op;
147- op.Init(x, y1, y2, workspace, tilingData, pipe);
148- op.Process();
149- }
150-}
151- 
152-__aicore__ inline void RunAxisOneCopyRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
153- KthValueTilingData* tilingData, TPipe* pipe)
154-{
155- KthValueAxisOneCopy<DTYPE_X> op;
156- op.Init(x, y1, y2, workspace, tilingData, pipe);
157- op.Process();
158-}
159- 
160-template <bool EnableMedian>
161-__aicore__ inline void RunRadixOneCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
162- TPipe* pipe)
163-{
164- constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>;
165- KthValueRadixOneCore<DTYPE_X, enableNanMode> op;
166- op.Init(x, y1, y2, tilingData, pipe);
167- op.Process();
168-}
169- 
170-template <bool EnableMedian>
171-__aicore__ inline void RunRadixSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace,
172- KthValueTilingData* tilingData, TPipe* pipe)
173-{
174- constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>;
175- if constexpr (sizeof(DTYPE_X) == 1) {
176- KthValueRadixSelect<DTYPE_X, uint8_t, enableNanMode> op;
177- op.Init(x, y1, y2, workspace, tilingData, pipe);
178- op.Process();
179- } else if constexpr (sizeof(DTYPE_X) == 2) {
180- KthValueRadixSelect<DTYPE_X, uint16_t, enableNanMode> op;
181- op.Init(x, y1, y2, workspace, tilingData, pipe);
182- op.Process();
183- } else if constexpr (sizeof(DTYPE_X) == 4) {
184- KthValueRadixSelect<DTYPE_X, uint32_t, enableNanMode> op;
185- op.Init(x, y1, y2, workspace, tilingData, pipe);
186- op.Process();
187- } else if constexpr (sizeof(DTYPE_X) == 8) {
188- KthValueRadixSelect<DTYPE_X, uint64_t, enableNanMode> op;
189- op.Init(x, y1, y2, workspace, tilingData, pipe);
190- op.Process();
191- }
192-}
193- 
194-template <bool EnableMedian>
195-__aicore__ inline void RunSmallAxisTwoStageRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
196- TPipe* pipe)
197-{
198- constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>;
199- KthValueSmallAxisTwoStage<DTYPE_X, enableNanMode> op;
200- op.Init(x, y1, y2, tilingData, pipe);
201- op.Process();
202-}
203- 
204-__aicore__ inline void RunSmallAxisShortRankSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2,
205- KthValueTilingData* tilingData, TPipe* pipe)
206-{
207- if constexpr (sizeof(DTYPE_X) == 8) {
208- KthValueSmallAxisShortRankSelect<DTYPE_X> op;
209- op.Init(x, y1, y2, tilingData, pipe);
210- op.Process();
211- }
212-}
213- 
214-template <bool EnableMedian, uint64_t schId>
215-__aicore__ inline bool TryRunSmallAxisRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData,
216- TPipe* pipe)
217-{
218- if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION) {
219- RunSmallAxisInsertionRoute<EnableMedian>(x, y1, y2, tilingData, pipe);
220- return true;
221- }
222- if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE) {
223- RunSmallAxisTwoStageRoute<EnableMedian>(x, y1, y2, tilingData, pipe);
224- return true;
225- }
226- if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT) {
227- RunSmallAxisShortRankSelectRoute(x, y1, y2, tilingData, pipe);
228- return true;
229- }
230- return false;
231-}
232- 
233-// EnableMedian selects the shared Median/NanMedian behavior. tilingData->medianMode
234-// distinguishes propagate-NaN Median from ignore-NaN NanMedian at runtime.
235-template <bool EnableMedian, uint64_t schId, uint64_t isInt32>
236-__aicore__ inline void Dispatch(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData,
237- TPipe* pipe)
238-{
239- GM_ADDR userWorkspace = AscendC::GetUserWorkspace(workspace);
240- if constexpr (schId == KTH_VALUE_SCHID_AXIS_ONE_COPY) {
241- RunAxisOneCopyRoute(x, y1, y2, userWorkspace, tilingData, pipe);
242- } else if constexpr (schId == KTH_VALUE_SCHID_MERGE_SORT || schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS) {
243- RunMergeSortRoute<EnableMedian, schId>(x, y1, y2, tilingData, pipe);
244- } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) {
245- RunRadixOneCoreRoute<EnableMedian>(x, y1, y2, tilingData, pipe);
246- } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_MORE_CORE) {
247- RunRadixMoreCoreRoute<EnableMedian, isInt32>(x, y1, y2, userWorkspace, tilingData, pipe);
248- } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_SELECT) {
249- RunRadixSelectRoute<EnableMedian>(x, y1, y2, userWorkspace, tilingData, pipe);
250- } else if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION ||
251- schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE ||
252- schId == KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT) {
253- TryRunSmallAxisRoute<EnableMedian, schId>(x, y1, y2, tilingData, pipe);
254- } else if constexpr (schId == KTH_VALUE_SCHID_MERGE_MORE_CORE) {
255- RunMergeMoreCoreRoute<EnableMedian>(x, y1, y2, userWorkspace, tilingData, pipe);
256- } else if constexpr (schId == KTH_VALUE_SCHID_MERGE_INTRA_CORE) {
257- RunMergeIntraCoreRoute<EnableMedian>(x, y1, y2, userWorkspace, tilingData, pipe);
258- } else if constexpr (schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS) {
259- RunNonLastSmallAxisRoute<EnableMedian, true>(x, y1, y2, userWorkspace, tilingData, pipe);
260- } else if constexpr (schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX) {
261- RunNonLastSmallAxisRoute<EnableMedian, false>(x, y1, y2, userWorkspace, tilingData, pipe);
262- }
263-}
264- 
265-} // namespace KthValue
266- 
267-#endif
Dmath/kth_value/op_kernel/arch35/kth_value_median_utils.h+0-232
@@ -1,232 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/**
12- * @file kth_value_median_utils.h
13- * @brief Median-mode helper primitives shared by kth_value schedule kernels
14- *
15- * @details This file provides the common primitives required by median semantics
16- * (NaN-aware median), reused by all schedule kernels:
17- * 1. Sort-value canonicalization: replaces every NaN with the canonical quiet NaN
18- * bit pattern so that NaNs sort identically into a deterministic suffix. The
19- * scalar CanonicalizeMedianSortValue additionally normalizes both signed zeros
20- * to +0 so stable sorting preserves their original order.
21- * 2. Non-NaN counting: vector reduction of the non-NaN element count of a UB
22- * segment into a scalar slot, with cross-chunk accumulation support
23- * (AccumulateNonNanCount / CountNonNan)
24- * 3. Median K resolution: maps the median to a sorted-array index according to
25- * medianMode (PROPAGATE_NAN/IGNORE_NAN/STATIC); the non-NaN count is either
26- * passed in by the caller or derived by binary search over sorted rows
27- * (ResolveMedianK / ResolveMedianKFromSorted / IsNanValue)
28- *
29- * All primitives operate on contiguous UB segments only and are axis/row agnostic;
30- * row iteration, multi-core splitting and tiling orchestration are owned by the
31- * callers (radix/merge/small_axis/non_last schedule files).
32- */
33- 
34-#ifndef KTH_VALUE_MEDIAN_UTILS_H
35-#define KTH_VALUE_MEDIAN_UTILS_H
36- 
37-#include <cmath>
38- 
39-#include "kernel_operator.h"
40-#include "op_kernel/platform_util.h"
41-#include "../../sort/arch35/common/util_type_simd.h"
42- 
43-namespace KthValue {
44-using namespace AscendC;
45- 
46-constexpr Reg::CastTrait MEDIAN_CAST_TO_FLOAT = {Reg::RegLayout::ZERO, Reg::SatMode::UNKNOWN,
47- Reg::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
48-constexpr Reg::CastTrait MEDIAN_CAST_FLOAT_TO_INT32 = {Reg::RegLayout::UNKNOWN, Reg::SatMode::NO_SAT,
49- Reg::MaskMergeMode::ZEROING, RoundMode::CAST_TRUNC};
50- 
51-template <typename T>
52-constexpr bool IS_MEDIAN_FLOAT_TYPE = IsSameType<T, float>::value || IsSameType<T, half>::value ||
53- IsSameType<T, bfloat16_t>::value;
54- 
55-// Replaces every NaN in-place with the canonical quiet NaN bit pattern so that all
56-// NaNs compare/sort identically and form a deterministic suffix after sorting.
57-template <typename T>
58-__aicore__ inline void CanonicalizeNanValues(LocalTensor<T> input, uint32_t count)
59-{
60- static_assert(IS_MEDIAN_FLOAT_TYPE<T>, "CanonicalizeNanValues only supports floating-point types");
61- if constexpr (IsSameType<float, T>::value) {
62- constexpr uint32_t VECTOR_LEN = Ops::Base::GetVRegSize() / sizeof(uint32_t);
63- constexpr uint32_t ABS_MASK = 0x7FFFFFFFU;
64- constexpr uint32_t INF_BITS = 0x7F800000U;
65- constexpr uint32_t CANONICAL_NAN_BITS = 0x7FC00000U;
66- uint16_t repeatTime = static_cast<uint16_t>(CeilDivision(count, VECTOR_LEN));
67- uint32_t remaining = count;
68- __ubuf__ uint32_t* inputPtr = reinterpret_cast<__ubuf__ uint32_t*>(input.GetPhyAddr());
69- __VEC_SCOPE__
70- {
71- Reg::RegTensor<uint32_t> rawBits, absBits, absMask, canonicalNan, result;
72- Reg::MaskReg validMask, nanMask;
73- Reg::Duplicate(absMask, ABS_MASK);
74- Reg::Duplicate(canonicalNan, CANONICAL_NAN_BITS);
75- for (uint16_t i = 0; i < repeatTime; ++i) {
76- validMask = Reg::UpdateMask<uint32_t>(remaining);
77- Reg::LoadAlign<uint32_t, Reg::LoadDist::DIST_NORM>(rawBits, inputPtr + i * VECTOR_LEN);
78- Reg::And(absBits, rawBits, absMask, validMask);
79- Reg::Compares<uint32_t, CMPMODE::GT>(nanMask, absBits, INF_BITS, validMask);
80- Reg::Select(result, canonicalNan, rawBits, nanMask);
81- Reg::StoreAlign<uint32_t, Reg::StoreDist::DIST_NORM_B32>(inputPtr + i * VECTOR_LEN, result, validMask);
82- }
83- }
84- } else {
85- constexpr uint32_t VECTOR_LEN = Ops::Base::GetVRegSize() / sizeof(uint16_t);
86- constexpr uint16_t ABS_MASK = 0x7FFFU;
87- constexpr uint16_t INF_BITS = IsSameType<half, T>::value ? 0x7C00U : 0x7F80U;
88- constexpr uint16_t CANONICAL_NAN_BITS = IsSameType<half, T>::value ? 0x7E00U : 0x7FC0U;
89- uint16_t repeatTime = static_cast<uint16_t>(CeilDivision(count, VECTOR_LEN));
90- uint32_t remaining = count;
91- __ubuf__ uint16_t* inputPtr = reinterpret_cast<__ubuf__ uint16_t*>(input.GetPhyAddr());
92- __VEC_SCOPE__
93- {
94- Reg::RegTensor<uint16_t> rawBits, absBits, absMask, canonicalNan, result;
95- Reg::MaskReg validMask, nanMask;
96- Reg::Duplicate(absMask, ABS_MASK);
97- Reg::Duplicate(canonicalNan, CANONICAL_NAN_BITS);
98- for (uint16_t i = 0; i < repeatTime; ++i) {
99- validMask = Reg::UpdateMask<uint16_t>(remaining);
100- Reg::LoadAlign<uint16_t, Reg::LoadDist::DIST_NORM>(rawBits, inputPtr + i * VECTOR_LEN);
101- Reg::And(absBits, rawBits, absMask, validMask);
102- Reg::Compares<uint16_t, CMPMODE::GT>(nanMask, absBits, INF_BITS, validMask);
103- Reg::Select(result, canonicalNan, rawBits, nanMask);
104- Reg::StoreAlign<uint16_t, Reg::StoreDist::DIST_NORM_B16>(inputPtr + i * VECTOR_LEN, result, validMask);
105- }
106- }
107- }
108-}
109- 
110-__aicore__ inline uint32_t ResolveMedianK(uint32_t staticK, uint32_t validCount, uint32_t nonNanCount,
111- uint32_t medianMode)
112-{
113- if (medianMode == MEDIAN_MODE_PROPAGATE_NAN && nonNanCount < validCount) {
114- return nonNanCount;
115- }
116- if (medianMode == MEDIAN_MODE_IGNORE_NAN) {
117- return nonNanCount == 0U ? 0U : (nonNanCount - 1U) / 2U;
118- }
119- return staticK;
120-}
121- 
122-template <typename T>
123-__simt_callee__ __aicore__ inline bool IsNanValue(T value)
124-{
125- float valueFp32 = static_cast<float>(value);
126- return valueFp32 != valueFp32;
127-}
128- 
129-template <typename T>
130-__simt_callee__ __aicore__ inline T CanonicalizeMedianSortValue(T value)
131-{
132- float valueFp32 = static_cast<float>(value);
133- if (valueFp32 != valueFp32) {
134- return static_cast<T>(NAN);
135- }
136- if (valueFp32 == 0.0f) {
137- return static_cast<T>(0.0f);
138- }
139- return value;
140-}
141- 
142-// Resolves the sorted-array index of the median for the given mode: derives the non-NaN
143-// count by lower-bound binary search (NaNs form a suffix) and maps it per medianMode.
144-template <typename T>
145-__simt_callee__ __aicore__ inline uint32_t ResolveMedianKFromSorted(const __ubuf__ T* sortedValues, uint32_t rowOffset,
146- uint32_t validCount, uint32_t staticK,
147- uint32_t medianMode)
148-{
149- if (medianMode == MEDIAN_MODE_STATIC) {
150- return staticK;
151- }
152- 
153- // NaNs are canonicalized before sorting and form a suffix. Find the first
154- // NaN (which is also the number of non-NaN values) with a lower-bound search.
155- uint32_t left = 0U;
156- uint32_t right = validCount;
157- while (left < right) {
158- uint32_t mid = left + (right - left) / 2U;
159- if (IsNanValue(sortedValues[rowOffset + mid])) {
160- right = mid;
161- } else {
162- left = mid + 1U;
163- }
164- }
165- uint32_t nonNanCount = left;
166- if (medianMode == MEDIAN_MODE_PROPAGATE_NAN && nonNanCount < validCount) {
167- return nonNanCount;
168- }
169- if (medianMode == MEDIAN_MODE_IGNORE_NAN) {
170- return nonNanCount == 0U ? 0U : (nonNanCount - 1U) / 2U;
171- }
172- return staticK;
173-}
174- 
175-// Vector-reduces the number of non-NaN elements in input[count] into a single UB scalar slot;
176-// with Reset=false the result is accumulated on top of the slot's existing value.
177-template <bool Reset, typename T>
178-__aicore__ inline void AccumulateNonNanCount(LocalTensor<T> input, uint32_t count, LocalTensor<uint32_t> scalar)
179-{
180- constexpr uint32_t FLOAT_VL = Ops::Base::GetVRegSize() / sizeof(float);
181- uint16_t repeatTime = static_cast<uint16_t>((count + FLOAT_VL - 1U) / FLOAT_VL);
182- uint32_t remaining = count;
183- __ubuf__ T* inputPtr = reinterpret_cast<__ubuf__ T*>(input.GetPhyAddr());
184- __ubuf__ int32_t* scalarPtr = reinterpret_cast<__ubuf__ int32_t*>(scalar.GetPhyAddr());
185- __VEC_SCOPE__
186- {
187- Reg::RegTensor<float> value, one, zero, nonNanFlag, reduced, total;
188- Reg::RegTensor<int32_t> totalInt;
189- Reg::MaskReg validMask, nanMask;
190- Reg::MaskReg oneMask = Reg::CreateMask<float, Reg::MaskPattern::VL1>();
191- Reg::Duplicate(one, 1.0f);
192- Reg::Duplicate(zero, 0.0f);
193- Reg::Duplicate(total, 0.0f, oneMask);
194- for (uint16_t i = 0; i < repeatTime; ++i) {
195- validMask = Reg::UpdateMask<float>(remaining);
196- if constexpr (IsSameType<float, T>::value) {
197- Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(
198- value, reinterpret_cast<__ubuf__ float*>(inputPtr) + i * FLOAT_VL);
199- } else {
200- Reg::RegTensor<T> valueB16;
201- Reg::LoadAlign<T, Reg::LoadDist::DIST_UNPACK_B16>(valueB16, inputPtr + i * FLOAT_VL);
202- Reg::Cast<float, T, MEDIAN_CAST_TO_FLOAT>(value, valueB16, validMask);
203- }
204- Reg::Compare<float, CMPMODE::NE>(nanMask, value, value, validMask);
205- Reg::Select(nonNanFlag, zero, one, nanMask);
206- Reg::Reduce<Reg::ReduceType::SUM>(reduced, nonNanFlag, validMask);
207- Reg::Add(total, total, reduced, oneMask);
208- }
209- Reg::Cast<int32_t, float, MEDIAN_CAST_FLOAT_TO_INT32>(totalInt, total, oneMask);
210- if constexpr (!Reset) {
211- Reg::RegTensor<int32_t> accumulated;
212- Reg::LoadAlign<int32_t, Reg::LoadDist::DIST_NORM>(accumulated, scalarPtr);
213- Reg::Add(totalInt, totalInt, accumulated, oneMask);
214- }
215- Reg::StoreAlign<int32_t, Reg::StoreDist::DIST_NORM_B32>(scalarPtr, totalInt, oneMask);
216- }
217-}
218- 
219-template <typename T>
220-__aicore__ inline uint32_t CountNonNan(LocalTensor<T> input, uint32_t count, LocalTensor<float> scalar, TPipe* pipe)
221-{
222- LocalTensor<uint32_t> scalarU32 = scalar.template ReinterpretCast<uint32_t>();
223- AccumulateNonNanCount<true>(input, count, scalarU32);
224- event_t countEvent = static_cast<event_t>(pipe->FetchEventID(HardEvent::V_S));
225- SetFlag<HardEvent::V_S>(countEvent);
226- WaitFlag<HardEvent::V_S>(countEvent);
227- return scalarU32.GetValue(0);
228-}
229- 
230-} // namespace KthValue
231- 
232-#endif
Dmath/kth_value/op_kernel/arch35/kth_value_merge_intra_core.h+0-301
@@ -1,301 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file kth_value_merge_intra_core.h
13- * \brief Intra-core block merge sort for kth_value (fp32)
14- * \details Each core independently sorts blocks and merges them via 4-way MrgSort.
15- * No cross-core synchronization needed in merge phase.
16- * Phase 3 extracts only the kth element instead of the full sorted output.
17- */
18- 
19-#ifndef KTH_VALUE_MERGE_INTRA_CORE_H
20-#define KTH_VALUE_MERGE_INTRA_CORE_H
21- 
22-#include <cmath>
23-#include "kernel_operator.h"
24-#include "op_kernel/math_util.h"
25-#include "op_kernel/platform_util.h"
26-#include "kth_value_median_utils.h"
27-#include "kth_value_tiling_data.h"
28-#include "../../sort/arch35/common/merge_sort_constants.h"
29-#include "../../sort/arch35/common/merge_intra_core_base.h"
30- 
31-namespace KthValue {
32-using namespace AscendC;
33- 
34-// Import shared constants from MergeSortConstants namespace
35-using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;
36-using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE;
37-using MergeSortConstants::DEALING_SORT_NUM_ONCE;
38-using MergeSortConstants::MERGE_INTRA_BUFFER_NUM;
39-using MergeSortConstants::MERGE_LIST_MAX_NUM;
40-using MergeSortConstants::UB_BLOCK_BYTES;
41- 
42-/**
43- * @brief Intra-core block merge sort for kth_value: independent per-core sort+merge,
44- * then extracts only the kth element from the sorted result.
45- * @tparam ValueType Input data type (float)
46- * @tparam IndexType Index data type (int32_t or int64_t)
47- * @tparam IsDescend Sort order: true for descending, false for ascending
48- * @tparam EnableMedian Enables NaN canonicalization and runtime median-k resolution for Median/NanMedian
49- */
50-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian = false>
51-class KthValueMergeIntraCore
52- : public MergeIntraCoreCommon::MergeIntraCoreBase<
53- KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>, ValueType, IndexType, IsDescend> {
54- using Base = MergeIntraCoreCommon::MergeIntraCoreBase<
55- KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>, ValueType, IndexType, IsDescend>;
56- friend Base;
57- 
58-public:
59- __aicore__ inline KthValueMergeIntraCore() {}
60- __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace,
61- const KthValueTilingData* tilingData, TPipe* pipe);
62- 
63-private:
64- // KthValue-specific members
65- uint32_t kthIndex_ = 0;
66- uint32_t staticKthIndex_ = 0;
67- uint32_t medianMode_ = 0;
68- TQue<QuePosition::VECOUT, 1> kthValueQueue_, kthIdxQueue_;
69- 
70- __aicore__ inline void InitPhase1Buffers();
71- __aicore__ inline void InitPhase2Buffers();
72- __aicore__ inline void InitPhase3Buffers();
73- __aicore__ inline void PrepareInputForSort(LocalTensor<ValueType> inputLocal, uint32_t alignedCount);
74- __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion);
75- __aicore__ inline void ExtractAndCopyChunk(int64_t cacheBatchOffset, uint32_t cacheOffset, int64_t outputOffset,
76- uint32_t elemProcessed, uint32_t elemCount);
77- __aicore__ inline uint32_t ResolveBatchK(int64_t batchIdx);
78-};
79- 
80-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
81-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::Init(
82- GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData, TPipe* pipe)
83-{
84- if (tilingData == nullptr || pipe == nullptr) {
85- return;
86- }
87- this->blockIdx_ = GetBlockIdx();
88- this->pipe_ = pipe;
89- 
90- // Parse tiling data
91- this->batchNum_ = tilingData->unsortedDimNum;
92- this->sortAxisNum_ = tilingData->lastAxisNum;
93- this->batchPerCore_ = tilingData->keyParams0;
94- this->blockSortSize_ = tilingData->numTileDataSize;
95- kthIndex_ = tilingData->kthIndex;
96- staticKthIndex_ = kthIndex_;
97- medianMode_ = tilingData->medianMode;
98- this->extractChunkSize_ = tilingData->keyParams4;
99- this->blocksPerRow_ = tilingData->lastDimTileNum;
100- this->maxCoreNum_ = tilingData->lastDimNeedCore;
101- this->alignNum_ = tilingData->keyParams3;
102- this->maxMergeIterations_ = tilingData->keyParams5;
103- 
104- if (this->blockSortSize_ == 0 || this->extractChunkSize_ == 0 || this->blocksPerRow_ == 0 ||
105- this->sortAxisNum_ <= 0) {
106- return;
107- }
108- 
109- // Set GM buffers before deriving cache sizes. This keeps KthValue init local while avoiding
110- // a long identical block with SortMergeIntraCore.
111- this->inputXGm_.SetGlobalBuffer((__gm__ ValueType*)x);
112- this->outValueGm_.SetGlobalBuffer((__gm__ ValueType*)value);
113- this->outIdxGm_.SetGlobalBuffer((__gm__ IndexType*)indices);
114- 
115- // Precompute constants
116- this->blockSortLen_ = AscendC::GetSortLen<ValueType>(this->blockSortSize_);
117- this->batchSortLen_ = AscendC::GetSortLen<ValueType>(this->alignNum_);
118- this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType);
119- this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE;
120- this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE;
121- this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ -
122- static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_);
123- 
124- // Cache stores sort struct data (8 bytes per element: index + value)
125- // Each core has its own cache region, reused across batches
126- // perBatchCacheLen: sort struct length for one batch (with ping-pong, in ValueType units)
127- int64_t perCoreCacheLen = static_cast<int64_t>(this->batchSortLen_) * 2; // ping-pong, reused per batch
128- 
129- this->cacheGm_.SetGlobalBuffer((__gm__ ValueType*)workspace +
130- static_cast<int64_t>(this->blockIdx_) * perCoreCacheLen);
131-}
132- 
133-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
134-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::InitPhase1Buffers()
135-{
136- this->pipe_->InitBuffer(this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType));
137- this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_);
138- this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_);
139- this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_);
140- this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t));
141-}
142- 
143-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
144-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::PrepareInputForSort(
145- LocalTensor<ValueType> inputLocal, uint32_t alignedCount)
146-{
147- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<ValueType>) {
148- if (medianMode_ != MEDIAN_MODE_STATIC) {
149- CanonicalizeNanValues(inputLocal, alignedCount);
150- }
151- }
152-}
153- 
154-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
155-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::InitPhase2Buffers()
156-{
157- uint32_t mergeBufferSize = MERGE_LIST_MAX_NUM * this->sortBufferSize_;
158- this->pipe_->InitBuffer(this->mergeInQueue_, 1, mergeBufferSize);
159- this->pipe_->InitBuffer(this->mergeOutQueue_, 1, mergeBufferSize);
160-}
161- 
162-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
163-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::InitPhase3Buffers()
164-{
165- uint32_t extractInSize = AscendC::GetSortLen<ValueType>(this->extractChunkSize_) * sizeof(ValueType);
166- this->pipe_->InitBuffer(this->extractInQueue_, MERGE_INTRA_BUFFER_NUM, extractInSize);
167- this->pipe_->InitBuffer(this->outValueQueue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(ValueType));
168- this->pipe_->InitBuffer(this->outIdxQueue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(uint32_t));
169- this->pipe_->InitBuffer(kthValueQueue_, 1, UB_BLOCK_BYTES);
170- if constexpr (IsSameType<int64_t, IndexType>::value) {
171- this->pipe_->InitBuffer(kthIdxQueue_, 1, UB_BLOCK_BYTES);
172- }
173-}
174- 
175-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
176-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::ExtractAndCopyOut(
177- int64_t batchIdx, uint32_t resultRegion)
178-{
179- int64_t outputOffset = batchIdx;
180- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<ValueType>) {
181- kthIndex_ = ResolveBatchK(batchIdx);
182- }
183- 
184- // resultRegion: 0 = Ping (offset 0), 1 = Pong (offset batchSortLen_)
185- int64_t cacheBatchOffset = (resultRegion == 1) ? this->batchSortLen_ : 0;
186- 
187- uint32_t elemProcessed = 0;
188- uint32_t cacheOffset = 0;
189- 
190- // The merged cache is full sorted [N], but KthValue only needs one element. Skip extract chunks
191- // until the chunk containing kthIndex_, then write exactly one value/index pair.
192- while (elemProcessed <= kthIndex_) {
193- uint32_t remainElem = static_cast<uint32_t>(this->sortAxisNum_) - elemProcessed;
194- uint32_t elemCount = (this->extractChunkSize_ <= remainElem) ? this->extractChunkSize_ : remainElem;
195- if (elemCount == 0)
196- break;
197- 
198- if (kthIndex_ < elemProcessed + elemCount) {
199- ExtractAndCopyChunk(cacheBatchOffset, cacheOffset, outputOffset, elemProcessed, elemCount);
200- break;
201- }
202- 
203- elemProcessed += elemCount;
204- cacheOffset += AscendC::GetSortLen<ValueType>(elemCount);
205- }
206-}
207- 
208-// Resolves the effective k for one batch row: recounts non-NaN elements by re-reading the
209-// raw GM row in chunks (the merged cache is sign-flip encoded, unusable for NaN tests),
210-// then maps it to a sorted index via ResolveMedianK.
211-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
212-__aicore__ inline uint32_t KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::ResolveBatchK(
213- int64_t batchIdx)
214-{
215- uint32_t nonNanCount = 0U;
216- uint32_t elemOffset = 0U;
217- uint32_t axisLen = static_cast<uint32_t>(this->sortAxisNum_);
218- LocalTensor<ValueType> scalarLocal = kthValueQueue_.template AllocTensor<ValueType>();
219- while (elemOffset < axisLen) {
220- uint32_t remain = axisLen - elemOffset;
221- uint32_t elemCount = this->extractChunkSize_ < remain ? this->extractChunkSize_ : remain;
222- LocalTensor<ValueType> inputLocal = this->extractInQueue_.template AllocTensor<ValueType>();
223- DataCopyExtParams loadParams{1, static_cast<uint32_t>(elemCount * sizeof(ValueType)), 0, 0, 0};
224- DataCopyPad(inputLocal, this->inputXGm_[batchIdx * this->sortAxisNum_ + elemOffset], loadParams,
225- {false, 0, 0, 0});
226- this->extractInQueue_.EnQue(inputLocal);
227- inputLocal = this->extractInQueue_.template DeQue<ValueType>();
228- nonNanCount += CountNonNan(inputLocal, elemCount, scalarLocal.template ReinterpretCast<float>(), this->pipe_);
229- this->extractInQueue_.FreeTensor(inputLocal);
230- elemOffset += elemCount;
231- }
232- kthValueQueue_.FreeTensor(scalarLocal);
233- return ResolveMedianK(staticKthIndex_, axisLen, nonNanCount, medianMode_);
234-}
235- 
236-template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMedian>
237-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::ExtractAndCopyChunk(
238- int64_t cacheBatchOffset, uint32_t cacheOffset, int64_t outputOffset, uint32_t elemProcessed, uint32_t elemCount)
239-{
240- LocalTensor<ValueType> kthCacheLocal = this->extractInQueue_.template AllocTensor<ValueType>();
241- DataCopyExtParams loadParams;
242- loadParams.blockCount = 1;
243- loadParams.blockLen = static_cast<uint32_t>(AscendC::GetSortLen<ValueType>(elemCount) * sizeof(ValueType));
244- loadParams.srcStride = 0;
245- loadParams.dstStride = 0;
246- DataCopyPad(kthCacheLocal, this->cacheGm_[cacheBatchOffset + cacheOffset], loadParams, {false, 0, 0, 0});
247- this->extractInQueue_.EnQue(kthCacheLocal);
248- 
249- kthCacheLocal = this->extractInQueue_.template DeQue<ValueType>();
250- LocalTensor<ValueType> valueLocal = this->outValueQueue_.template AllocTensor<ValueType>();
251- LocalTensor<uint32_t> indexLocal = this->outIdxQueue_.template AllocTensor<uint32_t>();
252- Extract(valueLocal, indexLocal, kthCacheLocal, Ops::Base::CeilDiv(elemCount, DEALING_EXTRACT_NUM_ONCE));
253- 
254- // Flip back sign bit for ascending order (was flipped in SortBlockToStruct)
255- if constexpr (!IsDescend) {
256- Adds(valueLocal.template ReinterpretCast<int32_t>(), valueLocal.template ReinterpretCast<int32_t>(), 0x80000000,
257- elemCount);
258- }
259- 
260- this->outValueQueue_.EnQue(valueLocal);
261- this->outIdxQueue_.EnQue(indexLocal);
262- this->extractInQueue_.FreeTensor(kthCacheLocal);
263- valueLocal = this->outValueQueue_.template DeQue<ValueType>();
264- indexLocal = this->outIdxQueue_.template DeQue<uint32_t>();
265- 
266- event_t eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
267- SetFlag<HardEvent::V_S>(eventIdVToS);
268- WaitFlag<HardEvent::V_S>(eventIdVToS);
269- uint32_t localK = kthIndex_ - elemProcessed;
270- // localK is relative to the extracted chunk, not the whole row.
271- LocalTensor<ValueType> outValue = kthValueQueue_.AllocTensor<ValueType>();
272- outValue.SetValue(0, valueLocal.GetValue(localK));
273- LocalTensor<int32_t> indexInt32 = indexLocal.template ReinterpretCast<int32_t>();
274- LocalTensor<int64_t> indexInt64 = kthIdxQueue_.AllocTensor<int64_t>();
275- indexInt64.SetValue(0, static_cast<int64_t>(indexInt32.GetValue(localK)));
276- kthValueQueue_.EnQue(outValue);
277- kthIdxQueue_.EnQue(indexInt64);
278- outValue = kthValueQueue_.DeQue<ValueType>();
279- indexInt64 = kthIdxQueue_.DeQue<int64_t>();
280- 
281- event_t eventIdSToMte3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_MTE3));
282- SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
283- WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3);
284- DataCopyExtParams outParams{1, static_cast<uint32_t>(sizeof(ValueType)), 0, 0, 0};
285- DataCopyPad(this->outValueGm_[outputOffset], outValue, outParams);
286- if constexpr (IsSameType<int64_t, IndexType>::value) {
287- outParams.blockLen = static_cast<uint32_t>(sizeof(int64_t));
288- DataCopyPad(this->outIdxGm_[outputOffset], indexInt64, outParams);
289- }
290- event_t eventIdMte3ToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_S));
291- SetFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
292- WaitFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
293- kthIdxQueue_.FreeTensor(indexInt64);
294- kthValueQueue_.FreeTensor(outValue);
295- this->outIdxQueue_.FreeTensor(indexLocal);
296- this->outValueQueue_.FreeTensor(valueLocal);
297-}
298- 
299-} // namespace KthValue
300- 
301-#endif // KTH_VALUE_MERGE_INTRA_CORE_H
Dmath/kth_value/op_kernel/arch35/kth_value_merge_sort_more_core.h+0-245
@@ -1,245 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file merge_sort_simd.h
13- * \brief merge_sort kernel entry
14- */
15-#ifndef KTH_VALUE_MERGE_SORT_MORE_CORE_H
16-#define KTH_VALUE_MERGE_SORT_MORE_CORE_H
17-#include <cmath>
18-#include "op_kernel/platform_util.h"
19-#include "kernel_operator.h"
20-#include "kth_value_median_utils.h"
21-#include "kth_value_tiling_data.h"
22-#include "../../sort/arch35/common/merge_sort_constants.h"
23-#include "../../sort/arch35/common/merge_more_core_base.h"
24- 
25-namespace KthValue {
26-using namespace AscendC;
27- 
28-// Import shared constants from MergeSortConstants namespace
29-using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;
30-using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE;
31-using MergeSortConstants::DEALING_SORT_NUM_ONCE;
32-using MergeSortConstants::MERGE_LIST_MAX_NUM;
33-using MergeSortConstants::MERGE_MORE_BUFFER_NUM;
34-using MergeSortConstants::MERGE_SORT_WORKSPACE_PARAM;
35-using MergeSortConstants::UB_BLOCK_BYTES;
36-using MergeSortConstants::XOR_OP_VALUE_FP;
37-using MergeSortConstants::XOR_OP_VALUE_HALF;
38- 
39-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian = false>
40-struct KthValueMergeSortMoreCore : public MergeMoreCoreCommon::MergeMoreCoreBase<
41- KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>,
42- T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> {
43- using Base = MergeMoreCoreCommon::MergeMoreCoreBase<
44- KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>, T, CONVERT_TYPE, IS_DESCEND,
45- INDEX_TYPE>;
46- friend Base;
47- 
48- __aicore__ inline KthValueMergeSortMoreCore() {}
49- __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace,
50- const KthValueTilingData* tilingData, TPipe* pipe);
51- __aicore__ inline void InitMergeBuffers();
52- __aicore__ inline void ExtractAndCopyOut();
53- __aicore__ inline void OnInputLoaded(LocalTensor<T> inputLocal, uint32_t tileNum);
54- __aicore__ inline void PrepareInputForSort(LocalTensor<T> inputLocal, uint32_t alignedCount);
55- __aicore__ inline uint32_t ResolveRowK();
56- 
57- // KthValue-specific member
58- uint32_t kthIndex_ = 0;
59- uint32_t medianMode_ = 0;
60- GlobalTensor<uint32_t> medianCountGm_;
61-};
62- 
63-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian>
64-__aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::Init(
65- GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const KthValueTilingData* tilingData,
66- TPipe* pipe)
67-{
68- if (tilingData == nullptr || pipe == nullptr) {
69- return;
70- }
71- this->blockIdx_ = GetBlockIdx();
72- this->pipe_ = pipe;
73- this->outputLastDimValue_ = tilingData->lastDimTileNum;
74- this->numTileData_ = tilingData->numTileDataSize;
75- this->frontCoreNum_ = tilingData->lastDimNeedCore;
76- kthIndex_ = tilingData->kthIndex;
77- medianMode_ = tilingData->medianMode;
78- if (this->frontCoreNum_ == 0U) {
79- return;
80- }
81- uint32_t sortBufferSize = 8;
82- this->rowIdx_ = this->blockIdx_ / this->frontCoreNum_;
83- this->rowCoreIdx_ = this->blockIdx_ % this->frontCoreNum_;
84- this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_);
85- // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per
86- // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.
87- uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) *
88- sortBufferSize);
89- uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE);
90- this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2;
91- this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE;
92- 
93- this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));
94- this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value));
95- this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));
96- this->workspaceGm_[0].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_,
97- rowWorkspaceElements);
98- this->workspaceGm_[1].SetGlobalBuffer(
99- (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements);
100- uint64_t medianCountOffset = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(tilingData->lastAxisNum) *
101- static_cast<uint64_t>(tilingData->unsortedDimNum) *
102- MERGE_SORT_WORKSPACE_PARAM * sizeof(uint32_t));
103- medianCountGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(workSpace + medianCountOffset),
104- tilingData->lastDimNeedCore * tilingData->unsortedDimNum);
105- 
106- uint32_t tailNum = this->outputLastDimValue_ - (this->frontCoreNum_ - 1) * this->numTileData_;
107- uint32_t alignTile = ROUND_UP_AGLIN(tailNum);
108- this->pipe_->InitBuffer(this->inputQueue_, MERGE_MORE_BUFFER_NUM, alignTile * sizeof(T));
109- 
110- this->pipe_->InitBuffer(this->sortedValueUb_, alignTile * sortBufferSize);
111- this->pipe_->InitBuffer(this->sortedValueIndexUb_, alignTile * sizeof(uint32_t));
112- this->pipe_->InitBuffer(this->concatTempBuf_, alignTile * sortBufferSize);
113- this->pipe_->InitBuffer(this->sortTempBuf_, alignTile * sortBufferSize);
114- this->pipe_->InitBuffer(this->sortedValueLocalCastTbuf_, alignTile * sortBufferSize);
115-}
116- 
117-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian>
118-__aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::OnInputLoaded(
119- LocalTensor<T> inputLocal, uint32_t tileNum)
120-{
121- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
122- LocalTensor<float> scratch = this->concatTempBuf_.template Get<float>();
123- uint32_t nonNanCount = CountNonNan(inputLocal, tileNum, scratch, this->pipe_);
124- LocalTensor<uint32_t> countLocal = scratch.template ReinterpretCast<uint32_t>();
125- countLocal.SetValue(0, nonNanCount);
126- event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_MTE3));
127- SetFlag<HardEvent::S_MTE3>(eventId);
128- WaitFlag<HardEvent::S_MTE3>(eventId);
129- DataCopyExtParams copyParams{1, static_cast<uint32_t>(sizeof(uint32_t)), 0, 0, 0};
130- DataCopyPad(medianCountGm_[this->blockIdx_], countLocal, copyParams);
131- }
132-}
133- 
134-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian>
135-__aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE,
136- EnableMedian>::PrepareInputForSort(LocalTensor<T> inputLocal,
137- uint32_t alignedCount)
138-{
139- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
140- if (medianMode_ != MEDIAN_MODE_STATIC) {
141- CanonicalizeNanValues(inputLocal, alignedCount);
142- }
143- }
144-}
145- 
146-// Resolves the effective k for the current row: gathers the per-core non-NaN counts
147-// published to the median count workspace by all frontCoreNum_ cores of this row, sums
148-// them up, then maps the row total to a sorted index via ResolveMedianK.
149-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian>
150-__aicore__ inline uint32_t
151-KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::ResolveRowK()
152-{
153- LocalTensor<uint32_t> countLocal = this->concatTempBuf_.template Get<uint32_t>();
154- uint32_t countOffset = this->rowIdx_ * this->frontCoreNum_;
155- DataCopyExtParams copyParams{1, this->frontCoreNum_ * static_cast<uint32_t>(sizeof(uint32_t)), 0, 0, 0};
156- DataCopyPad(countLocal, medianCountGm_[countOffset], copyParams, {false, 0, 0, 0});
157- event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE2_S));
158- SetFlag<HardEvent::MTE2_S>(eventId);
159- WaitFlag<HardEvent::MTE2_S>(eventId);
160- uint32_t nonNanCount = 0U;
161- for (uint32_t core = 0U; core < this->frontCoreNum_; ++core) {
162- nonNanCount += countLocal.GetValue(core);
163- }
164- return ResolveMedianK(kthIndex_, this->outputLastDimValue_, nonNanCount, medianMode_);
165-}
166- 
167-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian>
168-__aicore__ inline void
169-KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::InitMergeBuffers()
170-{
171- uint32_t sortBufferSize = 8;
172- this->pipe_->InitBuffer(this->sortedQueue_, MERGE_MORE_BUFFER_NUM,
173- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);
174- this->pipe_->InitBuffer(this->copyInQueue_, MERGE_MORE_BUFFER_NUM,
175- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);
176- this->pipe_->InitBuffer(this->castValueQueue_, MERGE_MORE_BUFFER_NUM,
177- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE));
178- this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM,
179- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(uint32_t));
180- this->pipe_->InitBuffer(this->outValueQueue_, 1, UB_BLOCK_BYTES);
181- this->pipe_->InitBuffer(this->outIndexQueue_, 1, UB_BLOCK_BYTES);
182-}
183- 
184-template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian>
185-__aicore__ inline void
186-KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::ExtractAndCopyOut()
187-{
188- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
189- if (this->outOffset_ == 0) {
190- kthIndex_ = ResolveRowK();
191- }
192- }
193- LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template DeQue<CONVERT_TYPE>();
194- LocalTensor<CONVERT_TYPE> castValue = this->castValueQueue_.template AllocTensor<CONVERT_TYPE>();
195- LocalTensor<uint32_t> castIndex = this->castIndexQueue_.template AllocTensor<uint32_t>();
196- AscendC::Extract(castValue, castIndex, sortTempBuffer,
197- ((this->curLoopSortedNum_ + DEALING_EXTRACT_NUM_ONCE - 1) / DEALING_EXTRACT_NUM_ONCE));
198- if constexpr (!IS_DESCEND) {
199- this->FlipSignBit(castValue, ROUND_UP_AGLIN(this->curLoopSortedNum_));
200- }
201- this->castValueQueue_.EnQue(castValue);
202- this->castIndexQueue_.EnQue(castIndex);
203- castValue = this->castValueQueue_.template DeQue<CONVERT_TYPE>();
204- castIndex = this->castIndexQueue_.template DeQue<uint32_t>();
205- 
206- uint32_t sortedBase = static_cast<uint32_t>(this->outOffset_);
207- uint32_t sortedEnd = sortedBase + static_cast<uint32_t>(this->curLoopSortedNum_);
208- // Extract chunks are scanned in sorted order. Only the chunk covering kthIndex_ writes GM output;
209- // other chunks are merge intermediates and are discarded for KthValue.
210- if (kthIndex_ >= sortedBase && kthIndex_ < sortedEnd) {
211- uint32_t localK = kthIndex_ - sortedBase;
212- event_t eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
213- SetFlag<HardEvent::V_S>(eventIdVToS);
214- WaitFlag<HardEvent::V_S>(eventIdVToS);
215- 
216- LocalTensor<T> outValue = this->outValueQueue_.template AllocTensor<T>();
217- LocalTensor<INDEX_TYPE> outIndex = this->outIndexQueue_.template AllocTensor<INDEX_TYPE>();
218- outValue.SetValue(0, castValue.GetValue(localK));
219- outIndex.SetValue(0, static_cast<INDEX_TYPE>(castIndex.GetValue(localK)));
220- this->outValueQueue_.EnQue(outValue);
221- this->outIndexQueue_.EnQue(outIndex);
222- 
223- outValue = this->outValueQueue_.template DeQue<T>();
224- outIndex = this->outIndexQueue_.template DeQue<INDEX_TYPE>();
225- event_t eventIdSToMte3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_MTE3));
226- SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
227- WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3);
228- DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(sizeof(T)), 0, 0, 0};
229- DataCopyPad(this->outValueGm_[this->rowIdx_], outValue, valueCopyParam);
230- DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(sizeof(INDEX_TYPE)), 0, 0, 0};
231- DataCopyPad(this->outIndexGm_[this->rowIdx_], outIndex, indexCopyParam);
232- event_t eventIdMte3ToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_S));
233- SetFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
234- WaitFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
235- this->outValueQueue_.FreeTensor(outValue);
236- this->outIndexQueue_.FreeTensor(outIndex);
237- this->allRemainElements_ = 0;
238- }
239- this->castIndexQueue_.FreeTensor(castIndex);
240- this->castValueQueue_.FreeTensor(castValue);
241- this->sortedQueue_.FreeTensor(sortTempBuffer);
242- this->outOffset_ += this->curLoopSortedNum_;
243-}
244-} // namespace KthValue
245-#endif // KTH_VALUE_MERGE_SORT_MORE_CORE_H
Dmath/kth_value/op_kernel/arch35/kth_value_merge_sort_one_core.h+0-308
@@ -1,308 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#ifndef KTH_VALUE_MERGE_SORT_ONE_CORE_H
12-#define KTH_VALUE_MERGE_SORT_ONE_CORE_H
13- 
14-#include <cmath>
15-#include "kernel_operator.h"
16-#include "op_kernel/platform_util.h"
17-#include "kth_value_median_utils.h"
18-#include "kth_value_tiling_data.h"
19-#include "../../sort/arch35/common/merge_sort_constants.h"
20-#include "../../sort/arch35/common/util_type_simd.h"
21- 
22-namespace KthValue {
23-using namespace AscendC;
24- 
25-using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;
26-using MergeSortConstants::DEALING_SORT_NUM_ONCE;
27-using MergeSortConstants::XOR_OP_VALUE_FP;
28-using MergeSortConstants::XOR_OP_VALUE_HALF;
29- 
30-constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize();
31- 
32-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis = 0, bool EnableMedian = false>
33-class KthValueMergeSortOneCore {
34-public:
35- __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
36- TPipe* pipe);
37- __aicore__ inline void Process();
38- 
39-private:
40- __aicore__ inline void ParseTilingData();
41- __aicore__ inline void InitIndexLocal();
42- __aicore__ inline void CopyDataIn(uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum);
43- __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offset, uint32_t count);
44- __aicore__ inline void SortRows(LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal,
45- LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum);
46- __aicore__ inline void CopyKthToGm(uint64_t outputOffset, uint32_t rowNum);
47- __aicore__ inline void ProcessSingleRound(uint32_t round);
48- 
49- static constexpr uint32_t SORT_STRUCT_BYTES = 8;
50- 
51- GlobalTensor<T> xGm_;
52- GlobalTensor<T> valuesGm_;
53- GlobalTensor<int64_t> indicesGm_;
54- 
55- TPipe* pipe_{nullptr};
56- const KthValueTilingData* tiling_{nullptr};
57- TQue<QuePosition::VECIN, 1> inQueueX_;
58- TQue<QuePosition::VECOUT, 1> outValueQueue_;
59- TQue<QuePosition::VECOUT, 1> outIndexQueue_;
60- TQue<QuePosition::VECOUT, 1> compactValueQueue_;
61- TQue<QuePosition::VECOUT, 1> compactIndexQueue_;
62- TBuf<TPosition::VECCALC> concatTmpBuf_;
63- TBuf<TPosition::VECCALC> sortTmpBuf_;
64- TBuf<TPosition::VECCALC> sortedLocalBuf_;
65- TBuf<TPosition::VECCALC> xCastBuf_;
66- TBuf<TPosition::VECCALC> sortedValueCastBuf_;
67- TBuf<TPosition::VECCALC> indexLocalBuf_;
68- LocalTensor<uint32_t> indexLocal_;
69- 
70- uint32_t blockIdx_ = 0;
71- uint32_t oneCoreRowNum_ = 0;
72- uint32_t numTileData_ = 0;
73- uint32_t alignSize_ = 0;
74- uint32_t sortLoopTimes_ = 0;
75- uint32_t unsortedDimParallel_ = 0;
76- uint32_t kthIndex_ = 0;
77- uint32_t medianMode_ = 0;
78- int64_t unsortedDimNum_ = 0;
79-};
80- 
81-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
82-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::Init(
83- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe)
84-{
85- if (tiling == nullptr || pipe == nullptr) {
86- return;
87- }
88- blockIdx_ = GetBlockIdx();
89- pipe_ = pipe;
90- tiling_ = tiling;
91- ParseTilingData();
92- 
93- xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
94- valuesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(values));
95- indicesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(indices));
96- 
97- uint32_t bufferNum = tiling_->keyParams4 == 0U ? 1U : tiling_->keyParams4;
98- pipe_->InitBuffer(inQueueX_, bufferNum, tiling_->keyParams1);
99- pipe_->InitBuffer(outValueQueue_, bufferNum, tiling_->keyParams1);
100- pipe_->InitBuffer(outIndexQueue_, bufferNum, tiling_->keyParams2);
101- pipe_->InitBuffer(compactValueQueue_, bufferNum, ROUND_UP_AGLIN(oneCoreRowNum_ * sizeof(T)));
102- pipe_->InitBuffer(compactIndexQueue_, bufferNum, ROUND_UP_AGLIN(oneCoreRowNum_ * sizeof(int64_t)));
103- 
104- uint32_t sortStructElems = alignSize_ * SORT_STRUCT_BYTES;
105- pipe_->InitBuffer(indexLocalBuf_, alignSize_ * sizeof(uint32_t));
106- pipe_->InitBuffer(concatTmpBuf_, tiling_->tmpUbSize);
107- pipe_->InitBuffer(sortTmpBuf_, sortStructElems * sizeof(CONVERT_TYPE));
108- pipe_->InitBuffer(sortedLocalBuf_, sortStructElems * sizeof(CONVERT_TYPE));
109- if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
110- pipe_->InitBuffer(xCastBuf_, alignSize_ * oneCoreRowNum_ * sizeof(CONVERT_TYPE));
111- pipe_->InitBuffer(sortedValueCastBuf_, alignSize_ * oneCoreRowNum_ * sizeof(CONVERT_TYPE));
112- }
113- indexLocal_ = indexLocalBuf_.AllocTensor<uint32_t>();
114- InitIndexLocal();
115-}
116- 
117-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
118-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::ParseTilingData()
119-{
120- oneCoreRowNum_ = tiling_->keyParams0;
121- numTileData_ = tiling_->numTileDataSize;
122- alignSize_ = tiling_->keyParams3;
123- sortLoopTimes_ = tiling_->sortLoopTimes;
124- unsortedDimParallel_ = tiling_->unsortedDimParallel;
125- kthIndex_ = tiling_->kthIndex;
126- medianMode_ = tiling_->medianMode;
127- unsortedDimNum_ = tiling_->unsortedDimNum;
128-}
129- 
130-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
131-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::InitIndexLocal()
132-{
133- __ubuf__ int32_t* indexValuePtr = reinterpret_cast<__ubuf__ int32_t*>(indexLocal_.GetPhyAddr());
134- uint32_t vfLenB32 = Ops::Base::GetVRegSize() / sizeof(int32_t);
135- uint16_t repeatTime = Ops::Base::CeilDiv(alignSize_, vfLenB32);
136- uint32_t alignSizeCopy = alignSize_;
137- __VEC_SCOPE__
138- {
139- Reg::RegTensor<int32_t> vciTensor;
140- Reg::RegTensor<int32_t> indexTensor;
141- Reg::Arange(vciTensor, 0);
142- for (uint16_t i = 0; i < repeatTime; ++i) {
143- Reg::MaskReg mask = Reg::UpdateMask<uint32_t>(alignSizeCopy);
144- Reg::Adds(indexTensor, vciTensor, i * vfLenB32, mask);
145- Reg::StoreAlign<int32_t, Reg::PostLiteral::POST_MODE_UPDATE>(indexValuePtr, indexTensor, vfLenB32, mask);
146- }
147- }
148-}
149- 
150-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
151-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::CopyDataIn(
152- uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum)
153-{
154- LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();
155- Duplicate(xLocal, static_cast<T>(NAN), alignSize_ * rowNum);
156- event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2));
157- SetFlag<HardEvent::V_MTE2>(eventId);
158- WaitFlag<HardEvent::V_MTE2>(eventId);
159- 
160- uint32_t currTileSizeAlign = (currTileSize * sizeof(T) + UB_BLOCK_SIZE - 1U) / UB_BLOCK_SIZE * UB_BLOCK_SIZE /
161- sizeof(T);
162- uint32_t dstStride = ((alignSize_ - currTileSizeAlign) * sizeof(T)) / UB_BLOCK_SIZE;
163- DataCopyPadExtParams<T> padParams{true, 0, static_cast<uint8_t>(currTileSizeAlign - currTileSize),
164- static_cast<T>(NAN)};
165- DataCopyExtParams copyParam{static_cast<uint16_t>(rowNum), static_cast<uint32_t>(currTileSize * sizeof(T)), 0,
166- dstStride, 0};
167- DataCopyPad(xLocal, xGm_[tileOffset], copyParam, padParams);
168- inQueueX_.EnQue<T>(xLocal);
169-}
170- 
171-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
172-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::FlipSignBit(
173- LocalTensor<CONVERT_TYPE> xLocal, uint32_t offset, uint32_t count)
174-{
175- if constexpr (IsSameType<float, CONVERT_TYPE>::value) {
176- LocalTensor<int32_t> castTensor = xLocal[offset].template ReinterpretCast<int32_t>();
177- Adds(castTensor, castTensor, XOR_OP_VALUE_FP, count);
178- } else if constexpr (IsSameType<half, CONVERT_TYPE>::value) {
179- LocalTensor<int16_t> castTensor = xLocal[offset].template ReinterpretCast<int16_t>();
180- Adds(castTensor, castTensor, XOR_OP_VALUE_HALF, count);
181- }
182-}
183- 
184-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
185-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::SortRows(
186- LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal, LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum)
187-{
188- uint32_t sortRepeatTimes = alignSize_ / DEALING_SORT_NUM_ONCE;
189- uint32_t concatRepeatTimes = alignSize_ / DEALING_CONCAT_NUM_ONCE;
190- LocalTensor<CONVERT_TYPE> concatTmp = concatTmpBuf_.Get<CONVERT_TYPE>();
191- LocalTensor<CONVERT_TYPE> sortedLocal = sortedLocalBuf_.Get<CONVERT_TYPE>();
192- LocalTensor<CONVERT_TYPE> sortTmp = sortTmpBuf_.Get<CONVERT_TYPE>();
193- LocalTensor<CONVERT_TYPE> xSortLocal;
194- LocalTensor<CONVERT_TYPE> sortedValueCast;
195- if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
196- xSortLocal = xCastBuf_.Get<CONVERT_TYPE>();
197- sortedValueCast = sortedValueCastBuf_.Get<CONVERT_TYPE>();
198- Cast(xSortLocal, xLocal, RoundMode::CAST_NONE, alignSize_ * rowNum);
199- } else {
200- xSortLocal = xLocal;
201- sortedValueCast = sortedValueLocal;
202- }
203- for (uint32_t row = 0; row < rowNum; ++row) {
204- uint32_t offset = row * alignSize_;
205- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<CONVERT_TYPE>) {
206- if (medianMode_ != MEDIAN_MODE_STATIC) {
207- CanonicalizeNanValues(xSortLocal[offset], alignSize_);
208- }
209- }
210- // MergeSort sorts sign-flipped ascending keys. Flip before Sort/Sort32 and flip back after Extract.
211- FlipSignBit(xSortLocal, offset, alignSize_);
212- if constexpr (isSort32SmallAxis == 1) {
213- Sort32<CONVERT_TYPE>(sortedLocal, xSortLocal[offset], indexLocal_, 1);
214- } else {
215- LocalTensor<CONVERT_TYPE> concatLocal;
216- Concat(concatLocal, xSortLocal[offset], concatTmp, concatRepeatTimes);
217- AscendC::Sort<CONVERT_TYPE, true>(sortedLocal, concatLocal, indexLocal_, sortTmp, sortRepeatTimes);
218- }
219- Extract(sortedValueCast[offset], sortedIndexLocal[offset], sortedLocal,
220- isSort32SmallAxis == 1 ? 1 : sortRepeatTimes);
221- FlipSignBit(sortedValueCast, offset, alignSize_);
222- }
223- if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
224- Cast(sortedValueLocal, sortedValueCast, RoundMode::CAST_RINT, alignSize_ * rowNum);
225- }
226-}
227- 
228-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
229-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::CopyKthToGm(
230- uint64_t outputOffset, uint32_t rowNum)
231-{
232- LocalTensor<T> sortedValueLocal = outValueQueue_.DeQue<T>();
233- LocalTensor<uint32_t> sortedIndexLocal = outIndexQueue_.DeQue<uint32_t>();
234- 
235- event_t eventIdVToS = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
236- SetFlag<HardEvent::V_S>(eventIdVToS);
237- WaitFlag<HardEvent::V_S>(eventIdVToS);
238- 
239- LocalTensor<T> compactValue = compactValueQueue_.AllocTensor<T>();
240- LocalTensor<int64_t> compactIndex = compactIndexQueue_.AllocTensor<int64_t>();
241- for (uint32_t row = 0; row < rowNum; ++row) {
242- uint32_t rowOffset = row * alignSize_;
243- uint32_t selectedK = kthIndex_;
244- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
245- uint32_t nonNanCount = CountNonNan(sortedValueLocal[rowOffset], numTileData_, concatTmpBuf_.Get<float>(),
246- pipe_);
247- selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_);
248- }
249- uint32_t srcOffset = rowOffset + selectedK;
250- compactValue.SetValue(row, sortedValueLocal.GetValue(srcOffset));
251- compactIndex.SetValue(row, static_cast<int64_t>(sortedIndexLocal.GetValue(srcOffset)));
252- }
253- compactValueQueue_.EnQue<T>(compactValue);
254- compactIndexQueue_.EnQue<int64_t>(compactIndex);
255- compactValue = compactValueQueue_.DeQue<T>();
256- compactIndex = compactIndexQueue_.DeQue<int64_t>();
257- event_t eventIdSToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));
258- SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
259- WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3);
260- DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(rowNum * sizeof(T)), 0, 0, 0};
261- DataCopyPad(valuesGm_[outputOffset], compactValue, valueCopyParam);
262- DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(rowNum * sizeof(int64_t)), 0, 0, 0};
263- DataCopyPad(indicesGm_[outputOffset], compactIndex, indexCopyParam);
264- compactValueQueue_.FreeTensor(compactValue);
265- compactIndexQueue_.FreeTensor(compactIndex);
266- outValueQueue_.FreeTensor(sortedValueLocal);
267- outIndexQueue_.FreeTensor(sortedIndexLocal);
268-}
269- 
270-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
271-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::ProcessSingleRound(
272- uint32_t round)
273-{
274- int64_t rowStart = (static_cast<int64_t>(blockIdx_) + static_cast<int64_t>(round) * unsortedDimParallel_) *
275- oneCoreRowNum_;
276- if (rowStart >= unsortedDimNum_) {
277- return;
278- }
279- uint32_t rowNum = oneCoreRowNum_;
280- int64_t remain = unsortedDimNum_ - rowStart;
281- if (remain < static_cast<int64_t>(rowNum)) {
282- rowNum = static_cast<uint32_t>(remain);
283- }
284- uint64_t inputOffset = static_cast<uint64_t>(rowStart) * numTileData_;
285- CopyDataIn(inputOffset, numTileData_, rowNum);
286- LocalTensor<T> xLocal = inQueueX_.DeQue<T>();
287- LocalTensor<T> sortedValueLocal = outValueQueue_.AllocTensor<T>();
288- LocalTensor<uint32_t> sortedIndexLocal = outIndexQueue_.AllocTensor<uint32_t>();
289- SortRows(xLocal, sortedValueLocal, sortedIndexLocal, rowNum);
290- outValueQueue_.EnQue<T>(sortedValueLocal);
291- outIndexQueue_.EnQue<uint32_t>(sortedIndexLocal);
292- inQueueX_.FreeTensor(xLocal);
293- CopyKthToGm(static_cast<uint64_t>(rowStart), rowNum);
294-}
295- 
296-template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian>
297-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::Process()
298-{
299- if (blockIdx_ >= GetBlockNum()) {
300- return;
301- }
302- for (uint32_t round = 0; round < sortLoopTimes_; ++round) {
303- ProcessSingleRound(round);
304- }
305-}
306-} // namespace KthValue
307- 
308-#endif
Dmath/kth_value/op_kernel/arch35/kth_value_non_last_small_axis.h+0-332
@@ -1,332 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#ifndef KTH_VALUE_NON_LAST_SMALL_AXIS_H
12-#define KTH_VALUE_NON_LAST_SMALL_AXIS_H
13- 
14-#include <type_traits>
15- 
16-#include "kernel_operator.h"
17-#include "kernel_tiling/kernel_tiling.h"
18-#include "op_kernel/platform_util.h"
19-#include "simt_api/asc_simt.h"
20-#include "kth_value_median_utils.h"
21-#include "kth_value_tiling_data.h"
22-#include "../../sort/arch35/common/non_last_small_axis_base.h"
23- 
24-namespace KthValue {
25-using namespace AscendC;
26- 
27-// Median hook (pre-sort): canonicalizes NaN/signed-zero in place over the transposed
28-// sort-major layout, which is strided (valueAxisElems), hence a scalar SIMT pass.
29-template <typename T>
30-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__
31- void SimtCanonicalizeNonLastMedianRows(uint32_t validElems, uint32_t axisLen, uint32_t valueAxisElems,
32- __ubuf__ T* sortInput)
33-{
34- for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < validElems;
35- idx += SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) {
36- uint32_t inner = idx / axisLen;
37- uint32_t axis = idx - inner * axisLen;
38- uint32_t offset = inner * valueAxisElems + axis;
39- sortInput[offset] = CanonicalizeMedianSortValue(sortInput[offset]);
40- }
41-}
42- 
43-// Median hook (post-sort, per inner lane): binary-searches the NaN boundary on the sorted
44-// row, resolves the effective k per medianMode, then traces back the selected element to the
45-// original input tile (raw T layout), preserving precision and the original NaN payload.
46-template <typename T, typename SortT>
47-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__
48- void SimtStoreNonLastMedian(uint32_t curInnerChunk, uint32_t axisLen, uint32_t kthIndex, uint32_t medianMode,
49- uint32_t inputRowElems, uint32_t valueAxisElems, uint32_t indexAxisElems,
50- __ubuf__ T* inputTile, __ubuf__ SortT* sortedValue, __ubuf__ uint32_t* sortedIndex,
51- __ubuf__ T* compactValue, __ubuf__ int64_t* compactIndex)
52-{
53- for (uint32_t inner = static_cast<uint32_t>(threadIdx.x); inner < curInnerChunk;
54- inner += SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) {
55- uint32_t rowOffset = inner * valueAxisElems;
56- uint32_t left = 0U;
57- uint32_t right = axisLen;
58- while (left < right) {
59- uint32_t middle = left + (right - left) / 2U;
60- if (IsNanValue(sortedValue[rowOffset + middle])) {
61- right = middle;
62- } else {
63- left = middle + 1U;
64- }
65- }
66- uint32_t selectedK = kthIndex;
67- if (medianMode == MEDIAN_MODE_PROPAGATE_NAN && left < axisLen) {
68- selectedK = left;
69- } else if (medianMode == MEDIAN_MODE_IGNORE_NAN) {
70- selectedK = left == 0U ? 0U : (left - 1U) / 2U;
71- }
72- uint32_t selectedIndex = sortedIndex[inner * indexAxisElems + selectedK];
73- compactValue[inner] = inputTile[selectedIndex * inputRowElems + inner];
74- compactIndex[inner] = static_cast<int64_t>(selectedIndex);
75- }
76-}
77- 
78-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian = false>
79-class KthValueNonLastSmallAxis
80- : public SmallAxisCommon::NonLastSmallAxisBase<
81- KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>, T,
82- std::conditional_t<UseMergeSort && std::is_same_v<T, bfloat16_t>, float, T>,
83- std::conditional_t<sizeof(T) <= sizeof(int16_t), int16_t, int32_t>,
84- std::conditional_t<sizeof(T) <= sizeof(int16_t), uint16_t, uint32_t>,
85- std::conditional_t<sizeof(T) == 1, std::conditional_t<std::is_same_v<T, uint8_t>, uint16_t, int16_t>, T>,
86- IsDescend, UseMergeSort, UseMergeSort && std::is_same_v<T, bfloat16_t>,
87- !UseMergeSort && SignedZeroSortCommon::IS_FLOATING_POINT_V<T>> {
88- using Base = SmallAxisCommon::NonLastSmallAxisBase<
89- KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>, T,
90- std::conditional_t<UseMergeSort && std::is_same_v<T, bfloat16_t>, float, T>,
91- std::conditional_t<sizeof(T) <= sizeof(int16_t), int16_t, int32_t>,
92- std::conditional_t<sizeof(T) <= sizeof(int16_t), uint16_t, uint32_t>,
93- std::conditional_t<sizeof(T) == 1, std::conditional_t<std::is_same_v<T, uint8_t>, uint16_t, int16_t>, T>,
94- IsDescend, UseMergeSort, UseMergeSort && std::is_same_v<T, bfloat16_t>,
95- !UseMergeSort && SignedZeroSortCommon::IS_FLOATING_POINT_V<T>>;
96- 
97- using SortT = std::conditional_t<UseMergeSort && std::is_same_v<T, bfloat16_t>, float, T>;
98- static constexpr bool IS_BF16_MERGE = UseMergeSort && std::is_same_v<T, bfloat16_t>;
99- static constexpr bool NORMALIZE_SIGNED_ZERO = !UseMergeSort && SignedZeroSortCommon::IS_FLOATING_POINT_V<T>;
100- 
101-public:
102- __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace,
103- const KthValueTilingData* tilingData, TPipe* pipe);
104- 
105- __aicore__ inline void Process();
106- __aicore__ inline void StoreTile(int64_t inputOffset, int64_t outputOffset, uint32_t curInnerChunk);
107- 
108-private:
109- __aicore__ inline void ParseTilingData();
110- __aicore__ inline void PrepareSortRows(uint32_t curInnerChunk);
111- __aicore__ inline void CopyKthToOutput(uint32_t curInnerChunk, int64_t outputOffset);
112- 
113- const KthValueTilingData* tilingData_ = nullptr;
114- GlobalTensor<T> valueGm_;
115- GlobalTensor<int64_t> indexGm_;
116- 
117- TBuf<TPosition::VECCALC> compactValueBuf_;
118- TBuf<TPosition::VECCALC> compactCastValueBuf_;
119- TBuf<TPosition::VECCALC> compactIndexBuf_;
120- TBuf<TPosition::VECCALC> sourceIndexBuf_;
121- 
122- LocalTensor<T> compactValue_;
123- LocalTensor<SortT> compactCastValue_;
124- LocalTensor<int64_t> compactIndex_;
125- 
126- uint32_t kthIndex_ = 0;
127- uint32_t medianMode_ = MEDIAN_MODE_STATIC;
128-};
129- 
130-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian>
131-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>::Init(
132- GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData, TPipe* pipe)
133-{
134- (void)workspace;
135- if (tilingData == nullptr || pipe == nullptr) {
136- return;
137- }
138- this->pipe_ = pipe;
139- tilingData_ = tilingData;
140- this->blockIdx_ = GetBlockIdx();
141- this->blockDim_ = GetBlockNum();
142- ParseTilingData();
143- 
144- this->inputGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
145- valueGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(values));
146- indexGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(indices));
147- 
148- if (this->axisLen_ == 0 || this->innerChunk_ == 0 || this->innerLoopNum_ == 0 || this->outerSize_ <= 0 ||
149- this->innerSize_ <= 0) {
150- return;
151- }
152- this->pipe_->InitBuffer(this->inputTileBuf_, this->axisLen_ * this->inputRowBytes_);
153- if constexpr (IS_BF16_MERGE) {
154- this->pipe_->InitBuffer(this->inputCastBuf_, this->innerChunk_ * this->inputValueAxisBytes_);
155- this->inputCast_ = this->inputCastBuf_.template Get<T>();
156- }
157- this->pipe_->InitBuffer(this->sortInputBuf_, this->innerChunk_ * this->valueAxisBytes_);
158- this->pipe_->InitBuffer(this->sortedValueBuf_, this->innerChunk_ * this->valueAxisBytes_);
159- this->pipe_->InitBuffer(this->sortedIndexBuf_, this->innerChunk_ * this->indexAxisBytes_);
160- this->pipe_->InitBuffer(compactValueBuf_, ROUND_UP_AGLIN(this->innerChunk_ * sizeof(T)));
161- if constexpr (IS_BF16_MERGE) {
162- this->pipe_->InitBuffer(compactCastValueBuf_, ROUND_UP_AGLIN(this->innerChunk_ * sizeof(SortT)));
163- compactCastValue_ = compactCastValueBuf_.template Get<SortT>();
164- }
165- this->pipe_->InitBuffer(compactIndexBuf_, ROUND_UP_AGLIN(this->innerChunk_ * sizeof(int64_t)));
166- if constexpr (NORMALIZE_SIGNED_ZERO) {
167- this->pipe_->InitBuffer(sourceIndexBuf_, this->indexAxisBytes_);
168- }
169- this->pipe_->InitBuffer(this->tmpBuf_, this->tmpUbSize_);
170- 
171- this->inputTile_ = this->inputTileBuf_.template Get<T>();
172- this->sortInput_ = this->sortInputBuf_.template Get<SortT>();
173- this->sortedValue_ = this->sortedValueBuf_.template Get<SortT>();
174- this->sortedIndex_ = this->sortedIndexBuf_.template Get<uint32_t>();
175- compactValue_ = compactValueBuf_.template Get<T>();
176- compactIndex_ = compactIndexBuf_.template Get<int64_t>();
177- this->tmp_ = this->tmpBuf_.template Get<uint8_t>();
178- if constexpr (NORMALIZE_SIGNED_ZERO) {
179- this->sourceIndex_ = sourceIndexBuf_.template Get<uint32_t>();
180- }
181-}
182- 
183-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian>
184-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>::Process()
185-{
186- if (this->blockIdx_ >= this->blockDim_ || this->axisLen_ == 0 || this->innerChunk_ == 0 ||
187- this->innerLoopNum_ == 0) {
188- return;
189- }
190- uint64_t tileCount = static_cast<uint64_t>(this->outerSize_) * static_cast<uint64_t>(this->innerLoopNum_);
191- uint64_t tilesPerCore = Ops::Base::CeilDiv(tileCount, static_cast<uint64_t>(this->blockDim_));
192- uint64_t startTile = static_cast<uint64_t>(this->blockIdx_) * tilesPerCore;
193- uint64_t endTile = startTile + tilesPerCore;
194- if (endTile > tileCount) {
195- endTile = tileCount;
196- }
197- for (uint64_t tileId = startTile; tileId < endTile; ++tileId) {
198- uint64_t outerId = tileId / this->innerLoopNum_;
199- uint32_t innerTileId = static_cast<uint32_t>(tileId - outerId * static_cast<uint64_t>(this->innerLoopNum_));
200- uint32_t curInnerChunk = this->GetCurrentInnerChunk(innerTileId);
201- if (curInnerChunk == 0U) {
202- continue;
203- }
204- int64_t innerStart = static_cast<int64_t>(innerTileId) * static_cast<int64_t>(this->innerChunk_);
205- int64_t inputOffset = static_cast<int64_t>(outerId) * static_cast<int64_t>(this->axisLen_) * this->innerSize_ +
206- innerStart;
207- int64_t outputOffset = static_cast<int64_t>(outerId) * this->innerSize_ + innerStart;
208- this->LoadTile(inputOffset, curInnerChunk);
209- this->TransposeToSortMajor(curInnerChunk);
210- PrepareSortRows(curInnerChunk);
211- this->SortRows(curInnerChunk);
212- StoreTile(inputOffset, outputOffset, curInnerChunk);
213- }
214-}
215- 
216-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian>
217-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>::ParseTilingData()
218-{
219- this->axisLen_ = static_cast<uint32_t>(tilingData_->lastAxisNum);
220- kthIndex_ = static_cast<uint32_t>(tilingData_->kthIndex);
221- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
222- medianMode_ = tilingData_->medianMode;
223- }
224- this->outerSize_ = tilingData_->outerSize;
225- this->innerSize_ = tilingData_->innerSize;
226- this->innerLoopNum_ = tilingData_->innerLoopNum;
227- this->innerChunk_ = tilingData_->innerChunk;
228- this->inputRowBytes_ = tilingData_->inputRowBytes;
229- this->valueAxisBytes_ = tilingData_->valueAxisBytes;
230- this->indexAxisBytes_ = tilingData_->indexAxisBytes;
231- this->inputRowElems_ = this->inputRowBytes_ / sizeof(T);
232- this->valueAxisElems_ = this->valueAxisBytes_ / sizeof(SortT);
233- this->indexAxisElems_ = this->indexAxisBytes_ / sizeof(uint32_t);
234- this->sortCount_ = tilingData_->keyParams0 == 0U ? this->axisLen_ : tilingData_->keyParams0;
235- if constexpr (IS_BF16_MERGE) {
236- this->inputValueAxisBytes_ = tilingData_->keyParams1;
237- this->inputValueAxisElems_ = this->inputValueAxisBytes_ / sizeof(T);
238- }
239- this->tmpUbSize_ = tilingData_->tmpUbSize;
240-}
241- 
242-// Median hook point: runs after transpose, before SortRows.
243-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian>
244-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>::PrepareSortRows(
245- uint32_t curInnerChunk)
246-{
247- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
248- if (medianMode_ != MEDIAN_MODE_STATIC) {
249- asc_vf_call<SimtCanonicalizeNonLastMedianRows<SortT>>(
250- dim3(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM), curInnerChunk * this->axisLen_, this->axisLen_,
251- this->valueAxisElems_, reinterpret_cast<__ubuf__ SortT*>(this->sortInput_.GetPhyAddr()));
252- event_t eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
253- SetFlag<HardEvent::V_S>(eventIdVToS);
254- WaitFlag<HardEvent::V_S>(eventIdVToS);
255- }
256- }
257-}
258- 
259-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian>
260-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>::CopyKthToOutput(
261- uint32_t curInnerChunk, int64_t outputOffset)
262-{
263- event_t eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
264- SetFlag<HardEvent::V_S>(eventIdVToS);
265- WaitFlag<HardEvent::V_S>(eventIdVToS);
266- uint32_t compactCastElems = 0U;
267- if constexpr (IS_BF16_MERGE && !EnableMedian) {
268- compactCastElems = ROUND_UP_AGLIN(curInnerChunk * sizeof(SortT)) / sizeof(SortT);
269- Duplicate(compactCastValue_, static_cast<SortT>(0), compactCastElems);
270- event_t eventIdVToSForInit = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
271- SetFlag<HardEvent::V_S>(eventIdVToSForInit);
272- WaitFlag<HardEvent::V_S>(eventIdVToSForInit);
273- }
274- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
275- // Median hook point: k resolution fused with value extraction, tracing back to the raw input tile.
276- asc_vf_call<SimtStoreNonLastMedian<T, SortT>>(
277- dim3(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM), curInnerChunk, this->axisLen_, kthIndex_, medianMode_,
278- this->inputRowElems_, this->valueAxisElems_, this->indexAxisElems_,
279- reinterpret_cast<__ubuf__ T*>(this->inputTile_.GetPhyAddr()),
280- reinterpret_cast<__ubuf__ SortT*>(this->sortedValue_.GetPhyAddr()),
281- reinterpret_cast<__ubuf__ uint32_t*>(this->sortedIndex_.GetPhyAddr()),
282- reinterpret_cast<__ubuf__ T*>(compactValue_.GetPhyAddr()),
283- reinterpret_cast<__ubuf__ int64_t*>(compactIndex_.GetPhyAddr()));
284- eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
285- SetFlag<HardEvent::V_S>(eventIdVToS);
286- WaitFlag<HardEvent::V_S>(eventIdVToS);
287- } else {
288- // Compact one kth element per inner lane. The GM output is [outer, 1, inner], so outputOffset
289- // points to a contiguous inner chunk.
290- for (uint32_t inner = 0; inner < curInnerChunk; ++inner) {
291- uint32_t srcOffset = inner * this->valueAxisElems_ + kthIndex_;
292- uint32_t idxOffset = inner * this->indexAxisElems_ + kthIndex_;
293- if constexpr (IS_BF16_MERGE) {
294- compactCastValue_.SetValue(inner, this->sortedValue_.GetValue(srcOffset));
295- } else {
296- compactValue_.SetValue(inner, this->sortedValue_.GetValue(srcOffset));
297- }
298- compactIndex_.SetValue(inner, static_cast<int64_t>(this->sortedIndex_.GetValue(idxOffset)));
299- }
300- }
301- if constexpr (IS_BF16_MERGE && !EnableMedian) {
302- event_t eventIdSToV = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_V));
303- SetFlag<HardEvent::S_V>(eventIdSToV);
304- WaitFlag<HardEvent::S_V>(eventIdSToV);
305- Cast(compactValue_, compactCastValue_, RoundMode::CAST_RINT, compactCastElems);
306- event_t eventIdVToMte3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_MTE3));
307- SetFlag<HardEvent::V_MTE3>(eventIdVToMte3);
308- WaitFlag<HardEvent::V_MTE3>(eventIdVToMte3);
309- }
310- event_t eventIdSToMte3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_MTE3));
311- SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
312- WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3);
313- DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(curInnerChunk * sizeof(T)), 0, 0, 0};
314- DataCopyPad(valueGm_[outputOffset], compactValue_, valueCopyParam);
315- DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(curInnerChunk * sizeof(int64_t)), 0, 0, 0};
316- DataCopyPad(indexGm_[outputOffset], compactIndex_, indexCopyParam);
317- event_t eventIdMte3ToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_S));
318- SetFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
319- WaitFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
320-}
321- 
322-template <typename T, bool IsDescend, bool UseMergeSort, bool EnableMedian>
323-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, EnableMedian>::StoreTile(
324- int64_t inputOffset, int64_t outputOffset, uint32_t curInnerChunk)
325-{
326- (void)inputOffset;
327- CopyKthToOutput(curInnerChunk, outputOffset);
328-}
329- 
330-} // namespace KthValue
331- 
332-#endif
Dmath/kth_value/op_kernel/arch35/kth_value_radix_one_core.h+0-237
@@ -1,237 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-#ifndef KTH_VALUE_RADIX_ONE_CORE_H
12-#define KTH_VALUE_RADIX_ONE_CORE_H
13- 
14-#include "kernel_tiling/kernel_tiling.h"
15-#include "kernel_operator.h"
16-#include "op_kernel/platform_util.h"
17-#include "kth_value_median_utils.h"
18-#include "kth_value_tiling_data.h"
19-#include "../../sort/arch35/common/signed_zero_sort_utils.h"
20-#include "../../sort/arch35/common/util_type_simd.h"
21- 
22-namespace KthValue {
23-using namespace AscendC;
24- 
25-template <typename T, bool EnableMedian = false>
26-class KthValueRadixOneCore {
27-public:
28- __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
29- TPipe* pipe);
30- __aicore__ inline void Process();
31- 
32-private:
33- __aicore__ inline void ParseTilingData();
34- __aicore__ inline void CopyInputToUb(LocalTensor<T>& xLocal, int64_t row);
35- __aicore__ inline void ProcessOneRow(int64_t row, uint32_t localOffset, LocalTensor<T>& compactValue,
36- LocalTensor<int64_t>& compactIndex);
37- __aicore__ inline void CopyOutputToGm(int64_t rowStart, uint32_t rowCount, LocalTensor<T>& compactValue,
38- LocalTensor<int64_t>& compactIndex);
39- __aicore__ inline void ProcessRows(int64_t rowStart, uint32_t rowCount);
40- 
41- static constexpr SortConfig SORT_CONFIG{SortType::RADIX_SORT, false};
42- static constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize();
43- 
44- GlobalTensor<T> xGm_;
45- GlobalTensor<T> valuesGm_;
46- GlobalTensor<int64_t> indicesGm_;
47- 
48- TPipe* pipe_{nullptr};
49- const KthValueTilingData* tiling_{nullptr};
50- TQue<QuePosition::VECIN, 1> inQueueX_;
51- TQue<QuePosition::VECOUT, 1> outValueQueue_;
52- TQue<QuePosition::VECOUT, 1> outIndexQueue_;
53- TQue<QuePosition::VECOUT, 1> compactValueQueue_;
54- TQue<QuePosition::VECOUT, 1> compactIndexQueue_;
55- TBuf<TPosition::VECCALC> tmpUb_;
56- TBuf<TPosition::VECCALC> sourceOrderBuf_;
57- LocalTensor<uint32_t> sourceOrder_;
58- 
59- uint32_t blockIdx_ = 0;
60- uint32_t blockNum_ = 0;
61- uint32_t numTileData_ = 0;
62- uint32_t kthIndex_ = 0;
63- uint32_t medianMode_ = 0;
64- uint32_t unsortedDimParallel_ = 0;
65- uint32_t xUbSize_ = 0;
66- uint32_t idxUbSize_ = 0;
67- uint32_t bufferNum_ = 1;
68- uint32_t tmpUbSize_ = 0;
69- uint32_t outputRowsPerLoop_ = 1;
70- int64_t lastAxisNum_ = 0;
71- int64_t unsortedDimNum_ = 0;
72-};
73- 
74-template <typename T, bool EnableMedian>
75-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
76- const KthValueTilingData* tiling, TPipe* pipe)
77-{
78- if (tiling == nullptr || pipe == nullptr) {
79- return;
80- }
81- blockIdx_ = GetBlockIdx();
82- blockNum_ = GetBlockNum();
83- pipe_ = pipe;
84- tiling_ = tiling;
85- ParseTilingData();
86- 
87- xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
88- valuesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(values));
89- indicesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(indices));
90- 
91- pipe_->InitBuffer(inQueueX_, bufferNum_, xUbSize_);
92- pipe_->InitBuffer(outValueQueue_, bufferNum_, xUbSize_);
93- pipe_->InitBuffer(outIndexQueue_, bufferNum_, idxUbSize_);
94- pipe_->InitBuffer(compactValueQueue_, bufferNum_, ROUND_UP_AGLIN(outputRowsPerLoop_ * sizeof(T)));
95- pipe_->InitBuffer(compactIndexQueue_, bufferNum_, ROUND_UP_AGLIN(outputRowsPerLoop_ * sizeof(int64_t)));
96- pipe_->InitBuffer(tmpUb_, tmpUbSize_);
97- if constexpr (SignedZeroSortCommon::IS_FLOATING_POINT_V<T>) {
98- pipe_->InitBuffer(sourceOrderBuf_, idxUbSize_);
99- sourceOrder_ = sourceOrderBuf_.Get<uint32_t>();
100- }
101-}
102- 
103-template <typename T, bool EnableMedian>
104-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::ParseTilingData()
105-{
106- numTileData_ = tiling_->numTileDataSize;
107- kthIndex_ = tiling_->kthIndex;
108- medianMode_ = tiling_->medianMode;
109- unsortedDimParallel_ = tiling_->unsortedDimParallel;
110- xUbSize_ = tiling_->keyParams0;
111- idxUbSize_ = tiling_->keyParams1;
112- bufferNum_ = tiling_->keyParams3 == 2U ? 2U : 1U;
113- outputRowsPerLoop_ = tiling_->keyParams4 == 0U ? 1U : tiling_->keyParams4;
114- tmpUbSize_ = tiling_->tmpUbSize;
115- lastAxisNum_ = tiling_->lastAxisNum;
116- unsortedDimNum_ = tiling_->unsortedDimNum;
117-}
118- 
119-template <typename T, bool EnableMedian>
120-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::CopyInputToUb(LocalTensor<T>& xLocal, int64_t row)
121-{
122- int64_t inputOffset = row * lastAxisNum_;
123- uint32_t copyBytes = static_cast<uint32_t>(numTileData_ * sizeof(T));
124- 
125- DataCopyExtParams inputCopyParam{1, copyBytes, 0, 0, 0};
126- DataCopyPadExtParams<T> inputPadParam{true, 0, 0, 0};
127- DataCopyPad(xLocal, xGm_[inputOffset], inputCopyParam, inputPadParam);
128-}
129- 
130-template <typename T, bool EnableMedian>
131-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::ProcessOneRow(int64_t row, uint32_t localOffset,
132- LocalTensor<T>& compactValue,
133- LocalTensor<int64_t>& compactIndex)
134-{
135- LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();
136- CopyInputToUb(xLocal, row);
137- inQueueX_.EnQue<T>(xLocal);
138- xLocal = inQueueX_.DeQue<T>();
139- LocalTensor<T> sortedValue = outValueQueue_.AllocTensor<T>();
140- LocalTensor<uint32_t> sortedIndex = outIndexQueue_.AllocTensor<uint32_t>();
141- LocalTensor<uint8_t> tmpUb = tmpUb_.Get<uint8_t>();
142- uint32_t selectedK = kthIndex_;
143- if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) {
144- if (medianMode_ != MEDIAN_MODE_STATIC) {
145- CanonicalizeNanValues(xLocal, numTileData_);
146- }
147- uint32_t nonNanCount = CountNonNan(xLocal, numTileData_, tmpUb_.Get<float>(), pipe_);
148- selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_);
149- }
150- 
151- bool hasNegativeZero = false;
152- if constexpr (SignedZeroSortCommon::IS_FLOATING_POINT_V<T>) {
153- hasNegativeZero = SignedZeroSortCommon::HasNegativeZeroVec(xLocal, sourceOrder_, numTileData_);
154- if (hasNegativeZero) {
155- SignedZeroSortCommon::PrepareSignedZeroKeysVec(xLocal, sourceOrder_, numTileData_);
156- }
157- AscendC::Sort<T, false, SORT_CONFIG>(sortedValue, sortedIndex, xLocal, tmpUb, numTileData_);
158- if (hasNegativeZero) {
159- SignedZeroSortCommon::RestoreSignedZeroValuesByIndexVec(sortedValue, sortedIndex, sourceOrder_,
160- numTileData_);
161- }
162- } else {
163- AscendC::Sort<T, false, SORT_CONFIG>(sortedValue, sortedIndex, xLocal, tmpUb, numTileData_);
164- }
165- inQueueX_.FreeTensor(xLocal);
166- outValueQueue_.EnQue<T>(sortedValue);
167- outIndexQueue_.EnQue<uint32_t>(sortedIndex);
168- sortedValue = outValueQueue_.DeQue<T>();
169- sortedIndex = outIndexQueue_.DeQue<uint32_t>();
170- event_t eventIdVToS = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
171- SetFlag<HardEvent::V_S>(eventIdVToS);
172- WaitFlag<HardEvent::V_S>(eventIdVToS);
173- 
174- compactValue.SetValue(localOffset, sortedValue.GetValue(selectedK));
175- compactIndex.SetValue(localOffset, static_cast<int64_t>(sortedIndex.GetValue(selectedK)));
176- 
177- outValueQueue_.FreeTensor(sortedValue);
178- outIndexQueue_.FreeTensor(sortedIndex);
179-}
180- 
181-template <typename T, bool EnableMedian>
182-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::CopyOutputToGm(int64_t rowStart, uint32_t rowCount,
183- LocalTensor<T>& compactValue,
184- LocalTensor<int64_t>& compactIndex)
185-{
186- event_t eventIdSToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));
187- SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
188- WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3);
189- 
190- DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(rowCount * sizeof(T)), 0, 0, 0};
191- DataCopyPad(valuesGm_[rowStart], compactValue, valueCopyParam);
192- DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(rowCount * sizeof(int64_t)), 0, 0, 0};
193- DataCopyPad(indicesGm_[rowStart], compactIndex, indexCopyParam);
194-}
195- 
196-template <typename T, bool EnableMedian>
197-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::ProcessRows(int64_t rowStart, uint32_t rowCount)
198-{
199- LocalTensor<T> compactValue = compactValueQueue_.AllocTensor<T>();
200- LocalTensor<int64_t> compactIndex = compactIndexQueue_.AllocTensor<int64_t>();
201- for (uint32_t i = 0; i < rowCount; ++i) {
202- ProcessOneRow(rowStart + i, i, compactValue, compactIndex);
203- }
204- compactValueQueue_.EnQue<T>(compactValue);
205- compactIndexQueue_.EnQue<int64_t>(compactIndex);
206- compactValue = compactValueQueue_.DeQue<T>();
207- compactIndex = compactIndexQueue_.DeQue<int64_t>();
208- CopyOutputToGm(rowStart, rowCount, compactValue, compactIndex);
209- compactValueQueue_.FreeTensor(compactValue);
210- compactIndexQueue_.FreeTensor(compactIndex);
211-}
212- 
213-template <typename T, bool EnableMedian>
214-__aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::Process()
215-{
216- if (blockIdx_ >= blockNum_ || blockIdx_ >= unsortedDimParallel_) {
217- return;
218- }
219- int64_t rowsPerCore = unsortedDimNum_ / static_cast<int64_t>(unsortedDimParallel_);
220- uint32_t tailRows = static_cast<uint32_t>(unsortedDimNum_ % static_cast<int64_t>(unsortedDimParallel_));
221- int64_t blockRows = rowsPerCore + (blockIdx_ < tailRows ? 1 : 0);
222- uint32_t priorTailRows = blockIdx_ < tailRows ? blockIdx_ : tailRows;
223- int64_t blockOffset = static_cast<int64_t>(blockIdx_) * rowsPerCore + static_cast<int64_t>(priorTailRows);
224- 
225- uint64_t blockRows64 = static_cast<uint64_t>(blockRows);
226- uint64_t outputRowsPerLoop64 = static_cast<uint64_t>(outputRowsPerLoop_);
227- uint64_t loopTimes = Ops::Base::CeilDiv(blockRows64, outputRowsPerLoop64);
228- for (uint64_t loopIdx = 0; loopIdx < loopTimes; ++loopIdx) {
229- uint64_t localOffset = loopIdx * outputRowsPerLoop64;
230- uint64_t remainRows = blockRows64 - localOffset;
231- uint32_t rowCount = remainRows < outputRowsPerLoop64 ? static_cast<uint32_t>(remainRows) : outputRowsPerLoop_;
232- ProcessRows(blockOffset + static_cast<int64_t>(localOffset), rowCount);
233- }
234-}
235-} // namespace KthValue
236- 
237-#endif