Star1309
1430
代码
代码
Issues119
Pull Requests249
Actions
项目讨论
Wiki
分析
Star1309
1430
已合并
sort cleancode处理 #4135
黄晓彬创建于 7月17日
sort cleancode处理 #4135
已合并
黄晓彬创建于 7月17日
从ConanHuang/ops-mathcodecheck合入到cann/ops-mathmaster
共 45 个文件变更+1540-1573
Mmath/kth_value/op_api/kth_value.cpp+6-7
@@ -79,8 +79,8 @@ static op::Shape MakeKthShape(const op::Shape& shape, int64_t dim)
79 return outShape;79 return outShape;
80}80}
81 81 
82-static std::tuple<aclTensor*, aclTensor*> KthValueAiCore(82+static std::tuple<aclTensor*, aclTensor*> KthValueAiCore(const aclTensor* self, int64_t k, int64_t dim,
83- const aclTensor* self, int64_t k, int64_t dim, aclOpExecutor* executor)83+ aclOpExecutor* executor)
84{84{
85 auto rank = static_cast<int64_t>(self->GetViewShape().GetDimNum());85 auto rank = static_cast<int64_t>(self->GetViewShape().GetDimNum());
86 auto normDim = dim < 0 ? dim + rank : dim;86 auto normDim = dim < 0 ? dim + rank : dim;
@@ -91,14 +91,13 @@ static std::tuple<aclTensor*, aclTensor*> KthValueAiCore(
91 OP_CHECK_NULL(indices, return {});91 OP_CHECK_NULL(indices, return {});
92 92 
93 auto ret = ADD_TO_LAUNCHER_LIST_AICORE(KthValue, OP_INPUT(self), OP_OUTPUT(values, indices), OP_ATTR(k, dim));93 auto ret = ADD_TO_LAUNCHER_LIST_AICORE(KthValue, OP_INPUT(self), OP_OUTPUT(values, indices), OP_ATTR(k, dim));
94- OP_CHECK(94+ OP_CHECK(ret == ACLNN_SUCCESS, OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "KthValue ADD_TO_LAUNCHER_LIST_AICORE failed."),
95- ret == ACLNN_SUCCESS, OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "KthValue ADD_TO_LAUNCHER_LIST_AICORE failed."),95+ return {});
96- return {});
97 return std::tie(values, indices);96 return std::tie(values, indices);
98}97}
99 98 
100-const std::tuple<aclTensor*, aclTensor*> KthValue(99+const std::tuple<aclTensor*, aclTensor*> KthValue(const aclTensor* self, int64_t k, int64_t dim,
101- const aclTensor* self, int64_t k, int64_t dim, aclOpExecutor* executor)100+ aclOpExecutor* executor)
102{101{
103 L0_DFX(KthValue, self, k, dim);102 L0_DFX(KthValue, self, k, dim);
104 if (!CheckParams(self, k, dim)) {103 if (!CheckParams(self, k, dim)) {
Mmath/kth_value/op_host/arch35/kth_value_tiling_arch35.cpp+31-23
@@ -174,7 +174,8 @@ static void InitKthValueBaseTiling(KthValueTilingData* tilingData, const SortKth
174// =============================================================================174// =============================================================================
175static bool CheckNonLastSmallAxisInput(int64_t axisLen, int64_t outerSize, int64_t innerSize, uint32_t& axisLen32)175static bool CheckNonLastSmallAxisInput(int64_t axisLen, int64_t outerSize, int64_t innerSize, uint32_t& axisLen32)
176{176{
177- if (axisLen < 2 || axisLen > NON_LAST_SMALL_AXIS_THRESHOLD || outerSize <= 0 || innerSize <= 0) {177+ if (axisLen < NON_LAST_SMALL_AXIS_MIN_AXIS_LEN || axisLen > NON_LAST_SMALL_AXIS_THRESHOLD || outerSize <= 0 ||
178+ innerSize <= 0) {
178 return false;179 return false;
179 }180 }
180 axisLen32 = static_cast<uint32_t>(axisLen);181 axisLen32 = static_cast<uint32_t>(axisLen);
@@ -283,12 +284,12 @@ static ge::graphStatus SetRadixOneCoreTiling(gert::TilingContext* context, const
283 std::to_string(tilingData->tmpUbSize).c_str(),284 std::to_string(tilingData->tmpUbSize).c_str(),
284 "The value of tmpUbSize must be less than or equal to remainUb."),285 "The value of tmpUbSize must be less than or equal to remainUb."),
285 return ge::GRAPH_FAILED);286 return ge::GRAPH_FAILED);
286- uint64_t doubleBufferRemainUb = info.ubSize > info.oneBufferQueSize * 2 ?287+ uint64_t doubleBufferRemainUb = info.ubSize > info.oneBufferQueSize * DOUBLE_BUFFER_NUM ?
287- (info.ubSize - info.oneBufferQueSize * 2) / info.blockUbSize *288+ (info.ubSize - info.oneBufferQueSize * DOUBLE_BUFFER_NUM) / info.blockUbSize *
288 info.blockUbSize :289 info.blockUbSize :
289 0;290 0;
290 if (static_cast<uint64_t>(tilingData->tmpUbSize) <= doubleBufferRemainUb) {291 if (static_cast<uint64_t>(tilingData->tmpUbSize) <= doubleBufferRemainUb) {
291- tilingData->keyParams3 = 2;292+ tilingData->keyParams3 = DOUBLE_BUFFER_NUM;
292 }293 }
293 return ge::GRAPH_SUCCESS;294 return ge::GRAPH_SUCCESS;
294}295}
@@ -924,6 +925,31 @@ static void SetKthValueTilingContext(gert::TilingContext* context, uint64_t schI
924 }925 }
925}926}
926 927 
928+static ge::graphStatus SelectAndFinalizeKthValueRoute(gert::TilingContext* context,
929+ const platform_ascendc::PlatformAscendC& ascendcPlatform,
930+ SortKthTileInfo& info, KthValueTilingData* tilingData)
931+{
932+ KthValueTilingData candidateTilingData = *tilingData;
933+ uint64_t schId = KTH_VALUE_SCHID_RADIX_MORE_CORE;
934+ uint32_t blockDim = 1;
935+ OP_CHECK_IF((SelectKthValueRoute(context, info, &candidateTilingData, blockDim, schId) != ge::GRAPH_SUCCESS),
936+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "SelectKthValueRoute", "GRAPH_FAILED",
937+ "The value of SelectKthValueRoute must be GRAPH_SUCCESS."),
938+ return ge::GRAPH_FAILED);
939+ OP_CHECK_IF((FinalizeKthValueRoute(context, ascendcPlatform, info, &candidateTilingData, schId, blockDim) !=
940+ ge::GRAPH_SUCCESS),
941+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "FinalizeKthValueRoute", "GRAPH_FAILED",
942+ "The value of FinalizeKthValueRoute must be GRAPH_SUCCESS."),
943+ return ge::GRAPH_FAILED);
944+ *tilingData = candidateTilingData;
945+ OP_LOGI(context->GetNodeName(),
946+ "KthValueTiling: schId=%lu, blockDim=%u, lastAxis=%ld, unsortedDim=%ld, "
947+ "isNonLastAxis=%d, dtypeSize=%u",
948+ schId, blockDim, info.lastAxis, info.unsortedDim, static_cast<int>(info.isNonLastAxis), info.dtypeSize);
949+ SetKthValueTilingContext(context, schId, info, blockDim);
950+ return ge::GRAPH_SUCCESS;
951+}
952+ 
927// =============================================================================953// =============================================================================
928// Main entry954// Main entry
929// =============================================================================955// =============================================================================
@@ -968,25 +994,7 @@ static ge::graphStatus Tiling4KthValue(gert::TilingContext* context)
968 "The value of ComputeKthValueUbInfo must be GRAPH_SUCCESS."),994 "The value of ComputeKthValueUbInfo must be GRAPH_SUCCESS."),
969 return ge::GRAPH_FAILED);995 return ge::GRAPH_FAILED);
970 InitKthValueBaseTiling(tilingData, info, oneCoreUbValid, *kAttr - 1);996 InitKthValueBaseTiling(tilingData, info, oneCoreUbValid, *kAttr - 1);
971- KthValueTilingData candidateTilingData = *tilingData;997+ return SelectAndFinalizeKthValueRoute(context, ascendcPlatform, info, tilingData);
972- uint64_t schId = KTH_VALUE_SCHID_RADIX_MORE_CORE;
973- uint32_t blockDim = 1;
974- OP_CHECK_IF((SelectKthValueRoute(context, info, &candidateTilingData, blockDim, schId) != ge::GRAPH_SUCCESS),
975- OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "SelectKthValueRoute", "GRAPH_FAILED",
976- "The value of SelectKthValueRoute must be GRAPH_SUCCESS."),
977- return ge::GRAPH_FAILED);
978- OP_CHECK_IF((FinalizeKthValueRoute(context, ascendcPlatform, info, &candidateTilingData, schId, blockDim) !=
979- ge::GRAPH_SUCCESS),
980- OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "FinalizeKthValueRoute", "GRAPH_FAILED",
981- "The value of FinalizeKthValueRoute must be GRAPH_SUCCESS."),
982- return ge::GRAPH_FAILED);
983- *tilingData = candidateTilingData;
984- OP_LOGI(context->GetNodeName(),
985- "KthValueTiling: schId=%lu, blockDim=%u, lastAxis=%ld, unsortedDim=%ld, "
986- "isNonLastAxis=%d, dtypeSize=%u",
987- schId, blockDim, info.lastAxis, info.unsortedDim, static_cast<int>(info.isNonLastAxis), info.dtypeSize);
988- SetKthValueTilingContext(context, schId, info, blockDim);
989- return ge::GRAPH_SUCCESS;
990}998}
991 999 
992static ge::graphStatus TilingPrepare4KthValue(gert::TilingParseContext* context)1000static ge::graphStatus TilingPrepare4KthValue(gert::TilingParseContext* context)
Mmath/kth_value/op_kernel/arch35/kth_value_merge_intra_core.h+16-19
@@ -47,17 +47,16 @@ using MergeSortConstants::UB_BLOCK_BYTES;
47 */47 */
48template <typename ValueType, typename IndexType, bool IsDescend>48template <typename ValueType, typename IndexType, bool IsDescend>
49class KthValueMergeIntraCore49class KthValueMergeIntraCore
50- : public MergeIntraCoreCommon::MergeIntraCoreBase<50+ : public MergeIntraCoreCommon::MergeIntraCoreBase<KthValueMergeIntraCore<ValueType, IndexType, IsDescend>,
51- KthValueMergeIntraCore<ValueType, IndexType, IsDescend>, ValueType, IndexType, IsDescend> {51+ ValueType, IndexType, IsDescend> {
52- using Base = MergeIntraCoreCommon::MergeIntraCoreBase<52+ using Base = MergeIntraCoreCommon::MergeIntraCoreBase<KthValueMergeIntraCore<ValueType, IndexType, IsDescend>,
53- KthValueMergeIntraCore<ValueType, IndexType, IsDescend>, ValueType, IndexType, IsDescend>;53+ ValueType, IndexType, IsDescend>;
54 friend Base;54 friend Base;
55 55 
56public:56public:
57 __aicore__ inline KthValueMergeIntraCore() {}57 __aicore__ inline KthValueMergeIntraCore() {}
58- __aicore__ inline void Init(58+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace,
59- GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData,59+ const KthValueTilingData* tilingData, TPipe* pipe);
60- TPipe* pipe);
61 60 
62private:61private:
63 // KthValue-specific members62 // KthValue-specific members
@@ -68,9 +67,8 @@ private:
68 __aicore__ inline void InitPhase2Buffers();67 __aicore__ inline void InitPhase2Buffers();
69 __aicore__ inline void InitPhase3Buffers();68 __aicore__ inline void InitPhase3Buffers();
70 __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion);69 __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion);
71- __aicore__ inline void ExtractAndCopyChunk(70+ __aicore__ inline void ExtractAndCopyChunk(int64_t cacheBatchOffset, uint32_t cacheOffset, int64_t outputOffset,
72- int64_t cacheBatchOffset, uint32_t cacheOffset, int64_t outputOffset, uint32_t elemProcessed,71+ uint32_t elemProcessed, uint32_t elemCount);
73- uint32_t elemCount);
74};72};
75 73 
76template <typename ValueType, typename IndexType, bool IsDescend>74template <typename ValueType, typename IndexType, bool IsDescend>
@@ -112,16 +110,16 @@ __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend>::
112 this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType);110 this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType);
113 this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE;111 this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE;
114 this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE;112 this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE;
115- this->lastBlockSize_ = static_cast<uint32_t>(113+ this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ -
116- this->sortAxisNum_ - static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_);114+ static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_);
117 115 
118 // Cache stores sort struct data (8 bytes per element: index + value)116 // Cache stores sort struct data (8 bytes per element: index + value)
119 // Each core has its own cache region, reused across batches117 // Each core has its own cache region, reused across batches
120 // perBatchCacheLen: sort struct length for one batch (with ping-pong, in ValueType units)118 // perBatchCacheLen: sort struct length for one batch (with ping-pong, in ValueType units)
121 int64_t perCoreCacheLen = static_cast<int64_t>(this->batchSortLen_) * 2; // ping-pong, reused per batch119 int64_t perCoreCacheLen = static_cast<int64_t>(this->batchSortLen_) * 2; // ping-pong, reused per batch
122 120 
123- this->cacheGm_.SetGlobalBuffer(121+ this->cacheGm_.SetGlobalBuffer((__gm__ ValueType*)workspace +
124- (__gm__ ValueType*)workspace + static_cast<int64_t>(this->blockIdx_) * perCoreCacheLen);122+ static_cast<int64_t>(this->blockIdx_) * perCoreCacheLen);
125}123}
126 124 
127template <typename ValueType, typename IndexType, bool IsDescend>125template <typename ValueType, typename IndexType, bool IsDescend>
@@ -156,8 +154,8 @@ __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend>::
156}154}
157 155 
158template <typename ValueType, typename IndexType, bool IsDescend>156template <typename ValueType, typename IndexType, bool IsDescend>
159-__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut(157+__aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut(int64_t batchIdx,
160- int64_t batchIdx, uint32_t resultRegion)158+ uint32_t resultRegion)
161{159{
162 int64_t outputOffset = batchIdx;160 int64_t outputOffset = batchIdx;
163 161 
@@ -205,9 +203,8 @@ __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend>::
205 203 
206 // Flip back sign bit for ascending order (was flipped in SortBlockToStruct)204 // Flip back sign bit for ascending order (was flipped in SortBlockToStruct)
207 if constexpr (!IsDescend) {205 if constexpr (!IsDescend) {
208- Adds(206+ Adds(valueLocal.template ReinterpretCast<int32_t>(), valueLocal.template ReinterpretCast<int32_t>(), 0x80000000,
209- valueLocal.template ReinterpretCast<int32_t>(), valueLocal.template ReinterpretCast<int32_t>(), 0x80000000,207+ elemCount);
210- elemCount);
211 }208 }
212 209 
213 this->outValueQueue_.EnQue(valueLocal);210 this->outValueQueue_.EnQue(valueLocal);
Mmath/kth_value/op_kernel/arch35/kth_value_merge_sort_more_core.h+19-26
@@ -36,16 +36,15 @@ using MergeSortConstants::XOR_OP_VALUE_HALF;
36 36 
37template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>37template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
38struct KthValueMergeSortMoreCore38struct KthValueMergeSortMoreCore
39- : public MergeMoreCoreCommon::MergeMoreCoreBase<39+ : public MergeMoreCoreCommon::MergeMoreCoreBase<KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>,
40- KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> {40+ T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> {
41 using Base = MergeMoreCoreCommon::MergeMoreCoreBase<41 using Base = MergeMoreCoreCommon::MergeMoreCoreBase<
42 KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>;42 KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>;
43 friend Base;43 friend Base;
44 44 
45 __aicore__ inline KthValueMergeSortMoreCore() {}45 __aicore__ inline KthValueMergeSortMoreCore() {}
46- __aicore__ inline void Init(46+ __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace,
47- GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const KthValueTilingData* tilingData,47+ const KthValueTilingData* tilingData, TPipe* pipe);
48- TPipe* pipe);
49 __aicore__ inline void InitMergeBuffers();48 __aicore__ inline void InitMergeBuffers();
50 __aicore__ inline void ExtractAndCopyOut();49 __aicore__ inline void ExtractAndCopyOut();
51 50 
@@ -76,18 +75,17 @@ __aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, IN
76 this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_);75 this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_);
77 // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per76 // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per
78 // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.77 // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.
79- uint64_t rowWorkspaceBytes =78+ uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) *
80- ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) * sortBufferSize);79+ sortBufferSize);
81 uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE);80 uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE);
82- this->rowWorkspaceOffset_ =81+ this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2;
83- static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2;
84 this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE;82 this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE;
85 83 
86 this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));84 this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));
87 this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value));85 this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value));
88 this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));86 this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));
89- this->workspaceGm_[0].SetGlobalBuffer(87+ this->workspaceGm_[0].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_,
90- (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_, rowWorkspaceElements);88+ rowWorkspaceElements);
91 this->workspaceGm_[1].SetGlobalBuffer(89 this->workspaceGm_[1].SetGlobalBuffer(
92 (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements);90 (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements);
93 91 
@@ -106,18 +104,14 @@ template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYP
106__aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::InitMergeBuffers()104__aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::InitMergeBuffers()
107{105{
108 uint32_t sortBufferSize = 8;106 uint32_t sortBufferSize = 8;
109- this->pipe_->InitBuffer(107+ this->pipe_->InitBuffer(this->sortedQueue_, MERGE_MORE_BUFFER_NUM,
110- this->sortedQueue_, MERGE_MORE_BUFFER_NUM,108+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);
111- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);109+ this->pipe_->InitBuffer(this->copyInQueue_, MERGE_MORE_BUFFER_NUM,
112- this->pipe_->InitBuffer(110+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);
113- this->copyInQueue_, MERGE_MORE_BUFFER_NUM,111+ this->pipe_->InitBuffer(this->castValueQueue_, MERGE_MORE_BUFFER_NUM,
114- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);112+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE));
115- this->pipe_->InitBuffer(113+ this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM,
116- this->castValueQueue_, MERGE_MORE_BUFFER_NUM,114+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(uint32_t));
117- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE));
118- this->pipe_->InitBuffer(
119- this->castIndexQueue_, MERGE_MORE_BUFFER_NUM,
120- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(uint32_t));
121 this->pipe_->InitBuffer(this->outValueQueue_, 1, UB_BLOCK_BYTES);115 this->pipe_->InitBuffer(this->outValueQueue_, 1, UB_BLOCK_BYTES);
122 this->pipe_->InitBuffer(this->outIndexQueue_, 1, UB_BLOCK_BYTES);116 this->pipe_->InitBuffer(this->outIndexQueue_, 1, UB_BLOCK_BYTES);
123}117}
@@ -128,9 +122,8 @@ __aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, IN
128 LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template DeQue<CONVERT_TYPE>();122 LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template DeQue<CONVERT_TYPE>();
129 LocalTensor<CONVERT_TYPE> castValue = this->castValueQueue_.template AllocTensor<CONVERT_TYPE>();123 LocalTensor<CONVERT_TYPE> castValue = this->castValueQueue_.template AllocTensor<CONVERT_TYPE>();
130 LocalTensor<uint32_t> castIndex = this->castIndexQueue_.template AllocTensor<uint32_t>();124 LocalTensor<uint32_t> castIndex = this->castIndexQueue_.template AllocTensor<uint32_t>();
131- AscendC::Extract(125+ AscendC::Extract(castValue, castIndex, sortTempBuffer,
132- castValue, castIndex, sortTempBuffer,126+ ((this->curLoopSortedNum_ + DEALING_EXTRACT_NUM_ONCE - 1) / DEALING_EXTRACT_NUM_ONCE));
133- ((this->curLoopSortedNum_ + DEALING_EXTRACT_NUM_ONCE - 1) / DEALING_EXTRACT_NUM_ONCE));
134 if constexpr (!IS_DESCEND) {127 if constexpr (!IS_DESCEND) {
135 this->FlipSignBit(castValue, ROUND_UP_AGLIN(this->curLoopSortedNum_));128 this->FlipSignBit(castValue, ROUND_UP_AGLIN(this->curLoopSortedNum_));
136 }129 }
Mmath/kth_value/op_kernel/arch35/kth_value_merge_sort_one_core.h+21-22
@@ -31,8 +31,8 @@ constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize();
31template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis = 0>31template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis = 0>
32class KthValueMergeSortOneCore {32class KthValueMergeSortOneCore {
33public:33public:
34- __aicore__ inline void Init(34+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
35- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe);35+ TPipe* pipe);
36 __aicore__ inline void Process();36 __aicore__ inline void Process();
37 37 
38private:38private:
@@ -40,9 +40,8 @@ private:
40 __aicore__ inline void InitIndexLocal();40 __aicore__ inline void InitIndexLocal();
41 __aicore__ inline void CopyDataIn(uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum);41 __aicore__ inline void CopyDataIn(uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum);
42 __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offset, uint32_t count);42 __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offset, uint32_t count);
43- __aicore__ inline void SortRows(43+ __aicore__ inline void SortRows(LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal,
44- LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal, LocalTensor<uint32_t> sortedIndexLocal,44+ LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum);
45- uint32_t rowNum);
46 __aicore__ inline void CopyKthToGm(uint64_t outputOffset, uint32_t rowNum);45 __aicore__ inline void CopyKthToGm(uint64_t outputOffset, uint32_t rowNum);
47 __aicore__ inline void ProcessSingleRound(uint32_t round);46 __aicore__ inline void ProcessSingleRound(uint32_t round);
48 47 
@@ -140,15 +139,16 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx
140 for (uint16_t i = 0; i < repeatTime; ++i) {139 for (uint16_t i = 0; i < repeatTime; ++i) {
141 MicroAPI::MaskReg mask = MicroAPI::UpdateMask<uint32_t>(alignSizeCopy);140 MicroAPI::MaskReg mask = MicroAPI::UpdateMask<uint32_t>(alignSizeCopy);
142 MicroAPI::Adds(indexTensor, vciTensor, i * vfLenB32, mask);141 MicroAPI::Adds(indexTensor, vciTensor, i * vfLenB32, mask);
143- MicroAPI::DataCopy<int32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(142+ MicroAPI::DataCopy<int32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(indexValuePtr, indexTensor, vfLenB32,
144- indexValuePtr, indexTensor, vfLenB32, mask);143+ mask);
145 }144 }
146 }145 }
147}146}
148 147 
149template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis>148template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis>
150-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis>::CopyDataIn(149+__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis>::CopyDataIn(uint64_t tileOffset,
151- uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum)150+ uint32_t currTileSize,
151+ uint32_t rowNum)
152{152{
153 LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();153 LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();
154 Duplicate(xLocal, static_cast<T>(NAN), alignSize_ * rowNum);154 Duplicate(xLocal, static_cast<T>(NAN), alignSize_ * rowNum);
@@ -156,13 +156,13 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx
156 SetFlag<HardEvent::V_MTE2>(eventId);156 SetFlag<HardEvent::V_MTE2>(eventId);
157 WaitFlag<HardEvent::V_MTE2>(eventId);157 WaitFlag<HardEvent::V_MTE2>(eventId);
158 158 
159- uint32_t currTileSizeAlign =159+ uint32_t currTileSizeAlign = (currTileSize * sizeof(T) + UB_BLOCK_SIZE - 1U) / UB_BLOCK_SIZE * UB_BLOCK_SIZE /
160- (currTileSize * sizeof(T) + UB_BLOCK_SIZE - 1U) / UB_BLOCK_SIZE * UB_BLOCK_SIZE / sizeof(T);160+ sizeof(T);
161 uint32_t dstStride = ((alignSize_ - currTileSizeAlign) * sizeof(T)) / UB_BLOCK_SIZE;161 uint32_t dstStride = ((alignSize_ - currTileSizeAlign) * sizeof(T)) / UB_BLOCK_SIZE;
162- DataCopyPadExtParams<T> padParams{162+ DataCopyPadExtParams<T> padParams{true, 0, static_cast<uint8_t>(currTileSizeAlign - currTileSize),
163- true, 0, static_cast<uint8_t>(currTileSizeAlign - currTileSize), static_cast<T>(NAN)};163+ static_cast<T>(NAN)};
164- DataCopyExtParams copyParam{164+ DataCopyExtParams copyParam{static_cast<uint16_t>(rowNum), static_cast<uint32_t>(currTileSize * sizeof(T)), 0,
165- static_cast<uint16_t>(rowNum), static_cast<uint32_t>(currTileSize * sizeof(T)), 0, dstStride, 0};165+ dstStride, 0};
166 DataCopyPad(xLocal, xGm_[tileOffset], copyParam, padParams);166 DataCopyPad(xLocal, xGm_[tileOffset], copyParam, padParams);
167 inQueueX_.EnQue<T>(xLocal);167 inQueueX_.EnQue<T>(xLocal);
168}168}
@@ -210,9 +210,8 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx
210 Concat(concatLocal, xSortLocal[offset], concatTmp, concatRepeatTimes);210 Concat(concatLocal, xSortLocal[offset], concatTmp, concatRepeatTimes);
211 AscendC::Sort<CONVERT_TYPE, true>(sortedLocal, concatLocal, indexLocal_, sortTmp, sortRepeatTimes);211 AscendC::Sort<CONVERT_TYPE, true>(sortedLocal, concatLocal, indexLocal_, sortTmp, sortRepeatTimes);
212 }212 }
213- Extract(213+ Extract(sortedValueCast[offset], sortedIndexLocal[offset], sortedLocal,
214- sortedValueCast[offset], sortedIndexLocal[offset], sortedLocal,214+ isSort32SmallAxis == 1 ? 1 : sortRepeatTimes);
215- isSort32SmallAxis == 1 ? 1 : sortRepeatTimes);
216 FlipSignBit(sortedValueCast, offset, alignSize_);215 FlipSignBit(sortedValueCast, offset, alignSize_);
217 }216 }
218 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {217 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
@@ -221,8 +220,8 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx
221}220}
222 221 
223template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis>222template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis>
224-__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis>::CopyKthToGm(223+__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis>::CopyKthToGm(uint64_t outputOffset,
225- uint64_t outputOffset, uint32_t rowNum)224+ uint32_t rowNum)
226{225{
227 LocalTensor<T> sortedValueLocal = outValueQueue_.DeQue<T>();226 LocalTensor<T> sortedValueLocal = outValueQueue_.DeQue<T>();
228 LocalTensor<uint32_t> sortedIndexLocal = outIndexQueue_.DeQue<uint32_t>();227 LocalTensor<uint32_t> sortedIndexLocal = outIndexQueue_.DeQue<uint32_t>();
@@ -258,8 +257,8 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx
258template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis>257template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis>
259__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis>::ProcessSingleRound(uint32_t round)258__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis>::ProcessSingleRound(uint32_t round)
260{259{
261- int64_t rowStart =260+ int64_t rowStart = (static_cast<int64_t>(blockIdx_) + static_cast<int64_t>(round) * unsortedDimParallel_) *
262- (static_cast<int64_t>(blockIdx_) + static_cast<int64_t>(round) * unsortedDimParallel_) * oneCoreRowNum_;261+ oneCoreRowNum_;
263 if (rowStart >= unsortedDimNum_) {262 if (rowStart >= unsortedDimNum_) {
264 return;263 return;
265 }264 }
Mmath/kth_value/op_kernel/arch35/kth_value_non_last_small_axis.h+11-9
@@ -44,9 +44,8 @@ class KthValueNonLastSmallAxis
44 static constexpr bool IS_BF16_MERGE = UseMergeSort && std::is_same_v<T, bfloat16_t>;44 static constexpr bool IS_BF16_MERGE = UseMergeSort && std::is_same_v<T, bfloat16_t>;
45 45 
46public:46public:
47- __aicore__ inline void Init(47+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace,
48- GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData,48+ const KthValueTilingData* tilingData, TPipe* pipe);
49- TPipe* pipe);
50 49 
51 __aicore__ inline void StoreTile(int64_t inputOffset, int64_t outputOffset, uint32_t curInnerChunk);50 __aicore__ inline void StoreTile(int64_t inputOffset, int64_t outputOffset, uint32_t curInnerChunk);
52 51 
@@ -70,8 +69,10 @@ private:
70};69};
71 70 
72template <typename T, bool IsDescend, bool UseMergeSort>71template <typename T, bool IsDescend, bool UseMergeSort>
73-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::Init(72+__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::Init(GM_ADDR x, GM_ADDR values,
74- GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData, TPipe* pipe)73+ GM_ADDR indices, GM_ADDR workspace,
74+ const KthValueTilingData* tilingData,
75+ TPipe* pipe)
75{76{
76 (void)workspace;77 (void)workspace;
77 if (tilingData == nullptr || pipe == nullptr) {78 if (tilingData == nullptr || pipe == nullptr) {
@@ -140,8 +141,8 @@ __aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::Par
140}141}
141 142 
142template <typename T, bool IsDescend, bool UseMergeSort>143template <typename T, bool IsDescend, bool UseMergeSort>
143-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::CopyKthToOutput(144+__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::CopyKthToOutput(uint32_t curInnerChunk,
144- uint32_t curInnerChunk, int64_t outputOffset)145+ int64_t outputOffset)
145{146{
146 event_t eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));147 event_t eventIdVToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S));
147 SetFlag<HardEvent::V_S>(eventIdVToS);148 SetFlag<HardEvent::V_S>(eventIdVToS);
@@ -188,8 +189,9 @@ __aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::Cop
188}189}
189 190 
190template <typename T, bool IsDescend, bool UseMergeSort>191template <typename T, bool IsDescend, bool UseMergeSort>
191-__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::StoreTile(192+__aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort>::StoreTile(int64_t inputOffset,
192- int64_t inputOffset, int64_t outputOffset, uint32_t curInnerChunk)193+ int64_t outputOffset,
194+ uint32_t curInnerChunk)
193{195{
194 (void)inputOffset;196 (void)inputOffset;
195 CopyKthToOutput(curInnerChunk, outputOffset);197 CopyKthToOutput(curInnerChunk, outputOffset);
Mmath/kth_value/op_kernel/arch35/kth_value_radix_more_core.h+54-51
@@ -49,45 +49,47 @@ struct KthValueRadixMoreInnerTilingData {
49// T1输入x dtype T2输出Idx dtype UT无符号的数据类型49// T1输入x dtype T2输出Idx dtype UT无符号的数据类型
50template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>50template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
51class KthValueRadixMoreInnerCore51class KthValueRadixMoreInnerCore
52- : public RadixSortCommon::RadixMoreCoreBase<52+ : public RadixSortCommon::RadixMoreCoreBase<KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>, T1, T2, UT, T3,
53- KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>, T1, T2, UT, T3, isDescend> {53+ isDescend> {
54- using Base = RadixSortCommon::RadixMoreCoreBase<54+ using Base = RadixSortCommon::RadixMoreCoreBase<KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>, T1, T2, UT,
55- KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>, T1, T2, UT, T3, isDescend>;55+ T3, isDescend>;
56 friend Base;56 friend Base;
57 57 
58public:58public:
59 __aicore__ inline KthValueRadixMoreInnerCore(){};59 __aicore__ inline KthValueRadixMoreInnerCore(){};
60- __aicore__ inline void Init(60+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,
61- GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,61+ const KthValueRadixMoreInnerTilingData* __restrict tilingData, TPipe* pipe);
62- const KthValueRadixMoreInnerTilingData* __restrict tilingData, TPipe* pipe);
63 __aicore__ inline void SetKthOutput(GM_ADDR value, GM_ADDR index, T3 kthIndex);62 __aicore__ inline void SetKthOutput(GM_ADDR value, GM_ADDR index, T3 kthIndex);
64 63 
65protected:64protected:
66 __aicore__ inline void ParserTilingData();65 __aicore__ inline void ParserTilingData();
67- __aicore__ inline void ScatterKeysGlobal(66+ __aicore__ inline void ScatterKeysGlobal(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
68- LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,67+ LocalTensor<uint32_t> xInputIndexLocal,
69- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,68+ LocalTensor<uint8_t> sortedValueLocal,
70- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,69+ LocalTensor<uint16_t> blockExcusiveSum,
71- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,70+ LocalTensor<T3> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag,
72- uint32_t cureTileSize, uint32_t sortLoopRound);71+ LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,
73- __aicore__ inline void ScatterOutInt32(72+ uint32_t cureTileSize, uint32_t sortLoopRound);
74- LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,73+ __aicore__ inline void ScatterOutInt32(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
75- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,74+ LocalTensor<uint32_t> xInputIndexLocal,
76- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,75+ LocalTensor<uint8_t> sortedValueLocal,
77- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,76+ LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
78- uint32_t cureTileSize, uint32_t sortLoopRound);77+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
78+ uint32_t round, T3 tileDataStart, uint32_t cureTileSize,
79+ uint32_t sortLoopRound);
79 __aicore__ inline void ScatterOutInt32ToInt64(80 __aicore__ inline void ScatterOutInt32ToInt64(
80 LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,81 LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
81 LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,82 LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,
82 LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<uint32_t> blockDataInGlobalPos,83 LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<uint32_t> blockDataInGlobalPos,
83 LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,84 LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,
84 uint32_t cureTileSize, uint32_t sortLoopRound);85 uint32_t cureTileSize, uint32_t sortLoopRound);
85- __aicore__ inline void ScatterOutInt64(86+ __aicore__ inline void ScatterOutInt64(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
86- LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,87+ LocalTensor<uint32_t> xInputIndexLocal,
87- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,88+ LocalTensor<uint8_t> sortedValueLocal,
88- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,89+ LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
89- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,90+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
90- uint32_t cureTileSize, uint32_t sortLoopRound);91+ uint32_t round, T3 tileDataStart, uint32_t cureTileSize,
92+ uint32_t sortLoopRound);
91 93 
92 const KthValueRadixMoreInnerTilingData* tilingData_;94 const KthValueRadixMoreInnerTilingData* tilingData_;
93 GlobalTensor<T1> kthValueGm_;95 GlobalTensor<T1> kthValueGm_;
@@ -173,8 +175,8 @@ __aicore__ inline void KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>::In
173}175}
174 176 
175template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>177template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
176-__aicore__ inline void KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>::SetKthOutput(178+__aicore__ inline void KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>::SetKthOutput(GM_ADDR value, GM_ADDR index,
177- GM_ADDR value, GM_ADDR index, T3 kthIndex)179+ T3 kthIndex)
178{180{
179 kthValueGm_.SetGlobalBuffer((__gm__ T1*)value);181 kthValueGm_.SetGlobalBuffer((__gm__ T1*)value);
180 kthIndexGm_.SetGlobalBuffer((__gm__ T2*)index);182 kthIndexGm_.SetGlobalBuffer((__gm__ T2*)index);
@@ -223,12 +225,13 @@ __aicore__ inline void KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>::Pa
223 * @param xInputIndexLocalAddr Original indices carried from the previous radix round.225 * @param xInputIndexLocalAddr Original indices carried from the previous radix round.
224 */226 */
225template <typename T1, typename T2, typename T3, typename KthIdxT, int32_t round>227template <typename T1, typename T2, typename T3, typename KthIdxT, int32_t round>
226-__simt_vf__ LAUNCH_BOUND(THREAD_DIM_NUM) __aicore__ void CopyOutKthGm(228+__simt_vf__ LAUNCH_BOUND(THREAD_DIM_NUM) __aicore__
227- T3 tileDataStart, uint64_t unSortIdOffset, T3 kthIndex, uint32_t outputRow, __ubuf__ uint16_t* blockExcusiveSumAddr,229+ void CopyOutKthGm(T3 tileDataStart, uint64_t unSortIdOffset, T3 kthIndex, uint32_t outputRow,
228- __gm__ volatile T3* excusiveBinsGmAddr, __ubuf__ T3* blockDataInGlobalPosAddr,230+ __ubuf__ uint16_t* blockExcusiveSumAddr, __gm__ volatile T3* excusiveBinsGmAddr,
229- __ubuf__ uint32_t* sortedIndexLocalAddr, __ubuf__ T3* xInputIndexLocalAddr, __ubuf__ T1* xInputValueLocalAddr,231+ __ubuf__ T3* blockDataInGlobalPosAddr, __ubuf__ uint32_t* sortedIndexLocalAddr,
230- __ubuf__ T3* blockHistFlagAddr, __ubuf__ uint16_t* blockHistAddr, __gm__ volatile T1* kthValueGmAddr,232+ __ubuf__ T3* xInputIndexLocalAddr, __ubuf__ T1* xInputValueLocalAddr,
231- __gm__ volatile KthIdxT* kthIndexGmAddr)233+ __ubuf__ T3* blockHistFlagAddr, __ubuf__ uint16_t* blockHistAddr,
234+ __gm__ volatile T1* kthValueGmAddr, __gm__ volatile KthIdxT* kthIndexGmAddr)
232{235{
233 // Phase 1: compute the global scatter base for each bucket in this tile. The formula is identical to the236 // Phase 1: compute the global scatter base for each bucket in this tile. The formula is identical to the
234 // regular radix scatter path, but only the bucket containing kthIndex will be consumed below.237 // regular radix scatter path, but only the bucket containing kthIndex will be consumed below.
@@ -448,10 +451,10 @@ __aicore__ inline void KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>::Sc
448 if constexpr (sizeof(T1) == sizeof(int8_t)) {451 if constexpr (sizeof(T1) == sizeof(int8_t)) {
449 // int8时只循环一次,所以scatter时肯定要按照输出数据类型452 // int8时只循环一次,所以scatter时肯定要按照输出数据类型
450 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;453 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;
451- uint64_t outputXUnsortedAxisOffset =454+ uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) *
452- static_cast<uint64_t>(unSortId) * static_cast<uint64_t>(this->totalDataNum_);455+ static_cast<uint64_t>(this->totalDataNum_);
453- uint64_t unSortIdOffset =456+ uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) +
454- static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;457+ round * RADIX_SORT_NUM;
455 uint64_t outputRow = static_cast<uint64_t>(sortLoopRound) * this->unsortedDimParallel_ + unSortId;458 uint64_t outputRow = static_cast<uint64_t>(sortLoopRound) * this->unsortedDimParallel_ + unSortId;
456 bool writeKthOutput = writeKthOutput_ && outputRow < static_cast<uint64_t>(this->unsortedDimNum_);459 bool writeKthOutput = writeKthOutput_ && outputRow < static_cast<uint64_t>(this->unsortedDimNum_);
457 if constexpr (IsSameType<T3, uint32_t>::value) {460 if constexpr (IsSameType<T3, uint32_t>::value) {
@@ -504,19 +507,19 @@ __aicore__ inline void KthValueRadixMoreInnerCore<T1, T2, UT, T3, isDescend>::Sc
504 }507 }
505 } else if constexpr (sizeof(T2) == sizeof(int32_t)) {508 } else if constexpr (sizeof(T2) == sizeof(int32_t)) {
506 // 输出idx本省就是int32，无需cast509 // 输出idx本省就是int32，无需cast
507- ScatterOutInt32(510+ ScatterOutInt32(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,
508- xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,511+ blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize,
509- blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize, sortLoopRound);512+ sortLoopRound);
510 } else if constexpr (IsSameType<T3, uint32_t>::value) {513 } else if constexpr (IsSameType<T3, uint32_t>::value) {
511 // 输出idx是int64，需要在最后一次scatter时cast为int64514 // 输出idx是int64，需要在最后一次scatter时cast为int64
512- ScatterOutInt32ToInt64(515+ ScatterOutInt32ToInt64(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,
513- xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,516+ blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize,
514- blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize, sortLoopRound);517+ sortLoopRound);
515 } else {518 } else {
516 // 计算过程中idx使用int64519 // 计算过程中idx使用int64
517- ScatterOutInt64(520+ ScatterOutInt64(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,
518- xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,521+ blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize,
519- blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize, sortLoopRound);522+ sortLoopRound);
520 }523 }
521}524}
522 525 
@@ -533,9 +536,8 @@ template <typename T, typename T3, typename KeyT>
533class KthValueRadixMoreCore {536class KthValueRadixMoreCore {
534public:537public:
535 __aicore__ inline KthValueRadixMoreCore(){};538 __aicore__ inline KthValueRadixMoreCore(){};
536- __aicore__ inline void Init(539+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace,
537- GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData,540+ const KthValueTilingData* tilingData, TPipe* pipe);
538- TPipe* pipe);
539 __aicore__ inline void Process();541 __aicore__ inline void Process();
540 542 
541protected:543protected:
@@ -567,8 +569,9 @@ protected:
567};569};
568 570 
569template <typename T, typename T3, typename KeyT>571template <typename T, typename T3, typename KeyT>
570-__aicore__ inline void KthValueRadixMoreCore<T, T3, KeyT>::Init(572+__aicore__ inline void KthValueRadixMoreCore<T, T3, KeyT>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
571- GM_ADDR x, GM_ADDR values, GM_ADDR indices, GM_ADDR workspace, const KthValueTilingData* tilingData, TPipe* pipe)573+ GM_ADDR workspace, const KthValueTilingData* tilingData,
574+ TPipe* pipe)
572{575{
573 if (tilingData == nullptr || pipe == nullptr) {576 if (tilingData == nullptr || pipe == nullptr) {
574 return;577 return;
Mmath/kth_value/op_kernel/arch35/kth_value_radix_one_core.h+14-12
@@ -23,17 +23,17 @@ using namespace AscendC;
23template <typename T>23template <typename T>
24class KthValueRadixOneCore {24class KthValueRadixOneCore {
25public:25public:
26- __aicore__ inline void Init(26+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
27- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe);27+ TPipe* pipe);
28 __aicore__ inline void Process();28 __aicore__ inline void Process();
29 29 
30private:30private:
31 __aicore__ inline void ParseTilingData();31 __aicore__ inline void ParseTilingData();
32 __aicore__ inline void CopyInputToUb(LocalTensor<T>& xLocal, int64_t row);32 __aicore__ inline void CopyInputToUb(LocalTensor<T>& xLocal, int64_t row);
33- __aicore__ inline void ProcessOneRow(33+ __aicore__ inline void ProcessOneRow(int64_t row, uint32_t localOffset, LocalTensor<T>& compactValue,
34- int64_t row, uint32_t localOffset, LocalTensor<T>& compactValue, LocalTensor<int64_t>& compactIndex);34+ LocalTensor<int64_t>& compactIndex);
35- __aicore__ inline void CopyOutputToGm(35+ __aicore__ inline void CopyOutputToGm(int64_t rowStart, uint32_t rowCount, LocalTensor<T>& compactValue,
36- int64_t rowStart, uint32_t rowCount, LocalTensor<T>& compactValue, LocalTensor<int64_t>& compactIndex);36+ LocalTensor<int64_t>& compactIndex);
37 __aicore__ inline void ProcessRows(int64_t rowStart, uint32_t rowCount);37 __aicore__ inline void ProcessRows(int64_t rowStart, uint32_t rowCount);
38 38 
39 static constexpr SortConfig SORT_CONFIG{SortType::RADIX_SORT, false};39 static constexpr SortConfig SORT_CONFIG{SortType::RADIX_SORT, false};
@@ -67,8 +67,8 @@ private:
67};67};
68 68 
69template <typename T>69template <typename T>
70-__aicore__ inline void KthValueRadixOneCore<T>::Init(70+__aicore__ inline void KthValueRadixOneCore<T>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
71- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe)71+ const KthValueTilingData* tiling, TPipe* pipe)
72{72{
73 if (tiling == nullptr || pipe == nullptr) {73 if (tiling == nullptr || pipe == nullptr) {
74 return;74 return;
@@ -118,8 +118,9 @@ __aicore__ inline void KthValueRadixOneCore<T>::CopyInputToUb(LocalTensor<T>& xL
118}118}
119 119 
120template <typename T>120template <typename T>
121-__aicore__ inline void KthValueRadixOneCore<T>::ProcessOneRow(121+__aicore__ inline void KthValueRadixOneCore<T>::ProcessOneRow(int64_t row, uint32_t localOffset,
122- int64_t row, uint32_t localOffset, LocalTensor<T>& compactValue, LocalTensor<int64_t>& compactIndex)122+ LocalTensor<T>& compactValue,
123+ LocalTensor<int64_t>& compactIndex)
123{124{
124 LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();125 LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>();
125 CopyInputToUb(xLocal, row);126 CopyInputToUb(xLocal, row);
@@ -147,8 +148,9 @@ __aicore__ inline void KthValueRadixOneCore<T>::ProcessOneRow(
147}148}
148 149 
149template <typename T>150template <typename T>
150-__aicore__ inline void KthValueRadixOneCore<T>::CopyOutputToGm(151+__aicore__ inline void KthValueRadixOneCore<T>::CopyOutputToGm(int64_t rowStart, uint32_t rowCount,
151- int64_t rowStart, uint32_t rowCount, LocalTensor<T>& compactValue, LocalTensor<int64_t>& compactIndex)152+ LocalTensor<T>& compactValue,
153+ LocalTensor<int64_t>& compactIndex)
152{154{
153 event_t eventIdSToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));155 event_t eventIdSToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3));
154 SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);156 SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
Mmath/kth_value/op_kernel/arch35/kth_value_small_axis_insertion.h+15-13
@@ -23,10 +23,11 @@ namespace KthValue {
23using namespace AscendC;23using namespace AscendC;
24 24 
25template <typename T, typename CONVERT_TYPE>25template <typename T, typename CONVERT_TYPE>
26-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::INSERTION_THREAD_NUM) __aicore__ void SimtStoreKthInsertionBatch(26+__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::INSERTION_THREAD_NUM) __aicore__
27- uint32_t validSegs, uint32_t kthIndex, uint32_t valueRowElems, uint32_t indexRowElems, uint64_t outputStart,27+ void SimtStoreKthInsertionBatch(uint32_t validSegs, uint32_t kthIndex, uint32_t valueRowElems,
28- __ubuf__ CONVERT_TYPE* values, __ubuf__ uint32_t* indices, __gm__ volatile T* outputValue,28+ uint32_t indexRowElems, uint64_t outputStart, __ubuf__ CONVERT_TYPE* values,
29- __gm__ volatile int64_t* outputIndex)29+ __ubuf__ uint32_t* indices, __gm__ volatile T* outputValue,
30+ __gm__ volatile int64_t* outputIndex)
30{31{
31 for (uint32_t seg = static_cast<uint32_t>(threadIdx.x); seg < validSegs;32 for (uint32_t seg = static_cast<uint32_t>(threadIdx.x); seg < validSegs;
32 seg += SmallAxisCommon::INSERTION_THREAD_NUM) {33 seg += SmallAxisCommon::INSERTION_THREAD_NUM) {
@@ -38,14 +39,15 @@ __simt_vf__ LAUNCH_BOUND(SmallAxisCommon::INSERTION_THREAD_NUM) __aicore__ void
38}39}
39 40 
40template <typename T, typename CONVERT_TYPE>41template <typename T, typename CONVERT_TYPE>
41-class KthValueSmallAxisInsertion : public SmallAxisCommon::SmallAxisInsertionBase<42+class KthValueSmallAxisInsertion
42- KthValueSmallAxisInsertion<T, CONVERT_TYPE>, T, CONVERT_TYPE, uint32_t, false> {43+ : public SmallAxisCommon::SmallAxisInsertionBase<KthValueSmallAxisInsertion<T, CONVERT_TYPE>, T, CONVERT_TYPE,
43- using Base = SmallAxisCommon::SmallAxisInsertionBase<44+ uint32_t, false> {
44- KthValueSmallAxisInsertion<T, CONVERT_TYPE>, T, CONVERT_TYPE, uint32_t, false>;45+ using Base = SmallAxisCommon::SmallAxisInsertionBase<KthValueSmallAxisInsertion<T, CONVERT_TYPE>, T, CONVERT_TYPE,
46+ uint32_t, false>;
45 47 
46public:48public:
47- __aicore__ inline void Init(49+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
48- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe);50+ TPipe* pipe);
49 __aicore__ inline void Process() { Base::Process(); }51 __aicore__ inline void Process() { Base::Process(); }
50 52 
51 friend Base;53 friend Base;
@@ -84,8 +86,8 @@ private:
84};86};
85 87 
86template <typename T, typename CONVERT_TYPE>88template <typename T, typename CONVERT_TYPE>
87-__aicore__ inline void KthValueSmallAxisInsertion<T, CONVERT_TYPE>::Init(89+__aicore__ inline void KthValueSmallAxisInsertion<T, CONVERT_TYPE>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
88- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe)90+ const KthValueTilingData* tiling, TPipe* pipe)
89{91{
90 if (tiling == nullptr || pipe == nullptr) {92 if (tiling == nullptr || pipe == nullptr) {
91 return;93 return;
@@ -116,7 +118,7 @@ template <typename T, typename CONVERT_TYPE>
116__aicore__ inline bool KthValueSmallAxisInsertion<T, CONVERT_TYPE>::IsProcessInvalid() const118__aicore__ inline bool KthValueSmallAxisInsertion<T, CONVERT_TYPE>::IsProcessInvalid() const
117{119{
118 return blockIdx_ >= blockDim_ || segmentLen_ == 0U || segmentsPerBatch_ == 0U ||120 return blockIdx_ >= blockDim_ || segmentLen_ == 0U || segmentsPerBatch_ == 0U ||
119- (IsNonLastMode() && innerSize_ <= 0);121+ (IsNonLastMode() && innerSize_ <= 0);
120}122}
121 123 
122template <typename T, typename CONVERT_TYPE>124template <typename T, typename CONVERT_TYPE>
Mmath/kth_value/op_kernel/arch35/kth_value_small_axis_two_stage.h+12-12
@@ -23,9 +23,10 @@ namespace KthValue {
23using namespace AscendC;23using namespace AscendC;
24 24 
25template <typename T>25template <typename T>
26-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__ void SimtStoreKthTwoStageBatch(26+__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__
27- uint32_t validSegs, uint32_t segmentLen, uint32_t kthIndex, uint64_t outputStart, __ubuf__ T* finalValues,27+ void SimtStoreKthTwoStageBatch(uint32_t validSegs, uint32_t segmentLen, uint32_t kthIndex, uint64_t outputStart,
28- __ubuf__ uint32_t* finalIdx, __gm__ volatile T* outputValue, __gm__ volatile int64_t* outputIndex)28+ __ubuf__ T* finalValues, __ubuf__ uint32_t* finalIdx, __gm__ volatile T* outputValue,
29+ __gm__ volatile int64_t* outputIndex)
29{30{
30 for (uint32_t seg = static_cast<uint32_t>(threadIdx.x); seg < validSegs;31 for (uint32_t seg = static_cast<uint32_t>(threadIdx.x); seg < validSegs;
31 seg += SmallAxisCommon::TWO_STAGE_THREAD_NUM) {32 seg += SmallAxisCommon::TWO_STAGE_THREAD_NUM) {
@@ -41,9 +42,8 @@ class KthValueSmallAxisTwoStage
41 using Base = SmallAxisCommon::SmallAxisTwoStageBase<KthValueSmallAxisTwoStage<T>, T, uint32_t, false>;42 using Base = SmallAxisCommon::SmallAxisTwoStageBase<KthValueSmallAxisTwoStage<T>, T, uint32_t, false>;
42 43 
43public:44public:
44- __aicore__ inline void Init(45+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling,
45- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe);46+ TPipe* pipe);
46- __aicore__ inline void Process() { Base::Process(); }
47 47 
48 friend Base;48 friend Base;
49 49 
@@ -80,8 +80,8 @@ private:
80};80};
81 81 
82template <typename T>82template <typename T>
83-__aicore__ inline void KthValueSmallAxisTwoStage<T>::Init(83+__aicore__ inline void KthValueSmallAxisTwoStage<T>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices,
84- GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe)84+ const KthValueTilingData* tiling, TPipe* pipe)
85{85{
86 if (tiling == nullptr || pipe == nullptr) {86 if (tiling == nullptr || pipe == nullptr) {
87 return;87 return;
@@ -179,16 +179,16 @@ __aicore__ inline int64_t KthValueSmallAxisTwoStage<T>::GetInputStart(uint32_t b
179}179}
180 180 
181template <typename T>181template <typename T>
182-__aicore__ inline void KthValueSmallAxisTwoStage<T>::LoadBatch(182+__aicore__ inline void KthValueSmallAxisTwoStage<T>::LoadBatch(uint32_t batchId, uint32_t validSegs,
183- uint32_t batchId, uint32_t validSegs, uint32_t totalElems)183+ uint32_t totalElems)
184{184{
185 if (IsNonLastMode()) {185 if (IsNonLastMode()) {
186 uint64_t outerId = static_cast<uint64_t>(batchId / innerLoopNum_);186 uint64_t outerId = static_cast<uint64_t>(batchId / innerLoopNum_);
187 uint64_t innerTileId = static_cast<uint64_t>(batchId % innerLoopNum_);187 uint64_t innerTileId = static_cast<uint64_t>(batchId % innerLoopNum_);
188 uint64_t innerStart = innerTileId * static_cast<uint64_t>(batchSize_);188 uint64_t innerStart = innerTileId * static_cast<uint64_t>(batchSize_);
189 uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);189 uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);
190- Base::LoadNonLastBatch(190+ Base::LoadNonLastBatch(inputGm_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_), validSegs,
191- inputGm_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_), validSegs, totalElems);191+ totalElems);
192 } else {192 } else {
193 Base::LoadContiguousBatch(inputGm_, GetInputStart(batchId), totalElems);193 Base::LoadContiguousBatch(inputGm_, GetInputStart(batchId), totalElems);
194 }194 }
Mmath/kth_value/op_kernel/arch35/kth_value_tiling_data.h+13-13
@@ -14,11 +14,11 @@
14#include <cstdint>14#include <cstdint>
15 15 
16struct KthValueTilingData {16struct KthValueTilingData {
17- uint32_t numTileDataSize; // h轴ub一次处理个数17+ uint32_t numTileDataSize; // h轴ub一次处理个数
18 uint32_t unsortedDimParallel; // b轴使用的核数18 uint32_t unsortedDimParallel; // b轴使用的核数
19- uint32_t lastDimTileNum; // h轴循环次数19+ uint32_t lastDimTileNum; // h轴循环次数
20- uint32_t sortLoopTimes; // b轴循环次数20+ uint32_t sortLoopTimes; // b轴循环次数
21- uint32_t lastDimNeedCore; // h轴需要的核数21+ uint32_t lastDimNeedCore; // h轴需要的核数
22 // keyParamsxxx 预留参数22 // keyParamsxxx 预留参数
23 // radix: globalHistGmWk_ 使用核数23 // radix: globalHistGmWk_ 使用核数
24 // radix_one_core: inqueX 的 ub 大小24 // radix_one_core: inqueX 的 ub 大小
@@ -53,15 +53,15 @@ struct KthValueTilingData {
53 // radix: 清零 globalHistGmWk_ chunk 大小53 // radix: 清零 globalHistGmWk_ chunk 大小
54 // intra_core: 最大归并迭代次数54 // intra_core: 最大归并迭代次数
55 uint32_t keyParams5;55 uint32_t keyParams5;
56- uint32_t tmpUbSize; // 高级api需要的临时ub大小56+ uint32_t tmpUbSize; // 高级api需要的临时ub大小
57- int64_t kthIndex; // 零基kth偏移，避免按k展开binary57+ int64_t kthIndex; // 零基kth偏移，避免按k展开binary
58- int64_t lastAxisNum; // h轴大小58+ int64_t lastAxisNum; // h轴大小
59- int64_t unsortedDimNum; // b轴大小59+ int64_t unsortedDimNum; // b轴大小
60- int64_t outerSize; // non-last-axis: 外层维度大小60+ int64_t outerSize; // non-last-axis: 外层维度大小
61- int64_t innerSize; // non-last-axis: 内层维度大小61+ int64_t innerSize; // non-last-axis: 内层维度大小
62- uint32_t innerLoopNum; // non-last-axis: inner分块循环次数62+ uint32_t innerLoopNum; // non-last-axis: inner分块循环次数
63- uint32_t innerChunk; // non-last-axis: inner分块大小63+ uint32_t innerChunk; // non-last-axis: inner分块大小
64- uint32_t inputRowBytes; // non-last-axis: 输入一行字节数64+ uint32_t inputRowBytes; // non-last-axis: 输入一行字节数
65 uint32_t valueAxisBytes; // non-last-axis: value轴字节数65 uint32_t valueAxisBytes; // non-last-axis: value轴字节数
66 uint32_t indexAxisBytes; // non-last-axis: index轴字节数66 uint32_t indexAxisBytes; // non-last-axis: index轴字节数
67};67};
Mmath/kth_value/tests/ut/op_host/test_kth_value_infershape.cpp+120-130
@@ -22,190 +22,180 @@ protected:
22 22 
23TEST_F(KthValueInfershape, kthvalue_infershape_2d_last_axis)23TEST_F(KthValueInfershape, kthvalue_infershape_2d_last_axis)
24{24{
25- gert::InfershapeContextPara infershapeContextPara(25+ gert::InfershapeContextPara infershapeContextPara("KthValue",
26- "KthValue",26+ {
27- {27+ {{{3, 10}, {3, 10}}, ge::DT_FLOAT, ge::FORMAT_ND},
28- {{{3, 10}, {3, 10}}, ge::DT_FLOAT, ge::FORMAT_ND},28+ },
29- },29+ {
30- {30+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},
31- {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},31+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
32- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},32+ },
33- },33+ {
34- {34+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},
35- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},35+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},
36- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},36+ });
37- });
38 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 1}, {3, 1}};37 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 1}, {3, 1}};
39 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);38 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
40}39}
41 40 
42TEST_F(KthValueInfershape, kthvalue_infershape_3d_first_axis)41TEST_F(KthValueInfershape, kthvalue_infershape_3d_first_axis)
43{42{
44- gert::InfershapeContextPara infershapeContextPara(43+ gert::InfershapeContextPara infershapeContextPara("KthValue",
45- "KthValue",44+ {
46- {45+ {{{4, 5, 6}, {4, 5, 6}}, ge::DT_FLOAT16, ge::FORMAT_ND},
47- {{{4, 5, 6}, {4, 5, 6}}, ge::DT_FLOAT16, ge::FORMAT_ND},46+ },
48- },47+ {
49- {48+ {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND},
50- {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND},49+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
51- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},50+ },
52- },51+ {
53- {52+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(2)},
54- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(2)},53+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(0)},
55- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(0)},54+ });
56- });
57 std::vector<std::vector<int64_t>> expectOutputShape = {{1, 5, 6}, {1, 5, 6}};55 std::vector<std::vector<int64_t>> expectOutputShape = {{1, 5, 6}, {1, 5, 6}};
58 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);56 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
59}57}
60 58 
61TEST_F(KthValueInfershape, kthvalue_infershape_3d_middle_axis)59TEST_F(KthValueInfershape, kthvalue_infershape_3d_middle_axis)
62{60{
63- gert::InfershapeContextPara infershapeContextPara(61+ gert::InfershapeContextPara infershapeContextPara("KthValue",
64- "KthValue",62+ {
65- {63+ {{{2, 8, 3}, {2, 8, 3}}, ge::DT_INT32, ge::FORMAT_ND},
66- {{{2, 8, 3}, {2, 8, 3}}, ge::DT_INT32, ge::FORMAT_ND},64+ },
67- },65+ {
68- {66+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
69- {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},67+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
70- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},68+ },
71- },69+ {
72- {70+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(3)},
73- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(3)},71+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},
74- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},72+ });
75- });
76 std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 3}, {2, 1, 3}};73 std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 3}, {2, 1, 3}};
77 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);74 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
78}75}
79 76 
80TEST_F(KthValueInfershape, kthvalue_infershape_1d)77TEST_F(KthValueInfershape, kthvalue_infershape_1d)
81{78{
82- gert::InfershapeContextPara infershapeContextPara(79+ gert::InfershapeContextPara infershapeContextPara("KthValue",
83- "KthValue",80+ {
84- {81+ {{{100}, {100}}, ge::DT_FLOAT, ge::FORMAT_ND},
85- {{{100}, {100}}, ge::DT_FLOAT, ge::FORMAT_ND},82+ },
86- },83+ {
87- {84+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},
88- {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},85+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
89- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},86+ },
90- },87+ {
91- {88+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},
92- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},89+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},
93- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},90+ });
94- });
95 std::vector<std::vector<int64_t>> expectOutputShape = {{1}, {1}};91 std::vector<std::vector<int64_t>> expectOutputShape = {{1}, {1}};
96 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);92 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
97}93}
98 94 
99TEST_F(KthValueInfershape, kthvalue_infershape_negative_dim)95TEST_F(KthValueInfershape, kthvalue_infershape_negative_dim)
100{96{
101- gert::InfershapeContextPara infershapeContextPara(97+ gert::InfershapeContextPara infershapeContextPara("KthValue",
102- "KthValue",98+ {
103- {99+ {{{2, 3, 4}, {2, 3, 4}}, ge::DT_BF16, ge::FORMAT_ND},
104- {{{2, 3, 4}, {2, 3, 4}}, ge::DT_BF16, ge::FORMAT_ND},100+ },
105- },101+ {
106- {102+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},
107- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},103+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
108- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},104+ },
109- },105+ {
110- {106+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(2)},
111- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(2)},107+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-2)},
112- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-2)},108+ });
113- });
114 std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 4}, {2, 1, 4}};109 std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 4}, {2, 1, 4}};
115 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);110 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
116}111}
117 112 
118TEST_F(KthValueInfershape, kthvalue_infershape_invalid_dim_out_of_range)113TEST_F(KthValueInfershape, kthvalue_infershape_invalid_dim_out_of_range)
119{114{
120- gert::InfershapeContextPara infershapeContextPara(115+ gert::InfershapeContextPara infershapeContextPara("KthValue",
121- "KthValue",116+ {
122- {117+ {{{3, 10}, {3, 10}}, ge::DT_FLOAT, ge::FORMAT_ND},
123- {{{3, 10}, {3, 10}}, ge::DT_FLOAT, ge::FORMAT_ND},118+ },
124- },119+ {
125- {120+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},
126- {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},121+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
127- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},122+ },
128- },123+ {
129- {124+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},
130- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},125+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(3)},
131- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(3)},126+ });
132- });
133 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 1}, {3, 1}};127 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 1}, {3, 1}};
134 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);128 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);
135}129}
136 130 
137TEST_F(KthValueInfershape, kthvalue_infershape_invalid_negative_dim_out_of_range)131TEST_F(KthValueInfershape, kthvalue_infershape_invalid_negative_dim_out_of_range)
138{132{
139- gert::InfershapeContextPara infershapeContextPara(133+ gert::InfershapeContextPara infershapeContextPara("KthValue",
140- "KthValue",134+ {
141- {135+ {{{3, 10}, {3, 10}}, ge::DT_FLOAT, ge::FORMAT_ND},
142- {{{3, 10}, {3, 10}}, ge::DT_FLOAT, ge::FORMAT_ND},136+ },
143- },137+ {
144- {138+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},
145- {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},139+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
146- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},140+ },
147- },141+ {
148- {142+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},
149- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(5)},143+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-3)},
150- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-3)},144+ });
151- });
152 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 1}, {3, 1}};145 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 1}, {3, 1}};
153 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);146 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);
154}147}
155 148 
156TEST_F(KthValueInfershape, kthvalue_infershape_scalar_input_fails)149TEST_F(KthValueInfershape, kthvalue_infershape_scalar_input_fails)
157{150{
158- gert::InfershapeContextPara infershapeContextPara(151+ gert::InfershapeContextPara infershapeContextPara("KthValue",
159- "KthValue",152+ {
160- {153+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},
161- {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},154+ },
162- },155+ {
163- {156+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},
164- {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND},157+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
165- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},158+ },
166- },159+ {
167- {160+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},
168- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},161+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},
169- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},162+ });
170- });
171 std::vector<std::vector<int64_t>> expectOutputShape = {{}, {}};163 std::vector<std::vector<int64_t>> expectOutputShape = {{}, {}};
172 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);164 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);
173}165}
174 166 
175TEST_F(KthValueInfershape, kthvalue_infershape_int8_dtype)167TEST_F(KthValueInfershape, kthvalue_infershape_int8_dtype)
176{168{
177- gert::InfershapeContextPara infershapeContextPara(169+ gert::InfershapeContextPara infershapeContextPara("KthValue",
178- "KthValue",170+ {
179- {171+ {{{5, 20}, {5, 20}}, ge::DT_INT8, ge::FORMAT_ND},
180- {{{5, 20}, {5, 20}}, ge::DT_INT8, ge::FORMAT_ND},172+ },
181- },173+ {
182- {174+ {{{}, {}}, ge::DT_INT8, ge::FORMAT_ND},
183- {{{}, {}}, ge::DT_INT8, ge::FORMAT_ND},175+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
184- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},176+ },
185- },177+ {
186- {178+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(10)},
187- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(10)},179+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},
188- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(-1)},180+ });
189- });
190 std::vector<std::vector<int64_t>> expectOutputShape = {{5, 1}, {5, 1}};181 std::vector<std::vector<int64_t>> expectOutputShape = {{5, 1}, {5, 1}};
191 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);182 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
192}183}
193 184 
194TEST_F(KthValueInfershape, kthvalue_infershape_uint64_dtype)185TEST_F(KthValueInfershape, kthvalue_infershape_uint64_dtype)
195{186{
196- gert::InfershapeContextPara infershapeContextPara(187+ gert::InfershapeContextPara infershapeContextPara("KthValue",
197- "KthValue",188+ {
198- {189+ {{{4, 8}, {4, 8}}, ge::DT_UINT64, ge::FORMAT_ND},
199- {{{4, 8}, {4, 8}}, ge::DT_UINT64, ge::FORMAT_ND},190+ },
200- },191+ {
201- {192+ {{{}, {}}, ge::DT_UINT64, ge::FORMAT_ND},
202- {{{}, {}}, ge::DT_UINT64, ge::FORMAT_ND},193+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
203- {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},194+ },
204- },195+ {
205- {196+ {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},
206- {"k", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},197+ {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},
207- {"dim", Ops::Math::AnyValue::CreateFrom<int64_t>(1)},198+ });
208- });
209 std::vector<std::vector<int64_t>> expectOutputShape = {{4, 1}, {4, 1}};199 std::vector<std::vector<int64_t>> expectOutputShape = {{4, 1}, {4, 1}};
210 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);200 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
211}201}
Mmath/sort/CMakeLists.txt+1-1
@@ -12,4 +12,4 @@
12set(SUPPORT_COMPUTE_UNIT "ascend950")12set(SUPPORT_COMPUTE_UNIT "ascend950")
13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14set(SUPPORT_TILING_DIR "arch35")14set(SUPPORT_TILING_DIR "arch35")
15-add_all_modules_sources(OPTYPE sort ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)15+add_all_modules_sources(OPTYPE sort ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
Mmath/sort/examples/test_aclnn_argsort.cpp+147-142
@@ -1,142 +1,147 @@
1-/**1+/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * 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").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.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,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.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.8+ * See LICENSE in the root of the software repository for the full text of the License.
9- */9+ */
10- 10+ 
11-#include <iostream>11+#include <iostream>
12-#include <vector>12+#include <vector>
13-#include "acl/acl.h"13+#include "acl/acl.h"
14-#include "aclnnop/aclnn_argsort.h"14+#include "aclnnop/aclnn_argsort.h"
15- 15+ 
16-#define CHECK_RET(cond, return_expr) \16+#define CHECK_RET(cond, return_expr) \
17- do { \17+ do { \
18- if (!(cond)) { \18+ if (!(cond)) { \
19- return_expr; \19+ return_expr; \
20- } \20+ } \
21- } while (0)21+ } while (0)
22- 22+ 
23-#define LOG_PRINT(message, ...) \23+#define LOG_PRINT(message, ...) \
24- do { \24+ do { \
25- printf(message, ##__VA_ARGS__); \25+ printf(message, ##__VA_ARGS__); \
26- } while (0)26+ } while (0)
27- 27+ 
28-int64_t GetShapeSize(const std::vector<int64_t>& shape) {28+int64_t GetShapeSize(const std::vector<int64_t>& shape)
29- int64_t shapeSize = 1;29+{
30- for (auto i : shape) {30+ int64_t shapeSize = 1;
31- shapeSize *= i;31+ for (auto i : shape) {
32- }32+ shapeSize *= i;
33- return shapeSize;33+ }
34-}34+ return shapeSize;
35- 35+}
36-int Init(int32_t deviceId, aclrtStream* stream) {36+ 
37- // 固定写法，资源初始化37+int Init(int32_t deviceId, aclrtStream* stream)
38- auto ret = aclInit(nullptr);38+{
39- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);39+ // 固定写法，资源初始化
40- ret = aclrtSetDevice(deviceId);40+ auto ret = aclInit(nullptr);
41- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);41+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
42- ret = aclrtCreateStream(stream);42+ ret = aclrtSetDevice(deviceId);
43- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);43+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
44- return 0;44+ ret = aclrtCreateStream(stream);
45-}45+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
46- 46+ return 0;
47-template <typename T>47+}
48-int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,48+ 
49- aclDataType dataType, aclTensor** tensor) {49+template <typename T>
50- auto size = GetShapeSize(shape) * sizeof(T);50+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
51- // 调用aclrtMalloc申请device侧内存51+ aclDataType dataType, aclTensor** tensor)
52- auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);52+{
53- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);53+ auto size = GetShapeSize(shape) * sizeof(T);
54- // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上54+ // 调用aclrtMalloc申请device侧内存
55- ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);55+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
56- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);56+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
57- 57+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
58- // 计算连续tensor的strides58+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
59- std::vector<int64_t> strides(shape.size(), 1);59+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
60- for (int64_t i = shape.size() - 2; i >= 0; i--) {60+ 
61- strides[i] = shape[i + 1] * strides[i + 1];61+ // 计算连续tensor的strides
62- }62+ std::vector<int64_t> strides(shape.size(), 1);
63- 63+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
64- // 调用aclCreateTensor接口创建aclTensor64+ strides[i] = shape[i + 1] * strides[i + 1];
65- *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,65+ }
66- shape.data(), shape.size(), *deviceAddr);66+ 
67- return 0;67+ // 调用aclCreateTensor接口创建aclTensor
68-}68+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
69- 69+ shape.data(), shape.size(), *deviceAddr);
70-int main() {70+ return 0;
71- // 1. （固定写法）device/stream初始化，参考acl API手册71+}
72- // 根据自己的实际device填写deviceId72+ 
73- int32_t deviceId = 0;73+int main()
74- aclrtStream stream;74+{
75- auto ret = Init(deviceId, &stream);75+ // 1. （固定写法）device/stream初始化，参考acl API手册
76- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);76+ // 根据自己的实际device填写deviceId
77- 77+ int32_t deviceId = 0;
78- // 2. 构造输入与输出，需要根据API的接口自定义构造78+ aclrtStream stream;
79- int64_t dim = 0;79+ auto ret = Init(deviceId, &stream);
80- bool descending = false;80+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
81- std::vector<int64_t> selfShape = {3, 4};81+ 
82- std::vector<int64_t> outIndicesShape = {3, 4};82+ // 2. 构造输入与输出，需要根据API的接口自定义构造
83- void* selfDeviceAddr = nullptr;83+ int64_t dim = 0;
84- void* outIndicesDeviceAddr = nullptr;84+ bool descending = false;
85- aclTensor* self = nullptr;85+ std::vector<int64_t> selfShape = {3, 4};
86- aclTensor* outIndices = nullptr;86+ std::vector<int64_t> outIndicesShape = {3, 4};
87- std::vector<int64_t> selfHostData = {7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6};87+ void* selfDeviceAddr = nullptr;
88- std::vector<int64_t> outIndicesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};88+ void* outIndicesDeviceAddr = nullptr;
89- 89+ aclTensor* self = nullptr;
90- // 创建self aclTensor90+ aclTensor* outIndices = nullptr;
91- ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_INT64, &self);91+ std::vector<int64_t> selfHostData = {7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6};
92- CHECK_RET(ret == ACL_SUCCESS, return ret);92+ std::vector<int64_t> outIndicesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
93- // 创建outValues和outIndices aclTensor93+ 
94- ret = CreateAclTensor(outIndicesHostData, outIndicesShape, &outIndicesDeviceAddr, aclDataType::ACL_INT64, &outIndices);94+ // 创建self aclTensor
95- CHECK_RET(ret == ACL_SUCCESS, return ret);95+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_INT64, &self);
96- 96+ CHECK_RET(ret == ACL_SUCCESS, return ret);
97- // 3. 调用CANN算子库API，需要修改为具体的Api名称97+ // 创建outValues和outIndices aclTensor
98- uint64_t workspaceSize = 0;98+ ret = CreateAclTensor(outIndicesHostData, outIndicesShape, &outIndicesDeviceAddr, aclDataType::ACL_INT64,
99- aclOpExecutor* executor;99+ &outIndices);
100- // 调用aclnnArgsort第一段接口100+ CHECK_RET(ret == ACL_SUCCESS, return ret);
101- ret = aclnnArgsortGetWorkspaceSize(self, dim, descending, outIndices, &workspaceSize, &executor);101+ 
102- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnArgsortGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);102+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
103- // 根据第一段接口计算出的workspaceSize申请device内存103+ uint64_t workspaceSize = 0;
104- void* workspaceAddr = nullptr;104+ aclOpExecutor* executor;
105- if (workspaceSize > 0) {105+ // 调用aclnnArgsort第一段接口
106- ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);106+ ret = aclnnArgsortGetWorkspaceSize(self, dim, descending, outIndices, &workspaceSize, &executor);
107- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);107+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnArgsortGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
108- }108+ // 根据第一段接口计算出的workspaceSize申请device内存
109- // 调用aclnnArgsort第二段接口109+ void* workspaceAddr = nullptr;
110- ret = aclnnArgsort(workspaceAddr, workspaceSize, executor, stream);110+ if (workspaceSize > 0) {
111- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnArgsort failed. ERROR: %d\n", ret); return ret);111+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
112- 112+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
113- // 4. （固定写法）同步等待任务执行结束113+ }
114- ret = aclrtSynchronizeStream(stream);114+ // 调用aclnnArgsort第二段接口
115- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);115+ ret = aclnnArgsort(workspaceAddr, workspaceSize, executor, stream);
116- 116+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnArgsort failed. ERROR: %d\n", ret); return ret);
117- // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改117+ 
118- auto size2 = GetShapeSize(outIndicesShape);118+ // 4. （固定写法）同步等待任务执行结束
119- std::vector<int64_t> resultData2(size2, 0);119+ ret = aclrtSynchronizeStream(stream);
120- ret = aclrtMemcpy(resultData2.data(), resultData2.size() * sizeof(resultData2[0]), outIndicesDeviceAddr,120+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
121- size2 * sizeof(resultData2[0]), ACL_MEMCPY_DEVICE_TO_HOST);121+ 
122- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);122+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
123- for (int64_t i = 0; i < size2; i++) {123+ auto size2 = GetShapeSize(outIndicesShape);
124- LOG_PRINT("result indices [%ld] is: %ld\n", i, resultData2[i]);124+ std::vector<int64_t> resultData2(size2, 0);
125- }125+ ret = aclrtMemcpy(resultData2.data(), resultData2.size() * sizeof(resultData2[0]), outIndicesDeviceAddr,
126- 126+ size2 * sizeof(resultData2[0]), ACL_MEMCPY_DEVICE_TO_HOST);
127- // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改127+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
128- aclDestroyTensor(self);128+ for (int64_t i = 0; i < size2; i++) {
129- aclDestroyTensor(outIndices);129+ LOG_PRINT("result indices [%ld] is: %ld\n", i, resultData2[i]);
130- 130+ }
131- // 7. 释放device 资源131+ 
132- aclrtFree(selfDeviceAddr);132+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
133- aclrtFree(outIndicesDeviceAddr);133+ aclDestroyTensor(self);
134- if (workspaceSize > 0) {134+ aclDestroyTensor(outIndices);
135- aclrtFree(workspaceAddr);135+ 
136- }136+ // 7. 释放device 资源
137- aclrtDestroyStream(stream);137+ aclrtFree(selfDeviceAddr);
138- aclrtResetDevice(deviceId);138+ aclrtFree(outIndicesDeviceAddr);
139- aclFinalize();139+ if (workspaceSize > 0) {
140- 140+ aclrtFree(workspaceAddr);
141- return 0;141+ }
142-}142+ aclrtDestroyStream(stream);
143+ aclrtResetDevice(deviceId);
144+ aclFinalize();
145+ 
146+ return 0;
147+}
Mmath/sort/examples/test_aclnn_sort.cpp+164-158
@@ -1,158 +1,164 @@
1-/**1+/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * 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").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.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,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.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.8+ * See LICENSE in the root of the software repository for the full text of the License.
9- */9+ */
10- 10+ 
11-#include <iostream>11+#include <iostream>
12-#include <vector>12+#include <vector>
13-#include "acl/acl.h"13+#include "acl/acl.h"
14-#include "aclnnop/aclnn_sort.h"14+#include "aclnnop/aclnn_sort.h"
15- 15+ 
16-#define CHECK_RET(cond, return_expr) \16+#define CHECK_RET(cond, return_expr) \
17- do { \17+ do { \
18- if (!(cond)) { \18+ if (!(cond)) { \
19- return_expr; \19+ return_expr; \
20- } \20+ } \
21- } while (0)21+ } while (0)
22- 22+ 
23-#define LOG_PRINT(message, ...) \23+#define LOG_PRINT(message, ...) \
24- do { \24+ do { \
25- printf(message, ##__VA_ARGS__); \25+ printf(message, ##__VA_ARGS__); \
26- } while (0)26+ } while (0)
27- 27+ 
28-int64_t GetShapeSize(const std::vector<int64_t>& shape) {28+int64_t GetShapeSize(const std::vector<int64_t>& shape)
29- int64_t shapeSize = 1;29+{
30- for (auto i : shape) {30+ int64_t shapeSize = 1;
31- shapeSize *= i;31+ for (auto i : shape) {
32- }32+ shapeSize *= i;
33- return shapeSize;33+ }
34-}34+ return shapeSize;
35- 35+}
36-int Init(int32_t deviceId, aclrtStream* stream) {36+ 
37- // 固定写法，资源初始化37+int Init(int32_t deviceId, aclrtStream* stream)
38- auto ret = aclInit(nullptr);38+{
39- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);39+ // 固定写法，资源初始化
40- ret = aclrtSetDevice(deviceId);40+ auto ret = aclInit(nullptr);
41- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);41+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
42- ret = aclrtCreateStream(stream);42+ ret = aclrtSetDevice(deviceId);
43- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);43+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
44- return 0;44+ ret = aclrtCreateStream(stream);
45-}45+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
46- 46+ return 0;
47-template <typename T>47+}
48-int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,48+ 
49- aclDataType dataType, aclTensor** tensor) {49+template <typename T>
50- auto size = GetShapeSize(shape) * sizeof(T);50+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
51- // 调用aclrtMalloc申请device侧内存51+ aclDataType dataType, aclTensor** tensor)
52- auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);52+{
53- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);53+ auto size = GetShapeSize(shape) * sizeof(T);
54- // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上54+ // 调用aclrtMalloc申请device侧内存
55- ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);55+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
56- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);56+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
57- 57+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
58- // 计算连续tensor的strides58+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
59- std::vector<int64_t> strides(shape.size(), 1);59+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
60- for (int64_t i = shape.size() - 2; i >= 0; i--) {60+ 
61- strides[i] = shape[i + 1] * strides[i + 1];61+ // 计算连续tensor的strides
62- }62+ std::vector<int64_t> strides(shape.size(), 1);
63- 63+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
64- // 调用aclCreateTensor接口创建aclTensor64+ strides[i] = shape[i + 1] * strides[i + 1];
65- *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr);65+ }
66- return 0;66+ 
67-}67+ // 调用aclCreateTensor接口创建aclTensor
68- 68+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
69-int main() {69+ shape.data(), shape.size(), *deviceAddr);
70- // 1. （固定写法）device/stream初始化，参考acl API手册70+ return 0;
71- // 根据自己的实际device填写deviceId71+}
72- int32_t deviceId = 0;72+ 
73- aclrtStream stream;73+int main()
74- auto ret = Init(deviceId, &stream);74+{
75- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);75+ // 1. （固定写法）device/stream初始化，参考acl API手册
76- 76+ // 根据自己的实际device填写deviceId
77- // 2. 构造输入与输出，需要根据API的接口自定义构造77+ int32_t deviceId = 0;
78- bool stable = false;78+ aclrtStream stream;
79- int64_t dim = 0;79+ auto ret = Init(deviceId, &stream);
80- bool descending = false;80+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
81- std::vector<int64_t> selfShape = {3, 4};81+ 
82- std::vector<int64_t> outValuesShape = {3, 4};82+ // 2. 构造输入与输出，需要根据API的接口自定义构造
83- std::vector<int64_t> outIndicesShape = {3, 4};83+ bool stable = false;
84- void* selfDeviceAddr = nullptr;84+ int64_t dim = 0;
85- void* outValuesDeviceAddr = nullptr;85+ bool descending = false;
86- void* outIndicesDeviceAddr = nullptr;86+ std::vector<int64_t> selfShape = {3, 4};
87- aclTensor* self = nullptr;87+ std::vector<int64_t> outValuesShape = {3, 4};
88- aclTensor* outValues = nullptr;88+ std::vector<int64_t> outIndicesShape = {3, 4};
89- aclTensor* outIndices = nullptr;89+ void* selfDeviceAddr = nullptr;
90- std::vector<int64_t> selfHostData = {7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6};90+ void* outValuesDeviceAddr = nullptr;
91- std::vector<int64_t> outValuesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};91+ void* outIndicesDeviceAddr = nullptr;
92- std::vector<int64_t> outIndicesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};92+ aclTensor* self = nullptr;
93- 93+ aclTensor* outValues = nullptr;
94- // 创建self aclTensor94+ aclTensor* outIndices = nullptr;
95- ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_INT64, &self);95+ std::vector<int64_t> selfHostData = {7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6};
96- CHECK_RET(ret == ACL_SUCCESS, return ret);96+ std::vector<int64_t> outValuesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
97- // 创建outValues和outIndices aclTensor97+ std::vector<int64_t> outIndicesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
98- ret = CreateAclTensor(outValuesHostData, outValuesShape, &outValuesDeviceAddr, aclDataType::ACL_INT64, &outValues);98+ 
99- CHECK_RET(ret == ACL_SUCCESS, return ret);99+ // 创建self aclTensor
100- ret = CreateAclTensor(outIndicesHostData, outIndicesShape, &outIndicesDeviceAddr, aclDataType::ACL_INT64, &outIndices);100+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_INT64, &self);
101- CHECK_RET(ret == ACL_SUCCESS, return ret);101+ CHECK_RET(ret == ACL_SUCCESS, return ret);
102- 102+ // 创建outValues和outIndices aclTensor
103- // 3. 调用CANN算子库API，需要修改为具体的Api名称103+ ret = CreateAclTensor(outValuesHostData, outValuesShape, &outValuesDeviceAddr, aclDataType::ACL_INT64, &outValues);
104- uint64_t workspaceSize = 0;104+ CHECK_RET(ret == ACL_SUCCESS, return ret);
105- aclOpExecutor* executor;105+ ret = CreateAclTensor(outIndicesHostData, outIndicesShape, &outIndicesDeviceAddr, aclDataType::ACL_INT64,
106- // 调用aclnnSort第一段接口106+ &outIndices);
107- ret = aclnnSortGetWorkspaceSize(self, stable, dim, descending, outValues, outIndices, &workspaceSize, &executor);107+ CHECK_RET(ret == ACL_SUCCESS, return ret);
108- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSortGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);108+ 
109- // 根据第一段接口计算出的workspaceSize申请device内存109+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
110- void* workspaceAddr = nullptr;110+ uint64_t workspaceSize = 0;
111- if (workspaceSize > 0) {111+ aclOpExecutor* executor;
112- ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);112+ // 调用aclnnSort第一段接口
113- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);113+ ret = aclnnSortGetWorkspaceSize(self, stable, dim, descending, outValues, outIndices, &workspaceSize, &executor);
114- }114+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSortGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
115- // 调用aclnnSort第二段接口115+ // 根据第一段接口计算出的workspaceSize申请device内存
116- ret = aclnnSort(workspaceAddr, workspaceSize, executor, stream);116+ void* workspaceAddr = nullptr;
117- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSort failed. ERROR: %d\n", ret); return ret);117+ if (workspaceSize > 0) {
118- 118+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
119- // 4. （固定写法）同步等待任务执行结束119+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
120- ret = aclrtSynchronizeStream(stream);120+ }
121- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);121+ // 调用aclnnSort第二段接口
122- 122+ ret = aclnnSort(workspaceAddr, workspaceSize, executor, stream);
123- // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改123+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSort failed. ERROR: %d\n", ret); return ret);
124- auto size = GetShapeSize(outValuesShape);124+ 
125- std::vector<int64_t> resultData(size, 0);125+ // 4. （固定写法）同步等待任务执行结束
126- ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outValuesDeviceAddr,126+ ret = aclrtSynchronizeStream(stream);
127- size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);127+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
128- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);128+ 
129- for (int64_t i = 0; i < size; i++) {129+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
130- LOG_PRINT("result values [%ld] is: %ld\n", i, resultData[i]);130+ auto size = GetShapeSize(outValuesShape);
131- }131+ std::vector<int64_t> resultData(size, 0);
132- 132+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outValuesDeviceAddr,
133- auto size2 = GetShapeSize(outIndicesShape);133+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
134- std::vector<int64_t> resultData2(size2, 0);134+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
135- ret = aclrtMemcpy(resultData2.data(), resultData2.size() * sizeof(resultData2[0]), outIndicesDeviceAddr,135+ for (int64_t i = 0; i < size; i++) {
136- size * sizeof(resultData2[0]), ACL_MEMCPY_DEVICE_TO_HOST);136+ LOG_PRINT("result values [%ld] is: %ld\n", i, resultData[i]);
137- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);137+ }
138- for (int64_t i = 0; i < size2; i++) {138+ 
139- LOG_PRINT("result indices [%ld] is: %ld\n", i, resultData2[i]);139+ auto size2 = GetShapeSize(outIndicesShape);
140- }140+ std::vector<int64_t> resultData2(size2, 0);
141- 141+ ret = aclrtMemcpy(resultData2.data(), resultData2.size() * sizeof(resultData2[0]), outIndicesDeviceAddr,
142- // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改142+ size * sizeof(resultData2[0]), ACL_MEMCPY_DEVICE_TO_HOST);
143- aclDestroyTensor(self);143+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
144- aclDestroyTensor(outValues);144+ for (int64_t i = 0; i < size2; i++) {
145- aclDestroyTensor(outIndices);145+ LOG_PRINT("result indices [%ld] is: %ld\n", i, resultData2[i]);
146- 146+ }
147- // 7. 释放device资源，需要根据具体API的接口定义修改147+ 
148- aclrtFree(selfDeviceAddr);148+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
149- aclrtFree(outValuesDeviceAddr);149+ aclDestroyTensor(self);
150- aclrtFree(outIndicesDeviceAddr);150+ aclDestroyTensor(outValues);
151- if (workspaceSize > 0) {151+ aclDestroyTensor(outIndices);
152- aclrtFree(workspaceAddr);152+ 
153- }153+ // 7. 释放device资源，需要根据具体API的接口定义修改
154- aclrtDestroyStream(stream);154+ aclrtFree(selfDeviceAddr);
155- aclrtResetDevice(deviceId);155+ aclrtFree(outValuesDeviceAddr);
156- aclFinalize();156+ aclrtFree(outIndicesDeviceAddr);
157- return 0;157+ if (workspaceSize > 0) {
158-}158+ aclrtFree(workspaceAddr);
159+ }
160+ aclrtDestroyStream(stream);
161+ aclrtResetDevice(deviceId);
162+ aclFinalize();
163+ return 0;
164+}
Mmath/sort/op_api/aclnn_sort.cpp+43-44
@@ -242,7 +242,7 @@ static aclIntArray* updatePerm(int64_t dim, int64_t dimSize, aclOpExecutor* exec
242 if (dim != dimSize - 1) {242 if (dim != dimSize - 1) {
243 auto perm = GetPermResult(dim, dimSize, executor); // 不需改动时的版本243 auto perm = GetPermResult(dim, dimSize, executor); // 不需改动时的版本
244 if (dimSize > DIM_MAX) {244 if (dimSize > DIM_MAX) {
245- perm = GetPermResult(1, 3, executor); // 1为sort的维度，3为reshape以后总共有3维245+ perm = GetPermResult(1, 3, executor); // 1为sort的维度，3为reshape以后总共有3维
246 }246 }
247 return perm;247 return perm;
248 }248 }
@@ -277,6 +277,8 @@ static bool IsNoTransposeProfitable(const aclTensor* self, int64_t dim)
277 277 
278static bool UseNoTranspose(const aclTensor* self, int64_t dim)278static bool UseNoTranspose(const aclTensor* self, int64_t dim)
279{279{
280+ constexpr int64_t NON_LAST_AXIS_MIN_LEN = 2;
281+ constexpr int64_t NON_LAST_AXIS_MAX_LEN = 2048;
280 if (!IsRegBase()) {282 if (!IsRegBase()) {
281 return false;283 return false;
282 }284 }
@@ -286,8 +288,7 @@ static bool UseNoTranspose(const aclTensor* self, int64_t dim)
286 }288 }
287 auto selfShape = self->GetViewShape();289 auto selfShape = self->GetViewShape();
288 int64_t axisLen = selfShape[dim];290 int64_t axisLen = selfShape[dim];
289- int64_t axisThreshold = 2048;291+ if (axisLen < NON_LAST_AXIS_MIN_LEN || axisLen > NON_LAST_AXIS_MAX_LEN) {
290- if (axisLen < 2 || axisLen > axisThreshold) {
291 return false;292 return false;
292 }293 }
293 return IsNoTransposeProfitable(self, dim);294 return IsNoTransposeProfitable(self, dim);
@@ -370,6 +371,43 @@ static const aclTensor* GetTensorWithValueZero(aclTensor* out, aclOpExecutor* ex
370 return viewCopyResult;371 return viewCopyResult;
371}372}
372 373 
374+static aclnnStatus BuildSortGraph(const aclTensor* self, int64_t dimPositive, int64_t dimSize, bool stable,
375+ bool descending, aclTensor* valuesOut, aclTensor* indicesOut, aclOpExecutor* executor)
376+{
377+ auto selfContiguous = l0op::Contiguous(self, executor);
378+ CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
379+ if (self->GetDataType() == op::DataType::DT_BF16 && !IsRegBase()) {
380+ selfContiguous = l0op::Cast(selfContiguous, op::DataType::DT_FLOAT, executor);
381+ CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
382+ }
383+ if (self->GetDataType() == op::DataType::DT_BOOL) {
384+ selfContiguous = l0op::Cast(selfContiguous, op::DataType::DT_UINT8, executor);
385+ CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
386+ }
387+ 
388+ auto selfShapeDetail = GetTensorShape(selfContiguous, executor);
389+ auto indicesType = indicesOut->GetDataType();
390+ std::tuple<const aclTensor*, const aclTensor*> sortRes;
391+ if (UseNoTranspose(selfContiguous, dimPositive)) {
392+ sortRes = l0op::Sort(selfContiguous, dimPositive, descending, stable, indicesType, executor);
393+ } else {
394+ if (dimSize > DIM_MAX) {
395+ auto shapeNew = reshapeShape(selfContiguous, dimPositive, executor);
396+ selfContiguous = reshapeIfLargeTensor(selfContiguous, executor, dimSize, shapeNew);
397+ }
398+ auto perm = updatePerm(dimPositive, dimSize, executor);
399+ sortRes = SortProcess(selfContiguous, perm, stable, descending, indicesType, executor);
400+ }
401+ CHECK_RET(CheckTupleNullptr(sortRes), ACLNN_ERR_PARAM_NULLPTR);
402+ 
403+ auto castRes = reshapeCastRes(sortRes, std::tie(valuesOut, indicesOut), dimSize, selfShapeDetail, executor);
404+ CHECK_RET(CheckTupleNullptr(castRes), ACLNN_ERR_PARAM_NULLPTR);
405+ auto viewCopyValues = l0op::ViewCopy(std::get<0>(castRes), valuesOut, executor);
406+ auto viewCopyIndices = l0op::ViewCopy(std::get<1>(castRes), indicesOut, executor);
407+ CHECK_RET(viewCopyValues != nullptr && viewCopyIndices != nullptr, ACLNN_ERR_PARAM_NULLPTR);
408+ return ACLNN_SUCCESS;
409+}
410+ 
373aclnnStatus aclnnSortGetWorkspaceSize(const aclTensor* self, bool stable, int64_t dim, bool descending,411aclnnStatus aclnnSortGetWorkspaceSize(const aclTensor* self, bool stable, int64_t dim, bool descending,
374 aclTensor* valuesOut, aclTensor* indicesOut, uint64_t* workspaceSize,412 aclTensor* valuesOut, aclTensor* indicesOut, uint64_t* workspaceSize,
375 aclOpExecutor** executor)413 aclOpExecutor** executor)
@@ -424,47 +462,8 @@ aclnnStatus aclnnSortGetWorkspaceSize(const aclTensor* self, bool stable, int64_
424 return ACLNN_SUCCESS;462 return ACLNN_SUCCESS;
425 }463 }
426 464 
427- auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());465+ ret = BuildSortGraph(self, dimPositive, dimSize, stable, descending, valuesOut, indicesOut, uniqueExecutor.get());
428- CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);466+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
429- 
430- // kernel暂不支持bf16输入，转为fp32进行计算
431- if (self->GetDataType() == op::DataType::DT_BF16 && !IsRegBase()) {
432- selfContiguous = l0op::Cast(selfContiguous, op::DataType::DT_FLOAT, uniqueExecutor.get());
433- CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
434- }
435- 
436- // kernel暂不支持bool输入，转为uint8进行计算
437- if (self->GetDataType() == op::DataType::DT_BOOL) {
438- selfContiguous = l0op::Cast(selfContiguous, op::DataType::DT_UINT8, uniqueExecutor.get());
439- CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);
440- }
441- 
442- auto selfShapeDetail = GetTensorShape(selfContiguous, uniqueExecutor.get()); // self最原始的shape
443- 
444- auto indicesType = indicesOut->GetDataType();
445- std::tuple<const aclTensor*, const aclTensor*> sortRes;
446- if (UseNoTranspose(selfContiguous, dimPositive)) {
447- sortRes = l0op::Sort(selfContiguous, dimPositive, descending, stable, indicesType, uniqueExecutor.get());
448- } else {
449- // 如果大于8维，需要reshape
450- if (dimSize > DIM_MAX) {
451- auto shapeNew = reshapeShape(selfContiguous, dimPositive, uniqueExecutor.get());
452- selfContiguous = reshapeIfLargeTensor(selfContiguous, uniqueExecutor.get(), dimSize, shapeNew);
453- }
454- auto perm = updatePerm(dimPositive, dimSize, uniqueExecutor.get());
455- sortRes = SortProcess(selfContiguous, perm, stable, descending, indicesType, uniqueExecutor.get());
456- }
457- CHECK_RET(CheckTupleNullptr(sortRes), ACLNN_ERR_PARAM_NULLPTR);
458- 
459- auto expectedCastRes = std::tie(valuesOut, indicesOut);
460- auto castRes = reshapeCastRes(sortRes, expectedCastRes, dimSize, selfShapeDetail, uniqueExecutor.get());
461- CHECK_RET(CheckTupleNullptr(castRes), ACLNN_ERR_PARAM_NULLPTR);
462- auto valuesCast = std::get<0>(castRes);
463- auto indicesCast = std::get<1>(castRes);
464- 
465- auto viewCopyValues = l0op::ViewCopy(valuesCast, valuesOut, uniqueExecutor.get());
466- auto viewCopyIndices = l0op::ViewCopy(indicesCast, indicesOut, uniqueExecutor.get());
467- CHECK_RET(viewCopyValues != nullptr && viewCopyIndices != nullptr, ACLNN_ERR_PARAM_NULLPTR);
468 467 
469 *workspaceSize = uniqueExecutor->GetWorkspaceSize();468 *workspaceSize = uniqueExecutor->GetWorkspaceSize();
470 uniqueExecutor.ReleaseTo(executor);469 uniqueExecutor.ReleaseTo(executor);
Mmath/sort/op_api/sort.cpp+48-47
@@ -34,11 +34,11 @@ static const int64_t NON_LAST_SMALL_AXIS_MIN = 2;
34static const int64_t NON_LAST_SMALL_AXIS_MAX = 2048;34static const int64_t NON_LAST_SMALL_AXIS_MAX = 2048;
35 35 
36// 根据排序轴的数据量大小判断是否支持aicore36// 根据排序轴的数据量大小判断是否支持aicore
37-static bool SocSupportDimSize(const aclTensor *self)37+static bool SocSupportDimSize(const aclTensor* self)
38{38{
39 // 该维度数据量为1 或数据量>100000 走AICPU39 // 该维度数据量为1 或数据量>100000 走AICPU
40 auto shapeSize = (int64_t)(self->GetViewShape().GetDimNum());40 auto shapeSize = (int64_t)(self->GetViewShape().GetDimNum());
41- auto lastDimSize = (self->GetViewShape())[shapeSize-1];41+ auto lastDimSize = (self->GetViewShape())[shapeSize - 1];
42 auto socVersion = GetCurrentPlatformInfo().GetSocVersion();42 auto socVersion = GetCurrentPlatformInfo().GetSocVersion();
43 if (socVersion == SocVersion::ASCEND310 || socVersion == SocVersion::ASCEND310B) {43 if (socVersion == SocVersion::ASCEND310 || socVersion == SocVersion::ASCEND310B) {
44 // sort轴数据量大于100k44 // sort轴数据量大于100k
@@ -55,14 +55,14 @@ static bool SocSupportDimSize(const aclTensor *self)
55}55}
56 56 
57// 根据dtype判断是否支持aicore57// 根据dtype判断是否支持aicore
58-static bool SocSupportDtype(const aclTensor *self)58+static bool SocSupportDtype(const aclTensor* self)
59{59{
60 auto socVersion = GetCurrentPlatformInfo().GetSocVersion();60 auto socVersion = GetCurrentPlatformInfo().GetSocVersion();
61 // AiCore只支持FLOAT16 + FLOAT3261 // AiCore只支持FLOAT16 + FLOAT32
62 if (CheckType(self->GetDataType(), AICORE_DTYPE_SUPPORT_LIST)) {62 if (CheckType(self->GetDataType(), AICORE_DTYPE_SUPPORT_LIST)) {
63 // 910和310芯片 + tensor为FLOAT32 则不支持AiCore63 // 910和310芯片 + tensor为FLOAT32 则不支持AiCore
64 if (((socVersion == SocVersion::ASCEND910) || (socVersion == SocVersion::ASCEND310)) &&64 if (((socVersion == SocVersion::ASCEND910) || (socVersion == SocVersion::ASCEND310)) &&
65- (self->GetDataType()==op::DataType::DT_FLOAT || self->GetDataType()==op::DataType::DT_BF16)) {65+ (self->GetDataType() == op::DataType::DT_FLOAT || self->GetDataType() == op::DataType::DT_BF16)) {
66 return false;66 return false;
67 }67 }
68 return true;68 return true;
@@ -70,9 +70,9 @@ static bool SocSupportDtype(const aclTensor *self)
70 return false;70 return false;
71}71}
72 72 
73-static bool IsAiCoreSupport(const aclTensor *self, bool stable, bool descending)73+static bool IsAiCoreSupport(const aclTensor* self, bool stable, bool descending)
74{74{
75- if (IsRegBase()){75+ if (IsRegBase()) {
76 return true;76 return true;
77 } else if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND310B && stable && !descending) {77 } else if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND310B && stable && !descending) {
78 return false;78 return false;
@@ -81,7 +81,7 @@ static bool IsAiCoreSupport(const aclTensor *self, bool stable, bool descending)
81 }81 }
82}82}
83 83 
84-static bool IsLastAxisOrSupportedNonLastAxis(const aclTensor *self, int64_t dim)84+static bool IsLastAxisOrSupportedNonLastAxis(const aclTensor* self, int64_t dim)
85{85{
86 auto dimSize = static_cast<int64_t>(self->GetViewShape().GetDimNum());86 auto dimSize = static_cast<int64_t>(self->GetViewShape().GetDimNum());
87 auto dimValue = dim < 0 ? dim + dimSize : dim;87 auto dimValue = dim < 0 ? dim + dimSize : dim;
@@ -92,25 +92,24 @@ static bool IsLastAxisOrSupportedNonLastAxis(const aclTensor *self, int64_t dim)
92 return axisLen >= NON_LAST_SMALL_AXIS_MIN && axisLen <= NON_LAST_SMALL_AXIS_MAX;92 return axisLen >= NON_LAST_SMALL_AXIS_MIN && axisLen <= NON_LAST_SMALL_AXIS_MAX;
93}93}
94 94 
95-void SortAiCore(const aclTensor *self, bool stable, int64_t dim, bool descending, aclTensor *values, aclTensor *indices,95+void SortAiCore(const aclTensor* self, bool stable, int64_t dim, bool descending, aclTensor* values, aclTensor* indices,
96- aclOpExecutor* executor)96+ aclOpExecutor* executor)
97{97{
98 L0_DFX(SortAiCore, self, stable, dim, descending, values, indices);98 L0_DFX(SortAiCore, self, stable, dim, descending, values, indices);
99 99 
100 auto dimSize = (int64_t)(self->GetViewShape().GetDimNum());100 auto dimSize = (int64_t)(self->GetViewShape().GetDimNum());
101- if ((dimSize!= dim + 1) && (dim != -1)) {101+ if ((dimSize != dim + 1) && (dim != -1)) {
102 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "dim must equal to the (number of dimensions - 1 ) or -1.");102 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "dim must equal to the (number of dimensions - 1 ) or -1.");
103 }103 }
104 104 
105- ADD_TO_LAUNCHER_LIST_AICORE(Sort, OP_INPUT(self), OP_OUTPUT(values, indices),105+ ADD_TO_LAUNCHER_LIST_AICORE(Sort, OP_INPUT(self), OP_OUTPUT(values, indices), OP_ATTR(dim, descending, stable));
106- OP_ATTR(dim, descending, stable));
107}106}
108 107 
109-static void SortAiCoreForDavid(const aclTensor *self, bool stable, int64_t dim, bool descending, aclTensor *values,108+static void SortAiCoreForDavid(const aclTensor* self, bool stable, int64_t dim, bool descending, aclTensor* values,
110- aclTensor *indices, op::DataType indicesType, aclOpExecutor* executor)109+ aclTensor* indices, op::DataType indicesType, aclOpExecutor* executor)
111{110{
112 L0_DFX(SortAiCoreForDavid, self, stable, dim, descending, values, indices, indicesType);111 L0_DFX(SortAiCoreForDavid, self, stable, dim, descending, values, indices, indicesType);
113- 112+ 
114 auto dimSize = static_cast<int64_t>(self->GetViewShape().GetDimNum());113 auto dimSize = static_cast<int64_t>(self->GetViewShape().GetDimNum());
115 if ((dim >= dimSize) || (dim + dimSize < 0)) {114 if ((dim >= dimSize) || (dim + dimSize < 0)) {
116 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Dim must be in range [-N, N-1]. Current dim is %ld.", dim);115 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Dim must be in range [-N, N-1]. Current dim is %ld.", dim);
@@ -118,29 +117,29 @@ static void SortAiCoreForDavid(const aclTensor *self, bool stable, int64_t dim,
118 }117 }
119 if (!IsLastAxisOrSupportedNonLastAxis(self, dim)) {118 if (!IsLastAxisOrSupportedNonLastAxis(self, dim)) {
120 OP_LOGE(ACLNN_ERR_PARAM_INVALID,119 OP_LOGE(ACLNN_ERR_PARAM_INVALID,
121- "Dim must be the last axis or a supported non-last axis with axis size in [2, 2048]. "120+ "Dim must be the last axis or a supported non-last axis with axis size in [2, 2048]. "
122- "Current dim is %ld.",121+ "Current dim is %ld.",
123- dim);122+ dim);
124 return;123 return;
125 }124 }
126 125 
127 ADD_TO_LAUNCHER_LIST_AICORE(Sort, OP_INPUT(self), OP_OUTPUT(values, indices),126 ADD_TO_LAUNCHER_LIST_AICORE(Sort, OP_INPUT(self), OP_OUTPUT(values, indices),
128- OP_ATTR(dim, descending, stable, indicesType));127+ OP_ATTR(dim, descending, stable, indicesType));
129}128}
130 129 
131-std::tuple<aclTensor*, aclTensor*> SortAiCpu(const aclTensor *self, bool stable, int64_t dim, bool descending,130+std::tuple<aclTensor*, aclTensor*> SortAiCpu(const aclTensor* self, bool stable, int64_t dim, bool descending,
132- aclTensor *values, aclTensor *indices, aclOpExecutor* executor)131+ aclTensor* values, aclTensor* indices, aclOpExecutor* executor)
133{132{
134 L0_DFX(SortAiCpu, self, stable, dim, descending, values, indices);133 L0_DFX(SortAiCpu, self, stable, dim, descending, values, indices);
135 134 
136 auto dimSize = (int64_t)(self->GetViewShape().GetDimNum());135 auto dimSize = (int64_t)(self->GetViewShape().GetDimNum());
137- if ((dim > (dimSize-1)) || (dim + dimSize < 0)) {136+ if ((dim > (dimSize - 1)) || (dim + dimSize < 0)) {
138 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "dim must be in range [-N, N-1]. Current dim is %ld.", dim);137 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "dim must be in range [-N, N-1]. Current dim is %ld.", dim);
139 }138 }
140 139 
141 static internal::AicpuTaskSpace space("Sort");140 static internal::AicpuTaskSpace space("Sort");
142 auto ret = ADD_TO_LAUNCHER_LIST_AICPU(Sort, OP_ATTR_NAMES({"axis", "descending", "stable"}), OP_INPUT(self),141 auto ret = ADD_TO_LAUNCHER_LIST_AICPU(Sort, OP_ATTR_NAMES({"axis", "descending", "stable"}), OP_INPUT(self),
143- OP_OUTPUT(values, indices), OP_ATTR(dim, descending, stable));142+ OP_OUTPUT(values, indices), OP_ATTR(dim, descending, stable));
144 if (ret != ACLNN_SUCCESS) {143 if (ret != ACLNN_SUCCESS) {
145 return std::tuple<aclTensor*, aclTensor*>(nullptr, nullptr);144 return std::tuple<aclTensor*, aclTensor*>(nullptr, nullptr);
146 }145 }
@@ -148,8 +147,8 @@ std::tuple<aclTensor*, aclTensor*> SortAiCpu(const aclTensor *self, bool stable,
148}147}
149 148 
150// 处理0维场景：self cast以后viewcopy, indices是为0149// 处理0维场景：self cast以后viewcopy, indices是为0
151-static aclnnStatus HandleDimZeroTensor(const aclTensor *self, aclTensor *valuesOut, aclTensor *indicesOut,150+static aclnnStatus HandleDimZeroTensor(const aclTensor* self, aclTensor* valuesOut, aclTensor* indicesOut,
152- aclOpExecutor* executor)151+ aclOpExecutor* executor)
153{152{
154 auto selfContiguous = l0op::Contiguous(self, executor);153 auto selfContiguous = l0op::Contiguous(self, executor);
155 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);154 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
@@ -171,7 +170,17 @@ static aclnnStatus HandleDimZeroTensor(const aclTensor *self, aclTensor *valuesO
171 return ACLNN_SUCCESS;170 return ACLNN_SUCCESS;
172}171}
173 172 
174-const std::tuple<aclTensor*, aclTensor*> Sort(const aclTensor *self, int64_t dim, bool descending, bool stable,173+static std::tuple<aclTensor*, aclTensor*> AllocateSortOutputs(const aclTensor* self, const op::Shape& shape,
174+ op::Format format, op::DataType indicesType,
175+ aclOpExecutor* executor)
176+{
177+ auto values = executor->AllocTensor(shape, self->GetDataType(), format);
178+ auto outputIndicesType = IsRegBase() ? indicesType : op::DataType::DT_INT32;
179+ auto indices = executor->AllocTensor(shape, outputIndicesType, format);
180+ return std::tie(values, indices);
181+}
182+ 
183+const std::tuple<aclTensor*, aclTensor*> Sort(const aclTensor* self, int64_t dim, bool descending, bool stable,
175 op::DataType indicesType, aclOpExecutor* executor)184 op::DataType indicesType, aclOpExecutor* executor)
176{185{
177 L0_DFX(Sort, self, dim, descending, stable, indicesType);186 L0_DFX(Sort, self, dim, descending, stable, indicesType);
@@ -181,24 +190,20 @@ const std::tuple<aclTensor*, aclTensor*> Sort(const aclTensor *self, int64_t dim
181 auto dimSize = (int64_t)(selfShape.GetDimNum());190 auto dimSize = (int64_t)(selfShape.GetDimNum());
182 if (dimSize < 0 || dimSize > AXIS_LIMIT) {191 if (dimSize < 0 || dimSize > AXIS_LIMIT) {
183 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor self dimension size must be in range [0, 8]. Current size is [%ld].",192 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor self dimension size must be in range [0, 8]. Current size is [%ld].",
184- dimSize);193+ dimSize);
185 }194 }
186 // handle dimsize=0195 // handle dimsize=0
187- if (dimSize == 0) { 196+ if (dimSize == 0) {
188 if (dim == 0 || dim == -1) {197 if (dim == 0 || dim == -1) {
189- auto valuesOut = executor->AllocTensor(selfShape, self->GetDataType(), selfFormat);198+ auto outputs = AllocateSortOutputs(self, selfShape, selfFormat, indicesType, executor);
190- aclTensor* indicesOut = nullptr;199+ auto valuesOut = std::get<0>(outputs);
191- if (IsRegBase()) {200+ auto indicesOut = std::get<1>(outputs);
192- indicesOut = executor->AllocTensor(selfShape, indicesType, selfFormat);
193- } else {
194- indicesOut = executor->AllocTensor(selfShape, op::DataType::DT_INT32, selfFormat);
195- }
196 auto res = HandleDimZeroTensor(self, valuesOut, indicesOut, executor);201 auto res = HandleDimZeroTensor(self, valuesOut, indicesOut, executor);
197 if (res != ACLNN_SUCCESS) {202 if (res != ACLNN_SUCCESS) {
198 OP_LOGE(ACLNN_ERR_INNER, "HandleDimZeroTensor error.");203 OP_LOGE(ACLNN_ERR_INNER, "HandleDimZeroTensor error.");
199 return std::tuple<aclTensor*, aclTensor*>(nullptr, nullptr);204 return std::tuple<aclTensor*, aclTensor*>(nullptr, nullptr);
200 } else {205 } else {
201- return std::tie(valuesOut, indicesOut); 206+ return std::tie(valuesOut, indicesOut);
202 }207 }
203 } else {208 } else {
204 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When dimSize == 0 , dim must be 0 or -1.");209 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When dimSize == 0 , dim must be 0 or -1.");
@@ -207,18 +212,14 @@ const std::tuple<aclTensor*, aclTensor*> Sort(const aclTensor *self, int64_t dim
207 }212 }
208 auto lastDimSize = selfShape[dimSize - 1];213 auto lastDimSize = selfShape[dimSize - 1];
209 // The Sort Op not support sort axis is 1 when input type is BF16..214 // The Sort Op not support sort axis is 1 when input type is BF16..
210- bool isNotSupport = (1 == lastDimSize && op::DataType::DT_BF16 == self->GetDataType());215+ bool isNotSupport = (1 == lastDimSize && op::DataType::DT_BF16 == self->GetDataType());
211 if (isNotSupport && !IsRegBase()) {216 if (isNotSupport && !IsRegBase()) {
212- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The sort axis value is not support 1 when input type is BF16.");217+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The sort axis value is not support 1 when input type is BF16.");
213- return std::tuple<aclTensor*, aclTensor*>(nullptr, nullptr);218+ return std::tuple<aclTensor*, aclTensor*>(nullptr, nullptr);
214- }
215- auto values = executor->AllocTensor(selfShape, self->GetDataType(), selfFormat);
216- aclTensor* indices = nullptr;
217- if (IsRegBase()) {
218- indices = executor->AllocTensor(selfShape, indicesType, selfFormat);
219- } else {
220- indices = executor->AllocTensor(selfShape, op::DataType::DT_INT32, selfFormat);
221 }219 }
220+ auto outputs = AllocateSortOutputs(self, selfShape, selfFormat, indicesType, executor);
221+ auto values = std::get<0>(outputs);
222+ auto indices = std::get<1>(outputs);
222 if (IsAiCoreSupport(self, stable, descending)) {223 if (IsAiCoreSupport(self, stable, descending)) {
223 if (IsRegBase()) {224 if (IsRegBase()) {
224 SortAiCoreForDavid(self, stable, dim, descending, values, indices, indicesType, executor);225 SortAiCoreForDavid(self, stable, dim, descending, values, indices, indicesType, executor);
@@ -230,4 +231,4 @@ const std::tuple<aclTensor*, aclTensor*> Sort(const aclTensor *self, int64_t dim
230 }231 }
231 return std::tie(values, indices);232 return std::tie(values, indices);
232}233}
233-}234+} // namespace l0op
Mmath/sort/op_host/arch35/sort_tiling.h+8-8
@@ -18,15 +18,15 @@
18 18 
19namespace optiling {19namespace optiling {
20struct SortTilingData {20struct SortTilingData {
21- int32_t task_num;21+ int32_t task_num;
22- int32_t num_per_task;22+ int32_t num_per_task;
23};23};
24 24 
25struct SortCompileInfo {25struct SortCompileInfo {
26- int32_t core_num;26+ int32_t core_num;
27- int32_t num_block;27+ int32_t num_block;
28- int32_t num_offset;28+ int32_t num_offset;
29- int32_t float_bytes;29+ int32_t float_bytes;
30};30};
31-} // namespace optiling31+} // namespace optiling
32-#endif // OPS_BUILT_IN_OP_TILING_RUNTIME_SORT_TILING_H32+#endif // OPS_BUILT_IN_OP_TILING_RUNTIME_SORT_TILING_H
Mmath/sort/op_host/arch35/sort_tiling_arch35.cpp+4-3
@@ -255,11 +255,11 @@ bool IsRadixSortOneCore(SortKthTileInfo& sortTileInfo)
255 return false;255 return false;
256 }256 }
257 257 
258- int64_t doubleBufferRemainUb = static_cast<int64_t>(sortTileInfo.ubSize) - oneBufferQueSize * 2;258+ int64_t doubleBufferRemainUb = static_cast<int64_t>(sortTileInfo.ubSize) - oneBufferQueSize * DOUBLE_BUFFER_NUM;
259 doubleBufferRemainUb = (doubleBufferRemainUb / static_cast<int64_t>(sortTileInfo.blockUbSize)) *259 doubleBufferRemainUb = (doubleBufferRemainUb / static_cast<int64_t>(sortTileInfo.blockUbSize)) *
260 static_cast<int64_t>(sortTileInfo.blockUbSize);260 static_cast<int64_t>(sortTileInfo.blockUbSize);
261 if (tmpUb <= doubleBufferRemainUb) {261 if (tmpUb <= doubleBufferRemainUb) {
262- sortTileInfo.keyParams3 = 2;262+ sortTileInfo.keyParams3 = DOUBLE_BUFFER_NUM;
263 }263 }
264 OP_LOGI("RadixSortTiling", "radix one-core bufferNum is %u", sortTileInfo.keyParams3);264 OP_LOGI("RadixSortTiling", "radix one-core bufferNum is %u", sortTileInfo.keyParams3);
265 return true;265 return true;
@@ -623,7 +623,8 @@ bool TryNonLastSmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTile
623 if (!sortTileInfo.isNonLastAxis) {623 if (!sortTileInfo.isNonLastAxis) {
624 return false;624 return false;
625 }625 }
626- if (sortTileInfo.lastAxis < 2 || sortTileInfo.lastAxis > NON_LAST_SMALL_AXIS_THRESHOLD) {626+ if (sortTileInfo.lastAxis < NON_LAST_SMALL_AXIS_MIN_AXIS_LEN ||
627+ sortTileInfo.lastAxis > NON_LAST_SMALL_AXIS_THRESHOLD) {
627 return false;628 return false;
628 }629 }
629 if (sortTileInfo.innerSize <= 0 || sortTileInfo.outerSize <= 0) {630 if (sortTileInfo.innerSize <= 0 || sortTileInfo.outerSize <= 0) {
Mmath/sort/op_host/arch35/sort_tiling_arch35.h+2-2
@@ -18,6 +18,6 @@
18#include "register/tilingdata_base.h"18#include "register/tilingdata_base.h"
19 19 
20namespace optiling {20namespace optiling {
21-ge::graphStatus SortTilingSimt(gert::TilingContext* context, int32_t maxCoreNum);21+ge::graphStatus SortTilingSimt(gert::TilingContext* context, int32_t maxCoreNum);
22}22}
23-#endif // SORT_TILING_ARCH35_H23+#endif // SORT_TILING_ARCH35_H
Mmath/sort/op_host/arch35/sort_tiling_common.cpp+34-25
@@ -67,8 +67,10 @@ uint64_t ComputeUbAfterSimtReserve(uint32_t ubSize)
67// Returns 0 when index is out of range or dataType is unsupported.67// Returns 0 when index is out of range or dataType is unsupported.
68uint32_t GetPreferredInnerChunk(ge::DataType dataType, uint32_t index)68uint32_t GetPreferredInnerChunk(ge::DataType dataType, uint32_t index)
69{69{
70+ constexpr uint32_t kMaxChunkCandidates = 6;
71+ constexpr uint32_t kOneByteChunkGroup = 3;
70 // Each row: dtype group's chunk candidates (decreasing powers of 2, padded with 0).72 // Each row: dtype group's chunk candidates (decreasing powers of 2, padded with 0).
71- static constexpr uint32_t kChunkCandidates[][6] = {73+ static constexpr uint32_t kChunkCandidates[][kMaxChunkCandidates] = {
72 {4, 2, 1, 0, 0, 0},74 {4, 2, 1, 0, 0, 0},
73 {8, 4, 2, 1, 0, 0},75 {8, 4, 2, 1, 0, 0},
74 {16, 8, 4, 2, 1, 0},76 {16, 8, 4, 2, 1, 0},
@@ -84,7 +86,7 @@ uint32_t GetPreferredInnerChunk(ge::DataType dataType, uint32_t index)
84 dataType == ge::DT_UINT16) {86 dataType == ge::DT_UINT16) {
85 group = 2; // 2-byte types87 group = 2; // 2-byte types
86 } else if (dataType == ge::DT_INT8 || dataType == ge::DT_UINT8) {88 } else if (dataType == ge::DT_INT8 || dataType == ge::DT_UINT8) {
87- group = 3; // 1-byte types89+ group = kOneByteChunkGroup;
88 } else {90 } else {
89 return 0;91 return 0;
90 }92 }
@@ -267,7 +269,8 @@ bool SearchNonLastSmallAxisPlan(
267 std::function<bool(SortKthTileInfo&, uint32_t, uint64_t&, NonLastSmallAxisCandidate&)> estimateUb,269 std::function<bool(SortKthTileInfo&, uint32_t, uint64_t&, NonLastSmallAxisCandidate&)> estimateUb,
268 NonLastSmallAxisCandidate& best, SortKthTileInfo* selectedInfo)270 NonLastSmallAxisCandidate& best, SortKthTileInfo* selectedInfo)
269{271{
270- for (uint32_t i = 0; i < 6; ++i) {272+ constexpr uint32_t kMaxChunkCandidates = 6;
273+ for (uint32_t i = 0; i < kMaxChunkCandidates; ++i) {
271 uint32_t chunk = GetPreferredInnerChunk(info.dataType, i);274 uint32_t chunk = GetPreferredInnerChunk(info.dataType, i);
272 if (chunk == 0U) {275 if (chunk == 0U) {
273 break;276 break;
@@ -589,6 +592,30 @@ bool ComputeRadixSortWorkspace(int64_t axisLen, uint32_t dtypeSize, uint32_t ind
589 return true;592 return true;
590}593}
591 594 
595+static bool FillRadixKernelResources(SortKthTileInfo& info, uint32_t indexSize, uint32_t lastDimTileNum,
596+ uint32_t tmpUbSize)
597+{
598+ RadixClearParams clearParams;
599+ if (!FillRadixKernelParams(info.dtypeSize, indexSize, info.coreNumNeed, lastDimTileNum, info.unsortedDimParallel,
600+ info.blockUbSize, tmpUbSize, clearParams)) {
601+ return false;
602+ }
603+ info.keyParams0 = clearParams.keyParams0;
604+ info.keyParams1 = clearParams.keyParams1;
605+ info.keyParams2 = clearParams.keyParams2;
606+ info.keyParams3 = clearParams.keyParams3;
607+ info.keyParams4 = clearParams.keyParams4;
608+ info.keyParams5 = clearParams.keyParams5;
609+ uint64_t sortWorkspaceSize = 0;
610+ if (!ComputeRadixSortWorkspace(info.lastAxis, info.dtypeSize, indexSize, lastDimTileNum, info.numTileDataSize,
611+ info.unsortedDimParallel, info.keyParams0, info.keyParams1, info.keyParams2,
612+ info.keyParams3, info.keyParams4, info.blockUbSize, sortWorkspaceSize)) {
613+ return false;
614+ }
615+ info.workspaceSize = static_cast<size_t>(sortWorkspaceSize + WORK_SPACE_SIZE);
616+ return true;
617+}
618+ 
592bool FillRadixMoreCoreInfo(SortKthTileInfo& info)619bool FillRadixMoreCoreInfo(SortKthTileInfo& info)
593{620{
594 uint32_t usableUb = info.ubSize > SIMT_UB ? info.ubSize - SIMT_UB : 0;621 uint32_t usableUb = info.ubSize > SIMT_UB ? info.ubSize - SIMT_UB : 0;
@@ -623,25 +650,7 @@ bool FillRadixMoreCoreInfo(SortKthTileInfo& info)
623 info.lastDimNeedCore = std::min(info.maxCoreNum, lastDimTileNum);650 info.lastDimNeedCore = std::min(info.maxCoreNum, lastDimTileNum);
624 info.coreNumNeed = info.unsortedDimParallel * info.lastDimNeedCore;651 info.coreNumNeed = info.unsortedDimParallel * info.lastDimNeedCore;
625 info.lastDimTileNum = lastDimTileNum;652 info.lastDimTileNum = lastDimTileNum;
626- RadixClearParams clearParams;653+ return FillRadixKernelResources(info, indexSize, lastDimTileNum, tmpUbSize);
627- if (!FillRadixKernelParams(info.dtypeSize, indexSize, info.coreNumNeed, lastDimTileNum, info.unsortedDimParallel,
628- info.blockUbSize, tmpUbSize, clearParams)) {
629- return false;
630- }
631- info.keyParams0 = clearParams.keyParams0;
632- info.keyParams1 = clearParams.keyParams1;
633- info.keyParams2 = clearParams.keyParams2;
634- info.keyParams3 = clearParams.keyParams3;
635- info.keyParams4 = clearParams.keyParams4;
636- info.keyParams5 = clearParams.keyParams5;
637- uint64_t sortWorkspaceSize = 0;
638- if (!ComputeRadixSortWorkspace(info.lastAxis, info.dtypeSize, indexSize, lastDimTileNum, info.numTileDataSize,
639- info.unsortedDimParallel, info.keyParams0, info.keyParams1, info.keyParams2,
640- info.keyParams3, info.keyParams4, info.blockUbSize, sortWorkspaceSize)) {
641- return false;
642- }
643- info.workspaceSize = static_cast<size_t>(sortWorkspaceSize + WORK_SPACE_SIZE);
644- return true;
645}654}
646 655 
647bool ComputeRadixOneCoreUbSizes(int64_t lastAxis, uint32_t dtypeSize, uint32_t indexElemSize, uint32_t blockUbSize,656bool ComputeRadixOneCoreUbSizes(int64_t lastAxis, uint32_t dtypeSize, uint32_t indexElemSize, uint32_t blockUbSize,
@@ -704,7 +713,7 @@ bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t
704 info.keyParams1 = plan.alignNum * plan.oneCoreRowNum * info.dtypeSize;713 info.keyParams1 = plan.alignNum * plan.oneCoreRowNum * info.dtypeSize;
705 info.keyParams2 = plan.alignNum * plan.oneCoreRowNum * indexDtypeSize;714 info.keyParams2 = plan.alignNum * plan.oneCoreRowNum * indexDtypeSize;
706 info.keyParams3 = plan.alignNum;715 info.keyParams3 = plan.alignNum;
707- info.keyParams4 = info.lastAxis > 2048 ? 1 : 2;716+ info.keyParams4 = info.lastAxis > ONE_CORE_DATA_SIZE ? 1 : DOUBLE_BUFFER_NUM;
708 info.tmpUbSize = std::max(concatTmpSize, info.blockUbSize);717 info.tmpUbSize = std::max(concatTmpSize, info.blockUbSize);
709 info.workspaceSize = WORK_SPACE_SIZE;718 info.workspaceSize = WORK_SPACE_SIZE;
710 return true;719 return true;
@@ -826,7 +835,7 @@ bool IsMergeIntraCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t u
826 uint32_t ubSize)835 uint32_t ubSize)
827{836{
828 if (dataType != ge::DT_FLOAT || axisLen <= static_cast<int64_t>(MERGE_SORT_MAX_AXIS_FP32) ||837 if (dataType != ge::DT_FLOAT || axisLen <= static_cast<int64_t>(MERGE_SORT_MAX_AXIS_FP32) ||
829- unsortedDim < static_cast<int64_t>(maxCoreNum / 2)) {838+ unsortedDim < static_cast<int64_t>(maxCoreNum / DOUBLE_BUFFER_NUM)) {
830 return false;839 return false;
831 }840 }
832 uint64_t maxBatch = static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) * static_cast<uint64_t>(maxCoreNum);841 uint64_t maxBatch = static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) * static_cast<uint64_t>(maxCoreNum);
@@ -1156,7 +1165,7 @@ static bool EstimateSmallAxisTwoStageBatching(const SortKthTileInfo& info, uint3
1156 maxBatch = std::min(maxBatch, MaxTwoStageU16SafeBatch(axisLen));1165 maxBatch = std::min(maxBatch, MaxTwoStageU16SafeBatch(axisLen));
1157 }1166 }
1158 TwoStageBatchPlan result;1167 TwoStageBatchPlan result;
1159- auto tryCandidate = [&](uint32_t candidate, TwoStageBatchPlan& p) -> bool {1168+ auto tryCandidate = [&info, &computeBatchNum](uint32_t candidate, TwoStageBatchPlan& p) -> bool {
1160 SmallAxisRoutePlan candidatePlan;1169 SmallAxisRoutePlan candidatePlan;
1161 if (!TrySmallAxisTwoStageBatchCandidate(info, candidate, computeBatchNum, candidatePlan)) {1170 if (!TrySmallAxisTwoStageBatchCandidate(info, candidate, computeBatchNum, candidatePlan)) {
1162 return false;1171 return false;
Mmath/sort/op_host/arch35/sort_tiling_common.h+4-2
@@ -37,8 +37,10 @@ constexpr uint32_t SMALL_TILE_DATA_NUM = 1024; // 测试数据得出一
37constexpr uint32_t SIMT_UB = 32768; // 预留了32k给simt使用37constexpr uint32_t SIMT_UB = 32768; // 预留了32k给simt使用
38constexpr int64_t RADIX_UINT32_VALUE_MAX = 0x3fffffff; // uint32 radix counters reserve the top two bits for state38constexpr int64_t RADIX_UINT32_VALUE_MAX = 0x3fffffff; // uint32 radix counters reserve the top two bits for state
39constexpr uint32_t SMALL_AXIS_THRESHOLD = 512;39constexpr uint32_t SMALL_AXIS_THRESHOLD = 512;
40+constexpr int64_t NON_LAST_SMALL_AXIS_MIN_AXIS_LEN = 2;
40constexpr int64_t NON_LAST_SMALL_AXIS_THRESHOLD = 2048;41constexpr int64_t NON_LAST_SMALL_AXIS_THRESHOLD = 2048;
41constexpr int64_t ONE_CORE_DATA_SIZE = 2048;42constexpr int64_t ONE_CORE_DATA_SIZE = 2048;
43+constexpr uint32_t DOUBLE_BUFFER_NUM = 2;
42constexpr uint32_t SORT32_SMALL_AXIS_THRESHOLD = 32;44constexpr uint32_t SORT32_SMALL_AXIS_THRESHOLD = 32;
43constexpr uint32_t SMALL_AXIS_MAX_DATACOPY_BLOCK_COUNT = 4095; // DataCopy hardware limit for blockCount45constexpr uint32_t SMALL_AXIS_MAX_DATACOPY_BLOCK_COUNT = 4095; // DataCopy hardware limit for blockCount
44constexpr uint32_t SORT_STRUCT_BYTES = 8; // fp32 sort struct size (index + value)46constexpr uint32_t SORT_STRUCT_BYTES = 8; // fp32 sort struct size (index + value)
@@ -193,10 +195,10 @@ struct NonLastSmallAxisCandidate {
193// sorting axisLen elements at each (outer, inner) position independently.195// sorting axisLen elements at each (outer, inner) position independently.
194// innerChunk controls how many adjacent inner positions are batched into one196// innerChunk controls how many adjacent inner positions are batched into one
195// sort invocation, trading UB for fewer total tiles (tileCount).197// sort invocation, trading UB for fewer total tiles (tileCount).
196-// Candidate evaluation:198+// Candidate quality is evaluated with the following metrics.
197// tileCount = outerSize × ceil(innerSize / innerChunk)199// tileCount = outerSize × ceil(innerSize / innerChunk)
198// activeCore = min(maxCoreNum, tileCount)200// activeCore = min(maxCoreNum, tileCount)
199-// Selection criteria (in priority order):201+// Candidates are selected in the following priority order.
200// 1. Maximise activeCore — better core utilisation202// 1. Maximise activeCore — better core utilisation
201// 2. Tie-break: prefer larger innerChunk — fewer kernel invocations per core203// 2. Tie-break: prefer larger innerChunk — fewer kernel invocations per core
202// innerChunk candidates come from GetPreferredInnerChunk(), which returns204// innerChunk candidates come from GetPreferredInnerChunk(), which returns
Mmath/sort/op_host/sort_infershape.cpp+13-12
@@ -17,20 +17,21 @@
17 17 
18using namespace ge;18using namespace ge;
19namespace ops {19namespace ops {
20-static ge::graphStatus SortInferShapeFunc(gert::InferShapeContext *context) {20+static ge::graphStatus SortInferShapeFunc(gert::InferShapeContext* context)
21- const gert::Shape* x1_shape = context->GetInputShape(0);21+{
22- OP_CHECK_NULL_WITH_CONTEXT(context, x1_shape);22+ const gert::Shape* x1_shape = context->GetInputShape(0);
23+ OP_CHECK_NULL_WITH_CONTEXT(context, x1_shape);
23 24 
24- gert::Shape *output_shape_1 = context->GetOutputShape(0);25+ gert::Shape* output_shape_1 = context->GetOutputShape(0);
25- OP_CHECK_NULL_WITH_CONTEXT(context, output_shape_1);26+ OP_CHECK_NULL_WITH_CONTEXT(context, output_shape_1);
26- gert::Shape *output_shape_2 = context->GetOutputShape(1);27+ gert::Shape* output_shape_2 = context->GetOutputShape(1);
27- OP_CHECK_NULL_WITH_CONTEXT(context, output_shape_2);28+ OP_CHECK_NULL_WITH_CONTEXT(context, output_shape_2);
28-
29- *output_shape_1 = *x1_shape;
30- *output_shape_2 = *x1_shape;
31 29 
32- return ge::GRAPH_SUCCESS;30+ *output_shape_1 = *x1_shape;
31+ *output_shape_2 = *x1_shape;
32+ 
33+ return ge::GRAPH_SUCCESS;
33}34}
34 35 
35IMPL_OP_INFERSHAPE(Sort).InferShape(SortInferShapeFunc);36IMPL_OP_INFERSHAPE(Sort).InferShape(SortInferShapeFunc);
36-} // namespace ops37+} // namespace ops
Mmath/sort/op_kernel/arch35/common/merge_intra_core_base.h+42-41
@@ -102,27 +102,28 @@ public:
102 // ===== Phase 1: Sort blocks in UB =====102 // ===== Phase 1: Sort blocks in UB =====
103 __aicore__ inline void Process();103 __aicore__ inline void Process();
104 __aicore__ inline void SortSingleBatchInUb(GlobalTensor<ValueType> inputX, int64_t batchOffset);104 __aicore__ inline void SortSingleBatchInUb(GlobalTensor<ValueType> inputX, int64_t batchOffset);
105- __aicore__ inline void CopyInBlock(105+ __aicore__ inline void CopyInBlock(GlobalTensor<ValueType> inputX, LocalTensor<ValueType> xLocal,
106- GlobalTensor<ValueType> inputX, LocalTensor<ValueType> xLocal, uint32_t elemOffset, uint32_t actualElem);106+ uint32_t elemOffset, uint32_t actualElem);
107- __aicore__ inline void SortBlockToStruct(107+ __aicore__ inline void SortBlockToStruct(LocalTensor<ValueType> xLocal, LocalTensor<ValueType> sortedLocal,
108- LocalTensor<ValueType> xLocal, LocalTensor<ValueType> sortedLocal, uint32_t actualElem, uint32_t baseOffset);108+ uint32_t actualElem, uint32_t baseOffset);
109 __aicore__ inline void CopyOutToCache(LocalTensor<ValueType> sortedLocal, uint32_t blockId, uint32_t actualElem);109 __aicore__ inline void CopyOutToCache(LocalTensor<ValueType> sortedLocal, uint32_t blockId, uint32_t actualElem);
110 110 
111 // ===== Phase 2: Merge sorted blocks =====111 // ===== Phase 2: Merge sorted blocks =====
112 __aicore__ inline uint32_t MergeSingleBatch();112 __aicore__ inline uint32_t MergeSingleBatch();
113- __aicore__ inline void MergeOneGroup(113+ __aicore__ inline void MergeOneGroup(uint32_t groupStart, uint32_t groupBlockCount, uint32_t fullBlockElemCount,
114- uint32_t groupStart, uint32_t groupBlockCount, uint32_t fullBlockElemCount, uint32_t fullBlockSortLen,114+ uint32_t fullBlockSortLen, uint32_t lastBlockElemCount, uint32_t numBlocks,
115- uint32_t lastBlockElemCount, uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset,115+ uint32_t pingPongFlag, uint32_t& cumulativeOffset,
116- uint32_t& mergedGroupElemCount);116+ uint32_t& mergedGroupElemCount);
117 __aicore__ inline void CopyBlockChunk(int64_t srcAddr, int64_t dstAddr, uint32_t elemCount);117 __aicore__ inline void CopyBlockChunk(int64_t srcAddr, int64_t dstAddr, uint32_t elemCount);
118 __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, MergeListContext& ctx);118 __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, MergeListContext& ctx);
119- __aicore__ inline uint32_t LoadListsToUb(119+ __aicore__ inline uint32_t LoadListsToUb(LocalTensor<ValueType> ubMainInput,
120- LocalTensor<ValueType> ubMainInput, uint16_t elementCountList[MERGE_LIST_MAX_NUM], const MergeListContext& ctx);120+ uint16_t elementCountList[MERGE_LIST_MAX_NUM],
121- __aicore__ inline uint32_t ExecuteMrgSort(121+ const MergeListContext& ctx);
122- LocalTensor<ValueType> dstLocal, LocalTensor<ValueType> ubMainInput,122+ __aicore__ inline uint32_t ExecuteMrgSort(LocalTensor<ValueType> dstLocal, LocalTensor<ValueType> ubMainInput,
123- uint16_t elementCountList[MERGE_LIST_MAX_NUM], uint32_t listSortedNums[MERGE_LIST_MAX_NUM], uint32_t listCount);123+ uint16_t elementCountList[MERGE_LIST_MAX_NUM],
124- __aicore__ inline void CopyRemainingList(124+ uint32_t listSortedNums[MERGE_LIST_MAX_NUM], uint32_t listCount);
125- MergeListContext& ctx, int64_t dstOffsetBase, uint32_t& dstCumulativeOffset);125+ __aicore__ inline void CopyRemainingList(MergeListContext& ctx, int64_t dstOffsetBase,
126+ uint32_t& dstCumulativeOffset);
126 127 
127private:128private:
128 __aicore__ inline void PhaseBarrierAndReset();129 __aicore__ inline void PhaseBarrierAndReset();
@@ -134,8 +135,8 @@ template <typename Derived, typename ValueType, typename IndexType, bool IsDesce
134__aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::Process()135__aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::Process()
135{136{
136 int64_t startBatch = static_cast<int64_t>(this->blockIdx_) * this->batchPerCore_;137 int64_t startBatch = static_cast<int64_t>(this->blockIdx_) * this->batchPerCore_;
137- int64_t endBatch =138+ int64_t endBatch = (startBatch + this->batchPerCore_ < this->batchNum_) ? (startBatch + this->batchPerCore_) :
138- (startBatch + this->batchPerCore_ < this->batchNum_) ? (startBatch + this->batchPerCore_) : this->batchNum_;139+ this->batchNum_;
139 140 
140 if (startBatch >= this->batchNum_) {141 if (startBatch >= this->batchNum_) {
141 return;142 return;
@@ -213,10 +214,10 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce
213 }214 }
214 215 
215 DataCopyExtParams dataCopyParam{1, static_cast<uint32_t>(actualElem * sizeof(ValueType)), 0, 0, 0};216 DataCopyExtParams dataCopyParam{1, static_cast<uint32_t>(actualElem * sizeof(ValueType)), 0, 0, 0};
216- uint32_t currTileSizeAlign =217+ uint32_t currTileSizeAlign = Ops::Base::CeilAlign(actualElem,
217- Ops::Base::CeilAlign(actualElem, static_cast<uint32_t>(UB_BLOCK_BYTES / sizeof(ValueType)));218+ static_cast<uint32_t>(UB_BLOCK_BYTES / sizeof(ValueType)));
218- DataCopyPadExtParams<ValueType> padParams{219+ DataCopyPadExtParams<ValueType> padParams{true, 0, static_cast<uint8_t>(currTileSizeAlign - actualElem),
219- true, 0, static_cast<uint8_t>(currTileSizeAlign - actualElem), defaultValue};220+ defaultValue};
220 DataCopyPad(xLocal, inputX[elemOffset], dataCopyParam, padParams);221 DataCopyPad(xLocal, inputX[elemOffset], dataCopyParam, padParams);
221}222}
222 223 
@@ -224,8 +225,8 @@ template <typename Derived, typename ValueType, typename IndexType, bool IsDesce
224__aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::SortBlockToStruct(225__aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::SortBlockToStruct(
225 LocalTensor<ValueType> xLocal, LocalTensor<ValueType> sortedLocal, uint32_t actualElem, uint32_t baseOffset)226 LocalTensor<ValueType> xLocal, LocalTensor<ValueType> sortedLocal, uint32_t actualElem, uint32_t baseOffset)
226{227{
227- uint32_t alignSize =228+ uint32_t alignSize = (actualElem == this->blockSortSize_) ? this->blockSortSize_ :
228- (actualElem == this->blockSortSize_) ? this->blockSortSize_ : Ops::Base::CeilAlign(actualElem, UB_BLOCK_BYTES);229+ Ops::Base::CeilAlign(actualElem, UB_BLOCK_BYTES);
229 uint32_t sortRepeatTimes = (actualElem == this->blockSortSize_) ?230 uint32_t sortRepeatTimes = (actualElem == this->blockSortSize_) ?
230 this->sortRepeatTimes_ :231 this->sortRepeatTimes_ :
231 Ops::Base::CeilDiv(alignSize, DEALING_SORT_NUM_ONCE);232 Ops::Base::CeilDiv(alignSize, DEALING_SORT_NUM_ONCE);
@@ -257,8 +258,8 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce
257 // Cache is reused per batch (only 1 batch's cache allocated per core)258 // Cache is reused per batch (only 1 batch's cache allocated per core)
258 int64_t cacheOffset = static_cast<int64_t>(blockId) * this->blockSortLen_;259 int64_t cacheOffset = static_cast<int64_t>(blockId) * this->blockSortLen_;
259 260 
260- uint32_t sortLen =261+ uint32_t sortLen = (actualElem == this->blockSortSize_) ? this->blockSortLen_ :
261- (actualElem == this->blockSortSize_) ? this->blockSortLen_ : AscendC::GetSortLen<ValueType>(actualElem);262+ AscendC::GetSortLen<ValueType>(actualElem);
262 // {blockCount, blockLen, srcStride, dstStride, rsv}263 // {blockCount, blockLen, srcStride, dstStride, rsv}
263 DataCopyExtParams copyParams{1, static_cast<uint32_t>(sortLen * sizeof(ValueType)), 0, 0, 0};264 DataCopyExtParams copyParams{1, static_cast<uint32_t>(sortLen * sizeof(ValueType)), 0, 0, 0};
264 265 
@@ -296,9 +297,8 @@ __aicore__ inline uint32_t MergeIntraCoreBase<Derived, ValueType, IndexType, IsD
296 uint32_t groupBlockCount = (i + MERGE_LIST_MAX_NUM <= numBlocks) ? MERGE_LIST_MAX_NUM : (numBlocks - i);297 uint32_t groupBlockCount = (i + MERGE_LIST_MAX_NUM <= numBlocks) ? MERGE_LIST_MAX_NUM : (numBlocks - i);
297 298 
298 uint32_t mergedGroupElemCount = 0;299 uint32_t mergedGroupElemCount = 0;
299- MergeOneGroup(300+ MergeOneGroup(i, groupBlockCount, fullBlockElemCount, fullBlockSortLen, lastBlockElemCount, numBlocks,
300- i, groupBlockCount, fullBlockElemCount, fullBlockSortLen, lastBlockElemCount, numBlocks, pingPongFlag,301+ pingPongFlag, cumulativeOffset, mergedGroupElemCount);
301- cumulativeOffset, mergedGroupElemCount);
302 302 
303 // First full group establishes the new full block size303 // First full group establishes the new full block size
304 if (newNumBlocks == 0) {304 if (newNumBlocks == 0) {
@@ -366,9 +366,9 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce
366 uint32_t chunkSize = (remainElems > this->blockSortSize_) ? this->blockSortSize_ : remainElems;366 uint32_t chunkSize = (remainElems > this->blockSortSize_) ? this->blockSortSize_ : remainElems;
367 if (chunkSize == 0)367 if (chunkSize == 0)
368 break;368 break;
369- CopyBlockChunk(369+ CopyBlockChunk(srcOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset),
370- srcOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset),370+ dstRegionOffset + cumulativeOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset),
371- dstRegionOffset + cumulativeOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), chunkSize);371+ chunkSize);
372 remainElems -= chunkSize;372 remainElems -= chunkSize;
373 srcChunkOffset += chunkSize;373 srcChunkOffset += chunkSize;
374 }374 }
@@ -378,8 +378,9 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce
378}378}
379 379 
380template <typename Derived, typename ValueType, typename IndexType, bool IsDescend>380template <typename Derived, typename ValueType, typename IndexType, bool IsDescend>
381-__aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::CopyBlockChunk(381+__aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::CopyBlockChunk(int64_t srcAddr,
382- int64_t srcAddr, int64_t dstAddr, uint32_t elemCount)382+ int64_t dstAddr,
383+ uint32_t elemCount)
383{384{
384 uint32_t sortLen = AscendC::GetSortLen<ValueType>(elemCount);385 uint32_t sortLen = AscendC::GetSortLen<ValueType>(elemCount);
385 LocalTensor<ValueType> srcLocal = this->mergeInQueue_.template AllocTensor<ValueType>();386 LocalTensor<ValueType> srcLocal = this->mergeInQueue_.template AllocTensor<ValueType>();
@@ -429,8 +430,8 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce
429 LocalTensor<ValueType> dstLocal = this->mergeOutQueue_.template AllocTensor<ValueType>();430 LocalTensor<ValueType> dstLocal = this->mergeOutQueue_.template AllocTensor<ValueType>();
430 431 
431 uint32_t listSortedNums[MERGE_LIST_MAX_NUM] = {0, 0, 0, 0};432 uint32_t listSortedNums[MERGE_LIST_MAX_NUM] = {0, 0, 0, 0};
432- uint32_t mergedCount =433+ uint32_t mergedCount = ExecuteMrgSort(dstLocal, ubMainInputCalc, elementCountList, listSortedNums,
433- ExecuteMrgSort(dstLocal, ubMainInputCalc, elementCountList, listSortedNums, remainListNum);434+ remainListNum);
434 if (mergedCount == 0) {435 if (mergedCount == 0) {
435 // FreeTensor can release AllocTensor'd buffers without EnQue/DeQue;436 // FreeTensor can release AllocTensor'd buffers without EnQue/DeQue;
436 // it only returns the buffer handle to the idle pool.437 // it only returns the buffer handle to the idle pool.
@@ -478,8 +479,8 @@ __aicore__ inline uint32_t MergeIntraCoreBase<Derived, ValueType, IndexType, IsD
478 DataCopyExtParams copyParams{479 DataCopyExtParams copyParams{
479 1, static_cast<uint32_t>(AscendC::GetSortLen<ValueType>(loadCount) * sizeof(ValueType)), 0, 0, 0};480 1, static_cast<uint32_t>(AscendC::GetSortLen<ValueType>(loadCount) * sizeof(ValueType)), 0, 0, 0};
480 DataCopyPadExtParams<ValueType> padParams{false, 0, 0, 0};481 DataCopyPadExtParams<ValueType> padParams{false, 0, 0, 0};
481- DataCopyPad(482+ DataCopyPad(ubMainInput[this->blockSortLen_ * remainListNum], this->cacheGm_[srcAddr], copyParams,
482- ubMainInput[this->blockSortLen_ * remainListNum], this->cacheGm_[srcAddr], copyParams, padParams);483+ padParams);
483 484 
484 remainListNum++;485 remainListNum++;
485 }486 }
@@ -520,11 +521,11 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce
520{521{
521 for (uint32_t listIdx = 0; listIdx < ctx.listCount; listIdx++) {522 for (uint32_t listIdx = 0; listIdx < ctx.listCount; listIdx++) {
522 while (ctx.remains[listIdx] > 0) {523 while (ctx.remains[listIdx] > 0) {
523- uint32_t loadCount =524+ uint32_t loadCount = (ctx.remains[listIdx] > this->blockSortSize_) ? this->blockSortSize_ :
524- (ctx.remains[listIdx] > this->blockSortSize_) ? this->blockSortSize_ : ctx.remains[listIdx];525+ ctx.remains[listIdx];
525 526 
526- CopyBlockChunk(527+ CopyBlockChunk(ctx.srcOffsets[listIdx] + ctx.gmOffsets[listIdx], dstOffsetBase + dstCumulativeOffset,
527- ctx.srcOffsets[listIdx] + ctx.gmOffsets[listIdx], dstOffsetBase + dstCumulativeOffset, loadCount);528+ loadCount);
528 529 
529 ctx.gmOffsets[listIdx] += AscendC::GetSortLen<ValueType>(loadCount);530 ctx.gmOffsets[listIdx] += AscendC::GetSortLen<ValueType>(loadCount);
530 ctx.remains[listIdx] -= loadCount;531 ctx.remains[listIdx] -= loadCount;
Mmath/sort/op_kernel/arch35/common/merge_more_core_base.h+48-47
@@ -32,6 +32,9 @@ using MergeSortConstants::DEALING_SORT_NUM_ONCE;
32using MergeSortConstants::FP32_DTYPE_BYTES;32using MergeSortConstants::FP32_DTYPE_BYTES;
33using MergeSortConstants::MERGE_LIST_MAX_NUM;33using MergeSortConstants::MERGE_LIST_MAX_NUM;
34using MergeSortConstants::MERGE_MORE_BUFFER_NUM;34using MergeSortConstants::MERGE_MORE_BUFFER_NUM;
35+using MergeSortConstants::MERGE_WORKSPACE_BUFFER_NUM;
36+using MergeSortConstants::THREE_WAY_MERGE_LIST_NUM;
37+using MergeSortConstants::TWO_WAY_MERGE_LIST_NUM;
35using MergeSortConstants::UB_BLOCK_BYTES;38using MergeSortConstants::UB_BLOCK_BYTES;
36using MergeSortConstants::XOR_OP_VALUE_FP;39using MergeSortConstants::XOR_OP_VALUE_FP;
37using MergeSortConstants::XOR_OP_VALUE_HALF;40using MergeSortConstants::XOR_OP_VALUE_HALF;
@@ -147,7 +150,7 @@ public:
147 this->currentTailElements_ = this->currentElements_ * (remainListNum - 1) + this->currentTailElements_;150 this->currentTailElements_ = this->currentElements_ * (remainListNum - 1) + this->currentTailElements_;
148 this->listNum_ = currentCoreNum;151 this->listNum_ = currentCoreNum;
149 this->currentElements_ = this->currentElements_ * MERGE_LIST_MAX_NUM;152 this->currentElements_ = this->currentElements_ * MERGE_LIST_MAX_NUM;
150- this->workSpaceFlag_ = (this->workSpaceFlag_ + 1) % 2;153+ this->workSpaceFlag_ = (this->workSpaceFlag_ + 1) % MERGE_WORKSPACE_BUFFER_NUM;
151 }154 }
152 this->workspaceInput_ = this->workspaceGm_[this->workSpaceFlag_];155 this->workspaceInput_ = this->workspaceGm_[this->workSpaceFlag_];
153 this->workspaceOutput_ = this->workspaceGm_[1 - this->workSpaceFlag_];156 this->workspaceOutput_ = this->workspaceGm_[1 - this->workSpaceFlag_];
@@ -190,13 +193,13 @@ public:
190 padParams.rightPadding = currTileSizeAlign - tileNum;193 padParams.rightPadding = currTileSizeAlign - tileNum;
191 padParams.paddingValue = static_cast<T>(defaultValue);194 padParams.paddingValue = static_cast<T>(defaultValue);
192 195 
193- AscendC::DataCopyPad(196+ AscendC::DataCopyPad(inputLocal, this->inputValueGm_[this->rowDataOffset_ + offsetPerCore], copyParams,
194- inputLocal, this->inputValueGm_[this->rowDataOffset_ + offsetPerCore], copyParams, padParams);197+ padParams);
195 this->inputQueue_.EnQue(inputLocal);198 this->inputQueue_.EnQue(inputLocal);
196 }199 }
197 200 
198- __aicore__ inline void InitIndexLocal(201+ __aicore__ inline void InitIndexLocal(uint32_t tileNum, LocalTensor<uint32_t> sortedValueIndexLocal,
199- uint32_t tileNum, LocalTensor<uint32_t> sortedValueIndexLocal, int64_t offsetPerCore)202+ int64_t offsetPerCore)
200 {203 {
201 PipeBarrier<PIPE_ALL>();204 PipeBarrier<PIPE_ALL>();
202 LocalTensor<int32_t> tempIndexLocal = sortedValueIndexLocal.ReinterpretCast<int32_t>();205 LocalTensor<int32_t> tempIndexLocal = sortedValueIndexLocal.ReinterpretCast<int32_t>();
@@ -204,14 +207,14 @@ public:
204 PipeBarrier<PIPE_ALL>();207 PipeBarrier<PIPE_ALL>();
205 }208 }
206 209 
207- __aicore__ inline void DoSort(210+ __aicore__ inline void DoSort(uint32_t tileNum, LocalTensor<T> inputLocal,
208- uint32_t tileNum, LocalTensor<T> inputLocal, LocalTensor<CONVERT_TYPE> sortedValueLocal,211+ LocalTensor<CONVERT_TYPE> sortedValueLocal,
209- LocalTensor<uint32_t> sortedValueIndexLocal)212+ LocalTensor<uint32_t> sortedValueIndexLocal)
210 {213 {
211 AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = this->sortTempBuf_.template Get<CONVERT_TYPE>();214 AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = this->sortTempBuf_.template Get<CONVERT_TYPE>();
212 AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = this->concatTempBuf_.template Get<CONVERT_TYPE>();215 AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = this->concatTempBuf_.template Get<CONVERT_TYPE>();
213- AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast =216+ AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = this->sortedValueLocalCastTbuf_
214- this->sortedValueLocalCastTbuf_.template Get<CONVERT_TYPE>();217+ .template Get<CONVERT_TYPE>();
215 218 
216 uint32_t aglinTileNum = ROUND_UP_AGLIN(tileNum);219 uint32_t aglinTileNum = ROUND_UP_AGLIN(tileNum);
217 uint32_t sortRepeatTimes = Ops::Base::CeilDiv(aglinTileNum, DEALING_SORT_NUM_ONCE);220 uint32_t sortRepeatTimes = Ops::Base::CeilDiv(aglinTileNum, DEALING_SORT_NUM_ONCE);
@@ -222,8 +225,8 @@ public:
222 }225 }
223 AscendC::LocalTensor<CONVERT_TYPE> concatLocal;226 AscendC::LocalTensor<CONVERT_TYPE> concatLocal;
224 AscendC::Concat(concatLocal, inputLocal, concatTempLocal, concatRepeatTimes);227 AscendC::Concat(concatLocal, inputLocal, concatTempLocal, concatRepeatTimes);
225- AscendC::Sort<CONVERT_TYPE, true>(228+ AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal,
226- sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal, sortRepeatTimes);229+ sortRepeatTimes);
227 }230 }
228 231 
229 __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t aglinTileNum)232 __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t aglinTileNum)
@@ -237,8 +240,8 @@ public:
237 }240 }
238 }241 }
239 242 
240- __aicore__ inline void CopyOutWorkSpace(243+ __aicore__ inline void CopyOutWorkSpace(uint32_t tileNum, int64_t offsetPerCore,
241- uint32_t tileNum, int64_t offsetPerCore, LocalTensor<CONVERT_TYPE> sortedValueLocal)244+ LocalTensor<CONVERT_TYPE> sortedValueLocal)
242 {245 {
243 event_t eventIdVToMTE3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_MTE3));246 event_t eventIdVToMTE3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_MTE3));
244 SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);247 SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
@@ -350,9 +353,9 @@ public:
350 this->remainListNum_ = 0;353 this->remainListNum_ = 0;
351 for (int64_t i = 0, j = 0; i < MERGE_LIST_MAX_NUM; i++) {354 for (int64_t i = 0, j = 0; i < MERGE_LIST_MAX_NUM; i++) {
352 // Copy at most onceMaxElements_ from each non-empty run so one MrgSort call can merge the current window.355 // Copy at most onceMaxElements_ from each non-empty run so one MrgSort call can merge the current window.
353- this->dealLengths_[i] =356+ this->dealLengths_[i] = (this->onceMaxElements_ > this->listRemainElements_[i] ?
354- (this->onceMaxElements_ > this->listRemainElements_[i] ? this->listRemainElements_[i] :357+ this->listRemainElements_[i] :
355- this->onceMaxElements_);358+ this->onceMaxElements_);
356 if (this->dealLengths_[i] > 0) {359 if (this->dealLengths_[i] > 0) {
357 DataCopyExtParams copyParams;360 DataCopyExtParams copyParams;
358 copyParams.blockCount = 1;361 copyParams.blockCount = 1;
@@ -360,9 +363,8 @@ public:
360 copyParams.srcStride = 0;363 copyParams.srcStride = 0;
361 copyParams.dstStride = 0;364 copyParams.dstStride = 0;
362 DataCopyPadExtParams<CONVERT_TYPE> padParams{false, 0, 0, 0};365 DataCopyPadExtParams<CONVERT_TYPE> padParams{false, 0, 0, 0};
363- DataCopyPad(366+ DataCopyPad(ubMainInput[GetSortLen<CONVERT_TYPE>(this->onceMaxElements_) * i],
364- ubMainInput[GetSortLen<CONVERT_TYPE>(this->onceMaxElements_) * i],367+ this->workspaceInput_[this->offsets_[i]], copyParams, padParams);
365- this->workspaceInput_[this->offsets_[i]], copyParams, padParams);
366 this->elementCountList_[j] = this->dealLengths_[i];368 this->elementCountList_[j] = this->dealLengths_[i];
367 this->remainListNum_ += 1;369 this->remainListNum_ += 1;
368 j++;370 j++;
@@ -374,19 +376,19 @@ public:
374 __aicore__ inline void UpdateMrgParam()376 __aicore__ inline void UpdateMrgParam()
375 {377 {
376 // MrgSort accepts four source lists. validBitTail marks which of the four are real for this window.378 // MrgSort accepts four source lists. validBitTail marks which of the four are real for this window.
377- if (this->remainListNum_ == 2) {379+ if (this->remainListNum_ == TWO_WAY_MERGE_LIST_NUM) {
378- this->elementCountList_[2] = 0;380+ this->elementCountList_[TWO_WAY_MERGE_LIST_NUM] = 0;
379- this->elementCountList_[3] = 0;381+ this->elementCountList_[THREE_WAY_MERGE_LIST_NUM] = 0;
380 this->validBitTail_ = 0b0011;382 this->validBitTail_ = 0b0011;
381- } else if (this->remainListNum_ == 3) {383+ } else if (this->remainListNum_ == THREE_WAY_MERGE_LIST_NUM) {
382- this->elementCountList_[3] = 0;384+ this->elementCountList_[THREE_WAY_MERGE_LIST_NUM] = 0;
383 this->validBitTail_ = 0b0111;385 this->validBitTail_ = 0b0111;
384- } else if (this->remainListNum_ == 4) {386+ } else if (this->remainListNum_ == MERGE_LIST_MAX_NUM) {
385 this->validBitTail_ = 0b1111;387 this->validBitTail_ = 0b1111;
386 } else {388 } else {
387 this->elementCountList_[1] = 0;389 this->elementCountList_[1] = 0;
388- this->elementCountList_[2] = 0;390+ this->elementCountList_[TWO_WAY_MERGE_LIST_NUM] = 0;
389- this->elementCountList_[3] = 0;391+ this->elementCountList_[THREE_WAY_MERGE_LIST_NUM] = 0;
390 this->validBitTail_ = 0b0001;392 this->validBitTail_ = 0b0001;
391 }393 }
392 }394 }
@@ -402,27 +404,26 @@ public:
402 j++;404 j++;
403 }405 }
404 }406 }
405- if (this->remainListNum_ == 2) {407+ if (this->remainListNum_ == TWO_WAY_MERGE_LIST_NUM) {
406 // Unused source-list arguments are duplicated because validBitTail controls which lists participate.408 // Unused source-list arguments are duplicated because validBitTail controls which lists participate.
407- MrgSortSrcList sortListTail =409+ MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0],
408- MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);410+ tmpUbInputs[0]);
409- MrgSort<CONVERT_TYPE, true>(411+ MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_,
410- sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_, this->validBitTail_, 1);412+ this->validBitTail_, 1);
411- } else if (this->remainListNum_ == 3) {413+ } else if (this->remainListNum_ == THREE_WAY_MERGE_LIST_NUM) {
412- MrgSortSrcList sortListTail =414+ MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2],
413- MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[0]);415+ tmpUbInputs[0]);
414- MrgSort<CONVERT_TYPE, true>(416+ MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_,
415- sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_, this->validBitTail_, 1);417+ this->validBitTail_, 1);
416- } else if (this->remainListNum_ == 4) {418+ } else if (this->remainListNum_ == MERGE_LIST_MAX_NUM) {
417- MrgSortSrcList sortListTail =419+ MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2],
418- MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], tmpUbInputs[3]);420+ tmpUbInputs[3]);
419- MrgSort<CONVERT_TYPE, true>(421+ MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_,
420- sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_, this->validBitTail_, 1);422+ this->validBitTail_, 1);
421 } else {423 } else {
422- AscendC::Copy(424+ AscendC::Copy(sortTempBuffer, tmpUbInputs[0],
423- sortTempBuffer, tmpUbInputs[0],425+ ROUND_UP_AGLIN(GetSortLen<CONVERT_TYPE>(this->elementCountList_[0]) * sizeof(CONVERT_TYPE)) /
424- ROUND_UP_AGLIN(GetSortLen<CONVERT_TYPE>(this->elementCountList_[0]) * sizeof(CONVERT_TYPE)) /426+ sizeof(CONVERT_TYPE));
425- sizeof(CONVERT_TYPE));
426 this->listSortedNums_[0] = this->elementCountList_[0];427 this->listSortedNums_[0] = this->elementCountList_[0];
427 }428 }
428 this->sortedQueue_.EnQue(sortTempBuffer);429 this->sortedQueue_.EnQue(sortTempBuffer);
Mmath/sort/op_kernel/arch35/common/merge_sort_constants.h+3-0
@@ -20,8 +20,11 @@ constexpr uint32_t FP32_DTYPE_BYTES = 4;
20constexpr uint32_t UB_BLOCK_BYTES = Ops::Base::GetUbBlockSize();20constexpr uint32_t UB_BLOCK_BYTES = Ops::Base::GetUbBlockSize();
21 21 
22constexpr uint32_t MERGE_LIST_MAX_NUM = 4;22constexpr uint32_t MERGE_LIST_MAX_NUM = 4;
23+constexpr uint32_t TWO_WAY_MERGE_LIST_NUM = 2;
24+constexpr uint32_t THREE_WAY_MERGE_LIST_NUM = 3;
23constexpr uint32_t MERGE_INTRA_BUFFER_NUM = 2;25constexpr uint32_t MERGE_INTRA_BUFFER_NUM = 2;
24constexpr uint32_t MERGE_MORE_BUFFER_NUM = 1;26constexpr uint32_t MERGE_MORE_BUFFER_NUM = 1;
27+constexpr uint32_t MERGE_WORKSPACE_BUFFER_NUM = 2;
25 28 
26constexpr int32_t XOR_OP_VALUE_FP = 0x80000000;29constexpr int32_t XOR_OP_VALUE_FP = 0x80000000;
27constexpr int16_t XOR_OP_VALUE_HALF = 0x8000;30constexpr int16_t XOR_OP_VALUE_HALF = 0x8000;
Mmath/sort/op_kernel/arch35/common/non_last_small_axis_base.h+1-1
@@ -181,7 +181,7 @@ protected:
181 }181 }
182 182 
183 // Transpose a tile via gather: convert row-major layout to column-major (or vice versa).183 // Transpose a tile via gather: convert row-major layout to column-major (or vice versa).
184- // Example (4x4 tile, inputRowElems_=4):184+ // The following 4x4 tile illustrates the transpose when inputRowElems_ is 4.
185 // Input (row-major): Output (column-major):185 // Input (row-major): Output (column-major):
186 // [0 1 2 3 ] [0 4 8 12]186 // [0 1 2 3 ] [0 4 8 12]
187 // [4 5 6 7 ] ------> [1 5 9 13]187 // [4 5 6 7 ] ------> [1 5 9 13]
Mmath/sort/op_kernel/arch35/common/radix_more_core_base.h+3-3
@@ -631,7 +631,7 @@ public:
631 if constexpr (sizeof(T3) == sizeof(uint32_t)) {631 if constexpr (sizeof(T3) == sizeof(uint32_t)) {
632 Duplicate(blockExcusiveUb, static_cast<uint32_t>(0), RADIX_SORT_NUM);632 Duplicate(blockExcusiveUb, static_cast<uint32_t>(0), RADIX_SORT_NUM);
633 } else {633 } else {
634- Duplicate(blockExcusiveUb, static_cast<uint32_t>(0), RADIX_SORT_NUM * 2);634+ Duplicate(blockExcusiveUb, static_cast<uint32_t>(0), RADIX_SORT_NUM * RADIX_HIST_BUFFER_NUM);
635 }635 }
636 for (uint32_t tileId = startTileId; tileId < this->lastDimTileNum_; tileId += this->lastDimRealCore_) {636 for (uint32_t tileId = startTileId; tileId < this->lastDimTileNum_; tileId += this->lastDimRealCore_) {
637 // tileOffset may exceed int32 range for large rows, so keep the address arithmetic in uint64.637 // tileOffset may exceed int32 range for large rows, so keep the address arithmetic in uint64.
@@ -1180,10 +1180,10 @@ public:
1180 } else {1180 } else {
1181 if constexpr (sizeof(T1) == sizeof(int8_t)) {1181 if constexpr (sizeof(T1) == sizeof(int8_t)) {
1182 this->inputXDbGm_.SetDoubleBuffer(this->outValueDbWK_, this->outValueGm_[gmOffset]);1182 this->inputXDbGm_.SetDoubleBuffer(this->outValueDbWK_, this->outValueGm_[gmOffset]);
1183- this->idxDbGm_.SetDoubleBuffer(this->outIdxDbWK_, this->outIdxGm_[gmOffset * 2]);1183+ this->idxDbGm_.SetDoubleBuffer(this->outIdxDbWK_, this->outIdxGm_[gmOffset * INT64_INDEX_SCALE]);
1184 } else {1184 } else {
1185 this->inputXDbGm_.SetDoubleBuffer(this->outValueGm_[gmOffset], this->outValueDbWK_);1185 this->inputXDbGm_.SetDoubleBuffer(this->outValueGm_[gmOffset], this->outValueDbWK_);
1186- this->idxDbGm_.SetDoubleBuffer(this->outIdxGm_[gmOffset * 2], this->outIdxDbWK_);1186+ this->idxDbGm_.SetDoubleBuffer(this->outIdxGm_[gmOffset * INT64_INDEX_SCALE], this->outIdxDbWK_);
1187 }1187 }
1188 }1188 }
1189 for (uint32_t round = 0; round < static_cast<uint32_t>(sizeof(T1)); round++) {1189 for (uint32_t round = 0; round < static_cast<uint32_t>(sizeof(T1)); round++) {
Mmath/sort/op_kernel/arch35/common/radix_sort_constants.h+2-0
@@ -38,6 +38,8 @@ const uint32_t SHIFT_BIT_NUM = 8;
38const uint32_t HIST_MASK_OUT_LEN = 8;38const uint32_t HIST_MASK_OUT_LEN = 8;
39const uint64_t XOR_OP_VALUE_B64 = 0x8000000000000000;39const uint64_t XOR_OP_VALUE_B64 = 0x8000000000000000;
40const uint32_t RADIX_SORT_NUM = 256;40const uint32_t RADIX_SORT_NUM = 256;
41+const uint32_t RADIX_HIST_BUFFER_NUM = 2;
42+const uint32_t INT64_INDEX_SCALE = sizeof(int64_t) / sizeof(uint32_t);
41const uint8_t XOR_OP_VALUE_B8 = (uint8_t(1) << 7);43const uint8_t XOR_OP_VALUE_B8 = (uint8_t(1) << 7);
42const uint32_t XOR_OP_VALUE = 0x80000000;44const uint32_t XOR_OP_VALUE = 0x80000000;
43const int16_t STATE_BIT_SHF_VALUE = 30;45const int16_t STATE_BIT_SHF_VALUE = 30;
Mmath/sort/op_kernel/arch35/common/radix_sort_simd_utils.h+16-17
@@ -31,9 +31,8 @@ using AscendC::MicroAPI::UpdateMask;
31// round, so the normalization lives in TwiddleInFp16/TwiddleInFp32.31// round, so the normalization lives in TwiddleInFp16/TwiddleInFp32.
32 32 
33template <typename UT, uint64_t isDescend>33template <typename UT, uint64_t isDescend>
34-__aicore__ inline void Twiddle(34+__aicore__ inline void Twiddle(uint16_t repeatTime, uint32_t vfLen, uint32_t inputNum, RegTensor<UT>& xorReg,
35- uint16_t repeatTime, uint32_t vfLen, uint32_t inputNum, RegTensor<UT>& xorReg, __local_mem__ UT* xValuePtr,35+ __local_mem__ UT* xValuePtr, __local_mem__ UT* uXValuePtr)
36- __local_mem__ UT* uXValuePtr)
37{36{
38 MicroAPI::MaskReg xorMask;37 MicroAPI::MaskReg xorMask;
39 MicroAPI::RegTensor<UT> inputReg;38 MicroAPI::RegTensor<UT> inputReg;
@@ -156,19 +155,19 @@ __aicore__ inline void TwiddleInFp16(LocalTensor<T1> inputX, LocalTensor<UT> uin
156 // 如果后期需要改为单次twiddleIn，则需要在提取位数时将负0转换为正0155 // 如果后期需要改为单次twiddleIn，则需要在提取位数时将负0转换为正0
157 // get -0.0 mask156 // get -0.0 mask
158 MicroAPI::MaskReg minusZeroMask;157 MicroAPI::MaskReg minusZeroMask;
159- MicroAPI::CompareScalar<uint16_t, CMPMODE::EQ>(158+ MicroAPI::CompareScalar<uint16_t, CMPMODE::EQ>(minusZeroMask, xorVectorOne, TWIDDLED_MINUS_ZERO_BITS_FP16,
160- minusZeroMask, xorVectorOne, TWIDDLED_MINUS_ZERO_BITS_FP16, maskB16);159+ maskB16);
161 // change -0.0 to +0.0160 // change -0.0 to +0.0
162 MicroAPI::RegTensor<uint16_t> resultReg;161 MicroAPI::RegTensor<uint16_t> resultReg;
163 MicroAPI::Select(resultReg, twiddledZeroReg, xorVectorOne, minusZeroMask);162 MicroAPI::Select(resultReg, twiddledZeroReg, xorVectorOne, minusZeroMask);
164 163 
165 if constexpr (isDescend == 0) {164 if constexpr (isDescend == 0) {
166- MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(165+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(uXValuePtr, resultReg, VF_LEN_B16,
167- uXValuePtr, resultReg, VF_LEN_B16, xorMask);166+ xorMask);
168 } else {167 } else {
169 MicroAPI::Not(vnotReg, resultReg, xorMask);168 MicroAPI::Not(vnotReg, resultReg, xorMask);
170- MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(169+ MicroAPI::DataCopy<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(uXValuePtr, vnotReg, VF_LEN_B16,
171- uXValuePtr, vnotReg, VF_LEN_B16, xorMask);170+ xorMask);
172 }171 }
173 }172 }
174 }173 }
@@ -212,19 +211,19 @@ __aicore__ inline void TwiddleInFp32(LocalTensor<T1> inputX, LocalTensor<UT> uin
212 // 如果后期需要改为单次twiddleIn，则需要在提取位数时将负0转换为正0211 // 如果后期需要改为单次twiddleIn，则需要在提取位数时将负0转换为正0
213 // get -0.0 mask212 // get -0.0 mask
214 MicroAPI::MaskReg minusZeroMask;213 MicroAPI::MaskReg minusZeroMask;
215- MicroAPI::CompareScalar<uint32_t, CMPMODE::EQ>(214+ MicroAPI::CompareScalar<uint32_t, CMPMODE::EQ>(minusZeroMask, xorVectorZero, TWIDDLED_MINUS_ZERO_BITS_FP32,
216- minusZeroMask, xorVectorZero, TWIDDLED_MINUS_ZERO_BITS_FP32, maskB32);215+ maskB32);
217 // change -0.0 to +0.0216 // change -0.0 to +0.0
218 MicroAPI::RegTensor<uint32_t> resultReg;217 MicroAPI::RegTensor<uint32_t> resultReg;
219 MicroAPI::Select(resultReg, twiddledZeroReg, xorVectorZero, minusZeroMask);218 MicroAPI::Select(resultReg, twiddledZeroReg, xorVectorZero, minusZeroMask);
220 219 
221 if constexpr (isDescend == 0) {220 if constexpr (isDescend == 0) {
222- MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(221+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(uXValuePtr, resultReg, VF_LEN_B32,
223- uXValuePtr, resultReg, VF_LEN_B32, xorMask);222+ xorMask);
224 } else {223 } else {
225 MicroAPI::Not(vnotReg, resultReg, maskB32);224 MicroAPI::Not(vnotReg, resultReg, maskB32);
226- MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(225+ MicroAPI::DataCopy<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(uXValuePtr, vnotReg, VF_LEN_B32,
227- uXValuePtr, vnotReg, VF_LEN_B32, xorMask);226+ xorMask);
228 }227 }
229 }228 }
230 }229 }
@@ -247,8 +246,8 @@ __aicore__ inline void ReverseInputData(LocalTensor<T1> inputX, LocalTensor<UT>
247 MicroAPI::MaskReg vnotMask = MicroAPI::UpdateMask<UT>(inputElementNum);246 MicroAPI::MaskReg vnotMask = MicroAPI::UpdateMask<UT>(inputElementNum);
248 MicroAPI::DataCopy<UT, MicroAPI::PostLiteral::POST_MODE_UPDATE>(inputVectorOne, inputXValuePtr, vfLen);247 MicroAPI::DataCopy<UT, MicroAPI::PostLiteral::POST_MODE_UPDATE>(inputVectorOne, inputXValuePtr, vfLen);
249 MicroAPI::Not(vnotVectorZero, inputVectorOne, predicateDefaultB8);248 MicroAPI::Not(vnotVectorZero, inputVectorOne, predicateDefaultB8);
250- MicroAPI::DataCopy<UT, MicroAPI::PostLiteral::POST_MODE_UPDATE>(249+ MicroAPI::DataCopy<UT, MicroAPI::PostLiteral::POST_MODE_UPDATE>(reverseInputXPtr, vnotVectorZero, vfLen,
251- reverseInputXPtr, vnotVectorZero, vfLen, vnotMask);250+ vnotMask);
252 }251 }
253 }252 }
254}253}
Mmath/sort/op_kernel/arch35/common/small_axis_insertion_base.h+15-14
@@ -33,17 +33,18 @@ constexpr uint32_t MAX_DATACOPY_BLOCK_COUNT = 4095;
33 * @tparam T Input storage data type33 * @tparam T Input storage data type
34 */34 */
35template <typename T>35template <typename T>
36-__simt_vf__ LAUNCH_BOUND(INSERTION_THREAD_NUM) __aicore__ void SimtLoadNonLastInsertionBatch(36+__simt_vf__ LAUNCH_BOUND(INSERTION_THREAD_NUM) __aicore__
37- uint32_t totalElems, uint32_t validSegs, uint32_t segmentLen, uint32_t valueRowElems, uint64_t outerBaseOffset,37+ void SimtLoadNonLastInsertionBatch(uint32_t totalElems, uint32_t validSegs, uint32_t segmentLen,
38- uint64_t innerStart, uint64_t innerSize, __gm__ volatile T* input, __ubuf__ T* output)38+ uint32_t valueRowElems, uint64_t outerBaseOffset, uint64_t innerStart,
39+ uint64_t innerSize, __gm__ volatile T* input, __ubuf__ T* output)
39{40{
40 // Load a non-last-axis tile as sort-major segments:41 // Load a non-last-axis tile as sort-major segments:
41 // GM [axis, inner] -> UB [inner segment, axis].42 // GM [axis, inner] -> UB [inner segment, axis].
42 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems; idx += INSERTION_THREAD_NUM) {43 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems; idx += INSERTION_THREAD_NUM) {
43 uint32_t axis = idx / validSegs;44 uint32_t axis = idx / validSegs;
44 uint32_t seg = idx - axis * validSegs;45 uint32_t seg = idx - axis * validSegs;
45- output[seg * valueRowElems + axis] =46+ output[seg * valueRowElems + axis] = input[outerBaseOffset + static_cast<uint64_t>(axis) * innerSize +
46- input[outerBaseOffset + static_cast<uint64_t>(axis) * innerSize + innerStart + seg];47+ innerStart + seg];
47 }48 }
48}49}
49 50 
@@ -54,9 +55,9 @@ __simt_vf__ LAUNCH_BOUND(INSERTION_THREAD_NUM) __aicore__ void SimtLoadNonLastIn
54 * @tparam IsDescend Sort order: true for descending, false for ascending55 * @tparam IsDescend Sort order: true for descending, false for ascending
55 */56 */
56template <typename T, typename IDX_T, bool IsDescend>57template <typename T, typename IDX_T, bool IsDescend>
57-__simt_vf__ LAUNCH_BOUND(INSERTION_THREAD_NUM) __aicore__ void SimtInsertionSortSegments(58+__simt_vf__ LAUNCH_BOUND(INSERTION_THREAD_NUM) __aicore__
58- uint32_t validSegs, uint32_t segmentLen, uint32_t valueRowElems, uint32_t idxRowElems, __ubuf__ T* valueBase,59+ void SimtInsertionSortSegments(uint32_t validSegs, uint32_t segmentLen, uint32_t valueRowElems,
59- __ubuf__ IDX_T* idxBase)60+ uint32_t idxRowElems, __ubuf__ T* valueBase, __ubuf__ IDX_T* idxBase)
60{61{
61 for (int32_t seg = threadIdx.x; seg < static_cast<int32_t>(validSegs);62 for (int32_t seg = threadIdx.x; seg < static_cast<int32_t>(validSegs);
62 seg += static_cast<int32_t>(INSERTION_THREAD_NUM)) {63 seg += static_cast<int32_t>(INSERTION_THREAD_NUM)) {
@@ -144,8 +145,8 @@ public:
144 indices_ = idxBuf_.template Get<IDX_T>();145 indices_ = idxBuf_.template Get<IDX_T>();
145 }146 }
146 147 
147- __aicore__ inline void LoadContiguousBatch(148+ __aicore__ inline void LoadContiguousBatch(GlobalTensor<T>& inputGm, int64_t inputStart, uint32_t validSegs,
148- GlobalTensor<T>& inputGm, int64_t inputStart, uint32_t validSegs, bool padTail)149+ bool padTail)
149 {150 {
150 if (validSegs == 0U) {151 if (validSegs == 0U) {
151 return;152 return;
@@ -153,8 +154,8 @@ public:
153 uint32_t dstRowStride = IsSameType<T, CONVERT_TYPE>::value ? valueRowStride_ : castRowStride_;154 uint32_t dstRowStride = IsSameType<T, CONVERT_TYPE>::value ? valueRowStride_ : castRowStride_;
154 uint32_t copiedRowElems = padTail ? ROUND_UP_AGLIN(segmentLen_ * sizeof(T)) / sizeof(T) : segmentLen_;155 uint32_t copiedRowElems = padTail ? ROUND_UP_AGLIN(segmentLen_ * sizeof(T)) / sizeof(T) : segmentLen_;
155 uint32_t dstStrideBlocks = ((dstRowStride - copiedRowElems) * sizeof(T)) / UB_BLOCK_SIZE;156 uint32_t dstStrideBlocks = ((dstRowStride - copiedRowElems) * sizeof(T)) / UB_BLOCK_SIZE;
156- DataCopyExtParams copyParam{157+ DataCopyExtParams copyParam{static_cast<uint16_t>(validSegs), static_cast<uint32_t>(segmentLen_ * sizeof(T)), 0,
157- static_cast<uint16_t>(validSegs), static_cast<uint32_t>(segmentLen_ * sizeof(T)), 0, dstStrideBlocks, 0};158+ dstStrideBlocks, 0};
158 DataCopyPadExtParams<T> padParams{159 DataCopyPadExtParams<T> padParams{
159 padTail, 0, static_cast<uint8_t>(padTail ? (copiedRowElems - segmentLen_) : 0U), static_cast<T>(0)};160 padTail, 0, static_cast<uint8_t>(padTail ? (copiedRowElems - segmentLen_) : 0U), static_cast<T>(0)};
160 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {161 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
@@ -172,8 +173,8 @@ public:
172 }173 }
173 }174 }
174 175 
175- __aicore__ inline void LoadNonLastBatch(176+ __aicore__ inline void LoadNonLastBatch(GlobalTensor<T>& inputGm, uint64_t outerBaseOffset, uint64_t innerStart,
176- GlobalTensor<T>& inputGm, uint64_t outerBaseOffset, uint64_t innerStart, uint64_t innerSize, uint32_t validSegs)177+ uint64_t innerSize, uint32_t validSegs)
177 {178 {
178 uint32_t totalElems = validSegs * segmentLen_;179 uint32_t totalElems = validSegs * segmentLen_;
179 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {180 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
Mmath/sort/op_kernel/arch35/common/small_axis_two_stage_base.h+27-29
@@ -31,16 +31,16 @@ constexpr uint32_t TWO_STAGE_THREAD_NUM = 1024;
31 * @tparam T Input storage data type31 * @tparam T Input storage data type
32 */32 */
33template <typename T>33template <typename T>
34-__simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__ void LoadNonLastBatchSimt(34+__simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__
35- uint32_t totalElems, uint32_t segmentLen, uint32_t validSegs, uint64_t outerBaseOffset, uint64_t innerStart,35+ void LoadNonLastBatchSimt(uint32_t totalElems, uint32_t segmentLen, uint32_t validSegs, uint64_t outerBaseOffset,
36- uint64_t innerSize, __gm__ volatile T* input, __ubuf__ T* output)36+ uint64_t innerStart, uint64_t innerSize, __gm__ volatile T* input, __ubuf__ T* output)
37{37{
38 // Gather the original [axis, inner] tile into [inner segment, axis] order for two-stage sorting.38 // Gather the original [axis, inner] tile into [inner segment, axis] order for two-stage sorting.
39 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems; idx += TWO_STAGE_THREAD_NUM) {39 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems; idx += TWO_STAGE_THREAD_NUM) {
40 uint32_t axis = idx / validSegs;40 uint32_t axis = idx / validSegs;
41 uint32_t seg = idx - axis * validSegs;41 uint32_t seg = idx - axis * validSegs;
42- output[seg * segmentLen + axis] =42+ output[seg * segmentLen + axis] = input[outerBaseOffset + static_cast<uint64_t>(axis) * innerSize + innerStart +
43- input[outerBaseOffset + static_cast<uint64_t>(axis) * innerSize + innerStart + seg];43+ seg];
44 }44 }
45}45}
46 46 
@@ -58,9 +58,10 @@ __simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__ void LoadNonLastBatchS
58 * @tparam OutIdxT Final index data type stored in UB58 * @tparam OutIdxT Final index data type stored in UB
59 */59 */
60template <typename T, typename OutIdxT>60template <typename T, typename OutIdxT>
61-__simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__ void RankInverseScatter(61+__simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__
62- uint32_t totalElems, uint32_t segmentLen, __ubuf__ T* stage1ValuePtr, __ubuf__ uint32_t* stage1OrderPtr,62+ void RankInverseScatter(uint32_t totalElems, uint32_t segmentLen, __ubuf__ T* stage1ValuePtr,
63- __ubuf__ uint32_t* rankInversePtr, __ubuf__ T* finalValuePtr, __ubuf__ OutIdxT* finalIdxPtr)63+ __ubuf__ uint32_t* stage1OrderPtr, __ubuf__ uint32_t* rankInversePtr,
64+ __ubuf__ T* finalValuePtr, __ubuf__ OutIdxT* finalIdxPtr)
64{65{
65 // Build inverse mapping: for each flattened index, store its rank from stage-1 sort.66 // Build inverse mapping: for each flattened index, store its rank from stage-1 sort.
66 // rankInverse[stage1Order[rank]] = rank, e.g. stage1Order=[1, 4, 0, 3, 2, 5]67 // rankInverse[stage1Order[rank]] = rank, e.g. stage1Order=[1, 4, 0, 3, 2, 5]
@@ -119,8 +120,9 @@ __simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__ void RankInverseScatte
119/**120/**
120 * @brief SIMT helper that packs stage-2 sort keys and rewrites stage-1 order to row-local lane ids.121 * @brief SIMT helper that packs stage-2 sort keys and rewrites stage-1 order to row-local lane ids.
121 */122 */
122-inline __simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__ void BuildStage2KeysSimt(123+inline __simt_vf__ LAUNCH_BOUND(TWO_STAGE_THREAD_NUM) __aicore__
123- uint32_t totalElems, uint32_t segmentLen, __ubuf__ uint32_t* stage1OrderPtr, __ubuf__ uint16_t* stage2KeyPtr)124+ void BuildStage2KeysSimt(uint32_t totalElems, uint32_t segmentLen, __ubuf__ uint32_t* stage1OrderPtr,
125+ __ubuf__ uint16_t* stage2KeyPtr)
124{126{
125 for (uint32_t rank = static_cast<uint32_t>(threadIdx.x); rank < totalElems; rank += TWO_STAGE_THREAD_NUM) {127 for (uint32_t rank = static_cast<uint32_t>(threadIdx.x); rank < totalElems; rank += TWO_STAGE_THREAD_NUM) {
126 uint32_t flatIdx = stage1OrderPtr[rank];128 uint32_t flatIdx = stage1OrderPtr[rank];
@@ -158,8 +160,8 @@ public:
158 }160 }
159 }161 }
160 162 
161- __aicore__ inline void InitSortBuffers(163+ __aicore__ inline void InitSortBuffers(TPipe* pipe, uint32_t maxFlatElems, uint32_t tmpUbSize, bool useRankInverse,
162- TPipe* pipe, uint32_t maxFlatElems, uint32_t tmpUbSize, bool useRankInverse, uint32_t finalIdxElemBytes)164+ uint32_t finalIdxElemBytes)
163 {165 {
164 pipe_ = pipe;166 pipe_ = pipe;
165 maxFlatElems_ = maxFlatElems;167 maxFlatElems_ = maxFlatElems;
@@ -206,9 +208,8 @@ public:
206 WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);208 WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
207 }209 }
208 210 
209- __aicore__ inline void LoadNonLastBatch(211+ __aicore__ inline void LoadNonLastBatch(GlobalTensor<T>& inputGm, uint64_t outerBaseOffset, uint64_t innerStart,
210- GlobalTensor<T>& inputGm, uint64_t outerBaseOffset, uint64_t innerStart, uint64_t innerSize, uint32_t validSegs,212+ uint64_t innerSize, uint32_t validSegs, uint32_t totalElems)
211- uint32_t totalElems)
212 {213 {
213 asc_vf_call<LoadNonLastBatchSimt<T>>(214 asc_vf_call<LoadNonLastBatchSimt<T>>(
214 dim3(TWO_STAGE_THREAD_NUM), totalElems, segmentLen_, validSegs, outerBaseOffset, innerStart, innerSize,215 dim3(TWO_STAGE_THREAD_NUM), totalElems, segmentLen_, validSegs, outerBaseOffset, innerStart, innerSize,
@@ -282,15 +283,15 @@ private:
282 283 
283 __aicore__ inline void BuildStage2Keys(uint32_t totalElems)284 __aicore__ inline void BuildStage2Keys(uint32_t totalElems)
284 {285 {
285- asc_vf_call<BuildStage2KeysSimt>(286+ asc_vf_call<BuildStage2KeysSimt>(dim3(TWO_STAGE_THREAD_NUM), totalElems, segmentLen_,
286- dim3(TWO_STAGE_THREAD_NUM), totalElems, segmentLen_, (__ubuf__ uint32_t*)stage1Order_.GetPhyAddr(),287+ (__ubuf__ uint32_t*)stage1Order_.GetPhyAddr(),
287- (__ubuf__ uint16_t*)stage2KeysIn_.GetPhyAddr());288+ (__ubuf__ uint16_t*)stage2KeysIn_.GetPhyAddr());
288 }289 }
289 290 
290 __aicore__ inline void Stage2Sort(uint32_t totalElems)291 __aicore__ inline void Stage2Sort(uint32_t totalElems)
291 {292 {
292- AscendC::Sort<uint16_t, false, kStage2SortConfig>(293+ AscendC::Sort<uint16_t, false, kStage2SortConfig>(stage2KeysOut_, stage2Order_, stage2KeysIn_, tmp_,
293- stage2KeysOut_, stage2Order_, stage2KeysIn_, tmp_, totalElems);294+ totalElems);
294 }295 }
295 296 
296 __aicore__ inline void BuildOutputs(uint32_t totalElems)297 __aicore__ inline void BuildOutputs(uint32_t totalElems)
@@ -313,22 +314,19 @@ private:
313 LocalTensor<uint32_t> gatherOffsets = stage2KeysOut_.template ReinterpretCast<uint32_t>();314 LocalTensor<uint32_t> gatherOffsets = stage2KeysOut_.template ReinterpretCast<uint32_t>();
314 LocalTensor<int32_t> gatherOffsetsInt = gatherOffsets.template ReinterpretCast<int32_t>();315 LocalTensor<int32_t> gatherOffsetsInt = gatherOffsets.template ReinterpretCast<int32_t>();
315 316 
316- Muls(317+ Muls(gatherOffsetsInt, stage2Order_.template ReinterpretCast<int32_t>(), static_cast<int32_t>(sizeof(T)),
317- gatherOffsetsInt, stage2Order_.template ReinterpretCast<int32_t>(), static_cast<int32_t>(sizeof(T)),318+ totalElems);
318- totalElems);
319 Gather(finalValues_, stage1Values_, gatherOffsets, 0, totalElems);319 Gather(finalValues_, stage1Values_, gatherOffsets, 0, totalElems);
320 320 
321- Muls(321+ Muls(gatherOffsetsInt, stage2Order_.template ReinterpretCast<int32_t>(), static_cast<int32_t>(sizeof(uint32_t)),
322- gatherOffsetsInt, stage2Order_.template ReinterpretCast<int32_t>(), static_cast<int32_t>(sizeof(uint32_t)),322+ totalElems);
323- totalElems);
324 if constexpr (IsSameType<FinalIdxT, int64_t>::value) {323 if constexpr (IsSameType<FinalIdxT, int64_t>::value) {
325 LocalTensor<int32_t> gatheredIdxInt32 = stage2Order_.template ReinterpretCast<int32_t>();324 LocalTensor<int32_t> gatheredIdxInt32 = stage2Order_.template ReinterpretCast<int32_t>();
326 Gather(gatheredIdxInt32, stage1Order_.template ReinterpretCast<int32_t>(), gatherOffsets, 0, totalElems);325 Gather(gatheredIdxInt32, stage1Order_.template ReinterpretCast<int32_t>(), gatherOffsets, 0, totalElems);
327 Cast(finalIdx_, gatheredIdxInt32, RoundMode::CAST_NONE, Ops::Base::CeilAlign(totalElems, 4u));326 Cast(finalIdx_, gatheredIdxInt32, RoundMode::CAST_NONE, Ops::Base::CeilAlign(totalElems, 4u));
328 } else {327 } else {
329- Gather(328+ Gather(finalIdx_.template ReinterpretCast<int32_t>(), stage1Order_.template ReinterpretCast<int32_t>(),
330- finalIdx_.template ReinterpretCast<int32_t>(), stage1Order_.template ReinterpretCast<int32_t>(),329+ gatherOffsets, 0, totalElems);
331- gatherOffsets, 0, totalElems);
332 }330 }
333 }331 }
334};332};
Mmath/sort/op_kernel/arch35/common/util_type_simd.h+7-15
@@ -17,7 +17,7 @@
17#include "kernel_operator.h"17#include "kernel_operator.h"
18#include "op_kernel/platform_util.h"18#include "op_kernel/platform_util.h"
19 19 
20-constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize(); // UB block size in bytes20+constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize(); // UB block size in bytes
21 21 
22__aicore__ inline uint32_t ROUND_UP_AGLIN(uint32_t x)22__aicore__ inline uint32_t ROUND_UP_AGLIN(uint32_t x)
23{23{
@@ -35,28 +35,20 @@ __aicore__ inline uint64_t ROUND_UP_AGLIN_UINT64(uint64_t x)
35 return (x + UB_BLOCK_SIZE - 1ULL) / UB_BLOCK_SIZE * UB_BLOCK_SIZE;35 return (x + UB_BLOCK_SIZE - 1ULL) / UB_BLOCK_SIZE * UB_BLOCK_SIZE;
36}36}
37 37 
38-template <typename T> struct DoubleBufferSimd {38+template <typename T>
39+struct DoubleBufferSimd {
39 AscendC::GlobalTensor<T> doubleBuffer_[2];40 AscendC::GlobalTensor<T> doubleBuffer_[2];
40 int selector_ = 0;41 int selector_ = 0;
41 __aicore__ inline DoubleBufferSimd() {}42 __aicore__ inline DoubleBufferSimd() {}
42 __aicore__ inline void SetDoubleBuffer(AscendC::GlobalTensor<T> currentBuffer,43 __aicore__ inline void SetDoubleBuffer(AscendC::GlobalTensor<T> currentBuffer,
43- AscendC::GlobalTensor<T> alternateBuffer)44+ AscendC::GlobalTensor<T> alternateBuffer)
44 {45 {
45 selector_ = 0;46 selector_ = 0;
46 doubleBuffer_[0] = currentBuffer;47 doubleBuffer_[0] = currentBuffer;
47 doubleBuffer_[1] = alternateBuffer;48 doubleBuffer_[1] = alternateBuffer;
48 }49 }
49- __aicore__ inline AscendC::GlobalTensor<T> Current() const50+ __aicore__ inline AscendC::GlobalTensor<T> Current() const { return doubleBuffer_[selector_]; }
50- {51+ __aicore__ inline AscendC::GlobalTensor<T> Alternate() const { return doubleBuffer_[selector_ ^ 1]; }
51- return doubleBuffer_[selector_];52+ __aicore__ inline void UpdateSelect() { selector_ = selector_ ^ 1; }
52- }
53- __aicore__ inline AscendC::GlobalTensor<T> Alternate() const
54- {
55- return doubleBuffer_[selector_ ^ 1];
56- }
57- __aicore__ inline void UpdateSelect()
58- {
59- selector_ = selector_ ^ 1;
60- }
61};53};
62#endif54#endif
Mmath/sort/op_kernel/arch35/merge_sort_big_size.h+24-23
@@ -25,32 +25,34 @@
25using namespace AscendC;25using namespace AscendC;
26 26 
27// Import shared constants from MergeSortConstants namespace27// Import shared constants from MergeSortConstants namespace
28+using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;
29+using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE;
30+using MergeSortConstants::DEALING_SORT_NUM_ONCE;
28using MergeSortConstants::MERGE_LIST_MAX_NUM;31using MergeSortConstants::MERGE_LIST_MAX_NUM;
29using MergeSortConstants::MERGE_MORE_BUFFER_NUM;32using MergeSortConstants::MERGE_MORE_BUFFER_NUM;
30using MergeSortConstants::UB_BLOCK_BYTES;33using MergeSortConstants::UB_BLOCK_BYTES;
31-using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;
32-using MergeSortConstants::DEALING_SORT_NUM_ONCE;
33-using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE;
34using MergeSortConstants::XOR_OP_VALUE_FP;34using MergeSortConstants::XOR_OP_VALUE_FP;
35using MergeSortConstants::XOR_OP_VALUE_HALF;35using MergeSortConstants::XOR_OP_VALUE_HALF;
36 36 
37template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>37template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
38-struct MergeSortBigSize : public MergeMoreCoreCommon::MergeMoreCoreBase<38+struct MergeSortBigSize
39- MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> {39+ : public MergeMoreCoreCommon::MergeMoreCoreBase<MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T,
40- using Base = MergeMoreCoreCommon::MergeMoreCoreBase<40+ CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> {
41- MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>;41+ using Base = MergeMoreCoreCommon::MergeMoreCoreBase<MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T,
42+ CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>;
42 friend Base;43 friend Base;
43 44 
44 __aicore__ inline MergeSortBigSize() {}45 __aicore__ inline MergeSortBigSize() {}
45 __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace,46 __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace,
46- const SortRegBaseTilingData* tilingData, TPipe* pipe);47+ const SortRegBaseTilingData* tilingData, TPipe* pipe);
47 __aicore__ inline void InitMergeBuffers();48 __aicore__ inline void InitMergeBuffers();
48 __aicore__ inline void ExtractAndCopyOut();49 __aicore__ inline void ExtractAndCopyOut();
49};50};
50 51 
51template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>52template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE>
52-__aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::Init(GM_ADDR inputValue,53+__aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::Init(
53- GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const SortRegBaseTilingData* tilingData, TPipe* pipe)54+ GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const SortRegBaseTilingData* tilingData,
55+ TPipe* pipe)
54{56{
55 this->blockIdx_ = GetBlockIdx();57 this->blockIdx_ = GetBlockIdx();
56 this->pipe_ = pipe;58 this->pipe_ = pipe;
@@ -66,18 +68,17 @@ __aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>
66 this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_);68 this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_);
67 // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per69 // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per
68 // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.70 // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses.
69- uint64_t rowWorkspaceBytes =71+ uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) *
70- ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) * sortBufferSize);72+ sortBufferSize);
71 uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE);73 uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE);
72- this->rowWorkspaceOffset_ =74+ this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2;
73- static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2;
74 this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE;75 this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE;
75 76 
76 this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));77 this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue));
77 this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value));78 this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value));
78 this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));79 this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices));
79- this->workspaceGm_[0].SetGlobalBuffer(80+ this->workspaceGm_[0].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_,
80- (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_, rowWorkspaceElements);81+ rowWorkspaceElements);
81 this->workspaceGm_[1].SetGlobalBuffer(82 this->workspaceGm_[1].SetGlobalBuffer(
82 (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements);83 (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements);
83 84 
@@ -97,16 +98,16 @@ __aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>
97{98{
98 uint32_t sortBufferSize = 8;99 uint32_t sortBufferSize = 8;
99 this->pipe_->InitBuffer(this->sortedQueue_, MERGE_MORE_BUFFER_NUM,100 this->pipe_->InitBuffer(this->sortedQueue_, MERGE_MORE_BUFFER_NUM,
100- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);101+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);
101 this->pipe_->InitBuffer(this->copyInQueue_, MERGE_MORE_BUFFER_NUM,102 this->pipe_->InitBuffer(this->copyInQueue_, MERGE_MORE_BUFFER_NUM,
102- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);103+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize);
103 this->pipe_->InitBuffer(this->castValueQueue_, MERGE_MORE_BUFFER_NUM,104 this->pipe_->InitBuffer(this->castValueQueue_, MERGE_MORE_BUFFER_NUM,
104- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE));105+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE));
105 this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM,106 this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM,
106- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(uint32_t));107+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(uint32_t));
107 if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {108 if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) {
108 this->pipe_->InitBuffer(this->outIndexQueue_, MERGE_MORE_BUFFER_NUM,109 this->pipe_->InitBuffer(this->outIndexQueue_, MERGE_MORE_BUFFER_NUM,
109- MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(INDEX_TYPE));110+ MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(INDEX_TYPE));
110 }111 }
111}112}
112 113 
@@ -121,7 +122,7 @@ __aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>
121 LocalTensor<CONVERT_TYPE> castValue = this->castValueQueue_.template AllocTensor<CONVERT_TYPE>();122 LocalTensor<CONVERT_TYPE> castValue = this->castValueQueue_.template AllocTensor<CONVERT_TYPE>();
122 LocalTensor<uint32_t> castIndex = this->castIndexQueue_.template AllocTensor<uint32_t>();123 LocalTensor<uint32_t> castIndex = this->castIndexQueue_.template AllocTensor<uint32_t>();
123 AscendC::Extract(castValue, castIndex, sortTempBuffer,124 AscendC::Extract(castValue, castIndex, sortTempBuffer,
124- Ops::Base::CeilDiv(this->curLoopSortedNum_, static_cast<int64_t>(DEALING_EXTRACT_NUM_ONCE)));125+ Ops::Base::CeilDiv(this->curLoopSortedNum_, static_cast<int64_t>(DEALING_EXTRACT_NUM_ONCE)));
125 if constexpr (!IS_DESCEND) {126 if constexpr (!IS_DESCEND) {
126 this->FlipSignBit(castValue, ROUND_UP_AGLIN(this->curLoopSortedNum_));127 this->FlipSignBit(castValue, ROUND_UP_AGLIN(this->curLoopSortedNum_));
127 }128 }
@@ -160,4 +161,4 @@ __aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>
160 this->sortedQueue_.FreeTensor(sortTempBuffer);161 this->sortedQueue_.FreeTensor(sortTempBuffer);
161 this->outOffset_ += this->curLoopSortedNum_;162 this->outOffset_ += this->curLoopSortedNum_;
162}163}
163-#endif //MERGE_SORT_BIG_SIZE_H164+#endif // MERGE_SORT_BIG_SIZE_H
Mmath/sort/op_kernel/arch35/sort_merge_intra_core.h+29-30
@@ -32,11 +32,10 @@ using namespace AscendC;
32 32 
33// Import shared constants from MergeSortConstants namespace33// Import shared constants from MergeSortConstants namespace
34using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;34using MergeSortConstants::DEALING_CONCAT_NUM_ONCE;
35-using MergeSortConstants::DEALING_SORT_NUM_ONCE;
36using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE;35using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE;
37-using MergeSortConstants::MERGE_LIST_MAX_NUM;36+using MergeSortConstants::DEALING_SORT_NUM_ONCE;
38using MergeSortConstants::MERGE_INTRA_BUFFER_NUM;37using MergeSortConstants::MERGE_INTRA_BUFFER_NUM;
39- 38+using MergeSortConstants::MERGE_LIST_MAX_NUM;
40 39 
41/**40/**
42 * @brief Intra-core block merge sort: independent per-core sort+merge without inter-core coordination41 * @brief Intra-core block merge sort: independent per-core sort+merge without inter-core coordination
@@ -45,10 +44,11 @@ using MergeSortConstants::MERGE_INTRA_BUFFER_NUM;
45 * @tparam IsDescend Sort order: true for descending, false for ascending44 * @tparam IsDescend Sort order: true for descending, false for ascending
46 */45 */
47template <typename ValueType, typename IndexType, bool IsDescend>46template <typename ValueType, typename IndexType, bool IsDescend>
48-class SortMergeIntraCore : public MergeIntraCoreCommon::MergeIntraCoreBase<47+class SortMergeIntraCore
49- SortMergeIntraCore<ValueType, IndexType, IsDescend>, ValueType, IndexType, IsDescend> {48+ : public MergeIntraCoreCommon::MergeIntraCoreBase<SortMergeIntraCore<ValueType, IndexType, IsDescend>, ValueType,
50- using Base = MergeIntraCoreCommon::MergeIntraCoreBase<49+ IndexType, IsDescend> {
51- SortMergeIntraCore<ValueType, IndexType, IsDescend>, ValueType, IndexType, IsDescend>;50+ using Base = MergeIntraCoreCommon::MergeIntraCoreBase<SortMergeIntraCore<ValueType, IndexType, IsDescend>,
51+ ValueType, IndexType, IsDescend>;
52 friend Base;52 friend Base;
53 53 
54public:54public:
@@ -56,14 +56,14 @@ public:
56 56 
57 __aicore__ inline SortMergeIntraCore() {}57 __aicore__ inline SortMergeIntraCore() {}
58 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace,58 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace,
59- const SortRegBaseTilingData* tilingData, TPipe* pipe);59+ const SortRegBaseTilingData* tilingData, TPipe* pipe);
60 60 
61protected:61protected:
62 __aicore__ inline void InitPhase1Buffers();62 __aicore__ inline void InitPhase1Buffers();
63 __aicore__ inline void InitPhase2Buffers();63 __aicore__ inline void InitPhase2Buffers();
64 __aicore__ inline void InitPhase3Buffers();64 __aicore__ inline void InitPhase3Buffers();
65- __aicore__ inline void ExtractAndCopyChunk(int64_t cacheBatchOffset, uint32_t cacheOffset,65+ __aicore__ inline void ExtractAndCopyChunk(int64_t cacheBatchOffset, uint32_t cacheOffset, int64_t outputOffset,
66- int64_t outputOffset, uint32_t elemProcessed, uint32_t elemCount);66+ uint32_t elemProcessed, uint32_t elemCount);
67 67 
68private:68private:
69 // Phase 3 only69 // Phase 3 only
@@ -91,8 +91,8 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Init
91 this->alignNum_ = tilingData->keyParams3;91 this->alignNum_ = tilingData->keyParams3;
92 this->maxMergeIterations_ = tilingData->keyParams5;92 this->maxMergeIterations_ = tilingData->keyParams5;
93 93 
94- if (this->blockSortSize_ == 0 || this->extractChunkSize_ == 0 ||94+ if (this->blockSortSize_ == 0 || this->extractChunkSize_ == 0 || this->blocksPerRow_ == 0 ||
95- this->blocksPerRow_ == 0 || this->sortAxisNum_ <= 0) {95+ this->sortAxisNum_ <= 0) {
96 return;96 return;
97 }97 }
98 98 
@@ -102,8 +102,8 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Init
102 this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType);102 this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType);
103 this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE;103 this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE;
104 this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE;104 this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE;
105- this->lastBlockSize_ = static_cast<uint32_t>(105+ this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ -
106- this->sortAxisNum_ - static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_);106+ static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_);
107 107 
108 // Set GM buffers108 // Set GM buffers
109 this->inputXGm_.SetGlobalBuffer((__gm__ ValueType*)x);109 this->inputXGm_.SetGlobalBuffer((__gm__ ValueType*)x);
@@ -113,10 +113,10 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Init
113 // Cache stores sort struct data (8 bytes per element: index + value)113 // Cache stores sort struct data (8 bytes per element: index + value)
114 // Each core has its own cache region, reused across batches114 // Each core has its own cache region, reused across batches
115 // perBatchCacheLen: sort struct length for one batch (with ping-pong, in ValueType units)115 // perBatchCacheLen: sort struct length for one batch (with ping-pong, in ValueType units)
116- int64_t perCoreCacheLen = static_cast<int64_t>(this->batchSortLen_) * 2; // ping-pong, reused per batch116+ int64_t perCoreCacheLen = static_cast<int64_t>(this->batchSortLen_) * 2; // ping-pong, reused per batch
117 117 
118 this->cacheGm_.SetGlobalBuffer((__gm__ ValueType*)workspace +118 this->cacheGm_.SetGlobalBuffer((__gm__ ValueType*)workspace +
119- static_cast<int64_t>(this->blockIdx_) * perCoreCacheLen);119+ static_cast<int64_t>(this->blockIdx_) * perCoreCacheLen);
120 120 
121 // Note: Queue/Buffer initialization is deferred to base Process() per phase121 // Note: Queue/Buffer initialization is deferred to base Process() per phase
122 // to optimize UB usage. Each phase initializes only what it needs.122 // to optimize UB usage. Each phase initializes only what it needs.
@@ -125,8 +125,7 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Init
125template <typename ValueType, typename IndexType, bool IsDescend>125template <typename ValueType, typename IndexType, bool IsDescend>
126__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::InitPhase1Buffers()126__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::InitPhase1Buffers()
127{127{
128- this->pipe_->InitBuffer(128+ this->pipe_->InitBuffer(this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType));
129- this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType));
130 this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_);129 this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_);
131 this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_);130 this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_);
132 this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_);131 this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_);
@@ -146,18 +145,17 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Init
146{145{
147 uint32_t extractInSize = AscendC::GetSortLen<ValueType>(this->extractChunkSize_) * sizeof(ValueType);146 uint32_t extractInSize = AscendC::GetSortLen<ValueType>(this->extractChunkSize_) * sizeof(ValueType);
148 this->pipe_->InitBuffer(this->extractInQueue_, MERGE_INTRA_BUFFER_NUM, extractInSize);147 this->pipe_->InitBuffer(this->extractInQueue_, MERGE_INTRA_BUFFER_NUM, extractInSize);
149- this->pipe_->InitBuffer(148+ this->pipe_->InitBuffer(this->outValueQueue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(ValueType));
150- this->outValueQueue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(ValueType));
151 this->pipe_->InitBuffer(this->outIdxQueue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(uint32_t));149 this->pipe_->InitBuffer(this->outIdxQueue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(uint32_t));
152 if constexpr (IsSameType<int64_t, IndexType>::value) {150 if constexpr (IsSameType<int64_t, IndexType>::value) {
153- this->pipe_->InitBuffer(151+ this->pipe_->InitBuffer(this->outIdxInt64Queue_, MERGE_INTRA_BUFFER_NUM,
154- this->outIdxInt64Queue_, MERGE_INTRA_BUFFER_NUM, this->extractChunkSize_ * sizeof(int64_t));152+ this->extractChunkSize_ * sizeof(int64_t));
155 }153 }
156}154}
157 155 
158template <typename ValueType, typename IndexType, bool IsDescend>156template <typename ValueType, typename IndexType, bool IsDescend>
159-__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut(157+__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut(int64_t batchIdx,
160- int64_t batchIdx, uint32_t resultRegion)158+ uint32_t resultRegion)
161{159{
162 int64_t outputOffset = batchIdx * this->sortAxisNum_;160 int64_t outputOffset = batchIdx * this->sortAxisNum_;
163 161 
@@ -169,8 +167,10 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Extr
169 167 
170 while (elemProcessed < this->sortAxisNum_) {168 while (elemProcessed < this->sortAxisNum_) {
171 uint32_t elemCount = (elemProcessed + this->extractChunkSize_ <= this->sortAxisNum_) ?169 uint32_t elemCount = (elemProcessed + this->extractChunkSize_ <= this->sortAxisNum_) ?
172- this->extractChunkSize_ : (this->sortAxisNum_ - elemProcessed);170+ this->extractChunkSize_ :
173- if (elemCount == 0) break;171+ (this->sortAxisNum_ - elemProcessed);
172+ if (elemCount == 0)
173+ break;
174 174 
175 ExtractAndCopyChunk(cacheBatchOffset, cacheOffset, outputOffset, elemProcessed, elemCount);175 ExtractAndCopyChunk(cacheBatchOffset, cacheOffset, outputOffset, elemProcessed, elemCount);
176 176 
@@ -181,8 +181,7 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Extr
181 181 
182template <typename ValueType, typename IndexType, bool IsDescend>182template <typename ValueType, typename IndexType, bool IsDescend>
183__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyChunk(183__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyChunk(
184- int64_t cacheBatchOffset, uint32_t cacheOffset,184+ int64_t cacheBatchOffset, uint32_t cacheOffset, int64_t outputOffset, uint32_t elemProcessed, uint32_t elemCount)
185- int64_t outputOffset, uint32_t elemProcessed, uint32_t elemCount)
186{185{
187 LocalTensor<ValueType> cacheLocal = this->extractInQueue_.template AllocTensor<ValueType>();186 LocalTensor<ValueType> cacheLocal = this->extractInQueue_.template AllocTensor<ValueType>();
188 DataCopyExtParams loadParams{187 DataCopyExtParams loadParams{
@@ -197,8 +196,8 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Extr
197 196 
198 // Flip back sign bit for ascending order (was flipped in SortBlockToStruct)197 // Flip back sign bit for ascending order (was flipped in SortBlockToStruct)
199 if constexpr (!IsDescend) {198 if constexpr (!IsDescend) {
200- Adds(valueLocal.template ReinterpretCast<int32_t>(),199+ Adds(valueLocal.template ReinterpretCast<int32_t>(), valueLocal.template ReinterpretCast<int32_t>(), 0x80000000,
201- valueLocal.template ReinterpretCast<int32_t>(), 0x80000000, elemCount);200+ elemCount);
202 }201 }
203 202 
204 this->outValueQueue_.EnQue(valueLocal);203 this->outValueQueue_.EnQue(valueLocal);
Mmath/sort/op_kernel/arch35/sort_merge_sort.h+31-26
@@ -36,29 +36,31 @@ class MergeSort {
36public:36public:
37 __aicore__ inline MergeSort(){};37 __aicore__ inline MergeSort(){};
38 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,38 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,
39- const SortRegBaseTilingData *__restrict tilingData, TPipe *pipe);39+ const SortRegBaseTilingData* __restrict tilingData, TPipe* pipe);
40 __aicore__ inline void Process();40 __aicore__ inline void Process();
41 41 
42private:42private:
43 __aicore__ inline void ProcessSingleBlockSort(GlobalTensor<T1> inputX, int32_t sortLoopRound);43 __aicore__ inline void ProcessSingleBlockSort(GlobalTensor<T1> inputX, int32_t sortLoopRound);
44 __aicore__ inline void ParserTilingData();44 __aicore__ inline void ParserTilingData();
45 __aicore__ inline void CopyDataIn(GlobalTensor<T1> inputX, uint64_t tileOffset, uint32_t currTileSize,45 __aicore__ inline void CopyDataIn(GlobalTensor<T1> inputX, uint64_t tileOffset, uint32_t currTileSize,
46- uint32_t oneCoreRowNum);46+ uint32_t oneCoreRowNum);
47 __aicore__ inline void VbsMergeSortBf16(LocalTensor<bfloat16_t> xLocal, LocalTensor<T1> sortedValueLocal,47 __aicore__ inline void VbsMergeSortBf16(LocalTensor<bfloat16_t> xLocal, LocalTensor<T1> sortedValueLocal,
48- LocalTensor<uint32_t> sortedValueIndexLocal, uint32_t currTileSize, uint32_t nowCoreRealRowNum);48+ LocalTensor<uint32_t> sortedValueIndexLocal, uint32_t currTileSize,
49+ uint32_t nowCoreRealRowNum);
49 __aicore__ inline void VbsMergeSort(LocalTensor<T1> xLocal, LocalTensor<T1> sortedValueLocal,50 __aicore__ inline void VbsMergeSort(LocalTensor<T1> xLocal, LocalTensor<T1> sortedValueLocal,
50- LocalTensor<uint32_t> sortedValueIndexLocal, uint32_t currTileSize, uint32_t nowCoreRealRowNum);51+ LocalTensor<uint32_t> sortedValueIndexLocal, uint32_t currTileSize,
52+ uint32_t nowCoreRealRowNum);
51 __aicore__ inline void CopyValue2Gm(uint64_t gmOffset, uint64_t tileOffset, uint32_t outputLastDimValue,53 __aicore__ inline void CopyValue2Gm(uint64_t gmOffset, uint64_t tileOffset, uint32_t outputLastDimValue,
52- uint32_t oneCoreRowNum);54+ uint32_t oneCoreRowNum);
53 __aicore__ inline void flipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offsetOneRow, uint32_t aglinTileSize);55 __aicore__ inline void flipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offsetOneRow, uint32_t aglinTileSize);
54 56 
55 GlobalTensor<T1> inputXGm_;57 GlobalTensor<T1> inputXGm_;
56 GlobalTensor<T1> outValueGm_;58 GlobalTensor<T1> outValueGm_;
57 GlobalTensor<T2> outIdxGm_;59 GlobalTensor<T2> outIdxGm_;
58 60 
59- TPipe *pipe_;61+ TPipe* pipe_;
60 // 599040 5120tile62 // 599040 5120tile
61- const SortRegBaseTilingData *tilingData_;63+ const SortRegBaseTilingData* tilingData_;
62 TQue<QuePosition::VECIN, 1> inQueueX_;64 TQue<QuePosition::VECIN, 1> inQueueX_;
63 TBuf<TPosition::VECCALC> tmpUb_;65 TBuf<TPosition::VECCALC> tmpUb_;
64 TQue<QuePosition::VECOUT, 1> outIdxQueue_;66 TQue<QuePosition::VECOUT, 1> outIdxQueue_;
@@ -83,16 +85,17 @@ private:
83};85};
84 86 
85template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis>87template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis>
86-__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::Init(GM_ADDR x, GM_ADDR value,88+__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::Init(
87- GM_ADDR sortIndex, GM_ADDR workspace, const SortRegBaseTilingData *__restrict tilingData, TPipe *pipe)89+ GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace, const SortRegBaseTilingData* __restrict tilingData,
90+ TPipe* pipe)
88{91{
89 blockIdx_ = GetBlockIdx();92 blockIdx_ = GetBlockIdx();
90 pipe_ = pipe;93 pipe_ = pipe;
91 tilingData_ = tilingData;94 tilingData_ = tilingData;
92 ParserTilingData();95 ParserTilingData();
93- inputXGm_.SetGlobalBuffer((__gm__ T1 *)x);96+ inputXGm_.SetGlobalBuffer((__gm__ T1*)x);
94- outValueGm_.SetGlobalBuffer((__gm__ T1 *)value);97+ outValueGm_.SetGlobalBuffer((__gm__ T1*)value);
95- outIdxGm_.SetGlobalBuffer((__gm__ T2 *)sortIndex);98+ outIdxGm_.SetGlobalBuffer((__gm__ T2*)sortIndex);
96 99 
97 uint32_t bufferNum = tilingData_->keyParams4;100 uint32_t bufferNum = tilingData_->keyParams4;
98 pipe_->InitBuffer(inQueueX_, bufferNum, tilingData_->keyParams1);101 pipe_->InitBuffer(inQueueX_, bufferNum, tilingData_->keyParams1);
@@ -110,7 +113,7 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
110 pipe_->InitBuffer(sortedValueLocalCastTbuf_, alignSize_ * sizeof(CONVERT_TYPE) * oneCoreRowNum_);113 pipe_->InitBuffer(sortedValueLocalCastTbuf_, alignSize_ * sizeof(CONVERT_TYPE) * oneCoreRowNum_);
111 indexLocal_ = indeXLocalTbuf_.AllocTensor<uint32_t>();114 indexLocal_ = indeXLocalTbuf_.AllocTensor<uint32_t>();
112 // init indexLocal value115 // init indexLocal value
113- __local_mem__ int32_t *indexValuePtr = (__ubuf__ int32_t *)indexLocal_.GetPhyAddr();116+ __local_mem__ int32_t* indexValuePtr = (__ubuf__ int32_t*)indexLocal_.GetPhyAddr();
114 uint32_t vfLenB32 = Ops::Base::GetVRegSize() / sizeof(int32_t);117 uint32_t vfLenB32 = Ops::Base::GetVRegSize() / sizeof(int32_t);
115 uint16_t repeatTime = CeilDivision(alignSize_, vfLenB32);118 uint16_t repeatTime = CeilDivision(alignSize_, vfLenB32);
116 uint32_t aglinTileSizeCopy = alignSize_;119 uint32_t aglinTileSizeCopy = alignSize_;
@@ -123,7 +126,7 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
123 MicroAPI::MaskReg vciMask = MicroAPI::UpdateMask<uint32_t>(aglinTileSizeCopy);126 MicroAPI::MaskReg vciMask = MicroAPI::UpdateMask<uint32_t>(aglinTileSizeCopy);
124 MicroAPI::Adds(indexTensor, vciTensor, i * vfLenB32, vciMask);127 MicroAPI::Adds(indexTensor, vciTensor, i * vfLenB32, vciMask);
125 MicroAPI::DataCopy<int32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(indexValuePtr, indexTensor, vfLenB32,128 MicroAPI::DataCopy<int32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(indexValuePtr, indexTensor, vfLenB32,
126- vciMask);129+ vciMask);
127 }130 }
128 }131 }
129}132}
@@ -210,13 +213,13 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
210 }213 }
211 // isSort32SmallAxis enables compile-time repeatTimes=1 for hardware Extract fast path214 // isSort32SmallAxis enables compile-time repeatTimes=1 for hardware Extract fast path
212 AscendC::Extract(sortedValueLocalCast[offsetOneRow], sortedValueIndexLocal[offsetOneRow], sortedLocal,215 AscendC::Extract(sortedValueLocalCast[offsetOneRow], sortedValueIndexLocal[offsetOneRow], sortedLocal,
213- isSort32SmallAxis == 1 ? 1 : extractRepeatTimes);216+ isSort32SmallAxis == 1 ? 1 : extractRepeatTimes);
214 if constexpr (isDescend == 0) {217 if constexpr (isDescend == 0) {
215 flipSignBit(sortedValueLocalCast, offsetOneRow, aglinTileSize);218 flipSignBit(sortedValueLocalCast, offsetOneRow, aglinTileSize);
216 }219 }
217 }220 }
218 AscendC::Cast(sortedValueLocal, sortedValueLocalCast, AscendC::RoundMode::CAST_RINT,221 AscendC::Cast(sortedValueLocal, sortedValueLocalCast, AscendC::RoundMode::CAST_RINT,
219- aglinTileSize * nowCoreRealRowNum);222+ aglinTileSize * nowCoreRealRowNum);
220}223}
221template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis>224template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis>
222__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::VbsMergeSort(225__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::VbsMergeSort(
@@ -246,7 +249,7 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
246 // 处理sort后的结果数据，输出排序后的value和index249 // 处理sort后的结果数据，输出排序后的value和index
247 // isSort32SmallAxis enables compile-time repeatTimes=1 for hardware Extract fast path250 // isSort32SmallAxis enables compile-time repeatTimes=1 for hardware Extract fast path
248 AscendC::Extract(sortedValueLocal[offsetOneRow], sortedValueIndexLocal[offsetOneRow], sortedLocal,251 AscendC::Extract(sortedValueLocal[offsetOneRow], sortedValueIndexLocal[offsetOneRow], sortedLocal,
249- isSort32SmallAxis == 1 ? 1 : extractRepeatTimes);252+ isSort32SmallAxis == 1 ? 1 : extractRepeatTimes);
250 if constexpr (isDescend == 0) {253 if constexpr (isDescend == 0) {
251 flipSignBit(sortedValueLocal, offsetOneRow, alignSize_);254 flipSignBit(sortedValueLocal, offsetOneRow, alignSize_);
252 }255 }
@@ -254,8 +257,8 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
254}257}
255 258 
256template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis>259template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis>
257-__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::CopyValue2Gm(uint64_t gmOffset,260+__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::CopyValue2Gm(
258- uint64_t tileOffset, uint32_t outputLastDimValue, uint32_t oneCoreRowNum)261+ uint64_t gmOffset, uint64_t tileOffset, uint32_t outputLastDimValue, uint32_t oneCoreRowNum)
259{262{
260 // value stride263 // value stride
261 uint32_t currTileSizeAlign = ROUND_UP_AGLIN(outputLastDimValue * sizeof(T1)) / sizeof(T1);264 uint32_t currTileSizeAlign = ROUND_UP_AGLIN(outputLastDimValue * sizeof(T1)) / sizeof(T1);
@@ -269,15 +272,17 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
269 // copy result out272 // copy result out
270 AscendC::LocalTensor<T1> outValueLocal = outValueQueue_.DeQue<T1>();273 AscendC::LocalTensor<T1> outValueLocal = outValueQueue_.DeQue<T1>();
271 AscendC::LocalTensor<T2> outIndexLocal = outIdxQueue_.DeQue<T2>();274 AscendC::LocalTensor<T2> outIndexLocal = outIdxQueue_.DeQue<T2>();
272- AscendC::DataCopyExtParams dataCopyParamValue{ static_cast<uint16_t>(oneCoreRowNum),275+ AscendC::DataCopyExtParams dataCopyParamValue{static_cast<uint16_t>(oneCoreRowNum),
273- static_cast<uint32_t>(outputLastDimValue * sizeof(T1)), ubStrideValue, 0, 0 };276+ static_cast<uint32_t>(outputLastDimValue * sizeof(T1)), ubStrideValue,
277+ 0, 0};
274 AscendC::DataCopyPad(outValueGm_[gmOffset + tileOffset], outValueLocal, dataCopyParamValue);278 AscendC::DataCopyPad(outValueGm_[gmOffset + tileOffset], outValueLocal, dataCopyParamValue);
275 279 
276- AscendC::DataCopyExtParams dataCopyParamIndex{ static_cast<uint16_t>(oneCoreRowNum), // 连续数据块的个数280+ AscendC::DataCopyExtParams dataCopyParamIndex{
281+ static_cast<uint16_t>(oneCoreRowNum), // 连续数据块的个数
277 static_cast<uint32_t>(outputLastDimValue * sizeof(T2)), // 每个连续传输数据块的长度，长度为Byte282 static_cast<uint32_t>(outputLastDimValue * sizeof(T2)), // 每个连续传输数据块的长度，长度为Byte
278 ubStrideIndex, // 源操作数，相邻连续数据块的间隔283 ubStrideIndex, // 源操作数，相邻连续数据块的间隔
279 0, // 目的操作数，相邻连续数据块的间隔284 0, // 目的操作数，相邻连续数据块的间隔
280- 0 };285+ 0};
281 AscendC::DataCopyPad(outIdxGm_[gmOffset + tileOffset], outIndexLocal, dataCopyParamIndex);286 AscendC::DataCopyPad(outIdxGm_[gmOffset + tileOffset], outIndexLocal, dataCopyParamIndex);
282 outIdxQueue_.FreeTensor(outIndexLocal);287 outIdxQueue_.FreeTensor(outIndexLocal);
283 outValueQueue_.FreeTensor(outValueLocal);288 outValueQueue_.FreeTensor(outValueLocal);
@@ -318,10 +323,10 @@ __aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallA
318 VbsMergeSort(xLocal, sortedValueLocal, sortedValueIndexLocal, numTileData_, nowCoreRealRowNum);323 VbsMergeSort(xLocal, sortedValueLocal, sortedValueIndexLocal, numTileData_, nowCoreRealRowNum);
319 }324 }
320 if constexpr (IsSameType<int64_t, T2>::value) {325 if constexpr (IsSameType<int64_t, T2>::value) {
321- AscendC::LocalTensor<int32_t> sortedValueIndexInt32Local =326+ AscendC::LocalTensor<int32_t> sortedValueIndexInt32Local = sortedValueIndexLocal
322- sortedValueIndexLocal.template ReinterpretCast<int32_t>();327+ .template ReinterpretCast<int32_t>();
323 AscendC::Cast(sortedValueIndexInt64Local, sortedValueIndexInt32Local, AscendC::RoundMode::CAST_NONE,328 AscendC::Cast(sortedValueIndexInt64Local, sortedValueIndexInt32Local, AscendC::RoundMode::CAST_NONE,
324- nowCoreRealRowNum * alignSize_);329+ nowCoreRealRowNum * alignSize_);
325 outIdxQueue_.EnQue<int64_t>(sortedValueIndexInt64Local);330 outIdxQueue_.EnQue<int64_t>(sortedValueIndexInt64Local);
326 } else {331 } else {
327 outIdxQueue_.EnQue<uint32_t>(sortedValueIndexLocal);332 outIdxQueue_.EnQue<uint32_t>(sortedValueIndexLocal);
Mmath/sort/op_kernel/arch35/sort_non_last_small_axis.h+40-34
@@ -25,9 +25,10 @@ namespace Sort {
25using namespace AscendC;25using namespace AscendC;
26 26 
27template <typename T, typename OutIdxT>27template <typename T, typename OutIdxT>
28-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__ void StoreSingleInnerOutput(28+__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__
29- uint32_t axisLen, uint32_t threadNum, uint64_t baseOffset, uint64_t innerSize, __ubuf__ T* sortedValue,29+ void StoreSingleInnerOutput(uint32_t axisLen, uint32_t threadNum, uint64_t baseOffset, uint64_t innerSize,
30- __ubuf__ uint32_t* sortedIndex, __gm__ volatile T* outputValue, __gm__ volatile OutIdxT* outputIndex)30+ __ubuf__ T* sortedValue, __ubuf__ uint32_t* sortedIndex, __gm__ volatile T* outputValue,
31+ __gm__ volatile OutIdxT* outputIndex)
31{32{
32 for (uint32_t axis = static_cast<uint32_t>(threadIdx.x); axis < axisLen; axis += threadNum) {33 for (uint32_t axis = static_cast<uint32_t>(threadIdx.x); axis < axisLen; axis += threadNum) {
33 uint64_t gmOffset = baseOffset + static_cast<uint64_t>(axis) * innerSize;34 uint64_t gmOffset = baseOffset + static_cast<uint64_t>(axis) * innerSize;
@@ -37,24 +38,26 @@ __simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicor
37}38}
38 39 
39template <typename OutIdxT>40template <typename OutIdxT>
40-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__ void BuildOutputIndexTile(41+__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__
41- uint32_t axisLen, uint32_t innerChunk, uint32_t threadNum, uint32_t indexAxisElems, uint32_t outputIndexRowElems,42+ void BuildOutputIndexTile(uint32_t axisLen, uint32_t innerChunk, uint32_t threadNum, uint32_t indexAxisElems,
42- __ubuf__ uint32_t* sortedIndex, __ubuf__ OutIdxT* outputIndex)43+ uint32_t outputIndexRowElems, __ubuf__ uint32_t* sortedIndex,
44+ __ubuf__ OutIdxT* outputIndex)
43{45{
44 uint32_t total = axisLen * innerChunk;46 uint32_t total = axisLen * innerChunk;
45 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += threadNum) {47 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += threadNum) {
46 uint32_t inner = idx / axisLen;48 uint32_t inner = idx / axisLen;
47 uint32_t axis = idx - inner * axisLen;49 uint32_t axis = idx - inner * axisLen;
48- outputIndex[axis * outputIndexRowElems + inner] =50+ outputIndex[axis * outputIndexRowElems + inner] = static_cast<OutIdxT>(
49- static_cast<OutIdxT>(sortedIndex[inner * indexAxisElems + axis]);51+ sortedIndex[inner * indexAxisElems + axis]);
50 }52 }
51}53}
52 54 
53template <typename T, typename SortT, typename OutIdxT>55template <typename T, typename SortT, typename OutIdxT>
54-__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__ void StoreConvertedOutputTile(56+__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM) __aicore__
55- uint32_t axisLen, uint32_t innerChunk, uint32_t threadNum, uint32_t valueAxisElems, uint32_t indexAxisElems,57+ void StoreConvertedOutputTile(uint32_t axisLen, uint32_t innerChunk, uint32_t threadNum, uint32_t valueAxisElems,
56- uint64_t baseOffset, uint64_t innerSize, __ubuf__ SortT* sortedValue, __ubuf__ uint32_t* sortedIndex,58+ uint32_t indexAxisElems, uint64_t baseOffset, uint64_t innerSize,
57- __gm__ volatile T* outputValue, __gm__ volatile OutIdxT* outputIndex)59+ __ubuf__ SortT* sortedValue, __ubuf__ uint32_t* sortedIndex,
60+ __gm__ volatile T* outputValue, __gm__ volatile OutIdxT* outputIndex)
58{61{
59 uint32_t total = axisLen * innerChunk;62 uint32_t total = axisLen * innerChunk;
60 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += threadNum) {63 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < total; idx += threadNum) {
@@ -87,13 +90,13 @@ public:
87 using SortT_ = typename std::conditional_t<UseMergeSort && std::is_same_v<T, bfloat16_t>, float, T>;90 using SortT_ = typename std::conditional_t<UseMergeSort && std::is_same_v<T, bfloat16_t>, float, T>;
88 using RangeType_ = std::conditional_t<sizeof(T) <= sizeof(int16_t), int16_t, int32_t>;91 using RangeType_ = std::conditional_t<sizeof(T) <= sizeof(int16_t), int16_t, int32_t>;
89 using IdxType_ = std::conditional_t<sizeof(T) <= sizeof(int16_t), uint16_t, uint32_t>;92 using IdxType_ = std::conditional_t<sizeof(T) <= sizeof(int16_t), uint16_t, uint32_t>;
90- using CastType_ =93+ using CastType_ = std::conditional_t<sizeof(T) == 1,
91- std::conditional_t<sizeof(T) == 1, std::conditional_t<std::is_same_v<T, uint8_t>, uint16_t, int16_t>, T>;94+ std::conditional_t<std::is_same_v<T, uint8_t>, uint16_t, int16_t>, T>;
92 static constexpr bool IsBf16Merge_ = UseMergeSort && std::is_same_v<T, bfloat16_t>;95 static constexpr bool IsBf16Merge_ = UseMergeSort && std::is_same_v<T, bfloat16_t>;
93 96 
94 __aicore__ inline SortNonLastSmallAxis() {}97 __aicore__ inline SortNonLastSmallAxis() {}
95- __aicore__ inline void Init(98+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,
96- GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace, const SortRegBaseTilingData* tilingData, TPipe* pipe);99+ const SortRegBaseTilingData* tilingData, TPipe* pipe);
97 100 
98 __aicore__ inline void ParseTilingData();101 __aicore__ inline void ParseTilingData();
99 __aicore__ inline void BuildOutputs(uint32_t curInnerChunk);102 __aicore__ inline void BuildOutputs(uint32_t curInnerChunk);
@@ -217,19 +220,19 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>
217 AscendC::MicroAPI::RegTensor<CastType_> valueReg;220 AscendC::MicroAPI::RegTensor<CastType_> valueReg;
218 AscendC::MicroAPI::RegTensor<RangeType_> valueIdxReg;221 AscendC::MicroAPI::RegTensor<RangeType_> valueIdxReg;
219 AscendC::MicroAPI::MaskReg valueMask = AscendC::MicroAPI::UpdateMask<CastType_>(curInnerChunk);222 AscendC::MicroAPI::MaskReg valueMask = AscendC::MicroAPI::UpdateMask<CastType_>(curInnerChunk);
220- AscendC::MicroAPI::MaskReg valueIdxMask =223+ AscendC::MicroAPI::MaskReg
221- AscendC::MicroAPI::CreateMask<RangeType_, AscendC::MicroAPI::MaskPattern::ALL>();224+ valueIdxMask = AscendC::MicroAPI::CreateMask<RangeType_, AscendC::MicroAPI::MaskPattern::ALL>();
222 225 
223 AscendC::MicroAPI::Arange(valueIdxReg, 0);226 AscendC::MicroAPI::Arange(valueIdxReg, 0);
224 AscendC::MicroAPI::Muls(valueIdxReg, valueIdxReg, static_cast<RangeType_>(this->valueAxisElems_), valueIdxMask);227 AscendC::MicroAPI::Muls(valueIdxReg, valueIdxReg, static_cast<RangeType_>(this->valueAxisElems_), valueIdxMask);
225 for (uint16_t axis = 0; axis < this->axisLen_; ++axis) {228 for (uint16_t axis = 0; axis < this->axisLen_; ++axis) {
226- AscendC::MicroAPI::DataCopyGather(229+ AscendC::MicroAPI::DataCopyGather(valueReg, sortedValueAddr + axis,
227- valueReg, sortedValueAddr + axis, (AscendC::MicroAPI::RegTensor<IdxType_>&)valueIdxReg, valueMask);230+ (AscendC::MicroAPI::RegTensor<IdxType_>&)valueIdxReg, valueMask);
228 if constexpr (sizeof(T) != 1) {231 if constexpr (sizeof(T) != 1) {
229 AscendC::MicroAPI::DataCopy(outputValueAddr + axis * this->inputRowElems_, valueReg, valueMask);232 AscendC::MicroAPI::DataCopy(outputValueAddr + axis * this->inputRowElems_, valueReg, valueMask);
230 } else {233 } else {
231- __local_mem__ CastType_* outputValueAddrB16 =234+ __local_mem__ CastType_* outputValueAddrB16 = reinterpret_cast<__local_mem__ CastType_*>(
232- reinterpret_cast<__local_mem__ CastType_*>(outputValueAddr + axis * this->inputRowElems_);235+ outputValueAddr + axis * this->inputRowElems_);
233 AscendC::MicroAPI::DataCopy<CastType_, AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(236 AscendC::MicroAPI::DataCopy<CastType_, AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(
234 outputValueAddrB16, valueReg, valueMask);237 outputValueAddrB16, valueReg, valueMask);
235 }238 }
@@ -248,23 +251,25 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>
248 251 
249 uint32_t valueBytes = curInnerChunk * sizeof(T);252 uint32_t valueBytes = curInnerChunk * sizeof(T);
250 uint32_t valueAlignedBytes = ROUND_UP_AGLIN(valueBytes);253 uint32_t valueAlignedBytes = ROUND_UP_AGLIN(valueBytes);
251- uint32_t valueSrcStride =254+ uint32_t valueSrcStride = (this->inputRowBytes_ > valueAlignedBytes) ?
252- (this->inputRowBytes_ > valueAlignedBytes) ? (this->inputRowBytes_ - valueAlignedBytes) / UB_BLOCK_SIZE : 0;255+ (this->inputRowBytes_ - valueAlignedBytes) / UB_BLOCK_SIZE :
256+ 0;
253 int64_t valueDstStride = (this->innerSize_ - static_cast<int64_t>(curInnerChunk)) * static_cast<int64_t>(sizeof(T));257 int64_t valueDstStride = (this->innerSize_ - static_cast<int64_t>(curInnerChunk)) * static_cast<int64_t>(sizeof(T));
254 // Store the restored [axisLen, curInnerChunk] tile back into the original258 // Store the restored [axisLen, curInnerChunk] tile back into the original
255 // non-last-axis GM layout without materializing a full transposed tensor.259 // non-last-axis GM layout without materializing a full transposed tensor.
256- DataCopyExtParams valueCopyParam{260+ DataCopyExtParams valueCopyParam{static_cast<uint16_t>(this->axisLen_), valueBytes, valueSrcStride, valueDstStride,
257- static_cast<uint16_t>(this->axisLen_), valueBytes, valueSrcStride, valueDstStride, 0};261+ 0};
258 DataCopyPad(outValueGm_[baseOffset], this->inputTile_, valueCopyParam);262 DataCopyPad(outValueGm_[baseOffset], this->inputTile_, valueCopyParam);
259 263 
260 uint32_t indexBytes = curInnerChunk * sizeof(OutIdxT);264 uint32_t indexBytes = curInnerChunk * sizeof(OutIdxT);
261 uint32_t indexAlignedBytes = ROUND_UP_AGLIN(indexBytes);265 uint32_t indexAlignedBytes = ROUND_UP_AGLIN(indexBytes);
262- uint32_t indexSrcStride =266+ uint32_t indexSrcStride = (outputIndexRowBytes_ > indexAlignedBytes) ?
263- (outputIndexRowBytes_ > indexAlignedBytes) ? (outputIndexRowBytes_ - indexAlignedBytes) / UB_BLOCK_SIZE : 0;267+ (outputIndexRowBytes_ - indexAlignedBytes) / UB_BLOCK_SIZE :
264- int64_t indexDstStride =268+ 0;
265- (this->innerSize_ - static_cast<int64_t>(curInnerChunk)) * static_cast<int64_t>(sizeof(OutIdxT));269+ int64_t indexDstStride = (this->innerSize_ - static_cast<int64_t>(curInnerChunk)) *
266- DataCopyExtParams indexCopyParam{270+ static_cast<int64_t>(sizeof(OutIdxT));
267- static_cast<uint16_t>(this->axisLen_), indexBytes, indexSrcStride, indexDstStride, 0};271+ DataCopyExtParams indexCopyParam{static_cast<uint16_t>(this->axisLen_), indexBytes, indexSrcStride, indexDstStride,
272+ 0};
268 DataCopyPad(outIndexGm_[baseOffset], outputIndex_, indexCopyParam);273 DataCopyPad(outIndexGm_[baseOffset], outputIndex_, indexCopyParam);
269 274 
270 // Sync: ensure MTE3 GM writes complete before the next tile starts275 // Sync: ensure MTE3 GM writes complete before the next tile starts
@@ -318,8 +323,9 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>
318}323}
319 324 
320template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort>325template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort>
321-__aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::StoreTile(326+__aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::StoreTile(int64_t inputOffset,
322- int64_t inputOffset, int64_t outputOffset, uint32_t curInnerChunk)327+ int64_t outputOffset,
328+ uint32_t curInnerChunk)
323{329{
324 (void)outputOffset;330 (void)outputOffset;
325 if constexpr (IsBf16Merge_) {331 if constexpr (IsBf16Merge_) {
Mmath/sort/op_kernel/arch35/sort_radix_sort_more_core.h+129-123
@@ -37,46 +37,54 @@ using RadixSortCommon::THREAD_DIM_NUM;
37 37 
38// T1输入x dtype T2输出Idx dtype UT无符号的数据类型38// T1输入x dtype T2输出Idx dtype UT无符号的数据类型
39template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>39template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
40-class SortRadixMoreCore : public RadixSortCommon::RadixMoreCoreBase<40+class SortRadixMoreCore : public RadixSortCommon::RadixMoreCoreBase<SortRadixMoreCore<T1, T2, UT, T3, isDescend>, T1,
41- SortRadixMoreCore<T1, T2, UT, T3, isDescend>, T1, T2, UT, T3, isDescend> {41+ T2, UT, T3, isDescend> {
42- using Base = RadixSortCommon::RadixMoreCoreBase<42+ using Base = RadixSortCommon::RadixMoreCoreBase<SortRadixMoreCore<T1, T2, UT, T3, isDescend>, T1, T2, UT, T3,
43- SortRadixMoreCore<T1, T2, UT, T3, isDescend>, T1, T2, UT, T3, isDescend>;43+ isDescend>;
44 friend Base;44 friend Base;
45+ 
45public:46public:
46 __aicore__ inline SortRadixMoreCore(){};47 __aicore__ inline SortRadixMoreCore(){};
47 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,48 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,
48- const SortRegBaseTilingData *__restrict tilingData, TPipe *pipe);49+ const SortRegBaseTilingData* __restrict tilingData, TPipe* pipe);
49 50 
50protected:51protected:
51 __aicore__ inline void ParserTilingData();52 __aicore__ inline void ParserTilingData();
52 __aicore__ inline void ScatterKeysGlobal(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,53 __aicore__ inline void ScatterKeysGlobal(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
53- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,54+ LocalTensor<uint32_t> xInputIndexLocal,
54- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,55+ LocalTensor<uint8_t> sortedValueLocal,
55- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,56+ LocalTensor<uint16_t> blockExcusiveSum,
56- uint32_t cureTileSize, uint32_t sortLoopRound);57+ LocalTensor<T3> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag,
58+ LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,
59+ uint32_t cureTileSize, uint32_t sortLoopRound);
57 __aicore__ inline void ScatterOutInt32(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,60 __aicore__ inline void ScatterOutInt32(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
58- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,61+ LocalTensor<uint32_t> xInputIndexLocal,
59- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,62+ LocalTensor<uint8_t> sortedValueLocal,
60- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,63+ LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
61- uint32_t cureTileSize);64+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
65+ uint32_t round, T3 tileDataStart, uint32_t cureTileSize);
62 __aicore__ inline void ScatterOutInt32ToInt64(LocalTensor<T1> xInputValueLocal,66 __aicore__ inline void ScatterOutInt32ToInt64(LocalTensor<T1> xInputValueLocal,
63- LocalTensor<uint32_t> sortedIndexLocal,67+ LocalTensor<uint32_t> sortedIndexLocal,
64- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,68+ LocalTensor<uint32_t> xInputIndexLocal,
65- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<uint32_t> blockDataInGlobalPos,69+ LocalTensor<uint8_t> sortedValueLocal,
66- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,70+ LocalTensor<uint16_t> blockExcusiveSum,
67- uint32_t cureTileSize);71+ LocalTensor<uint32_t> blockDataInGlobalPos,
72+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
73+ uint32_t round, T3 tileDataStart, uint32_t cureTileSize);
68 __aicore__ inline void ScatterOutInt64(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,74 __aicore__ inline void ScatterOutInt64(LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal,
69- LocalTensor<uint32_t> xInputIndexLocal, LocalTensor<uint8_t> sortedValueLocal,75+ LocalTensor<uint32_t> xInputIndexLocal,
70- LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,76+ LocalTensor<uint8_t> sortedValueLocal,
71- LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,77+ LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
72- uint32_t cureTileSize);78+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
79+ uint32_t round, T3 tileDataStart, uint32_t cureTileSize);
73 80 
74- const SortRegBaseTilingData *tilingData_;81+ const SortRegBaseTilingData* tilingData_;
75};82};
76 83 
77template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>84template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
78-__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::Init(GM_ADDR x, GM_ADDR value,85+__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::Init(
79- GM_ADDR sortIndex, GM_ADDR workspace, const SortRegBaseTilingData *__restrict tilingData, TPipe *pipe)86+ GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace, const SortRegBaseTilingData* __restrict tilingData,
87+ TPipe* pipe)
80{88{
81 this->blockIdx_ = GetBlockIdx();89 this->blockIdx_ = GetBlockIdx();
82 this->pipe_ = pipe;90 this->pipe_ = pipe;
@@ -87,16 +95,16 @@ __aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::Init(GM_ADD
87 this->factor_ = 2;95 this->factor_ = 2;
88 }96 }
89 97 
90- this->inputXGm_.SetGlobalBuffer((__gm__ T1 *)x);98+ this->inputXGm_.SetGlobalBuffer((__gm__ T1*)x);
91- this->outValueGm_.SetGlobalBuffer((__gm__ T1 *)value);99+ this->outValueGm_.SetGlobalBuffer((__gm__ T1*)value);
92- this->outIdxGm_.SetGlobalBuffer((__gm__ uint32_t *)sortIndex);100+ this->outIdxGm_.SetGlobalBuffer((__gm__ uint32_t*)sortIndex);
93 uint64_t wkOffset = this->clearCoreSize0_ * this->clearCore0_;101 uint64_t wkOffset = this->clearCoreSize0_ * this->clearCore0_;
94 uint64_t oneBlockNumB32 = this->oneBlock_ / sizeof(int32_t);102 uint64_t oneBlockNumB32 = this->oneBlock_ / sizeof(int32_t);
95 if constexpr (sizeof(T3) == sizeof(int64_t)) {103 if constexpr (sizeof(T3) == sizeof(int64_t)) {
96 wkOffset = wkOffset * 2;104 wkOffset = wkOffset * 2;
97 }105 }
98 wkOffset = Ops::Base::CeilAlign(wkOffset, oneBlockNumB32);106 wkOffset = Ops::Base::CeilAlign(wkOffset, oneBlockNumB32);
99- this->excusiveBinsGmWk_.SetGlobalBuffer((__gm__ uint32_t *)workspace, wkOffset);107+ this->excusiveBinsGmWk_.SetGlobalBuffer((__gm__ uint32_t*)workspace, wkOffset);
100 wkOffset = wkOffset * sizeof(uint32_t);108 wkOffset = wkOffset * sizeof(uint32_t);
101 109 
102 uint64_t histOffset = this->clearCout_ * this->clearSize_ * this->clearCore1_;110 uint64_t histOffset = this->clearCout_ * this->clearSize_ * this->clearCore1_;
@@ -104,7 +112,7 @@ __aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::Init(GM_ADD
104 histOffset = histOffset * 2;112 histOffset = histOffset * 2;
105 }113 }
106 histOffset = Ops::Base::CeilAlign(histOffset, oneBlockNumB32);114 histOffset = Ops::Base::CeilAlign(histOffset, oneBlockNumB32);
107- this->globalHistGmWk_.SetGlobalBuffer((__gm__ uint32_t *)(workspace + wkOffset), histOffset);115+ this->globalHistGmWk_.SetGlobalBuffer((__gm__ uint32_t*)(workspace + wkOffset), histOffset);
108 wkOffset = wkOffset + histOffset * sizeof(uint32_t);116 wkOffset = wkOffset + histOffset * sizeof(uint32_t);
109 117 
110 uint64_t dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;118 uint64_t dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
@@ -112,23 +120,23 @@ __aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::Init(GM_ADD
112 dbOffset = dbOffset * 2;120 dbOffset = dbOffset * 2;
113 }121 }
114 dbOffset = Ops::Base::CeilAlign(dbOffset, oneBlockNumB32);122 dbOffset = Ops::Base::CeilAlign(dbOffset, oneBlockNumB32);
115- this->outIdxDbWK_.SetGlobalBuffer((__gm__ uint32_t *)(workspace + wkOffset), dbOffset);123+ this->outIdxDbWK_.SetGlobalBuffer((__gm__ uint32_t*)(workspace + wkOffset), dbOffset);
116 wkOffset = wkOffset + dbOffset * sizeof(uint32_t);124 wkOffset = wkOffset + dbOffset * sizeof(uint32_t);
117 125 
118 uint64_t histTileOffset = this->lastDimTileNum_ * RADIX_SORT_NUM * this->unsortedDimParallel_;126 uint64_t histTileOffset = this->lastDimTileNum_ * RADIX_SORT_NUM * this->unsortedDimParallel_;
119- this->histTileGmWk_.SetGlobalBuffer((__gm__ uint16_t *)(workspace + wkOffset), histTileOffset);127+ this->histTileGmWk_.SetGlobalBuffer((__gm__ uint16_t*)(workspace + wkOffset), histTileOffset);
120 wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);128 wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);
121- this->histCumsumTileGmWk_.SetGlobalBuffer((__gm__ uint16_t *)(workspace + wkOffset), histTileOffset);129+ this->histCumsumTileGmWk_.SetGlobalBuffer((__gm__ uint16_t*)(workspace + wkOffset), histTileOffset);
122 wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);130 wkOffset = wkOffset + histTileOffset * sizeof(uint16_t);
123 131 
124 uint64_t xB8Offset = static_cast<uint64_t>(this->lastDimTileNum_) * this->numTileData_ * this->unsortedDimParallel_;132 uint64_t xB8Offset = static_cast<uint64_t>(this->lastDimTileNum_) * this->numTileData_ * this->unsortedDimParallel_;
125 xB8Offset = Ops::Base::CeilAlign(xB8Offset, this->oneBlock_);133 xB8Offset = Ops::Base::CeilAlign(xB8Offset, this->oneBlock_);
126- this->xB8GmWk_.SetGlobalBuffer((__gm__ uint8_t *)(workspace + wkOffset), xB8Offset);134+ this->xB8GmWk_.SetGlobalBuffer((__gm__ uint8_t*)(workspace + wkOffset), xB8Offset);
127 wkOffset = wkOffset + xB8Offset * sizeof(uint8_t);135 wkOffset = wkOffset + xB8Offset * sizeof(uint8_t);
128 136 
129 dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;137 dbOffset = this->totalDataNum_ * this->unsortedDimParallel_;
130 dbOffset = Ops::Base::CeilAlign(dbOffset * sizeof(T1), this->oneBlock_) / sizeof(T1);138 dbOffset = Ops::Base::CeilAlign(dbOffset * sizeof(T1), this->oneBlock_) / sizeof(T1);
131- this->outValueDbWK_.SetGlobalBuffer((__gm__ T1 *)(workspace + wkOffset), dbOffset);139+ this->outValueDbWK_.SetGlobalBuffer((__gm__ T1*)(workspace + wkOffset), dbOffset);
132 140 
133 this->pipe_->InitBuffer(this->inQueueX_, 1, this->numTileData_ * sizeof(T1));141 this->pipe_->InitBuffer(this->inQueueX_, 1, this->numTileData_ * sizeof(T1));
134 this->pipe_->InitBuffer(this->inQueueIndex_, 1, this->numTileData_ * sizeof(T3));142 this->pipe_->InitBuffer(this->inQueueIndex_, 1, this->numTileData_ * sizeof(T3));
@@ -168,39 +176,37 @@ __aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ParserTilin
168 176 
169template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>177template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
170__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutInt32(178__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutInt32(
171- LocalTensor<T1> xInputValueLocal,179+ LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,
172- LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,180+ LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
173- LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum,181+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,
174- LocalTensor<T3> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,182+ uint32_t cureTileSize)
175- uint32_t round, T3 tileDataStart, uint32_t cureTileSize)
176{183{
177 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;184 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;
178 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;185 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;
179 uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;186 uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;
180 if (round == 0) {187 if (round == 0) {
181- asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 0>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,188+ asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 0>>(
182- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),189+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
183- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),190+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
184- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ T3 *)(xInputIndexLocal.GetPhyAddr()),191+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
185- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),192+ (__ubuf__ T3*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
186- (__ubuf__ T3 *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),193+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ T3*)(blockHistFlag.GetPhyAddr()),
187- (__gm__ uint32_t *)(this->idxDbGm_.Alternate().GetPhyAddr()),194+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ uint32_t*)(this->idxDbGm_.Alternate().GetPhyAddr()),
188- (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));195+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
189 } else {196 } else {
190- asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 1>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,197+ asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 1>>(
191- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),198+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
192- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),199+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
193- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ T3 *)(xInputIndexLocal.GetPhyAddr()),200+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
194- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),201+ (__ubuf__ T3*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
195- (__ubuf__ T3 *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),202+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ T3*)(blockHistFlag.GetPhyAddr()),
196- (__gm__ uint32_t *)(this->idxDbGm_.Alternate().GetPhyAddr()),203+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ uint32_t*)(this->idxDbGm_.Alternate().GetPhyAddr()),
197- (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));204+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
198 }205 }
199}206}
200template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>207template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
201__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutInt32ToInt64(208__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutInt32ToInt64(
202- LocalTensor<T1> xInputValueLocal,209+ LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,
203- LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,
204 LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum,210 LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum,
205 LocalTensor<uint32_t> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,211 LocalTensor<uint32_t> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,
206 uint32_t round, T3 tileDataStart, uint32_t cureTileSize)212 uint32_t round, T3 tileDataStart, uint32_t cureTileSize)
@@ -209,101 +215,101 @@ __aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutI
209 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;215 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;
210 uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;216 uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;
211 if (round == 0) {217 if (round == 0) {
212- asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 0>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,218+ asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 0>>(
213- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),219+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
214- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),220+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
215- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ uint32_t *)(xInputIndexLocal.GetPhyAddr()),221+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
216- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),222+ (__ubuf__ uint32_t*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
217- (__ubuf__ uint32_t *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),223+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ uint32_t*)(blockHistFlag.GetPhyAddr()),
218- (__gm__ uint32_t *)(this->idxDbGm_.Alternate().GetPhyAddr()),224+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ uint32_t*)(this->idxDbGm_.Alternate().GetPhyAddr()),
219- (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));225+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
220 } else if (round < static_cast<uint32_t>(sizeof(T1) - 1)) {226 } else if (round < static_cast<uint32_t>(sizeof(T1) - 1)) {
221- asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 1>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,227+ asc_vf_call<CopyOutGm<T1, uint32_t, T3, uint32_t, 1>>(
222- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),228+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
223- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),229+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
224- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ uint32_t *)(xInputIndexLocal.GetPhyAddr()),230+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
225- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),231+ (__ubuf__ uint32_t*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
226- (__ubuf__ uint32_t *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),232+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ uint32_t*)(blockHistFlag.GetPhyAddr()),
227- (__gm__ uint32_t *)(this->idxDbGm_.Alternate().GetPhyAddr()),233+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ uint32_t*)(this->idxDbGm_.Alternate().GetPhyAddr()),
228- (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));234+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
229 } else {235 } else {
230 GlobalTensor<T2> outIdxT2 = (this->idxDbGm_.Alternate()).template ReinterpretCast<T2>();236 GlobalTensor<T2> outIdxT2 = (this->idxDbGm_.Alternate()).template ReinterpretCast<T2>();
231- asc_vf_call<CopyOutGm<T1, T2, T3, T2, 1>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,237+ asc_vf_call<CopyOutGm<T1, T2, T3, T2, 1>>(
232- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),238+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
233- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),239+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
234- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ uint32_t *)(xInputIndexLocal.GetPhyAddr()),240+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
235- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),241+ (__ubuf__ uint32_t*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
236- (__ubuf__ uint32_t *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),242+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ uint32_t*)(blockHistFlag.GetPhyAddr()),
237- (__gm__ T2 *)(outIdxT2.GetPhyAddr()), (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));243+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ T2*)(outIdxT2.GetPhyAddr()),
244+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
238 }245 }
239}246}
240template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>247template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
241__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutInt64(248__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterOutInt64(
242- LocalTensor<T1> xInputValueLocal,249+ LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,
243- LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,250+ LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
244- LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum,251+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,
245- LocalTensor<T3> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,252+ uint32_t cureTileSize)
246- uint32_t round, T3 tileDataStart, uint32_t cureTileSize)
247{253{
248 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;254 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;
249 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;255 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;
250 uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;256 uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;
251 257 
252 if (round == 0) {258 if (round == 0) {
253- asc_vf_call<CopyOutGm<T1, T2, T3, T2, 0>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,259+ asc_vf_call<CopyOutGm<T1, T2, T3, T2, 0>>(
254- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),260+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
255- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),261+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
256- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ T3 *)(xInputIndexLocal.GetPhyAddr()),262+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
257- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),263+ (__ubuf__ T3*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
258- (__ubuf__ T3 *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),264+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ T3*)(blockHistFlag.GetPhyAddr()),
259- (__gm__ T2 *)(this->idxDbGm_.Alternate().GetPhyAddr()),265+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ T2*)(this->idxDbGm_.Alternate().GetPhyAddr()),
260- (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));266+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
261 } else {267 } else {
262- asc_vf_call<CopyOutGm<T1, T2, T3, T2, 1>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,268+ asc_vf_call<CopyOutGm<T1, T2, T3, T2, 1>>(
263- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),269+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
264- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),270+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
265- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ T3 *)(xInputIndexLocal.GetPhyAddr()),271+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
266- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),272+ (__ubuf__ T3*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
267- (__ubuf__ T3 *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),273+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ T3*)(blockHistFlag.GetPhyAddr()),
268- (__gm__ T2 *)(this->idxDbGm_.Alternate().GetPhyAddr()),274+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ T2*)(this->idxDbGm_.Alternate().GetPhyAddr()),
269- (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));275+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
270 }276 }
271}277}
272template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>278template <typename T1, typename T2, typename UT, typename T3, uint64_t isDescend>
273__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterKeysGlobal(279__aicore__ inline void SortRadixMoreCore<T1, T2, UT, T3, isDescend>::ScatterKeysGlobal(
274- LocalTensor<T1> xInputValueLocal,280+ LocalTensor<T1> xInputValueLocal, LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,
275- LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<uint32_t> xInputIndexLocal,281+ LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum, LocalTensor<T3> blockDataInGlobalPos,
276- LocalTensor<uint8_t> sortedValueLocal, LocalTensor<uint16_t> blockExcusiveSum,282+ LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist, uint32_t round, T3 tileDataStart,
277- LocalTensor<T3> blockDataInGlobalPos, LocalTensor<uint32_t> blockHistFlag, LocalTensor<uint16_t> blockHist,283+ uint32_t cureTileSize, uint32_t sortLoopRound)
278- uint32_t round, T3 tileDataStart, uint32_t cureTileSize, uint32_t sortLoopRound)
279{284{
280 (void)sortLoopRound;285 (void)sortLoopRound;
281 if constexpr (sizeof(T1) == sizeof(int8_t)) {286 if constexpr (sizeof(T1) == sizeof(int8_t)) {
282 // int8时只循环一次,所以scatter时肯定要按照输出数据类型287 // int8时只循环一次,所以scatter时肯定要按照输出数据类型
283 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;288 uint32_t unSortId = this->blockIdx_ / this->lastDimRealCore_;
284 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;289 uint64_t outputXUnsortedAxisOffset = static_cast<uint64_t>(unSortId) * this->totalDataNum_;
285- uint64_t unSortIdOffset =290+ uint64_t unSortIdOffset = static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) +
286- static_cast<uint64_t>(unSortId) * RADIX_SORT_NUM * sizeof(T1) + round * RADIX_SORT_NUM;291+ round * RADIX_SORT_NUM;
287 GlobalTensor<T2> outIdxT2 = (this->idxDbGm_.Alternate()).template ReinterpretCast<T2>();292 GlobalTensor<T2> outIdxT2 = (this->idxDbGm_.Alternate()).template ReinterpretCast<T2>();
288- asc_vf_call<CopyOutGm<T1, T2, T3, T2, 0>>(dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize,293+ asc_vf_call<CopyOutGm<T1, T2, T3, T2, 0>>(
289- outputXUnsortedAxisOffset, unSortIdOffset, (__ubuf__ uint16_t *)(blockExcusiveSum.GetPhyAddr()),294+ dim3(THREAD_DIM_NUM), tileDataStart, cureTileSize, outputXUnsortedAxisOffset, unSortIdOffset,
290- (__gm__ T3 *)(this->excusiveBinsGmWk_.GetPhyAddr()), (__ubuf__ T3 *)(blockDataInGlobalPos.GetPhyAddr()),295+ (__ubuf__ uint16_t*)(blockExcusiveSum.GetPhyAddr()), (__gm__ T3*)(this->excusiveBinsGmWk_.GetPhyAddr()),
291- (__ubuf__ uint32_t *)(sortedIndexLocal.GetPhyAddr()), (__ubuf__ T3 *)(xInputIndexLocal.GetPhyAddr()),296+ (__ubuf__ T3*)(blockDataInGlobalPos.GetPhyAddr()), (__ubuf__ uint32_t*)(sortedIndexLocal.GetPhyAddr()),
292- (__ubuf__ uint8_t *)(sortedValueLocal.GetPhyAddr()), (__ubuf__ T1 *)(xInputValueLocal.GetPhyAddr()),297+ (__ubuf__ T3*)(xInputIndexLocal.GetPhyAddr()), (__ubuf__ uint8_t*)(sortedValueLocal.GetPhyAddr()),
293- (__ubuf__ T3 *)(blockHistFlag.GetPhyAddr()), (__ubuf__ uint16_t *)(blockHist.GetPhyAddr()),298+ (__ubuf__ T1*)(xInputValueLocal.GetPhyAddr()), (__ubuf__ T3*)(blockHistFlag.GetPhyAddr()),
294- (__gm__ T2 *)(outIdxT2.GetPhyAddr()), (__gm__ T1 *)(this->inputXDbGm_.Alternate().GetPhyAddr()));299+ (__ubuf__ uint16_t*)(blockHist.GetPhyAddr()), (__gm__ T2*)(outIdxT2.GetPhyAddr()),
300+ (__gm__ T1*)(this->inputXDbGm_.Alternate().GetPhyAddr()));
295 } else if constexpr (sizeof(T2) == sizeof(int32_t)) {301 } else if constexpr (sizeof(T2) == sizeof(int32_t)) {
296 // 输出idx本省就是int32，无需cast302 // 输出idx本省就是int32，无需cast
297- ScatterOutInt32(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal,303+ ScatterOutInt32(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,
298- blockExcusiveSum, blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize);304+ blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize);
299- } else if constexpr (IsSameType<T3, uint32_t>::value){305+ } else if constexpr (IsSameType<T3, uint32_t>::value) {
300 // 输出idx是int64，需要在最后一次scatter时cast为int64306 // 输出idx是int64，需要在最后一次scatter时cast为int64
301- ScatterOutInt32ToInt64(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal,307+ ScatterOutInt32ToInt64(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,
302- blockExcusiveSum, blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize);308+ blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize);
303 } else {309 } else {
304 // 计算过程中idx使用int64310 // 计算过程中idx使用int64
305- ScatterOutInt64(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal,311+ ScatterOutInt64(xInputValueLocal, sortedIndexLocal, xInputIndexLocal, sortedValueLocal, blockExcusiveSum,
306- blockExcusiveSum, blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize);312+ blockDataInGlobalPos, blockHistFlag, blockHist, round, tileDataStart, cureTileSize);
307 }313 }
308}314}
309} // namespace Sort315} // namespace Sort
Mmath/sort/op_kernel/arch35/sort_radix_sort_one_core.h+15-12
@@ -28,7 +28,7 @@ class SortRadixOneCore {
28public:28public:
29 __aicore__ inline SortRadixOneCore(){};29 __aicore__ inline SortRadixOneCore(){};
30 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,30 __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace,
31- const SortRegBaseTilingData *__restrict tilingData, TPipe *pipe);31+ const SortRegBaseTilingData* __restrict tilingData, TPipe* pipe);
32 __aicore__ inline void Process();32 __aicore__ inline void Process();
33 33 
34private:34private:
@@ -39,13 +39,13 @@ private:
39 GlobalTensor<T1> outValueGm_;39 GlobalTensor<T1> outValueGm_;
40 GlobalTensor<uint32_t> outIdxGm_;40 GlobalTensor<uint32_t> outIdxGm_;
41 41 
42- TPipe *pipe_;42+ TPipe* pipe_;
43- const SortRegBaseTilingData *tilingData_;43+ const SortRegBaseTilingData* tilingData_;
44 TQue<QuePosition::VECIN, 1> inQueueX_;44 TQue<QuePosition::VECIN, 1> inQueueX_;
45 TBuf<TPosition::VECCALC> tmpUb_;45 TBuf<TPosition::VECCALC> tmpUb_;
46 TQue<QuePosition::VECOUT, 1> outIdxQueue_;46 TQue<QuePosition::VECOUT, 1> outIdxQueue_;
47 TQue<QuePosition::VECOUT, 1> outValueQueue_;47 TQue<QuePosition::VECOUT, 1> outValueQueue_;
48- static constexpr SortConfig sortConfigMuti{ SortType::RADIX_SORT, isDescend };48+ static constexpr SortConfig sortConfigMuti{SortType::RADIX_SORT, isDescend};
49 49 
50 int64_t totalDataNum_ = 0;50 int64_t totalDataNum_ = 0;
51 uint32_t numTileData_ = 0;51 uint32_t numTileData_ = 0;
@@ -63,15 +63,17 @@ private:
63 63 
64template <typename T1, typename T2, bool isDescend>64template <typename T1, typename T2, bool isDescend>
65__aicore__ inline void SortRadixOneCore<T1, T2, isDescend>::Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex,65__aicore__ inline void SortRadixOneCore<T1, T2, isDescend>::Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex,
66- GM_ADDR workspace, const SortRegBaseTilingData *__restrict tilingData, TPipe *pipe)66+ GM_ADDR workspace,
67+ const SortRegBaseTilingData* __restrict tilingData,
68+ TPipe* pipe)
67{69{
68 blockIdx_ = GetBlockIdx();70 blockIdx_ = GetBlockIdx();
69 pipe_ = pipe;71 pipe_ = pipe;
70 tilingData_ = tilingData;72 tilingData_ = tilingData;
71 ParserTilingData();73 ParserTilingData();
72- inputXGm_.SetGlobalBuffer((__gm__ T1 *)x);74+ inputXGm_.SetGlobalBuffer((__gm__ T1*)x);
73- outValueGm_.SetGlobalBuffer((__gm__ T1 *)value);75+ outValueGm_.SetGlobalBuffer((__gm__ T1*)value);
74- outIdxGm_.SetGlobalBuffer((__gm__ uint32_t *)sortIndex);76+ outIdxGm_.SetGlobalBuffer((__gm__ uint32_t*)sortIndex);
75 realCoreNum_ = GetBlockNum();77 realCoreNum_ = GetBlockNum();
76 78 
77 pipe_->InitBuffer(inQueueX_, bufferNum_, xUbSize_);79 pipe_->InitBuffer(inQueueX_, bufferNum_, xUbSize_);
@@ -91,13 +93,14 @@ __aicore__ inline void SortRadixOneCore<T1, T2, isDescend>::ParserTilingData()
91 tmpUbSize_ = tilingData_->tmpUbSize; // 高级api需要用的ub大小93 tmpUbSize_ = tilingData_->tmpUbSize; // 高级api需要用的ub大小
92 xUbSize_ = tilingData_->keyParams0; // xInQue需要的ub大小94 xUbSize_ = tilingData_->keyParams0; // xInQue需要的ub大小
93 yUbSize_ = tilingData_->keyParams1; // y2OutQue需要的ub大小95 yUbSize_ = tilingData_->keyParams1; // y2OutQue需要的ub大小
94- halfIndex_ = tilingData_->keyParams2; // 输出idx如果是int64，cast时需要从ub的一半开始96+ halfIndex_ = tilingData_->keyParams2; // 输出idx如果是int64，cast时需要从ub的一半开始
95- bufferNum_ = tilingData_->keyParams3 == 2 ? 2 : 1; // one-core队列buffer数97+ bufferNum_ = tilingData_->keyParams3 == 2 ? 2 : 1; // one-core队列buffer数
96}98}
97 99 
98template <typename T1, typename T2, bool isDescend>100template <typename T1, typename T2, bool isDescend>
99__aicore__ inline void SortRadixOneCore<T1, T2, isDescend>::ProcessRadixSortOneCore(GlobalTensor<T1> inputXGm,101__aicore__ inline void SortRadixOneCore<T1, T2, isDescend>::ProcessRadixSortOneCore(GlobalTensor<T1> inputXGm,
100- int64_t gmOffset, uint32_t sortLoopRound)102+ int64_t gmOffset,
103+ uint32_t sortLoopRound)
101{104{
102 uint32_t unsortDimIndex = blockIdx_ + sortLoopRound * unsortedDimParallel_;105 uint32_t unsortDimIndex = blockIdx_ + sortLoopRound * unsortedDimParallel_;
103 if (unsortDimIndex >= unsortedDimNum_) {106 if (unsortDimIndex >= unsortedDimNum_) {
@@ -105,7 +108,7 @@ __aicore__ inline void SortRadixOneCore<T1, T2, isDescend>::ProcessRadixSortOneC
105 }108 }
106 LocalTensor<T1> xLocal = inQueueX_.AllocTensor<T1>();109 LocalTensor<T1> xLocal = inQueueX_.AllocTensor<T1>();
107 uint32_t tileOffset = blockIdx_ * numTileData_;110 uint32_t tileOffset = blockIdx_ * numTileData_;
108- DataCopyPadExtParams<T1> padParams{ false, 0, 0, 0 };111+ DataCopyPadExtParams<T1> padParams{false, 0, 0, 0};
109 DataCopyExtParams dataCopyParam;112 DataCopyExtParams dataCopyParam;
110 dataCopyParam.blockCount = 1;113 dataCopyParam.blockCount = 1;
111 dataCopyParam.blockLen = numTileData_ * sizeof(T1);114 dataCopyParam.blockLen = numTileData_ * sizeof(T1);
Mmath/sort/op_kernel/arch35/sort_small_axis_insertion.h+37-44
@@ -24,10 +24,11 @@ using namespace AscendC;
24 24 
25template <typename T, typename OUT_IDX_T>25template <typename T, typename OUT_IDX_T>
26__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::INSERTION_THREAD_NUM) __aicore__26__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::INSERTION_THREAD_NUM) __aicore__
27-void SimtStoreNonLastInsertionBatch(uint32_t totalElems, uint32_t validSegs, uint32_t segmentLen,27+ void SimtStoreNonLastInsertionBatch(uint32_t totalElems, uint32_t validSegs, uint32_t segmentLen,
28- uint32_t valueRowElems, uint32_t indexRowElems, uint64_t outerBaseOffset, uint64_t innerStart,28+ uint32_t valueRowElems, uint32_t indexRowElems, uint64_t outerBaseOffset,
29- uint64_t innerSize, __ubuf__ T *inputValue, __ubuf__ OUT_IDX_T *inputIdx,29+ uint64_t innerStart, uint64_t innerSize, __ubuf__ T* inputValue,
30- __gm__ volatile T *outputValue, __gm__ volatile OUT_IDX_T *outputIdx)30+ __ubuf__ OUT_IDX_T* inputIdx, __gm__ volatile T* outputValue,
31+ __gm__ volatile OUT_IDX_T* outputIdx)
31{32{
32 // Store sorted sort-major segments back to the original non-last-axis GM layout.33 // Store sorted sort-major segments back to the original non-last-axis GM layout.
33 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems;34 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems;
@@ -41,19 +42,17 @@ void SimtStoreNonLastInsertionBatch(uint32_t totalElems, uint32_t validSegs, uin
41}42}
42 43 
43template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>44template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
44-class SortSmallAxisInsertion : public SmallAxisCommon::SmallAxisInsertionBase<45+class SortSmallAxisInsertion
45- SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>, T, CONVERT_TYPE, OUT_IDX_T, IsDescend> {46+ : public SmallAxisCommon::SmallAxisInsertionBase<SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>, T,
46- using Base = SmallAxisCommon::SmallAxisInsertionBase<47+ CONVERT_TYPE, OUT_IDX_T, IsDescend> {
47- SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>, T, CONVERT_TYPE, OUT_IDX_T, IsDescend>;48+ using Base = SmallAxisCommon::SmallAxisInsertionBase<SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>,
49+ T, CONVERT_TYPE, OUT_IDX_T, IsDescend>;
48 50 
49public:51public:
50 __aicore__ inline SortSmallAxisInsertion() {}52 __aicore__ inline SortSmallAxisInsertion() {}
51 __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,53 __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,
52- const SortRegBaseTilingData *tilingData, TPipe *pipe);54+ const SortRegBaseTilingData* tilingData, TPipe* pipe);
53- __aicore__ inline void Process()55+ __aicore__ inline void Process() { Base::Process(); }
54- {
55- Base::Process();
56- }
57 56 
58 friend Base;57 friend Base;
59 58 
@@ -78,7 +77,7 @@ private:
78 __aicore__ inline void StoreBatch(int64_t segStart, uint32_t validSegs);77 __aicore__ inline void StoreBatch(int64_t segStart, uint32_t validSegs);
79 __aicore__ inline void StoreNonLastBatch(uint64_t outerId, uint64_t innerStart, uint32_t validSegs);78 __aicore__ inline void StoreNonLastBatch(uint64_t outerId, uint64_t innerStart, uint32_t validSegs);
80 79 
81- const SortRegBaseTilingData *tilingData_ = nullptr;80+ const SortRegBaseTilingData* tilingData_ = nullptr;
82 81 
83 GlobalTensor<T> inputXGm_;82 GlobalTensor<T> inputXGm_;
84 GlobalTensor<T> outValueGm_;83 GlobalTensor<T> outValueGm_;
@@ -93,7 +92,7 @@ private:
93 92 
94template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>93template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
95__aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::Init(94__aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::Init(
96- GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace, const SortRegBaseTilingData *tilingData, TPipe *pipe)95+ GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace, const SortRegBaseTilingData* tilingData, TPipe* pipe)
97{96{
98 (void)workspace;97 (void)workspace;
99 if (tilingData == nullptr || pipe == nullptr) {98 if (tilingData == nullptr || pipe == nullptr) {
@@ -110,12 +109,11 @@ __aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
110 outerSize_ = tilingData_->outerSize;109 outerSize_ = tilingData_->outerSize;
111 innerSize_ = tilingData_->innerSize;110 innerSize_ = tilingData_->innerSize;
112 innerLoopNum_ = tilingData_->innerLoopNum;111 innerLoopNum_ = tilingData_->innerLoopNum;
113- inputXGm_.SetGlobalBuffer((__gm__ T *)x);112+ inputXGm_.SetGlobalBuffer((__gm__ T*)x);
114- outValueGm_.SetGlobalBuffer((__gm__ T *)y);113+ outValueGm_.SetGlobalBuffer((__gm__ T*)y);
115- outIdxGm_.SetGlobalBuffer((__gm__ OUT_IDX_T *)idx);114+ outIdxGm_.SetGlobalBuffer((__gm__ OUT_IDX_T*)idx);
116 115 
117- if (segmentLen_ == 0 || segmentsPerBatch_ == 0 ||116+ if (segmentLen_ == 0 || segmentsPerBatch_ == 0 || segmentsPerBatch_ > SmallAxisCommon::MAX_DATACOPY_BLOCK_COUNT) {
118- segmentsPerBatch_ > SmallAxisCommon::MAX_DATACOPY_BLOCK_COUNT) {
119 return;117 return;
120 }118 }
121 119 
@@ -126,7 +124,7 @@ template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
126__aicore__ inline bool SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::IsProcessInvalid() const124__aicore__ inline bool SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::IsProcessInvalid() const
127{125{
128 return blockIdx_ >= blockDim_ || segmentLen_ == 0 || segmentsPerBatch_ == 0 ||126 return blockIdx_ >= blockDim_ || segmentLen_ == 0 || segmentsPerBatch_ == 0 ||
129- (isNonLastAxis_ && (innerLoopNum_ == 0 || outerSize_ <= 0 || innerSize_ <= 0));127+ (isNonLastAxis_ && (innerLoopNum_ == 0 || outerSize_ <= 0 || innerSize_ <= 0));
130}128}
131 129 
132template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>130template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
@@ -143,7 +141,7 @@ __aicore__ inline uint32_t SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, Is
143 return 0;141 return 0;
144 }142 }
145 return innerRemain >= static_cast<int64_t>(segmentsPerBatch_) ? segmentsPerBatch_ :143 return innerRemain >= static_cast<int64_t>(segmentsPerBatch_) ? segmentsPerBatch_ :
146- static_cast<uint32_t>(innerRemain);144+ static_cast<uint32_t>(innerRemain);
147 }145 }
148 int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(segmentsPerBatch_);146 int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(segmentsPerBatch_);
149 int64_t segRemain = totalSegs_ - segStart;147 int64_t segRemain = totalSegs_ - segStart;
@@ -157,8 +155,8 @@ __aicore__ inline uint32_t SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, Is
157}155}
158 156 
159template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>157template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
160-__aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::ProcessBatch(158+__aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::ProcessBatch(uint32_t batchId,
161- uint32_t batchId, uint32_t validSegs)159+ uint32_t validSegs)
162{160{
163 int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(segmentsPerBatch_);161 int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(segmentsPerBatch_);
164 uint64_t outerId = 0;162 uint64_t outerId = 0;
@@ -191,18 +189,17 @@ __aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
191 // Convert widened sorted values back to the output storage dtype before SIMT GM stores.189 // Convert widened sorted values back to the output storage dtype before SIMT GM stores.
192 LocalTensor<T> storeBuf = castBuf_.template Get<T>();190 LocalTensor<T> storeBuf = castBuf_.template Get<T>();
193 Cast(storeBuf, values_, RoundMode::CAST_RINT, validSegs * valueRowStride_);191 Cast(storeBuf, values_, RoundMode::CAST_RINT, validSegs * valueRowStride_);
194- asc_vf_call<SimtStoreNonLastInsertionBatch<T, OUT_IDX_T>>(dim3(SmallAxisCommon::INSERTION_THREAD_NUM),192+ asc_vf_call<SimtStoreNonLastInsertionBatch<T, OUT_IDX_T>>(
195- totalElems, validSegs, segmentLen_, castRowStride_, indexRowStride_, outerBaseOffset, innerStart,193+ dim3(SmallAxisCommon::INSERTION_THREAD_NUM), totalElems, validSegs, segmentLen_, castRowStride_,
196- static_cast<uint64_t>(innerSize_), (__ubuf__ T *)storeBuf.GetPhyAddr(),194+ indexRowStride_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_),
197- (__ubuf__ OUT_IDX_T *)indices_.GetPhyAddr(), (__gm__ volatile T *)outValueGm_.GetPhyAddr(),195+ (__ubuf__ T*)storeBuf.GetPhyAddr(), (__ubuf__ OUT_IDX_T*)indices_.GetPhyAddr(),
198- (__gm__ volatile OUT_IDX_T *)outIdxGm_.GetPhyAddr());196+ (__gm__ volatile T*)outValueGm_.GetPhyAddr(), (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());
199 } else {197 } else {
200- asc_vf_call<SimtStoreNonLastInsertionBatch<T, OUT_IDX_T>>(dim3(SmallAxisCommon::INSERTION_THREAD_NUM),198+ asc_vf_call<SimtStoreNonLastInsertionBatch<T, OUT_IDX_T>>(
201- totalElems,199+ dim3(SmallAxisCommon::INSERTION_THREAD_NUM), totalElems, validSegs, segmentLen_, valueRowStride_,
202- validSegs, segmentLen_, valueRowStride_, indexRowStride_, outerBaseOffset, innerStart,200+ indexRowStride_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_),
203- static_cast<uint64_t>(innerSize_), (__ubuf__ T *)values_.GetPhyAddr(),201+ (__ubuf__ T*)values_.GetPhyAddr(), (__ubuf__ OUT_IDX_T*)indices_.GetPhyAddr(),
204- (__ubuf__ OUT_IDX_T *)indices_.GetPhyAddr(), (__gm__ volatile T *)outValueGm_.GetPhyAddr(),202+ (__gm__ volatile T*)outValueGm_.GetPhyAddr(), (__gm__ volatile OUT_IDX_T*)outIdxGm_.GetPhyAddr());
205- (__gm__ volatile OUT_IDX_T *)outIdxGm_.GetPhyAddr());
206 }203 }
207 event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));204 event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S));
208 SetFlag<HardEvent::V_S>(eventId);205 SetFlag<HardEvent::V_S>(eventId);
@@ -211,7 +208,7 @@ __aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
211 208 
212template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>209template <typename T, typename CONVERT_TYPE, typename OUT_IDX_T, bool IsDescend>
213__aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::StoreBatch(int64_t segStart,210__aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDescend>::StoreBatch(int64_t segStart,
214- uint32_t validSegs)211+ uint32_t validSegs)
215{212{
216 if (validSegs == 0) {213 if (validSegs == 0) {
217 return;214 return;
@@ -225,8 +222,8 @@ __aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
225 // For diff-type, valueRowStride_ is intentionally equal to castRowStride_, so the222 // For diff-type, valueRowStride_ is intentionally equal to castRowStride_, so the
226 // GM-transfer T buffer and the CONVERT_TYPE sort buffer share the same row pitch.223 // GM-transfer T buffer and the CONVERT_TYPE sort buffer share the same row pitch.
227 uint32_t valueSrcStrideBlocks = ((valueRowStride_ - segmentLen_) * sizeof(T)) / UB_BLOCK_SIZE;224 uint32_t valueSrcStrideBlocks = ((valueRowStride_ - segmentLen_) * sizeof(T)) / UB_BLOCK_SIZE;
228- DataCopyExtParams valueCopyParam{225+ DataCopyExtParams valueCopyParam{static_cast<uint16_t>(validSegs), static_cast<uint32_t>(segmentLen_ * sizeof(T)),
229- static_cast<uint16_t>(validSegs), static_cast<uint32_t>(segmentLen_ * sizeof(T)), valueSrcStrideBlocks, 0, 0 };226+ valueSrcStrideBlocks, 0, 0};
230 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {227 if constexpr (!IsSameType<T, CONVERT_TYPE>::value) {
231 LocalTensor<T> storeBuf = castBuf_.template Get<T>();228 LocalTensor<T> storeBuf = castBuf_.template Get<T>();
232 Cast(storeBuf, values_, RoundMode::CAST_RINT, validSegs * valueRowStride_);229 Cast(storeBuf, values_, RoundMode::CAST_RINT, validSegs * valueRowStride_);
@@ -240,12 +237,8 @@ __aicore__ inline void SortSmallAxisInsertion<T, CONVERT_TYPE, OUT_IDX_T, IsDesc
240 }237 }
241 238 
242 uint32_t idxStrideBlocks = ((indexRowStride_ - segmentLen_) * sizeof(OUT_IDX_T)) / UB_BLOCK_SIZE;239 uint32_t idxStrideBlocks = ((indexRowStride_ - segmentLen_) * sizeof(OUT_IDX_T)) / UB_BLOCK_SIZE;
243- DataCopyExtParams idxCopyParam{240+ DataCopyExtParams idxCopyParam{static_cast<uint16_t>(validSegs),
244- static_cast<uint16_t>(validSegs),241+ static_cast<uint32_t>(segmentLen_ * sizeof(OUT_IDX_T)), idxStrideBlocks, 0, 0};
245- static_cast<uint32_t>(segmentLen_ * sizeof(OUT_IDX_T)),
246- idxStrideBlocks,
247- 0,
248- 0 };
249 DataCopyPad(outIdxGm_[gmOffset], indices_, idxCopyParam);242 DataCopyPad(outIdxGm_[gmOffset], indices_, idxCopyParam);
250 243 
251 // The next batch is issued by scalar/control flow after these MTE3 writebacks complete.244 // The next batch is issued by scalar/control flow after these MTE3 writebacks complete.
Mmath/sort/op_kernel/arch35/sort_small_axis_two_stage.h+39-39
@@ -23,10 +23,10 @@ using namespace AscendC;
23 23 
24template <typename T, typename OutIdxT>24template <typename T, typename OutIdxT>
25__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__25__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__
26-void StoreNonLastBatchSimt(uint32_t totalElems, uint32_t segmentLen, uint32_t validSegs,26+ void StoreNonLastBatchSimt(uint32_t totalElems, uint32_t segmentLen, uint32_t validSegs, uint64_t outerBaseOffset,
27- uint64_t outerBaseOffset, uint64_t innerStart, uint64_t innerSize,27+ uint64_t innerStart, uint64_t innerSize, __ubuf__ T* inputValue,
28- __ubuf__ T *inputValue, __ubuf__ OutIdxT *inputIdx,28+ __ubuf__ OutIdxT* inputIdx, __gm__ volatile T* outputValue,
29- __gm__ volatile T *outputValue, __gm__ volatile OutIdxT *outputIdx)29+ __gm__ volatile OutIdxT* outputIdx)
30{30{
31 // Scatter final sorted rows from [inner segment, axis] order back to original GM offsets.31 // Scatter final sorted rows from [inner segment, axis] order back to original GM offsets.
32 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems;32 for (uint32_t idx = static_cast<uint32_t>(threadIdx.x); idx < totalElems;
@@ -41,19 +41,16 @@ void StoreNonLastBatchSimt(uint32_t totalElems, uint32_t segmentLen, uint32_t va
41}41}
42 42 
43template <typename T, typename OutIdxT, bool IsDescend>43template <typename T, typename OutIdxT, bool IsDescend>
44-class SortSmallAxisTwoStage : public SmallAxisCommon::SmallAxisTwoStageBase<44+class SortSmallAxisTwoStage
45- SortSmallAxisTwoStage<T, OutIdxT, IsDescend>, T, OutIdxT, IsDescend> {45+ : public SmallAxisCommon::SmallAxisTwoStageBase<SortSmallAxisTwoStage<T, OutIdxT, IsDescend>, T, OutIdxT,
46- using Base = SmallAxisCommon::SmallAxisTwoStageBase<46+ IsDescend> {
47- SortSmallAxisTwoStage<T, OutIdxT, IsDescend>, T, OutIdxT, IsDescend>;47+ using Base = SmallAxisCommon::SmallAxisTwoStageBase<SortSmallAxisTwoStage<T, OutIdxT, IsDescend>, T, OutIdxT,
48+ IsDescend>;
48 49 
49public:50public:
50 __aicore__ inline SortSmallAxisTwoStage() {}51 __aicore__ inline SortSmallAxisTwoStage() {}
51 __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,52 __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx, GM_ADDR workspace,
52- const SortRegBaseTilingData *tilingData, TPipe *pipe);53+ const SortRegBaseTilingData* tilingData, TPipe* pipe);
53- __aicore__ inline void Process()
54- {
55- Base::Process();
56- }
57 54 
58 friend Base;55 friend Base;
59 56 
@@ -71,10 +68,10 @@ private:
71 __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchId) const;68 __aicore__ inline uint32_t ComputeValidSegs(uint32_t batchId) const;
72 __aicore__ inline void ProcessBatch(uint32_t batchId, uint32_t validSegs);69 __aicore__ inline void ProcessBatch(uint32_t batchId, uint32_t validSegs);
73 __aicore__ inline void StoreBatch(int64_t segStart, uint32_t totalElems);70 __aicore__ inline void StoreBatch(int64_t segStart, uint32_t totalElems);
74- __aicore__ inline void StoreNonLastBatch(uint64_t outerId, uint64_t innerStart,71+ __aicore__ inline void StoreNonLastBatch(uint64_t outerId, uint64_t innerStart, uint32_t validSegs,
75- uint32_t validSegs, uint32_t totalElems);72+ uint32_t totalElems);
76 73 
77- const SortRegBaseTilingData *tilingData_ = nullptr;74+ const SortRegBaseTilingData* tilingData_ = nullptr;
78 75 
79 GlobalTensor<T> inputXGm_;76 GlobalTensor<T> inputXGm_;
80 GlobalTensor<T> outValueGm_;77 GlobalTensor<T> outValueGm_;
@@ -88,8 +85,10 @@ private:
88};85};
89 86 
90template <typename T, typename OutIdxT, bool IsDescend>87template <typename T, typename OutIdxT, bool IsDescend>
91-__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::Init(GM_ADDR x, GM_ADDR y,88+__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR idx,
92- GM_ADDR idx, GM_ADDR workspace, const SortRegBaseTilingData *tilingData, TPipe *pipe)89+ GM_ADDR workspace,
90+ const SortRegBaseTilingData* tilingData,
91+ TPipe* pipe)
93{92{
94 (void)workspace;93 (void)workspace;
95 if (tilingData == nullptr || pipe == nullptr) {94 if (tilingData == nullptr || pipe == nullptr) {
@@ -108,24 +107,23 @@ __aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::Init(GM_ADD
108 innerSize_ = tilingData_->innerSize;107 innerSize_ = tilingData_->innerSize;
109 innerLoopNum_ = tilingData_->innerLoopNum;108 innerLoopNum_ = tilingData_->innerLoopNum;
110 109 
111- inputXGm_.SetGlobalBuffer((__gm__ T *)x);110+ inputXGm_.SetGlobalBuffer((__gm__ T*)x);
112- outValueGm_.SetGlobalBuffer((__gm__ T *)y);111+ outValueGm_.SetGlobalBuffer((__gm__ T*)y);
113- outIdxGm_.SetGlobalBuffer((__gm__ OutIdxT *)idx);112+ outIdxGm_.SetGlobalBuffer((__gm__ OutIdxT*)idx);
114 113 
115 if (batchSize_ == 0 || segmentLen_ == 0 || maxFlatElems_ == 0) {114 if (batchSize_ == 0 || segmentLen_ == 0 || maxFlatElems_ == 0) {
116 return;115 return;
117 }116 }
118 117 
119 constexpr uint32_t kAliasElemBytes = static_cast<uint32_t>(sizeof(OutIdxT));118 constexpr uint32_t kAliasElemBytes = static_cast<uint32_t>(sizeof(OutIdxT));
120- Base::InitSortBuffers(pipe, maxFlatElems_, tilingData_->tmpUbSize,119+ Base::InitSortBuffers(pipe, maxFlatElems_, tilingData_->tmpUbSize, tilingData_->keyParams2 != 0U, kAliasElemBytes);
121- tilingData_->keyParams2 != 0U, kAliasElemBytes);
122}120}
123 121 
124template <typename T, typename OutIdxT, bool IsDescend>122template <typename T, typename OutIdxT, bool IsDescend>
125__aicore__ inline bool SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::IsProcessInvalid() const123__aicore__ inline bool SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::IsProcessInvalid() const
126{124{
127 return blockIdx_ >= blockDim_ || batchSize_ == 0 || segmentLen_ == 0 ||125 return blockIdx_ >= blockDim_ || batchSize_ == 0 || segmentLen_ == 0 ||
128- (isNonLastAxis_ && (innerLoopNum_ == 0 || outerSize_ <= 0 || innerSize_ <= 0));126+ (isNonLastAxis_ && (innerLoopNum_ == 0 || outerSize_ <= 0 || innerSize_ <= 0));
129}127}
130 128 
131template <typename T, typename OutIdxT, bool IsDescend>129template <typename T, typename OutIdxT, bool IsDescend>
@@ -141,7 +139,8 @@ __aicore__ inline uint32_t SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::Compute
141 return 0;139 return 0;
142 }140 }
143 uint32_t validInnerSegs = remainingInTile >= static_cast<int64_t>(batchSize_) ?141 uint32_t validInnerSegs = remainingInTile >= static_cast<int64_t>(batchSize_) ?
144- batchSize_ : static_cast<uint32_t>(remainingInTile);142+ batchSize_ :
143+ static_cast<uint32_t>(remainingInTile);
145 return validInnerSegs;144 return validInnerSegs;
146 }145 }
147 int64_t batchStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);146 int64_t batchStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);
@@ -156,8 +155,7 @@ __aicore__ inline uint32_t SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::Compute
156}155}
157 156 
158template <typename T, typename OutIdxT, bool IsDescend>157template <typename T, typename OutIdxT, bool IsDescend>
159-__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::ProcessBatch(158+__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::ProcessBatch(uint32_t batchId, uint32_t validSegs)
160- uint32_t batchId, uint32_t validSegs)
161{159{
162 uint32_t totalElems = validSegs * segmentLen_;160 uint32_t totalElems = validSegs * segmentLen_;
163 int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);161 int64_t segStart = static_cast<int64_t>(batchId) * static_cast<int64_t>(batchSize_);
@@ -170,8 +168,8 @@ __aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::ProcessBatc
170 uint32_t innerTileId = batchId % innerLoopNum_;168 uint32_t innerTileId = batchId % innerLoopNum_;
171 innerStart = static_cast<uint64_t>(innerTileId) * static_cast<uint64_t>(batchSize_);169 innerStart = static_cast<uint64_t>(innerTileId) * static_cast<uint64_t>(batchSize_);
172 uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);170 uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);
173- Base::LoadNonLastBatch(inputXGm_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_),171+ Base::LoadNonLastBatch(inputXGm_, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_), validSegs,
174- validSegs, totalElems);172+ totalElems);
175 } else {173 } else {
176 Base::LoadContiguousBatch(inputXGm_, segStart * static_cast<int64_t>(segmentLen_), totalElems);174 Base::LoadContiguousBatch(inputXGm_, segStart * static_cast<int64_t>(segmentLen_), totalElems);
177 }175 }
@@ -184,27 +182,29 @@ __aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::ProcessBatc
184}182}
185 183 
186template <typename T, typename OutIdxT, bool IsDescend>184template <typename T, typename OutIdxT, bool IsDescend>
187-__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::StoreNonLastBatch(185+__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::StoreNonLastBatch(uint64_t outerId,
188- uint64_t outerId, uint64_t innerStart, uint32_t validSegs, uint32_t totalElems)186+ uint64_t innerStart,
187+ uint32_t validSegs,
188+ uint32_t totalElems)
189{189{
190 uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);190 uint64_t outerBaseOffset = outerId * static_cast<uint64_t>(segmentLen_) * static_cast<uint64_t>(innerSize_);
191 // finalValues_/finalIdx_ are already in per-segment order when this store runs.191 // finalValues_/finalIdx_ are already in per-segment order when this store runs.
192- asc_vf_call<StoreNonLastBatchSimt<T, OutIdxT>>(dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM),192+ asc_vf_call<StoreNonLastBatchSimt<T, OutIdxT>>(
193- totalElems, segmentLen_, validSegs, outerBaseOffset, innerStart, static_cast<uint64_t>(innerSize_),193+ dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), totalElems, segmentLen_, validSegs, outerBaseOffset, innerStart,
194- (__ubuf__ T *)finalValues_.GetPhyAddr(), (__ubuf__ OutIdxT *)finalIdx_.GetPhyAddr(),194+ static_cast<uint64_t>(innerSize_), (__ubuf__ T*)finalValues_.GetPhyAddr(),
195- (__gm__ volatile T *)outValueGm_.GetPhyAddr(), (__gm__ volatile OutIdxT *)outIdxGm_.GetPhyAddr());195+ (__ubuf__ OutIdxT*)finalIdx_.GetPhyAddr(), (__gm__ volatile T*)outValueGm_.GetPhyAddr(),
196+ (__gm__ volatile OutIdxT*)outIdxGm_.GetPhyAddr());
196 event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));197 event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
197 SetFlag<HardEvent::V_S>(eventId);198 SetFlag<HardEvent::V_S>(eventId);
198 WaitFlag<HardEvent::V_S>(eventId);199 WaitFlag<HardEvent::V_S>(eventId);
199}200}
200 201 
201template <typename T, typename OutIdxT, bool IsDescend>202template <typename T, typename OutIdxT, bool IsDescend>
202-__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::StoreBatch(int64_t segStart,203+__aicore__ inline void SortSmallAxisTwoStage<T, OutIdxT, IsDescend>::StoreBatch(int64_t segStart, uint32_t totalElems)
203- uint32_t totalElems)
204{204{
205 int64_t gmOffset = segStart * static_cast<int64_t>(segmentLen_);205 int64_t gmOffset = segStart * static_cast<int64_t>(segmentLen_);
206- DataCopyExtParams valueCopyParam{ 1, static_cast<uint32_t>(totalElems * sizeof(T)), 0, 0, 0 };206+ DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(totalElems * sizeof(T)), 0, 0, 0};
207- DataCopyExtParams idxCopyParam{ 1, static_cast<uint32_t>(totalElems * sizeof(OutIdxT)), 0, 0, 0 };207+ DataCopyExtParams idxCopyParam{1, static_cast<uint32_t>(totalElems * sizeof(OutIdxT)), 0, 0, 0};
208 // Rank-inverse writes final buffers via SIMT VF; BuildOutputs writes them via vector APIs.208 // Rank-inverse writes final buffers via SIMT VF; BuildOutputs writes them via vector APIs.
209 // Wait for the producing VF/vector work before GM writeback.209 // Wait for the producing VF/vector work before GM writeback.
210 event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));210 event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
Mmath/sort/op_kernel/arch35/sort_tiling_data.h+7-7
@@ -17,11 +17,11 @@
17#define _SORT_REGBASE_TILING_DATA_H_17#define _SORT_REGBASE_TILING_DATA_H_
18 18 
19struct SortRegBaseTilingData {19struct SortRegBaseTilingData {
20- uint32_t numTileDataSize; // h轴ub一次处理个数20+ uint32_t numTileDataSize; // h轴ub一次处理个数
21 uint32_t unsortedDimParallel; // b轴使用的核数21 uint32_t unsortedDimParallel; // b轴使用的核数
22- uint32_t lastDimTileNum; // h轴循环次数22+ uint32_t lastDimTileNum; // h轴循环次数
23- uint32_t sortLoopTimes; // b轴循环次数23+ uint32_t sortLoopTimes; // b轴循环次数
24- uint32_t lastDimNeedCore; // h轴需要的核数24+ uint32_t lastDimNeedCore; // h轴需要的核数
25 // keyParamsxxx 预留参数25 // keyParamsxxx 预留参数
26 // radix: globalHistGmWk_ 使用核数26 // radix: globalHistGmWk_ 使用核数
27 // radix_one_core: inqueX 的 ub 大小27 // radix_one_core: inqueX 的 ub 大小
@@ -45,9 +45,9 @@ struct SortRegBaseTilingData {
45 // merge_sort(sch0): 队列 buffer 数45 // merge_sort(sch0): 队列 buffer 数
46 // intra_core(sch4): extract chunk 大小46 // intra_core(sch4): extract chunk 大小
47 uint32_t keyParams4;47 uint32_t keyParams4;
48- uint32_t keyParams5; // radix：清零chunk大小；intra_core(sch4)：最大归并迭代次数48+ uint32_t keyParams5; // radix：清零chunk大小；intra_core(sch4)：最大归并迭代次数
49- uint32_t tmpUbSize; // sort高级api需要的临时ub大小49+ uint32_t tmpUbSize; // sort高级api需要的临时ub大小
50- int64_t lastAxisNum; // h轴大小50+ int64_t lastAxisNum; // h轴大小
51 int64_t unsortedDimNum; // b轴大小51 int64_t unsortedDimNum; // b轴大小
52 // Non-last small-axis fields describe the original [outer, axis, inner] GM layout52 // Non-last small-axis fields describe the original [outer, axis, inner] GM layout
53 // and the aligned UB row strides used by tile-local transpose schedules.53 // and the aligned UB row strides used by tile-local transpose schedules.
Mmath/sort/op_kernel/arch35/sort_tiling_key.h+58-57
@@ -17,6 +17,8 @@
17#define _SORT_TILING_KEY_DECL_H_17#define _SORT_TILING_KEY_DECL_H_
18#include "ascendc/host_api/tiling/template_argument.h"18#include "ascendc/host_api/tiling/template_argument.h"
19 19 
20+#define SORT_SCHID_0 0
21+#define SORT_SCHID_1 1
20#define SORT_SCHID_2 222#define SORT_SCHID_2 2
21#define SORT_SCHID_3 323#define SORT_SCHID_3 3
22#define SORT_SCHID_4 424#define SORT_SCHID_4 4
@@ -27,74 +29,73 @@
27#define SORT_SCHID_9 929#define SORT_SCHID_9 9
28#define SORT_SCHID_10 1030#define SORT_SCHID_10 10
29 31 
30-#define SORT_TPL_KEY_DECL() \32+#define SORT_TPL_KEY_DECL() \
31- ASCENDC_TPL_UINT_DECL(schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1, SORT_SCHID_2, SORT_SCHID_3, \33+ ASCENDC_TPL_UINT_DECL(schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, SORT_SCHID_0, SORT_SCHID_1, SORT_SCHID_2, \
32- SORT_SCHID_4, SORT_SCHID_5, SORT_SCHID_6, SORT_SCHID_7, SORT_SCHID_8, \34+ SORT_SCHID_3, SORT_SCHID_4, SORT_SCHID_5, SORT_SCHID_6, SORT_SCHID_7, SORT_SCHID_8, \
33- SORT_SCHID_9, SORT_SCHID_10), \35+ SORT_SCHID_9, SORT_SCHID_10), \
34- ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1), \36+ ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1), \
35- ASCENDC_TPL_UINT_DECL(isDescend, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1) \37+ ASCENDC_TPL_UINT_DECL(isDescend, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1)
36 38 
37-#define SORT_TPL_RADIX_MORE_CORE_KEY_SEL() \39+#define SORT_TPL_RADIX_MORE_CORE_KEY_SEL() \
38- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_2), \40+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_2), \
39- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \41+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \
40- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \42+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
41 43 
42-#define SORT_TPL_MERGE_SORT_KEY_SEL() \44+#define SORT_TPL_MERGE_SORT_KEY_SEL() \
43- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, 0, SORT_SCHID_8), \45+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_0, SORT_SCHID_8), \
44- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \46+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \
45- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \47+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
46 48 
47-#define SORT_TPL_RADIX_ONE_CORE_KEY_SEL() \49+#define SORT_TPL_RADIX_ONE_CORE_KEY_SEL() \
48- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, 1), \50+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_1), \
49- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \51+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \
50- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \52+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
51 53 
52-#define SORT_TPL_MERGE_BIG_SIZE_KEY_SEL() \54+#define SORT_TPL_MERGE_BIG_SIZE_KEY_SEL() \
53- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_3), \55+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_3), \
54- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \56+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \
55- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \57+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
56 58 
57-#define SORT_TPL_MERGE_INTRA_CORE_KEY_SEL() \59+#define SORT_TPL_MERGE_INTRA_CORE_KEY_SEL() \
58- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_4), \60+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_4), \
59- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \61+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \
60- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \62+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
61 63 
62-#define SORT_TPL_SMALL_AXIS_INSERTION_KEY_SEL() \64+#define SORT_TPL_SMALL_AXIS_INSERTION_KEY_SEL() \
63- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_5), \65+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_5), \
64- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \66+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \
65- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \67+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
66 68 
67-#define SORT_TPL_SMALL_AXIS_TWO_STAGE_KEY_SEL() \69+#define SORT_TPL_SMALL_AXIS_TWO_STAGE_KEY_SEL() \
68- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_6), \70+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_6), \
69- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \71+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \
70- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \72+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
71 73 
72-#define SORT_TPL_AXIS_ONE_COPY_KEY_SEL() \74+#define SORT_TPL_AXIS_ONE_COPY_KEY_SEL() \
73- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_7), \75+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_7), \
74- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \76+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 0, 1), \
75- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \77+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
76 78 
77-#define SORT_TPL_NON_LAST_SMALL_AXIS_KEY_SEL() \79+#define SORT_TPL_NON_LAST_SMALL_AXIS_KEY_SEL() \
78- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_9), \80+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_9), \
79- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \81+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \
80- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \82+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
81 83 
82-#define SORT_TPL_NON_LAST_SMALL_AXIS_RADIX_KEY_SEL() \84+#define SORT_TPL_NON_LAST_SMALL_AXIS_RADIX_KEY_SEL() \
83- ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_10), \85+ ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_10), \
84- ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \86+ ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \
85- ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) \87+ ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1)
86 88 
87ASCENDC_TPL_ARGS_DECL(Sort, SORT_TPL_KEY_DECL());89ASCENDC_TPL_ARGS_DECL(Sort, SORT_TPL_KEY_DECL());
88 90 
89-ASCENDC_TPL_SEL(91+ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(SORT_TPL_RADIX_MORE_CORE_KEY_SEL()),
90- ASCENDC_TPL_ARGS_SEL(SORT_TPL_RADIX_MORE_CORE_KEY_SEL()),92+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_MERGE_SORT_KEY_SEL()),
91- ASCENDC_TPL_ARGS_SEL(SORT_TPL_MERGE_SORT_KEY_SEL()),93+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_RADIX_ONE_CORE_KEY_SEL()),
92- ASCENDC_TPL_ARGS_SEL(SORT_TPL_RADIX_ONE_CORE_KEY_SEL()),94+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_MERGE_BIG_SIZE_KEY_SEL()),
93- ASCENDC_TPL_ARGS_SEL(SORT_TPL_MERGE_BIG_SIZE_KEY_SEL()),95+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_MERGE_INTRA_CORE_KEY_SEL()),
94- ASCENDC_TPL_ARGS_SEL(SORT_TPL_MERGE_INTRA_CORE_KEY_SEL()),96+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_SMALL_AXIS_INSERTION_KEY_SEL()),
95- ASCENDC_TPL_ARGS_SEL(SORT_TPL_SMALL_AXIS_INSERTION_KEY_SEL()),97+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_SMALL_AXIS_TWO_STAGE_KEY_SEL()),
96- ASCENDC_TPL_ARGS_SEL(SORT_TPL_SMALL_AXIS_TWO_STAGE_KEY_SEL()),98+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_AXIS_ONE_COPY_KEY_SEL()),
97- ASCENDC_TPL_ARGS_SEL(SORT_TPL_AXIS_ONE_COPY_KEY_SEL()),99+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_NON_LAST_SMALL_AXIS_KEY_SEL()),
98- ASCENDC_TPL_ARGS_SEL(SORT_TPL_NON_LAST_SMALL_AXIS_KEY_SEL()),100+ ASCENDC_TPL_ARGS_SEL(SORT_TPL_NON_LAST_SMALL_AXIS_RADIX_KEY_SEL()));
99- ASCENDC_TPL_ARGS_SEL(SORT_TPL_NON_LAST_SMALL_AXIS_RADIX_KEY_SEL()));
100#endif101#endif
Mmath/sort/op_kernel/sort_apt.cpp+84-127
@@ -12,6 +12,8 @@
12 * \file sort_apt.cpp12 * \file sort_apt.cpp
13 * \brief13 * \brief
14 */14 */
15+#include <type_traits>
16+ 
15#include "kernel_tiling/kernel_tiling.h"17#include "kernel_tiling/kernel_tiling.h"
16#include "basic_api/kernel_vec_intf.h"18#include "basic_api/kernel_vec_intf.h"
17#include "arch35/sort_tiling_key.h"19#include "arch35/sort_tiling_key.h"
@@ -29,134 +31,89 @@
29using namespace AscendC;31using namespace AscendC;
30using namespace Sort;32using namespace Sort;
31 33 
32-template <uint64_t schId, uint64_t isInt32, uint64_t isDescend>34+template <typename Op>
33-__global__ __aicore__ void sort(35+__aicore__ inline void LaunchSortKernel(GM_ADDR input, GM_ADDR values, GM_ADDR indices, GM_ADDR userWorkspace,
34- GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, GM_ADDR tiling)36+ const SortRegBaseTilingData* sortTiling, TPipe* pipeline)
35{37{
36- REGISTER_TILING_DEFAULT(SortRegBaseTilingData);38+ Op op;
37- GET_TILING_DATA_WITH_STRUCT(SortRegBaseTilingData, tilingData, tiling);39+ op.Init(input, values, indices, userWorkspace, sortTiling, pipeline);
40+ op.Process();
41+}
38 42 
39- GM_ADDR usrWorkspace = AscendC::GetUserWorkspace(workspace);43+template <uint64_t isInt32, uint64_t isDescend>
40- KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);44+__aicore__ inline void LaunchRadixMoreCore(GM_ADDR radixInput, GM_ADDR radixValues, GM_ADDR radixIndices,
41- TPipe pipe;45+ GM_ADDR radixWorkspace, const SortRegBaseTilingData* radixTiling,
42- constexpr bool isDescending = (isDescend != 0);46+ TPipe* radixPipeline)
43- if constexpr (schId == 7) {47+{
44- Sort::SortAxisOneCopy<DTYPE_X, DTYPE_Y2> op;48+ using IndexType = std::conditional_t<isInt32 == 1, uint32_t, int64_t>;
45- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);49+ using RadixType = std::conditional_t<
46- op.Process();50+ sizeof(DTYPE_X) == sizeof(uint8_t), uint8_t,
47- return;51+ std::conditional_t<sizeof(DTYPE_X) == sizeof(uint16_t), uint16_t,
48- }52+ std::conditional_t<sizeof(DTYPE_X) == sizeof(uint32_t), uint32_t, uint64_t>>>;
49- if constexpr (schId == 2) {53+ LaunchSortKernel<SortRadixMoreCore<DTYPE_X, DTYPE_Y2, RadixType, IndexType, isDescend>>(
50- if constexpr (sizeof(DTYPE_X) == 1) {54+ radixInput, radixValues, radixIndices, radixWorkspace, radixTiling, radixPipeline);
51- if constexpr (isInt32 == 1) {55+}
52- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint8_t, uint32_t, isDescend> op;56+ 
53- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);57+template <uint64_t schId, uint64_t isDescend>
54- op.Process();58+__aicore__ inline void LaunchMergeSortRoute(GM_ADDR mergeInput, GM_ADDR mergeValues, GM_ADDR mergeIndices,
55- } else {59+ GM_ADDR mergeWorkspace, const SortRegBaseTilingData* mergeTiling,
56- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint8_t, int64_t, isDescend> op;60+ TPipe* mergePipeline)
57- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);61+{
58- op.Process();62+ constexpr bool isSort32SmallAxis = (schId == SORT_SCHID_8);
59- }63+ if constexpr (std::is_same_v<bfloat16_t, DTYPE_X>) {
60- }64+ LaunchSortKernel<MergeSort<DTYPE_X, DTYPE_Y2, float, isDescend, isSort32SmallAxis>>(
61- if constexpr (sizeof(DTYPE_X) == 2) {65+ mergeInput, mergeValues, mergeIndices, mergeWorkspace, mergeTiling, mergePipeline);
62- if constexpr (isInt32 == 1) {66+ } else if constexpr (std::is_same_v<half, DTYPE_X> || std::is_same_v<float, DTYPE_X>) {
63- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint16_t, uint32_t, isDescend> op;67+ LaunchSortKernel<MergeSort<DTYPE_X, DTYPE_Y2, DTYPE_X, isDescend, isSort32SmallAxis>>(
64- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);68+ mergeInput, mergeValues, mergeIndices, mergeWorkspace, mergeTiling, mergePipeline);
65- op.Process();69+ }
66- } else {70+}
67- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint16_t, int64_t, isDescend> op;71+ 
68- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);72+template <uint64_t schId, uint64_t isInt32, uint64_t isDescend>
69- op.Process();73+__global__ __aicore__ void sort(GM_ADDR input, GM_ADDR sortedValues, GM_ADDR sortedIndices, GM_ADDR workBuffer,
70- }74+ GM_ADDR tilingAddress)
71- }75+{
72- if constexpr (sizeof(DTYPE_X) == 4) {76+ REGISTER_TILING_DEFAULT(SortRegBaseTilingData);
73- if constexpr (isInt32 == 1) {77+ GET_TILING_DATA_WITH_STRUCT(SortRegBaseTilingData, sortTilingData, tilingAddress);
74- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint32_t, uint32_t, isDescend> op;78+ 
75- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);79+ GM_ADDR sortWorkspace = AscendC::GetUserWorkspace(workBuffer);
76- op.Process();80+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
77- } else {81+ TPipe sortPipeline;
78- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint32_t, int64_t, isDescend> op;82+ constexpr bool isDescending = (isDescend != 0);
79- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);83+ if constexpr (schId == SORT_SCHID_7) {
80- op.Process();84+ LaunchSortKernel<Sort::SortAxisOneCopy<DTYPE_X, DTYPE_Y2>>(input, sortedValues, sortedIndices, sortWorkspace,
81- }85+ &sortTilingData, &sortPipeline);
82- }86+ } else if constexpr (schId == SORT_SCHID_2) {
83- if constexpr (sizeof(DTYPE_X) == 8) {87+ LaunchRadixMoreCore<isInt32, isDescend>(input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData,
84- if constexpr (isInt32 == 1) {88+ &sortPipeline);
85- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint64_t, uint32_t, isDescend> op;89+ } else if constexpr (schId == SORT_SCHID_1) {
86- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);90+ LaunchSortKernel<SortRadixOneCore<DTYPE_X, DTYPE_Y2, isDescending>>(
87- op.Process();91+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
88- } else {92+ } else if constexpr (schId == SORT_SCHID_0 || schId == SORT_SCHID_8) {
89- SortRadixMoreCore<DTYPE_X, DTYPE_Y2, uint64_t, int64_t, isDescend> op;93+ LaunchMergeSortRoute<schId, isDescend>(input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData,
90- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);94+ &sortPipeline);
91- op.Process();95+ } else if constexpr (schId == SORT_SCHID_3 && std::is_same_v<float, DTYPE_X>) {
92- }96+ LaunchSortKernel<MergeSortBigSize<DTYPE_X, DTYPE_X, isDescending, DTYPE_Y2>>(
93- }97+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
94- return;98+ } else if constexpr (schId == SORT_SCHID_4 && std::is_same_v<float, DTYPE_X>) {
95- }99+ LaunchSortKernel<Sort::SortMergeIntraCore<DTYPE_X, DTYPE_Y2, isDescending>>(
96- if constexpr (schId == 1) {100+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
97- SortRadixOneCore<DTYPE_X, DTYPE_Y2, isDescending> op;101+ } else if constexpr (schId == SORT_SCHID_5 && std::is_same_v<bfloat16_t, DTYPE_X>) {
98- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);102+ LaunchSortKernel<Sort::SortSmallAxisInsertion<DTYPE_X, float, DTYPE_Y2, isDescending>>(
99- op.Process();103+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
100- return;104+ } else if constexpr (schId == SORT_SCHID_5) {
101- }105+ LaunchSortKernel<Sort::SortSmallAxisInsertion<DTYPE_X, DTYPE_X, DTYPE_Y2, isDescending>>(
102- if constexpr (schId == 0 || schId == 8) {106+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
103- constexpr uint64_t isSort32SmallAxis = (schId == 8);107+ } else if constexpr (schId == SORT_SCHID_6) {
104- if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) {108+ LaunchSortKernel<Sort::SortSmallAxisTwoStage<DTYPE_X, DTYPE_Y2, isDescending>>(
105- MergeSort<DTYPE_X, DTYPE_Y2, float, isDescend, isSort32SmallAxis> op;109+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
106- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);110+ } else if constexpr (schId == SORT_SCHID_9 || schId == SORT_SCHID_10) {
107- op.Process();111+ constexpr bool useMergeSort = (schId == SORT_SCHID_9);
108- } else if constexpr (IsSameType<float, DTYPE_X>::value || IsSameType<half, DTYPE_X>::value) {112+ constexpr bool supportMergeSort = std::is_same_v<DTYPE_X, half> || std::is_same_v<DTYPE_X, float> ||
109- MergeSort<DTYPE_X, DTYPE_Y2, DTYPE_X, isDescend, isSort32SmallAxis> op;113+ std::is_same_v<DTYPE_X, bfloat16_t>;
110- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);114+ if constexpr (!useMergeSort || supportMergeSort) {
111- op.Process();115+ LaunchSortKernel<Sort::SortNonLastSmallAxis<DTYPE_X, DTYPE_Y2, isDescending, useMergeSort>>(
112- }116+ input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline);
113- return;117+ }
114- }
115- if constexpr (schId == 3) {
116- if constexpr (IsSameType<float, DTYPE_X>::value) {
117- MergeSortBigSize<DTYPE_X, DTYPE_X, isDescending, DTYPE_Y2> op;
118- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);
119- op.Process();
120- }
121- return;
122- }
123- if constexpr (schId == 4) {
124- if constexpr (IsSameType<float, DTYPE_X>::value) {
125- Sort::SortMergeIntraCore<DTYPE_X, DTYPE_Y2, isDescending> op;
126- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);
127- op.Process();
128- }
129- return;
130- }
131- if constexpr (schId == 5) {
132- if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) {
133- Sort::SortSmallAxisInsertion<DTYPE_X, float, DTYPE_Y2, isDescending> op;
134- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);
135- op.Process();
136- } else {
137- Sort::SortSmallAxisInsertion<DTYPE_X, DTYPE_X, DTYPE_Y2, isDescending> op;
138- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);
139- op.Process();
140- }
141- return;
142- }
143- if constexpr (schId == 6) {
144- Sort::SortSmallAxisTwoStage<DTYPE_X, DTYPE_Y2, isDescending> op;
145- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);
146- op.Process();
147- return;
148- }
149- if constexpr (schId == 9 || schId == 10) {
150- constexpr bool useMergeSort = (schId == 9);
151- constexpr bool supportMergeSort =
152- IsSameType<DTYPE_X, float>::value || IsSameType<DTYPE_X, half>::value ||
153- IsSameType<DTYPE_X, bfloat16_t>::value;
154- // Host tiling only assigns merge sort to supported dtypes; this guard prevents unsupported instantiation.
155- if constexpr (!useMergeSort || supportMergeSort) {
156- Sort::SortNonLastSmallAxis<DTYPE_X, DTYPE_Y2, isDescending, useMergeSort> op;
157- op.Init(x, y1, y2, usrWorkspace, &tilingData, &pipe);
158- op.Process();
159- }
160- return;
161 }118 }
162}119}
Mmath/sort/tests/ut/op_api/test_aclnn_sort.cpp+83-103
@@ -28,31 +28,31 @@ protected:
28 28 
29 vector<int64_t> GetViewDims(int num_of_dim)29 vector<int64_t> GetViewDims(int num_of_dim)
30 {30 {
31- vector<int64_t> view_dims= {1, 2, 2, 2}; // NCHW + NHWC + HWCN31+ vector<int64_t> view_dims = {1, 2, 2, 2}; // NCHW + NHWC + HWCN
32- if (num_of_dim == 5) // NDHWC + NCDHW32+ if (num_of_dim == 5) // NDHWC + NCDHW
33 view_dims = {1, 4, 1, 1, 2};33 view_dims = {1, 4, 1, 1, 2};
34- else if (num_of_dim == 0) // 空tensor34+ else if (num_of_dim == 0) // 空tensor
35 view_dims = {0};35 view_dims = {0};
36- else if (num_of_dim == 1) // 1维36+ else if (num_of_dim == 1) // 1维
37 view_dims = {10};37 view_dims = {10};
38- else if (num_of_dim == 8) // 8维38+ else if (num_of_dim == 8) // 8维
39 view_dims = {2, 1, 2, 1, 2, 4, 5, 7};39 view_dims = {2, 1, 2, 1, 2, 4, 5, 7};
40- else if (num_of_dim == 10) // > 8维40+ else if (num_of_dim == 10) // > 8维
41 view_dims = {2, 1, 4, 4, 3, 1, 2, 4, 5, 7};41 view_dims = {2, 1, 4, 4, 3, 1, 2, 4, 5, 7};
42- else if (num_of_dim != 4) // ND42+ else if (num_of_dim != 4) // ND
43 view_dims = {1, 4, 1, 1, 2, 1};43 view_dims = {1, 4, 1, 1, 2, 1};
44 44 
45 return view_dims;45 return view_dims;
46 }46 }
47 47 
48- vector<float> tensor_value = {1.3232, 2.123, 3.234, -234.4, 23.52, -1.6, 71.8, 823};48+ vector<float> tensor_value = {1.3232, 2.123, 3.234, -234.4, 23.52, -1.6, 71.8, 823};
49 49 
50- // num_of_dim: 这个tensor是几维的50+ // num_of_dim: 这个tensor是几维的
51- void test_run(aclDataType test_dtype, aclFormat test_format, int num_of_dim, bool stable=false, int64_t dim=1,51+ void test_run(aclDataType test_dtype, aclFormat test_format, int num_of_dim, bool stable = false, int64_t dim = 1,
52- bool descending=false)52+ bool descending = false)
53 {53 {
54 auto view_dims = GetViewDims(num_of_dim);54 auto view_dims = GetViewDims(num_of_dim);
55- const vector<int64_t>& view_dims_final= const_cast <vector<int64_t>&>(view_dims);55+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
56 56 
57 auto self = TensorDesc(view_dims_final, test_dtype, test_format).ValueRange(-2, 2).Value(tensor_value);57 auto self = TensorDesc(view_dims_final, test_dtype, test_format).ValueRange(-2, 2).Value(tensor_value);
58 auto values = TensorDesc(view_dims_final, test_dtype, test_format).Precision(0.00001, 0.00001);58 auto values = TensorDesc(view_dims_final, test_dtype, test_format).Precision(0.00001, 0.00001);
@@ -66,13 +66,15 @@ protected:
66 // ut.TestPrecision();66 // ut.TestPrecision();
67 }67 }
68 68 
69- void test_run_invalid(aclDataType test_dtype, aclFormat test_format, int num_of_dim, bool stable=false, int64_t dim=1, bool descending=false)69+ void test_run_invalid(aclDataType test_dtype, aclFormat test_format, int num_of_dim, bool stable = false,
70+ int64_t dim = 1, bool descending = false)
70 {71 {
71 auto view_dims = GetViewDims(num_of_dim);72 auto view_dims = GetViewDims(num_of_dim);
72- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);73+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
73 74 
74 auto self = TensorDesc(view_dims_final, test_dtype, test_format)75 auto self = TensorDesc(view_dims_final, test_dtype, test_format)
75- .ValueRange(-2, 2).Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});76+ .ValueRange(-2, 2)
77+ .Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});
76 auto values = TensorDesc(view_dims_final, test_dtype, test_format);78 auto values = TensorDesc(view_dims_final, test_dtype, test_format);
77 auto indices = TensorDesc(view_dims_final, ACL_INT64, test_format);79 auto indices = TensorDesc(view_dims_final, ACL_INT64, test_format);
78 80 
@@ -84,14 +86,15 @@ protected:
84 }86 }
85 87 
86 void test_run_dtype_not_consistent(aclDataType self_dtype, aclDataType values_dtype, aclDataType indices_dtype,88 void test_run_dtype_not_consistent(aclDataType self_dtype, aclDataType values_dtype, aclDataType indices_dtype,
87- aclFormat test_format, int num_of_dim, bool stable=false, int64_t dim=1, bool descending=false,89+ aclFormat test_format, int num_of_dim, bool stable = false, int64_t dim = 1,
88- bool invalid=true)90+ bool descending = false, bool invalid = true)
89 {91 {
90 auto view_dims = GetViewDims(num_of_dim);92 auto view_dims = GetViewDims(num_of_dim);
91- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);93+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
92 94 
93 auto self = TensorDesc(view_dims_final, self_dtype, test_format)95 auto self = TensorDesc(view_dims_final, self_dtype, test_format)
94- .ValueRange(-2, 2).Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});96+ .ValueRange(-2, 2)
97+ .Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});
95 auto values = TensorDesc(view_dims_final, values_dtype, test_format);98 auto values = TensorDesc(view_dims_final, values_dtype, test_format);
96 auto indices = TensorDesc(view_dims_final, indices_dtype, test_format);99 auto indices = TensorDesc(view_dims_final, indices_dtype, test_format);
97 100 
@@ -107,14 +110,16 @@ protected:
107 }110 }
108 }111 }
109 112 
110- void test_run_format_not_consistent(aclDataType test_dtype, aclFormat self_format, aclFormat values_format, aclFormat indices_format,113+ void test_run_format_not_consistent(aclDataType test_dtype, aclFormat self_format, aclFormat values_format,
111- int num_of_dim, bool stable=false, int64_t dim=1, bool descending=false)114+ aclFormat indices_format, int num_of_dim, bool stable = false, int64_t dim = 1,
115+ bool descending = false)
112 {116 {
113 auto view_dims = GetViewDims(num_of_dim);117 auto view_dims = GetViewDims(num_of_dim);
114- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);118+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
115 119 
116 auto self = TensorDesc(view_dims_final, test_dtype, self_format)120 auto self = TensorDesc(view_dims_final, test_dtype, self_format)
117- .ValueRange(-2, 2).Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});121+ .ValueRange(-2, 2)
122+ .Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});
118 auto values = TensorDesc(view_dims_final, test_dtype, values_format);123 auto values = TensorDesc(view_dims_final, test_dtype, values_format);
119 auto indices = TensorDesc(view_dims_final, ACL_INT64, indices_format);124 auto indices = TensorDesc(view_dims_final, ACL_INT64, indices_format);
120 125 
@@ -132,10 +137,11 @@ protected:
132 bool descending = false;137 bool descending = false;
133 138 
134 auto view_dims = GetViewDims(num_of_dim);139 auto view_dims = GetViewDims(num_of_dim);
135- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);140+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
136 141 
137 auto self = TensorDesc(view_dims_final, test_dtype, test_format)142 auto self = TensorDesc(view_dims_final, test_dtype, test_format)
138- .ValueRange(-2, 2).Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});143+ .ValueRange(-2, 2)
144+ .Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});
139 auto values = TensorDesc(view_dims_final, test_dtype, test_format);145 auto values = TensorDesc(view_dims_final, test_dtype, test_format);
140 auto indices = TensorDesc(view_dims_final, ACL_INT64, test_format);146 auto indices = TensorDesc(view_dims_final, ACL_INT64, test_format);
141 147 
@@ -166,22 +172,23 @@ protected:
166 EXPECT_EQ(getWorkspaceResult6, ACLNN_ERR_PARAM_NULLPTR);172 EXPECT_EQ(getWorkspaceResult6, ACLNN_ERR_PARAM_NULLPTR);
167 173 
168 auto ut7 = OP_API_UT(aclnnSort, INPUT((aclTensor*)nullptr, stable, dim, descending),174 auto ut7 = OP_API_UT(aclnnSort, INPUT((aclTensor*)nullptr, stable, dim, descending),
169- OUTPUT((aclTensor*)nullptr, (aclTensor*)nullptr));175+ OUTPUT((aclTensor*)nullptr, (aclTensor*)nullptr));
170 aclnnStatus getWorkspaceResult7 = ut7.TestGetWorkspaceSize(&workspaceSize);176 aclnnStatus getWorkspaceResult7 = ut7.TestGetWorkspaceSize(&workspaceSize);
171 EXPECT_EQ(getWorkspaceResult7, ACLNN_ERR_PARAM_NULLPTR);177 EXPECT_EQ(getWorkspaceResult7, ACLNN_ERR_PARAM_NULLPTR);
172 }178 }
173 179 
174- void test_run_indices_longTensor (aclDataType test_dtype, aclFormat test_format, int num_of_dim)180+ void test_run_indices_longTensor(aclDataType test_dtype, aclFormat test_format, int num_of_dim)
175 {181 {
176 int64_t dim = 1;182 int64_t dim = 1;
177 bool stable = false;183 bool stable = false;
178 bool descending = false;184 bool descending = false;
179 185 
180 auto view_dims = GetViewDims(num_of_dim);186 auto view_dims = GetViewDims(num_of_dim);
181- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);187+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
182 188 
183 auto self = TensorDesc(view_dims_final, test_dtype, test_format)189 auto self = TensorDesc(view_dims_final, test_dtype, test_format)
184- .ValueRange(-2, 2).Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});190+ .ValueRange(-2, 2)
191+ .Value(vector<float>{1, 2, 3, 4, 5, 6, 7, 8});
185 auto values = TensorDesc(view_dims_final, test_dtype, test_format).Precision(0.00001, 0.00001);192 auto values = TensorDesc(view_dims_final, test_dtype, test_format).Precision(0.00001, 0.00001);
186 auto indices = TensorDesc(view_dims_final, ACL_INT64, test_format).Precision(0.00001, 0.00001);193 auto indices = TensorDesc(view_dims_final, ACL_INT64, test_format).Precision(0.00001, 0.00001);
187 194 
@@ -195,10 +202,10 @@ protected:
195 // ut.TestPrecision();202 // ut.TestPrecision();
196 }203 }
197 204 
198- void test_run_new(aclDataType test_dtype, vector<int64_t> view_dims, int dim=-1, bool stable=false,205+ void test_run_new(aclDataType test_dtype, vector<int64_t> view_dims, int dim = -1, bool stable = false,
199- bool descending=false)206+ bool descending = false)
200 {207 {
201- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);208+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
202 209 
203 auto self = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND).ValueRange(-2, 2);210 auto self = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND).ValueRange(-2, 2);
204 auto values = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND);211 auto values = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND);
@@ -212,10 +219,10 @@ protected:
212 // ut.TestPrecision();219 // ut.TestPrecision();
213 }220 }
214 221 
215- void test_run_aicpu_stable(aclDataType test_dtype, vector<int64_t> view_dims, int dim=-1, bool stable=false,222+ void test_run_aicpu_stable(aclDataType test_dtype, vector<int64_t> view_dims, int dim = -1, bool stable = false,
216- bool descending=false)223+ bool descending = false)
217 {224 {
218- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);225+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
219 226 
220 auto self = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND).ValueRange(-2, 2);227 auto self = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND).ValueRange(-2, 2);
221 auto values = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND);228 auto values = TensorDesc(view_dims_final, test_dtype, ACL_FORMAT_ND);
@@ -235,46 +242,25 @@ protected:
235///////////////////////////////////////242///////////////////////////////////////
236 243 
237// sort输入支持float16 aicore244// sort输入支持float16 aicore
238-TEST_F(l2_sort_test, l2_sort_test_01)245+TEST_F(l2_sort_test, l2_sort_test_01) { test_run(ACL_FLOAT16, ACL_FORMAT_ND, 6); }
239-{
240- test_run(ACL_FLOAT16, ACL_FORMAT_ND, 6);
241-}
242 246 
243// sort输入支持float32 1980 aicpu, 1971 aicore247// sort输入支持float32 1980 aicpu, 1971 aicore
244-TEST_F(l2_sort_test, l2_sort_test_02)248+TEST_F(l2_sort_test, l2_sort_test_02) { test_run(ACL_FLOAT, ACL_FORMAT_NCHW, 4); }
245-{
246- test_run(ACL_FLOAT, ACL_FORMAT_NCHW, 4);
247-}
248 249 
249// sort输入支持int8 aicpu250// sort输入支持int8 aicpu
250-TEST_F(l2_sort_test, l2_sort_test_03)251+TEST_F(l2_sort_test, l2_sort_test_03) { test_run(ACL_INT8, ACL_FORMAT_NHWC, 4); }
251-{
252- test_run(ACL_INT8, ACL_FORMAT_NHWC, 4);
253-}
254 252 
255// sort输入支持int16 aicpu253// sort输入支持int16 aicpu
256-TEST_F(l2_sort_test, l2_sort_test_04)254+TEST_F(l2_sort_test, l2_sort_test_04) { test_run(ACL_INT16, ACL_FORMAT_HWCN, 4); }
257-{
258- test_run(ACL_INT16, ACL_FORMAT_HWCN, 4);
259-}
260 255 
261// sort输入支持int32 aicpu256// sort输入支持int32 aicpu
262-TEST_F(l2_sort_test, l2_sort_test_05)257+TEST_F(l2_sort_test, l2_sort_test_05) { test_run(ACL_INT32, ACL_FORMAT_NDHWC, 5); }
263-{
264- test_run(ACL_INT32, ACL_FORMAT_NDHWC, 5);
265-}
266 258 
267-//sort输入支持int64 aicpu259+// sort输入支持int64 aicpu
268-TEST_F(l2_sort_test, l2_sort_test_06)260+TEST_F(l2_sort_test, l2_sort_test_06) { test_run(ACL_INT64, ACL_FORMAT_NCDHW, 5); }
269-{
270- test_run(ACL_INT64, ACL_FORMAT_NCDHW, 5);
271-}
272 261 
273-//sort输入支持uint8 aicpu262+// sort输入支持uint8 aicpu
274-TEST_F(l2_sort_test, l2_sort_test_07)263+TEST_F(l2_sort_test, l2_sort_test_07) { test_run(ACL_UINT8, ACL_FORMAT_ND, 6); }
275-{
276- test_run(ACL_UINT8, ACL_FORMAT_ND, 6);
277-}
278 264 
279// // NPU sort输入不支持bool, bfloat16, double, complex64, complex128265// // NPU sort输入不支持bool, bfloat16, double, complex64, complex128
280// TEST_F(l2_sort_test, l2_sort_test_08)266// TEST_F(l2_sort_test, l2_sort_test_08)
@@ -288,18 +274,15 @@ TEST_F(l2_sort_test, l2_sort_test_07)
288 274 
289// tensor的dtype不一致275// tensor的dtype不一致
290TEST_F(l2_sort_test, l2_sort_test_09)276TEST_F(l2_sort_test, l2_sort_test_09)
291-{ 277+{
292 // self + values dtype不一致是允许的278 // self + values dtype不一致是允许的
293 test_run_dtype_not_consistent(ACL_INT8, ACL_INT16, ACL_INT64, ACL_FORMAT_ND, 6, false, 1, false, false);279 test_run_dtype_not_consistent(ACL_INT8, ACL_INT16, ACL_INT64, ACL_FORMAT_ND, 6, false, 1, false, false);
294- test_run_dtype_not_consistent(ACL_INT8, ACL_INT8, ACL_FLOAT, ACL_FORMAT_ND, 6); // indices 非INT64280+ test_run_dtype_not_consistent(ACL_INT8, ACL_INT8, ACL_FLOAT, ACL_FORMAT_ND, 6); // indices 非INT64
295- test_run_dtype_not_consistent(ACL_INT8, ACL_INT8, ACL_INT32, ACL_FORMAT_ND, 6); // indices 非INT64281+ test_run_dtype_not_consistent(ACL_INT8, ACL_INT8, ACL_INT32, ACL_FORMAT_ND, 6); // indices 非INT64
296}282}
297 283 
298// tensor的dtype是未定义的284// tensor的dtype是未定义的
299-TEST_F(l2_sort_test, l2_sort_test_10)285+TEST_F(l2_sort_test, l2_sort_test_10) { test_run_invalid(ACL_DT_UNDEFINED, ACL_FORMAT_ND, 6); }
300-{
301- test_run_invalid(ACL_DT_UNDEFINED, ACL_FORMAT_ND, 6);
302-}
303 286 
304///////////////////////////////////////287///////////////////////////////////////
305///// 检查空指针 /////288///// 检查空指针 /////
@@ -328,7 +311,6 @@ TEST_F(l2_sort_test, l2_sort_test_12)
328 test_run_indices_longTensor(ACL_INT64, ACL_FORMAT_ND, 4);311 test_run_indices_longTensor(ACL_INT64, ACL_FORMAT_ND, 4);
329}312}
330 313 
331- 
332///////////////////////////////////////314///////////////////////////////////////
333///// 数据格式覆盖 /////315///// 数据格式覆盖 /////
334///////////////////////////////////////316///////////////////////////////////////
@@ -360,10 +342,9 @@ TEST_F(l2_sort_test, l2_sort_test_16)
360{342{
361 // aicore343 // aicore
362 auto view_dims = GetViewDims(0);344 auto view_dims = GetViewDims(0);
363- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);345+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
364 346 
365- auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND)347+ auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(-2, 2).Value(vector<float>{});
366- .ValueRange(-2, 2).Value(vector<float>{});
367 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);348 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);
368 auto indices = TensorDesc(view_dims_final, ACL_INT64, ACL_FORMAT_ND);349 auto indices = TensorDesc(view_dims_final, ACL_INT64, ACL_FORMAT_ND);
369 350 
@@ -385,7 +366,6 @@ TEST_F(l2_sort_test, l2_sort_test_16)
385 // ut.TestPrecision();366 // ut.TestPrecision();
386}367}
387 368 
388- 
389///////////////////////////////////////369///////////////////////////////////////
390///// 支持非连续tensor /////370///// 支持非连续tensor /////
391///////////////////////////////////////371///////////////////////////////////////
@@ -393,8 +373,8 @@ TEST_F(l2_sort_test, l2_sort_test_16)
393// 非连续tensor aicore + aicpu373// 非连续tensor aicore + aicpu
394TEST_F(l2_sort_test, l2_sort_test_17)374TEST_F(l2_sort_test, l2_sort_test_17)
395{375{
396- vector<int64_t> view_dims= {5, 4};376+ vector<int64_t> view_dims = {5, 4};
397- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);377+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
398 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND, {1, 5}, 0, {4, 5}).ValueRange(300, 350);378 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND, {1, 5}, 0, {4, 5}).ValueRange(300, 350);
399 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);379 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);
400 auto indices = TensorDesc(view_dims_final, ACL_INT64, ACL_FORMAT_ND);380 auto indices = TensorDesc(view_dims_final, ACL_INT64, ACL_FORMAT_ND);
@@ -436,13 +416,13 @@ TEST_F(l2_sort_test, l2_sort_test_19)
436// dim在范围 [-N, N-1] 范围之外 aicore + aicpu416// dim在范围 [-N, N-1] 范围之外 aicore + aicpu
437TEST_F(l2_sort_test, l2_sort_test_20)417TEST_F(l2_sort_test, l2_sort_test_20)
438{418{
439- test_run(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, -4, false); // -N419+ test_run(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, -4, false); // -N
440- test_run(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, 3, false); // N-1420+ test_run(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, 3, false); // N-1
441 test_run_invalid(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, -5, false);421 test_run_invalid(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, -5, false);
442 test_run_invalid(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, 4, false);422 test_run_invalid(ACL_FLOAT16, ACL_FORMAT_NCHW, 4, false, 4, false);
443 423 
444- test_run(ACL_INT8, ACL_FORMAT_NCHW, 4, false, -4, false); // -N424+ test_run(ACL_INT8, ACL_FORMAT_NCHW, 4, false, -4, false); // -N
445- test_run(ACL_INT16, ACL_FORMAT_NCHW, 4, false, 3, false); // N-1425+ test_run(ACL_INT16, ACL_FORMAT_NCHW, 4, false, 3, false); // N-1
446 test_run_invalid(ACL_INT64, ACL_FORMAT_NCHW, 4, false, -5, false);426 test_run_invalid(ACL_INT64, ACL_FORMAT_NCHW, 4, false, -5, false);
447 test_run_invalid(ACL_UINT8, ACL_FORMAT_NCHW, 4, false, 4, false);427 test_run_invalid(ACL_UINT8, ACL_FORMAT_NCHW, 4, false, 4, false);
448}428}
@@ -459,7 +439,8 @@ TEST_F(l2_sort_test, l2_sort_test_22)
459}439}
460 440 
461// sort的维度数据量为100001 aicore + aicpu441// sort的维度数据量为100001 aicore + aicpu
462-TEST_F(l2_sort_test, l2_sort_test_23) {442+TEST_F(l2_sort_test, l2_sort_test_23)
443+{
463 // test_run_new(ACL_FLOAT16, {3, 4, 4, 100001}, -1, true, false); // values正常, indices正常. st通过444 // test_run_new(ACL_FLOAT16, {3, 4, 4, 100001}, -1, true, false); // values正常, indices正常. st通过
464 // test_run_new(ACL_INT64, {3, 4, 4, 100001}, -1, false, false); // values正常, indices正常. st通过445 // test_run_new(ACL_INT64, {3, 4, 4, 100001}, -1, false, false); // values正常, indices正常. st通过
465}446}
@@ -489,7 +470,6 @@ TEST_F(l2_sort_test, l2_sort_test_26)
489 test_run_new(ACL_INT32, {2, 1, 4, 4, 3, 1, 2, 4, 5, 7}, -1, false, true);470 test_run_new(ACL_INT32, {2, 1, 4, 4, 3, 1, 2, 4, 5, 7}, -1, false, true);
490}471}
491 472 
492- 
493///////////////////////////////////////473///////////////////////////////////////
494///// sort stable为T / F /////474///// sort stable为T / F /////
495///////////////////////////////////////475///////////////////////////////////////
@@ -497,18 +477,18 @@ TEST_F(l2_sort_test, l2_sort_test_26)
497// stable为True/False aicore + aicpu477// stable为True/False aicore + aicpu
498TEST_F(l2_sort_test, l2_sort_test_27)478TEST_F(l2_sort_test, l2_sort_test_27)
499{479{
500- test_run_new(ACL_FLOAT16, {2, 1, 2, 1, 2, 4, 5, 7}, -1, true, true); // stable True480+ test_run_new(ACL_FLOAT16, {2, 1, 2, 1, 2, 4, 5, 7}, -1, true, true); // stable True
501- test_run_new(ACL_FLOAT16, {2, 1, 2, 1, 2, 4, 5, 7}, -1, false, true); // stable false481+ test_run_new(ACL_FLOAT16, {2, 1, 2, 1, 2, 4, 5, 7}, -1, false, true); // stable false
502 // aicpu不支持stable为True482 // aicpu不支持stable为True
503- test_run_new(ACL_INT64, {2, 1, 2, 1, 2, 4, 5, 7}, -1, false, true); // stable false483+ test_run_new(ACL_INT64, {2, 1, 2, 1, 2, 4, 5, 7}, -1, false, true); // stable false
504- test_run_aicpu_stable(ACL_INT64, {2, 1, 2, 1, 2, 4, 5, 7}, -1, true, false); // stable为True时，因为走AICPU所以报错484+ test_run_aicpu_stable(ACL_INT64, {2, 1, 2, 1, 2, 4, 5, 7}, -1, true, false); // stable为True时，因为走AICPU所以报错
505}485}
506 486 
507// value相等时的排序487// value相等时的排序
508TEST_F(l2_sort_test, l2_sort_test_28)488TEST_F(l2_sort_test, l2_sort_test_28)
509{489{
510- vector<int64_t> view_dims= GetViewDims(8);;490+ vector<int64_t> view_dims = GetViewDims(8);
511- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);491+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
512 492 
513 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(1, 1);493 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(1, 1);
514 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);494 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);
@@ -553,10 +533,10 @@ TEST_F(l2_sort_test, l2_sort_test_28)
553// shape必须一致533// shape必须一致
554TEST_F(l2_sort_test, l2_sort_test_29)534TEST_F(l2_sort_test, l2_sort_test_29)
555{535{
556- vector<int64_t> view_dims= {3, 4, 2};536+ vector<int64_t> view_dims = {3, 4, 2};
557- vector<int64_t> view_dims1= {3, 2, 4};537+ vector<int64_t> view_dims1 = {3, 2, 4};
558- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);538+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
559- const vector<int64_t>& view_dims_final1 = const_cast <vector<int64_t>&>(view_dims1);539+ const vector<int64_t>& view_dims_final1 = const_cast<vector<int64_t>&>(view_dims1);
560 540 
561 // indices ! = self541 // indices ! = self
562 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(1, 1);542 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(1, 1);
@@ -583,8 +563,8 @@ TEST_F(l2_sort_test, l2_sort_test_29)
583///////////////////////////////////////563///////////////////////////////////////
584TEST_F(l2_sort_test, l2_sort_test_34)564TEST_F(l2_sort_test, l2_sort_test_34)
585{565{
586- test_run_aicpu_stable(ACL_FLOAT16, {2, 3, 3, 1}, -1, true); // 因为sort为1，走aicpu所以报错566+ test_run_aicpu_stable(ACL_FLOAT16, {2, 3, 3, 1}, -1, true); // 因为sort为1，走aicpu所以报错
587- test_run_aicpu_stable(ACL_INT64, {2, 4, 5, 7}, -1, true); // dtype走aicpu，所以报错567+ test_run_aicpu_stable(ACL_INT64, {2, 4, 5, 7}, -1, true); // dtype走aicpu，所以报错
588}568}
589 569 
590///////////////////////////////////////570///////////////////////////////////////
@@ -594,15 +574,15 @@ TEST_F(l2_sort_test, l2_sort_test_34)
594// torch.tensor(a)能正常计算574// torch.tensor(a)能正常计算
595TEST_F(l2_sort_test, l2_sort_test_35)575TEST_F(l2_sort_test, l2_sort_test_35)
596{576{
597- vector<int64_t> view_dims= {};577+ vector<int64_t> view_dims = {};
598- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);578+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
599 579 
600 // indices ! = self580 // indices ! = self
601 auto self = TensorDesc(view_dims, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(10, 10);581 auto self = TensorDesc(view_dims, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(10, 10);
602 auto values = TensorDesc(view_dims, ACL_FLOAT16, ACL_FORMAT_ND);582 auto values = TensorDesc(view_dims, ACL_FLOAT16, ACL_FORMAT_ND);
603 auto indices = TensorDesc(view_dims, ACL_INT64, ACL_FORMAT_ND);583 auto indices = TensorDesc(view_dims, ACL_INT64, ACL_FORMAT_ND);
604 584 
605- auto ut = OP_API_UT(aclnnSort, INPUT(self, false, -1, false), OUTPUT(values, indices)); 585+ auto ut = OP_API_UT(aclnnSort, INPUT(self, false, -1, false), OUTPUT(values, indices));
606 uint64_t workspace_size = 0;586 uint64_t workspace_size = 0;
607 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);587 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
608 EXPECT_EQ(aclRet, ACL_SUCCESS);588 EXPECT_EQ(aclRet, ACL_SUCCESS);
@@ -611,15 +591,15 @@ TEST_F(l2_sort_test, l2_sort_test_35)
611 591 
612TEST_F(l2_sort_test, ascend310B_l2_sort_test_36)592TEST_F(l2_sort_test, ascend310B_l2_sort_test_36)
613{593{
614- vector<int64_t> view_dims= {20};594+ vector<int64_t> view_dims = {20};
615- const vector<int64_t>& view_dims_final = const_cast <vector<int64_t>&>(view_dims);595+ const vector<int64_t>& view_dims_final = const_cast<vector<int64_t>&>(view_dims);
616 596 
617 // indices ! = self597 // indices ! = self
618 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(-10, 10);598 auto self = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND).ValueRange(-10, 10);
619 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);599 auto values = TensorDesc(view_dims_final, ACL_FLOAT16, ACL_FORMAT_ND);
620 auto indices = TensorDesc(view_dims_final, ACL_INT64, ACL_FORMAT_ND);600 auto indices = TensorDesc(view_dims_final, ACL_INT64, ACL_FORMAT_ND);
621 601 
622- auto ut = OP_API_UT(aclnnSort, INPUT(self, true, -1, false), OUTPUT(values, indices)); 602+ auto ut = OP_API_UT(aclnnSort, INPUT(self, true, -1, false), OUTPUT(values, indices));
623 uint64_t workspace_size = 0;603 uint64_t workspace_size = 0;
624 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);604 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
625 EXPECT_EQ(aclRet, ACL_SUCCESS);605 EXPECT_EQ(aclRet, ACL_SUCCESS);