Star879
1779
代码
代码
Issues233
Pull Requests328
Actions
项目讨论
Wiki
分析
Star879
1779
已合并
usm排序模板性能优化 #2445
jinpenghe创建于 3月9日
usm排序模板性能优化 #2445
已合并
jinpenghe创建于 3月9日
从jinpenghe/ops-nnusmprof合入到cann/ops-nnmaster
共 9 个文件变更+346-104
Mindex/unsorted_segment_common/op_host/arch35/unsorted_segment_simd_dyn_sort_tiling.cpp+49-18
@@ -20,9 +20,9 @@ static constexpr uint64_t TEMPLATE_SIMD_DYN_SORT = 7000;
20static constexpr uint64_t LAST_DIM_SIMD_COND = 128;20static constexpr uint64_t LAST_DIM_SIMD_COND = 128;
21static constexpr uint64_t BUFFER_NUM = 1;21static constexpr uint64_t BUFFER_NUM = 1;
22static constexpr uint64_t SIMD_RESERVED_SIZE = 8192;22static constexpr uint64_t SIMD_RESERVED_SIZE = 8192;
23-static constexpr uint64_t BASE_A_SIZE = 512;23+static constexpr uint64_t BASE_A_SIZE = 4096;
24static constexpr uint64_t BASE_BLOCK_SIZE = 8192;24static constexpr uint64_t BASE_BLOCK_SIZE = 8192;
25-static constexpr uint64_t COL_LIMIT_SIZE = 1024;25+static constexpr uint64_t COL_LIMIT_SIZE = 2048;
26static constexpr uint64_t DOUBLE = 2;26static constexpr uint64_t DOUBLE = 2;
27static constexpr uint32_t SORT_STAT_PADDING = 64;27static constexpr uint32_t SORT_STAT_PADDING = 64;
28 28 
@@ -61,6 +61,7 @@ void UnsortedSegmentSimdDynSortTiling::SetTilingData()
61 tilingData->sortBaseS = sortBaseS_;61 tilingData->sortBaseS = sortBaseS_;
62 tilingData->sortBaseA = sortBaseA_;62 tilingData->sortBaseA = sortBaseA_;
63 tilingData->sortSharedBufSize = static_cast<uint64_t>(sortSharedBufSize_);63 tilingData->sortSharedBufSize = static_cast<uint64_t>(sortSharedBufSize_);
64+ tilingData->idCastMode = idCastMode_;
64}65}
65 66 
66void UnsortedSegmentSimdDynSortTiling::DoBlockTiling()67void UnsortedSegmentSimdDynSortTiling::DoBlockTiling()
@@ -69,7 +70,7 @@ void UnsortedSegmentSimdDynSortTiling::DoBlockTiling()
69 usedCoreNum_ =70 usedCoreNum_ =
70 std::min(totalCoreNum_, static_cast<uint64_t>(inputOuterDim_ * colNumAlign * BASE_A_SIZE / BASE_BLOCK_SIZE));71 std::min(totalCoreNum_, static_cast<uint64_t>(inputOuterDim_ * colNumAlign * BASE_A_SIZE / BASE_BLOCK_SIZE));
71 usedCoreNum_ = usedCoreNum_ == 0 ? 1 : usedCoreNum_;72 usedCoreNum_ = usedCoreNum_ == 0 ? 1 : usedCoreNum_;
72- std::tie(sTileNum_, aTileNum_) = AutoTiling(usedCoreNum_, colNumAlign, COL_LIMIT_SIZE);73+ std::tie(sTileNum_, aTileNum_) = AutoTiling(usedCoreNum_, colNumAlign, COL_LIMIT_SIZE, true);
73 74 
74 normBlockS_ = Ops::Base::CeilDiv(inputOuterDim_, sTileNum_);75 normBlockS_ = Ops::Base::CeilDiv(inputOuterDim_, sTileNum_);
75 sTileNum_ = Ops::Base::CeilDiv(inputOuterDim_, normBlockS_);76 sTileNum_ = Ops::Base::CeilDiv(inputOuterDim_, normBlockS_);
@@ -95,14 +96,26 @@ uint64_t UnsortedSegmentSimdDynSortTiling::CalBestBaseSize(uint64_t baseXoStart,
95 baseXoEnd = baseXoEnd + 1UL;96 baseXoEnd = baseXoEnd + 1UL;
96 while (baseXoEnd - baseXoStart > 1UL) {97 while (baseXoEnd - baseXoStart > 1UL) {
97 baseXoMid = (baseXoStart + baseXoEnd) / DOUBLE;98 baseXoMid = (baseXoStart + baseXoEnd) / DOUBLE;
98- uint64_t sortNeedTmpSize = static_cast<uint64_t>(GetSortTmpSize(baseXoMid, false));99+ if (idCastMode_ == 0) {
99- tmpTotalSize = Ops::Base::CeilAlign(baseXoMid * sortBaseA_ * dataTypeBytes_, ubBlockSize_) * BUFFER_NUM + // xQue100+ uint64_t sortNeedTmpSize = static_cast<uint64_t>(GetSortTmpSize(idType_, baseXoMid, false));
100- Ops::Base::CeilAlign(baseXoMid * idTypeBytes_, ubBlockSize_) * BUFFER_NUM + // idQue101+ tmpTotalSize = Ops::Base::CeilAlign(baseXoMid * sortBaseA_ * dataTypeBytes_, ubBlockSize_) * BUFFER_NUM + // xQue
101- Ops::Base::CeilAlign(baseXoMid * sortBaseA_ * dataTypeBytes_, ubBlockSize_) + // resBuf102+ Ops::Base::CeilAlign(baseXoMid * idTypeBytes_, ubBlockSize_) * BUFFER_NUM + // idQue
102- Ops::Base::CeilAlign(baseXoMid * idTypeBytes_, ubBlockSize_) + // sortedkeyBuf103+ Ops::Base::CeilAlign(baseXoMid * sortBaseA_ * dataTypeBytes_, ubBlockSize_) + // resBuf
103- Ops::Base::CeilAlign(baseXoMid * sizeof(uint32_t), ubBlockSize_) * idsSortBufCnt + // sortedIdxBuf104+ Ops::Base::CeilAlign(baseXoMid * idTypeBytes_, ubBlockSize_) + // sortedkeyBuf
104- SORT_STAT_PADDING + SORT_STAT_PADDING + // sort padding105+ Ops::Base::CeilAlign(baseXoMid * sizeof(uint32_t), ubBlockSize_) * idsSortBufCnt + // sortedIdxBuf
105- Ops::Base::CeilAlign(sortNeedTmpSize, ubBlockSize_); // sort shared buf size106+ SORT_STAT_PADDING + SORT_STAT_PADDING + // sort padding
107+ Ops::Base::CeilAlign(sortNeedTmpSize, ubBlockSize_); // sort shared buf size
108+ } else {
109+ uint64_t sortNeedTmpSize = static_cast<uint64_t>(GetSortTmpSize(idCastDtype_, baseXoMid, false));
110+ tmpTotalSize = Ops::Base::CeilAlign(baseXoMid * sortBaseA_ * dataTypeBytes_, ubBlockSize_) * BUFFER_NUM + // xQue
111+ Ops::Base::CeilAlign(baseXoMid * idTypeBytes_, ubBlockSize_) * BUFFER_NUM + // idQue
112+ Ops::Base::CeilAlign(baseXoMid * idCastDtypeSize_, ubBlockSize_) * BUFFER_NUM + // idCastQue
113+ Ops::Base::CeilAlign(baseXoMid * sortBaseA_ * dataTypeBytes_, ubBlockSize_) + // resBuf
114+ Ops::Base::CeilAlign(baseXoMid * idCastDtypeSize_, ubBlockSize_) + // sortedkeyBuf
115+ Ops::Base::CeilAlign(baseXoMid * sizeof(uint32_t), ubBlockSize_) * idsSortBufCnt + // sortedIdxBuf
116+ SORT_STAT_PADDING + SORT_STAT_PADDING + // sort padding
117+ Ops::Base::CeilAlign(sortNeedTmpSize, ubBlockSize_); // sort shared buf size
118+ }
106 if (tmpTotalSize <= ubSize_) {119 if (tmpTotalSize <= ubSize_) {
107 baseXoStart = baseXoMid;120 baseXoStart = baseXoMid;
108 } else {121 } else {
@@ -114,6 +127,7 @@ uint64_t UnsortedSegmentSimdDynSortTiling::CalBestBaseSize(uint64_t baseXoStart,
114 127 
115ge::graphStatus UnsortedSegmentSimdDynSortTiling::DoOpTiling()128ge::graphStatus UnsortedSegmentSimdDynSortTiling::DoOpTiling()
116{129{
130+ GetCastTypeForSort();
117 baseA_ = BASE_A_SIZE / dataTypeBytes_;131 baseA_ = BASE_A_SIZE / dataTypeBytes_;
118 DoBlockTiling();132 DoBlockTiling();
119 133 
@@ -139,17 +153,33 @@ ge::graphStatus UnsortedSegmentSimdDynSortTiling::DoOpTiling()
139 if (coreMaxS < sortBaseS_) {153 if (coreMaxS < sortBaseS_) {
140 sortBaseS_ = coreMaxS;154 sortBaseS_ = coreMaxS;
141 uint64_t idsSortBufCnt = 2;155 uint64_t idsSortBufCnt = 2;
142- uint64_t sortNeedTmpSize = static_cast<uint64_t>(GetSortTmpSize(sortBaseS_, false));156+ uint64_t remainSize = 0;
143- uint64_t remainSize = ubSize_ - Ops::Base::CeilAlign(sortBaseS_ * idTypeBytes_, ubBlockSize_) * BUFFER_NUM -157+ uint64_t sortNeedTmpSize = 0;
144- Ops::Base::CeilAlign(sortBaseS_ * idTypeBytes_, ubBlockSize_) -158+ if (idCastMode_ == 0) {
145- Ops::Base::CeilAlign(sortBaseS_ * sizeof(uint32_t), ubBlockSize_) * idsSortBufCnt -159+ sortNeedTmpSize = static_cast<uint64_t>(GetSortTmpSize(idType_, sortBaseS_, false));
146- Ops::Base::CeilAlign(sortNeedTmpSize, ubBlockSize_) - SORT_STAT_PADDING - SORT_STAT_PADDING;160+ remainSize = ubSize_ - Ops::Base::CeilAlign(sortBaseS_ * idTypeBytes_, ubBlockSize_) * BUFFER_NUM -
161+ Ops::Base::CeilAlign(sortBaseS_ * idTypeBytes_, ubBlockSize_) -
162+ Ops::Base::CeilAlign(sortBaseS_ * sizeof(uint32_t), ubBlockSize_) * idsSortBufCnt -
163+ Ops::Base::CeilAlign(sortNeedTmpSize, ubBlockSize_) - SORT_STAT_PADDING - SORT_STAT_PADDING;
164+ } else {
165+ sortNeedTmpSize = static_cast<uint64_t>(GetSortTmpSize(idCastDtype_, sortBaseS_, false));
166+ remainSize = ubSize_ - Ops::Base::CeilAlign(sortBaseS_ * idTypeBytes_, ubBlockSize_) * BUFFER_NUM -
167+ Ops::Base::CeilAlign(sortBaseS_ * idCastDtypeSize_, ubBlockSize_) * BUFFER_NUM -
168+ Ops::Base::CeilAlign(sortBaseS_ * idCastDtypeSize_, ubBlockSize_) -
169+ Ops::Base::CeilAlign(sortBaseS_ * sizeof(uint32_t), ubBlockSize_) * idsSortBufCnt -
170+ Ops::Base::CeilAlign(sortNeedTmpSize, ubBlockSize_) - SORT_STAT_PADDING - SORT_STAT_PADDING;
171+ }
147 sortBaseA_ =172 sortBaseA_ =
148 (remainSize - BUFFER_NUM * ubBlockSize_ - ubBlockSize_) / (BUFFER_NUM + 1) / (sortBaseS_ * dataTypeBytes_);173 (remainSize - BUFFER_NUM * ubBlockSize_ - ubBlockSize_) / (BUFFER_NUM + 1) / (sortBaseS_ * dataTypeBytes_);
149 sortBaseA_ = Ops::Base::FloorAlign(sortBaseA_, ubBlockSize_ / dataTypeBytes_);174 sortBaseA_ = Ops::Base::FloorAlign(sortBaseA_, ubBlockSize_ / dataTypeBytes_);
150 }175 }
151- sortSharedBufSize_ = GetSortTmpSize(sortBaseS_, false);176+ 
152- SetTilingData();177+ if (idCastMode_ == 0) {
178+ sortSharedBufSize_ = GetSortTmpSize(idType_, sortBaseS_, false);
179+ } else {
180+ sortSharedBufSize_ = GetSortTmpSize(idCastDtype_, sortBaseS_, false);
181+ }
182+ SetTilingData();
153 return ge::GRAPH_SUCCESS;183 return ge::GRAPH_SUCCESS;
154}184}
155 185 
@@ -182,6 +212,7 @@ void UnsortedSegmentSimdDynSortTiling::DumpTilingInfo()
182 info << ", sortBaseS: " << tilingData->sortBaseS;212 info << ", sortBaseS: " << tilingData->sortBaseS;
183 info << ", sortBaseA: " << tilingData->sortBaseA;213 info << ", sortBaseA: " << tilingData->sortBaseA;
184 info << ", sortSharedBufSize: " << tilingData->sortSharedBufSize;214 info << ", sortSharedBufSize: " << tilingData->sortSharedBufSize;
215+ info << ", idCastMode: " << tilingData->idCastMode;
185 OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());216 OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());
186}217}
187 218 
Mindex/unsorted_segment_common/op_host/arch35/unsorted_segment_simd_spilt_col_tiling.cpp+1-1
@@ -24,7 +24,7 @@ static constexpr uint64_t RATIO_BY_SORT = 5;
24 24 
25bool UnsortedSegmentSimdSplitColTiling::IsCapable()25bool UnsortedSegmentSimdSplitColTiling::IsCapable()
26{26{
27- if (innerDim_ * dataTypeBytes_ > totalCoreNum_ * LAST_DIM_SIMD_COND && ratio_ > RATIO_BY_SORT) {27+ if (innerDim_ * dataTypeBytes_ > totalCoreNum_ * LAST_DIM_SIMD_COND && ratio_ < RATIO_BY_SORT) {
28 return IsFullLoad();28 return IsFullLoad();
29 }29 }
30 return false;30 return false;
Mindex/unsorted_segment_common/op_host/arch35/unsorted_segment_sort_simt_tiling.cpp+1-1
@@ -52,7 +52,7 @@ ge::graphStatus UnsortedSegmentSortSimtTiling::CalcTiling()
52 int64_t totalIndexSize = Ops::Base::CeilAlign(mid * idTypeBytes_, ubBlockSize_) * DOUBLE +52 int64_t totalIndexSize = Ops::Base::CeilAlign(mid * idTypeBytes_, ubBlockSize_) * DOUBLE +
53 Ops::Base::CeilAlign(mid * idTypeBytes_, ubBlockSize_) + ubBlockSize_ * DOUBLE +53 Ops::Base::CeilAlign(mid * idTypeBytes_, ubBlockSize_) + ubBlockSize_ * DOUBLE +
54 Ops::Base::CeilAlign(mid * sizeof(uint32_t), ubBlockSize_);54 Ops::Base::CeilAlign(mid * sizeof(uint32_t), ubBlockSize_);
55- sortTmpSize = GetSortTmpSize(mid, false);55+ sortTmpSize = GetSortTmpSize(idType_, mid, false);
56 sortTmpSize = Ops::Base::CeilAlign(sortTmpSize, static_cast<int64_t>(ubBlockSize_));56 sortTmpSize = Ops::Base::CeilAlign(sortTmpSize, static_cast<int64_t>(ubBlockSize_));
57 int64_t tmpTotalSize =57 int64_t tmpTotalSize =
58 totalIndexSize + sortTmpSize + Ops::Base::CeilAlign(mid * innerDim_ * dataTypeBytes_, ubBlockSize_) * DOUBLE;58 totalIndexSize + sortTmpSize + Ops::Base::CeilAlign(mid * innerDim_ * dataTypeBytes_, ubBlockSize_) * DOUBLE;
Mindex/unsorted_segment_common/op_host/arch35/unsorted_segment_tiling_common.cpp+39-3
@@ -24,6 +24,11 @@ static constexpr uint32_t INPUT_SEGMENT_IDS_INDEX = 1;
24static constexpr uint32_t INPUT_NUM_SEGMENTS_INDEX = 2;24static constexpr uint32_t INPUT_NUM_SEGMENTS_INDEX = 2;
25static constexpr uint32_t OUTPUT_DATA_INDEX = 0;25static constexpr uint32_t OUTPUT_DATA_INDEX = 0;
26constexpr uint64_t ASCENDC_WORKSPACE = static_cast<uint64_t>(16) * 1024 * 1024;26constexpr uint64_t ASCENDC_WORKSPACE = static_cast<uint64_t>(16) * 1024 * 1024;
27+static constexpr uint64_t CAST_INT32_TO_INT16 = 1; // int32 Cast int16
28+static constexpr uint64_t CAST_INT64_TO_INT32 = 2; // int64 Cast int32
29+static constexpr uint64_t CAST_INT64_TO_INT16 = 3; // int64 Cast int16
30+static constexpr uint64_t CAST_INT32_TO_UINT8 = 4; // int32 Cast uint8
31+static constexpr uint64_t CAST_INT64_TO_UINT8 = 5; // int64 Cast uint8
27 32 
28static const std::set<ge::DataType> DATA_TYPE_SUPPORT = {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16,33static const std::set<ge::DataType> DATA_TYPE_SUPPORT = {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16,
29 ge::DT_INT32, ge::DT_INT64, ge::DT_UINT32,34 ge::DT_INT32, ge::DT_INT64, ge::DT_UINT32,
@@ -41,6 +46,37 @@ static std::string ToString(std::set<ge::DataType> supportDtypes)
41 return ss.str();46 return ss.str();
42}47}
43 48 
49+void UnsortedSegmentBaseTiling::GetCastTypeForSort()
50+{
51+ idCastDtype_ = idType_;
52+ 
53+ if (idType_ == ge::DT_INT32) {
54+ if (outputOuterDim_ < UINT8_MAX) {
55+ idCastMode_ = CAST_INT32_TO_UINT8; // int32 Cast uint8
56+ idCastDtype_ = ge::DT_UINT8;
57+ } else if (outputOuterDim_ < INT16_MAX) {
58+ idCastMode_ = CAST_INT32_TO_INT16; // int32 Cast int16
59+ idCastDtype_ = ge::DT_INT16;
60+ }
61+ } else {
62+ if (outputOuterDim_ < UINT8_MAX) {
63+ idCastMode_ = CAST_INT64_TO_UINT8; // int64 Cast uint8
64+ idCastDtype_ = ge::DT_UINT8;
65+ } else if (outputOuterDim_ < INT16_MAX) {
66+ idCastMode_ = CAST_INT64_TO_INT16; // int64 Cast int16
67+ idCastDtype_ = ge::DT_INT16;
68+ } else if (outputOuterDim_ < INT32_MAX) {
69+ idCastMode_ = CAST_INT64_TO_INT32; // int64 Cast int32
70+ idCastDtype_ = ge::DT_INT32;
71+ }
72+ }
73+ 
74+ if (idCastMode_ != 0) {
75+ idCastDtypeSize_ = ge::GetSizeByDataType(idCastDtype_);
76+ }
77+}
78+ 
79+ 
44ge::graphStatus UnsortedSegmentBaseTiling::GetPlatformInfo()80ge::graphStatus UnsortedSegmentBaseTiling::GetPlatformInfo()
45{81{
46 auto compileInfo = context_->GetCompileInfo<UnsortedSegmentCompileInfo>();82 auto compileInfo = context_->GetCompileInfo<UnsortedSegmentCompileInfo>();
@@ -80,7 +116,7 @@ std::tuple<int64_t, int64_t> UnsortedSegmentBaseTiling::FlatInput(const gert::Sh
80 return std::make_tuple(outer, inner);116 return std::make_tuple(outer, inner);
81}117}
82 118 
83-uint32_t UnsortedSegmentBaseTiling::GetSortTmpSize(uint32_t lastAxisNum, bool isDescend)119+uint32_t UnsortedSegmentBaseTiling::GetSortTmpSize(ge::DataType dataType, uint32_t lastAxisNum, bool isDescend)
84{120{
85 std::vector<int64_t> shapeVec = {lastAxisNum};121 std::vector<int64_t> shapeVec = {lastAxisNum};
86 ge::Shape srcShape(shapeVec);122 ge::Shape srcShape(shapeVec);
@@ -91,7 +127,7 @@ uint32_t UnsortedSegmentBaseTiling::GetSortTmpSize(uint32_t lastAxisNum, bool is
91 config.hasDstIndex = true;127 config.hasDstIndex = true;
92 uint32_t maxValue = 0;128 uint32_t maxValue = 0;
93 uint32_t minValue = 0;129 uint32_t minValue = 0;
94- AscendC::GetSortMaxMinTmpSize(srcShape, idType_, ge::DT_UINT32, false, config, maxValue, minValue);130+ AscendC::GetSortMaxMinTmpSize(srcShape, dataType, ge::DT_UINT32, false, config, maxValue, minValue);
95 return maxValue;131 return maxValue;
96}132}
97 133 
@@ -236,7 +272,7 @@ std::tuple<uint64_t, uint64_t> UnsortedSegmentBaseTiling::AutoTiling(
236 allTiling.begin(), allTiling.end(), [](const std::vector<uint64_t>& a, const std::vector<uint64_t>& b) {272 allTiling.begin(), allTiling.end(), [](const std::vector<uint64_t>& a, const std::vector<uint64_t>& b) {
237 constexpr int NIndex = 1;273 constexpr int NIndex = 1;
238 constexpr int DeltaIndex = 3;274 constexpr int DeltaIndex = 3;
239- return std::make_pair(a[DeltaIndex], a[NIndex]) < std::make_pair(b[DeltaIndex], b[NIndex]);275+ return std::make_pair(a[NIndex], a[DeltaIndex]) < std::make_pair(b[NIndex], b[DeltaIndex]);
240 });276 });
241 } else {277 } else {
242 std::sort(278 std::sort(
Mindex/unsorted_segment_common/op_host/arch35/unsorted_segment_tiling_common.h+7-2
@@ -66,7 +66,8 @@ protected:
66 ge::graphStatus CheckInputDtype();66 ge::graphStatus CheckInputDtype();
67 bool ShapeStartsWith(const gert::Shape shape, const gert::Shape prefix);67 bool ShapeStartsWith(const gert::Shape shape, const gert::Shape prefix);
68 std::tuple<int64_t, int64_t> FlatInput(const gert::Shape shape, const gert::Shape prefix);68 std::tuple<int64_t, int64_t> FlatInput(const gert::Shape shape, const gert::Shape prefix);
69- uint32_t GetSortTmpSize(uint32_t lastAxisNum, bool isDescend);69+ uint32_t GetSortTmpSize(ge::DataType dataType, uint32_t lastAxisNum, bool isDescend);
70+ void GetCastTypeForSort();
70 std::set<uint64_t> FindUniqueCut(uint64_t usedCoreNum);71 std::set<uint64_t> FindUniqueCut(uint64_t usedCoreNum);
71 std::tuple<uint64_t, uint64_t> AutoTiling(72 std::tuple<uint64_t, uint64_t> AutoTiling(
72 uint64_t usedCoreNum, uint64_t colNumAlign, uint64_t colLimitSize, bool colTileNumMin = false);73 uint64_t usedCoreNum, uint64_t colNumAlign, uint64_t colLimitSize, bool colTileNumMin = false);
@@ -86,8 +87,12 @@ public:
86 uint64_t idTypeBytes_ = 0;87 uint64_t idTypeBytes_ = 0;
87 uint64_t dataShapeSize_ = 0;88 uint64_t dataShapeSize_ = 0;
88 uint64_t ratio_ = 0;89 uint64_t ratio_ = 0;
90+ uint64_t idCastMode_ = 0; // 0: 不Cast; 1：int32 Cast int16; 2：int64 Cast int32; 3：int64 Cast int16; 4:int32 Cast uint8; 5:int64 Cast uint8.
91+ int64_t idCastDtypeSize_ = 0;
92+ 
89 ge::DataType dataType_ = ge::DT_UNDEFINED;93 ge::DataType dataType_ = ge::DT_UNDEFINED;
90- ge::DataType idType_ = ge::DT_UNDEFINED; 94+ ge::DataType idType_ = ge::DT_UNDEFINED;
95+ ge::DataType idCastDtype_ = ge::DT_UNDEFINED;
91};96};
92} // namespace optiling97} // namespace optiling
93#endif // UNSORTED_SEGMENT_COMMON_TILING_H98#endif // UNSORTED_SEGMENT_COMMON_TILING_H
Mindex/unsorted_segment_common/op_kernel/arch35/unsorted_segment_base.h+173-46
@@ -41,6 +41,7 @@ constexpr uint32_t ROW_NUM = 16;
41constexpr uint32_t COUNT = 64;41constexpr uint32_t COUNT = 64;
42constexpr uint32_t HALFTIME = 4;42constexpr uint32_t HALFTIME = 4;
43constexpr uint32_t TWO = 2;43constexpr uint32_t TWO = 2;
44+constexpr uint32_t THREE = 3;
44constexpr uint32_t VF_SIZE = platform::GetVRegSize();45constexpr uint32_t VF_SIZE = platform::GetVRegSize();
45constexpr uint32_t VF_B32 = VF_SIZE / sizeof(int32_t);46constexpr uint32_t VF_B32 = VF_SIZE / sizeof(int32_t);
46constexpr uint64_t MIN_FACTOR = 2 * 1024;47constexpr uint64_t MIN_FACTOR = 2 * 1024;
@@ -49,10 +50,28 @@ constexpr uint64_t GM_ALIGN = 512;
49constexpr float FLOAT32_MAX = 3.4028235e+38f;50constexpr float FLOAT32_MAX = 3.4028235e+38f;
50constexpr half FLOAT16_MAX = 65504.0f;51constexpr half FLOAT16_MAX = 65504.0f;
51constexpr bfloat16_t BFLOAT16_MAX = 3.3895314e+38f;52constexpr bfloat16_t BFLOAT16_MAX = 3.3895314e+38f;
52-// constexpr int32_t INT32_MAX = 2147483647;53+ 
53-// constexpr int64_t UINT32_MAX = 4294967295;54+constexpr uint32_t CAST_0 = 0;
54-// constexpr int64_t INT64_MAX = 9223372036854775807LL;55+constexpr uint32_t CAST_1 = 1;
55-// constexpr uint64_t UINT64_MAX = 18446744073709551615ULL;56+constexpr uint32_t CAST_2 = 2;
57+constexpr uint32_t CAST_3 = 3;
58+constexpr uint32_t CAST_4 = 4;
59+constexpr uint32_t CAST_5 = 5;
60+constexpr uint32_t MASK_UINT8 = 255;
61+constexpr int64_t VFLEN_INT64 = platform::GetVRegSize() / sizeof(int64_t);
62+constexpr int64_t VFLEN_INT32 = platform::GetVRegSize() / sizeof(int32_t);
63+constexpr int64_t VFLEN_INT16 = platform::GetVRegSize() / sizeof(int16_t);
64+constexpr int64_t VFLEN_INT16HALF = platform::GetVRegSize() / sizeof(int16_t) / TWO;
65+constexpr int64_t VFLEN_UINT8 = platform::GetVRegSize() / sizeof(uint8_t);
66+constexpr int64_t VFLEN_UINT8HALFHALF = platform::GetVRegSize() / sizeof(uint8_t) / HALFTIME;
67+ 
68+template <typename T, uint32_t CAST_MODE>
69+struct CastType {
70+ using type = typename std::conditional<
71+ CAST_MODE == CAST_1, int16_t, typename std::conditional<CAST_MODE == CAST_2, int32_t,
72+ typename std::conditional<CAST_MODE == CAST_3, int16_t, typename std::conditional<CAST_MODE == CAST_4, uint8_t,
73+ typename std::conditional<CAST_MODE == CAST_5, uint8_t, T>::type>::type>::type>::type>::type;
74+};
56 75 
57typedef struct {76typedef struct {
58 uint16_t segCount;77 uint16_t segCount;
@@ -228,61 +247,169 @@ __aicore__ inline void CopyOut(
228 DataCopyPad(dstGm[offset], srcLocal, dataCoptExtParams);247 DataCopyPad(dstGm[offset], srcLocal, dataCoptExtParams);
229}248}
230 249 
231-template <typename Index>250+template <typename IDX_T, typename CAST_T, uint32_t castType>
232-__aicore__ inline void UniqueGetElm(251+__aicore__ inline void IndicesSortCast(LocalTensor<IDX_T> indicesLocal, LocalTensor<CAST_T> indicesCastLocal,
233- const LocalTensor<Index>& sortedIndice, LocalTensor<int32_t>& noDupRes, uint32_t processIdx, uint32_t shiftOffset,252+ LocalTensor<int32_t> indicesCastTmpLocal, uint32_t indicesCount)
234- uint32_t vfIndicesNum, int64_t& arNum)
235{253{
236- __local_mem__ Index* sortedIndicesAddr = (__ubuf__ Index*)sortedIndice.GetPhyAddr();254+ if constexpr (castType == CAST_4) { // int32 Cast uint8
237- __local_mem__ int32_t* noDupResAddr = (__ubuf__ int32_t*)noDupRes.GetPhyAddr();255+ CompareScalar(indicesCastLocal, indicesLocal, static_cast<IDX_T>(0), CMPMODE::GE, indicesCount);
256+ Select(indicesLocal, indicesCastLocal, indicesLocal, static_cast<IDX_T>(MASK_UINT8), SELMODE::VSEL_TENSOR_SCALAR_MODE, indicesCount);
257+ Cast<CAST_T, IDX_T>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
258+ } else if constexpr (castType == CAST_3) { // int64 Cast int16
259+ Cast<int32_t, IDX_T>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
260+ Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
261+ } else if constexpr (castType == CAST_5) { // int64 Cast uint8
262+ CompareScalar(indicesCastLocal, indicesLocal, static_cast<IDX_T>(0), CMPMODE::GE, indicesCount);
263+ Select(indicesLocal, indicesCastLocal, indicesLocal, static_cast<IDX_T>(MASK_UINT8), SELMODE::VSEL_TENSOR_SCALAR_MODE, indicesCount);
264+ Cast<int32_t, IDX_T>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
265+ Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
266+ } else { // CAST_1 + CAST_2, int32 Cast int16 + int64 Cast int32
267+ Cast<CAST_T, IDX_T>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
268+ }
269+}
238 270 
239- uint16_t loopCnt = (uint16_t)((processIdx + vfIndicesNum) / vfIndicesNum);271+template<typename IDX_T>
272+__aicore__ inline void ComputeUniqueIdNumInt64(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
273+{
274+ uint32_t counter = dataLen + 1;
275+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, selReg;
276+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
277+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskHalf;
278+ AscendC::MicroAPI::UnalignReg u0, uOut;
279+ for (uint16_t i = 0; i < loopCnt; ++i) {
280+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT64);
281+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
282+ auto startAddr = indicesAddr + i * VFLEN_INT64;
283+ DataCopy(sortedIdxReg, startAddr);
284+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
285+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
286+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
287+ AscendC::MicroAPI::MaskPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskHalf, cmpMask);
288+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskHalf);
289+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
290+ uniqueIdCountsAddr, selReg, uOut);
291+ }
292+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
293+}
240 294 
241- int32_t scalar = 0;295+template<typename IDX_T>
242- uint32_t counter = processIdx + 1;296+__aicore__ inline void ComputeUniqueIdNumInt32(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
297+{
298+ uint32_t counter = dataLen + 1;
299+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, selReg;
300+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
301+ AscendC::MicroAPI::MaskReg cmpMask, maskReg;
302+ AscendC::MicroAPI::UnalignReg u0, uOut;
303+ for (uint16_t i = 0; i < loopCnt; ++i) {
304+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT32);
305+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
306+ auto startAddr = indicesAddr + i * VFLEN_INT32;
307+ DataCopy(sortedIdxReg, startAddr);
308+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
309+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
310+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
311+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, cmpMask);
312+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
313+ uniqueIdCountsAddr, selReg, uOut);
314+ }
315+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
316+}
317+ 
318+template<typename IDX_T>
319+__aicore__ inline void ComputeUniqueIdNumInt16(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
320+{
321+ uint32_t counter = dataLen + 1;
322+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, orderReg2, selReg, selReg2;
323+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
324+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskDouble1, maskDouble2;
325+ AscendC::MicroAPI::UnalignReg u0, uOut;
326+ for (uint16_t i = 0; i < loopCnt; ++i) {
327+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT16);
328+ AscendC::MicroAPI::Arange(orderReg2, i * VFLEN_INT16 + VFLEN_INT16HALF);
329+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
330+ auto startAddr = indicesAddr + i * VFLEN_INT16;
331+ DataCopy(sortedIdxReg, startAddr);
332+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
333+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
334+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
335+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskDouble1, cmpMask);
336+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskDouble2, cmpMask);
337+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskDouble1);
338+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskDouble2);
339+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg, uOut);
340+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg2, uOut);
341+ }
342+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
343+}
344+ 
345+template<typename IDX_T>
346+__aicore__ inline void ComputeUniqueIdNumUint8(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
347+{
348+ uint32_t counter = dataLen + 1;
349+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, orderReg2, orderReg3, orderReg4;
350+ AscendC::MicroAPI::RegTensor<int32_t> selReg, selReg2, selReg3, selReg4;
351+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
352+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskFour1, maskFour2, maskFour3, maskFour4;
353+ AscendC::MicroAPI::UnalignReg u0, uOut;
354+ for (uint16_t i = 0; i < loopCnt; ++i) {
355+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_UINT8);
356+ AscendC::MicroAPI::Arange(orderReg2, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF);
357+ AscendC::MicroAPI::Arange(orderReg3, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF * TWO);
358+ AscendC::MicroAPI::Arange(orderReg4, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF * THREE);
359+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
360+ auto startAddr = indicesAddr + i * VFLEN_UINT8;
361+ DataCopy(sortedIdxReg, startAddr);
362+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
363+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
364+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
365+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour3, cmpMask);
366+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour4, cmpMask);
367+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour1, maskFour3);
368+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour2, maskFour3);
369+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour3, maskFour4);
370+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour4, maskFour4);
371+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskFour1);
372+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskFour2);
373+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg3, orderReg3, maskFour3);
374+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg4, orderReg4, maskFour4);
375+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg, uOut);
376+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg2, uOut);
377+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg3, uOut);
378+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg4, uOut);
379+ }
380+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
381+}
382+ 
383+template <typename IDX_T>
384+__aicore__ inline int64_t UniqueGetElm(
385+ const LocalTensor<IDX_T>& sortedIndice, LocalTensor<int32_t>& noDupRes, int64_t dataLen)
386+{
387+ __local_mem__ IDX_T* indicesAddr = (__local_mem__ IDX_T*)sortedIndice[(ONE_BLOCK_SIZE / sizeof(IDX_T))].GetPhyAddr();
388+ __local_mem__ int32_t* uniqueIdCountsAddr = (__local_mem__ int32_t*)noDupRes.GetPhyAddr();
389+ 
390+ constexpr int64_t vfLen = platform::GetVRegSize() / sizeof(IDX_T);
391+ uint16_t loopCnt = ops::CeilDiv(dataLen + 1, vfLen);
243 __VEC_SCOPE__392 __VEC_SCOPE__
244 {393 {
245- AscendC::MicroAPI::RegTensor<int32_t> orderReg;
246- AscendC::MicroAPI::RegTensor<int32_t> selReg;
247- AscendC::MicroAPI::RegTensor<Index> indicesReg;
248- AscendC::MicroAPI::RegTensor<Index> indicesShiftOneReg;
249- 
250- AscendC::MicroAPI::MaskReg cmpMask;
251- AscendC::MicroAPI::MaskReg maskRegUpdate;
252- AscendC::MicroAPI::UnalignReg u0;
253- MicroAPI::UnalignReg ureg;
254 AscendC::MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();394 AscendC::MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
255 395 
256- for (uint16_t i = 0; i < loopCnt; ++i) {396+ if constexpr (std::is_same<int64_t, IDX_T>::value) {
257- scalar = i * vfIndicesNum;397+ ComputeUniqueIdNumInt64<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
258- auto sortedIndicesAddrUpdate = sortedIndicesAddr + shiftOffset + i * vfIndicesNum;398+ } else if constexpr (std::is_same<int32_t, IDX_T>::value) {
259- AscendC::MicroAPI::Arange(orderReg, scalar);399+ ComputeUniqueIdNumInt32<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
260- maskRegUpdate = AscendC::MicroAPI::UpdateMask<Index>(counter);400+ } else if constexpr (std::is_same<int16_t, IDX_T>::value) {
261- AscendC::MicroAPI::DataCopy(indicesReg, sortedIndicesAddrUpdate);401+ ComputeUniqueIdNumInt16<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
262- AscendC::MicroAPI::DataCopyUnAlignPre(u0, sortedIndicesAddrUpdate - 1);402+ } else { // uint8
263- AscendC::MicroAPI::DataCopyUnAlign<Index>(indicesShiftOneReg, u0, sortedIndicesAddrUpdate - 1);403+ ComputeUniqueIdNumUint8<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
264- AscendC::MicroAPI::Compare<Index, CMPMODE::NE>(cmpMask, indicesReg, indicesShiftOneReg, maskRegUpdate);
265- if constexpr (IsSameType<Index, int64_t>::value) {
266- AscendC::MicroAPI::MaskReg maskHalf;
267- AscendC::MicroAPI::MaskPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskHalf, cmpMask);
268- // vSQZ
269- AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(
270- selReg, orderReg, maskHalf);
271- } else {
272- // vSQZ
273- AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(
274- selReg, orderReg, cmpMask);
275- }
276- AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
277- noDupResAddr, selReg, ureg);
278- AscendC::MicroAPI::DataCopyUnAlignPost(noDupResAddr, ureg);
279 }404 }
280 }405 }
406+ int64_t uniqueIdNum = ((AscendC::MicroAPI::GetSpr<AscendC::SpecialPurposeReg::AR>()) / sizeof(int32_t)) - 1;
407+ return uniqueIdNum;
281}408}
282 409 
283__aicore__ inline void UniqueStat(LocalTensor<int32_t>& noDupRes, int64_t& arNum)410__aicore__ inline void UniqueStat(LocalTensor<int32_t>& noDupRes, int64_t& arNum)
284{411{
285- __local_mem__ int32_t* noDupResAddr = (__ubuf__ int32_t*)noDupRes.GetPhyAddr();412+ __local_mem__ int32_t* noDupResAddr = (__local_mem__ int32_t*)noDupRes.GetPhyAddr();
286 413 
287 uint16_t loopCntStatFre = (arNum + VF_B32 - 1) / VF_B32;414 uint16_t loopCntStatFre = (arNum + VF_B32 - 1) / VF_B32;
288 uint32_t counterStatFre = static_cast<uint32_t>(arNum);415 uint32_t counterStatFre = static_cast<uint32_t>(arNum);
Mindex/unsorted_segment_common/op_kernel/arch35/unsorted_segment_simd_dyn_sort.h+40-28
@@ -22,8 +22,9 @@ using namespace AscendC;
22constexpr uint32_t DYN_SORT_DB_BUF = 1;22constexpr uint32_t DYN_SORT_DB_BUF = 1;
23constexpr uint32_t SORT_PADDING = 64;23constexpr uint32_t SORT_PADDING = 64;
24constexpr uint32_t HELP_FRE = 2;24constexpr uint32_t HELP_FRE = 2;
25+static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};
25 26 
26-template <typename X_T, typename IDS_T, uint8_t Mode>27+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
27class KernelSimdDynSort28class KernelSimdDynSort
28{29{
29public:30public:
@@ -35,13 +36,14 @@ public:
35 __aicore__ inline void CopyResToGm(uint32_t cols, uint32_t colsAlign, uint64_t ubOffset, int64_t& arNum);36 __aicore__ inline void CopyResToGm(uint32_t cols, uint32_t colsAlign, uint64_t ubOffset, int64_t& arNum);
36 __aicore__ inline void Compute();37 __aicore__ inline void Compute();
37 __aicore__ inline void Process();38 __aicore__ inline void Process();
38- 39+ using CAST_T = typename CastType<IDS_T, CAST_MODE>::type;
39private:40private:
40 AscendC::GlobalTensor<X_T> xGm_;41 AscendC::GlobalTensor<X_T> xGm_;
41 AscendC::GlobalTensor<X_T> yGm_;42 AscendC::GlobalTensor<X_T> yGm_;
42 AscendC::GlobalTensor<IDS_T> idsGm_;43 AscendC::GlobalTensor<IDS_T> idsGm_;
43 TQue<QuePosition::VECIN, DYN_SORT_DB_BUF> xQue_;44 TQue<QuePosition::VECIN, DYN_SORT_DB_BUF> xQue_;
44 TQue<QuePosition::VECIN, DYN_SORT_DB_BUF> idsQue_;45 TQue<QuePosition::VECIN, DYN_SORT_DB_BUF> idsQue_;
46+ TQue<QuePosition::VECIN, DYN_SORT_DB_BUF> idsCastQue_;
45 TQue<QuePosition::VECOUT, 1> outQueueRes_;47 TQue<QuePosition::VECOUT, 1> outQueueRes_;
46 TBuf<QuePosition::VECCALC> noDupBuf_;48 TBuf<QuePosition::VECCALC> noDupBuf_;
47 TBuf<QuePosition::VECCALC> sortedIdxBuf_;49 TBuf<QuePosition::VECCALC> sortedIdxBuf_;
@@ -50,12 +52,11 @@ private:
50 TPipe* pipe_ = nullptr;52 TPipe* pipe_ = nullptr;
51 const UnsortedSegmentSimdDynSortTilingData* td_;53 const UnsortedSegmentSimdDynSortTilingData* td_;
52 static constexpr uint32_t vfLengthX_ = VF_SIZE / sizeof(X_T);54 static constexpr uint32_t vfLengthX_ = VF_SIZE / sizeof(X_T);
53- static constexpr uint32_t shiftOffset_ = ONE_BLOCK_SIZE / sizeof(IDS_T);55+ static constexpr uint32_t shiftOffset_ = ONE_BLOCK_SIZE / sizeof(CAST_T);
54- static constexpr uint32_t vfIndicesNum_ = VF_SIZE / sizeof(IDS_T);
55};56};
56 57 
57-template <typename X_T, typename IDS_T, uint8_t Mode>58+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
58-__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR output)59+__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode, CAST_MODE>::Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR output)
59{60{
60 InitGm<X_T, Mode>(output, td_->outputOuterDim * td_->innerDim);61 InitGm<X_T, Mode>(output, td_->outputOuterDim * td_->innerDim);
61 62 
@@ -67,17 +68,22 @@ __aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::Init(GM_ADDR x, GM_A
67 pipe_->InitBuffer(outQueueRes_, 1, td_->sortBaseS * td_->sortBaseA * sizeof(X_T));68 pipe_->InitBuffer(outQueueRes_, 1, td_->sortBaseS * td_->sortBaseA * sizeof(X_T));
68 69 
69 uint64_t idsAlignB32 = Aligned(static_cast<uint64_t>(td_->sortBaseS * sizeof(uint32_t)), ONE_BLOCK_SIZE);70 uint64_t idsAlignB32 = Aligned(static_cast<uint64_t>(td_->sortBaseS * sizeof(uint32_t)), ONE_BLOCK_SIZE);
70- uint64_t idsAlign = Aligned(static_cast<uint64_t>(td_->sortBaseS * sizeof(IDS_T)), ONE_BLOCK_SIZE);71+ uint64_t idsAlign = Aligned(static_cast<uint64_t>(td_->sortBaseS * sizeof(CAST_T)), ONE_BLOCK_SIZE);
71 72 
72 pipe_->InitBuffer(idsQue_, DYN_SORT_DB_BUF, idsAlign);73 pipe_->InitBuffer(idsQue_, DYN_SORT_DB_BUF, idsAlign);
73 pipe_->InitBuffer(noDupBuf_, idsAlignB32 + SORT_PADDING);74 pipe_->InitBuffer(noDupBuf_, idsAlignB32 + SORT_PADDING);
74- pipe_->InitBuffer(sortedKeyBuf_, idsAlign + SORT_PADDING);75+ if constexpr (CAST_MODE == CAST_0) {
76+ pipe_->InitBuffer(sortedKeyBuf_, idsAlign + SORT_PADDING);
77+ } else {
78+ pipe_->InitBuffer(idsCastQue_, DYN_SORT_DB_BUF, idsAlign);
79+ pipe_->InitBuffer(sortedKeyBuf_, idsAlign + SORT_PADDING);
80+ }
75 pipe_->InitBuffer(sortedIdxBuf_, idsAlignB32);81 pipe_->InitBuffer(sortedIdxBuf_, idsAlignB32);
76 pipe_->InitBuffer(sharedTmpBuf_, Aligned(static_cast<uint64_t>(td_->sortSharedBufSize), ONE_BLOCK_SIZE));82 pipe_->InitBuffer(sharedTmpBuf_, Aligned(static_cast<uint64_t>(td_->sortSharedBufSize), ONE_BLOCK_SIZE));
77}83}
78 84 
79-template <typename X_T, typename IDS_T, uint8_t Mode>85+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
80-__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::ProcessIndices(86+__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode, CAST_MODE>::ProcessIndices(
81 uint64_t blockOffsetIdx, uint64_t sLoop, uint32_t rows, int64_t& arNum)87 uint64_t blockOffsetIdx, uint64_t sLoop, uint32_t rows, int64_t& arNum)
82{88{
83 LocalTensor<IDS_T> idsLocal = idsQue_.AllocTensor<IDS_T>();89 LocalTensor<IDS_T> idsLocal = idsQue_.AllocTensor<IDS_T>();
@@ -86,24 +92,30 @@ __aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::ProcessIndices(
86 92 
87 idsLocal = idsQue_.DeQue<IDS_T>();93 idsLocal = idsQue_.DeQue<IDS_T>();
88 LocalTensor<uint32_t> sortedIdxLocal = sortedIdxBuf_.Get<uint32_t>();94 LocalTensor<uint32_t> sortedIdxLocal = sortedIdxBuf_.Get<uint32_t>();
89- LocalTensor<IDS_T> sortedKeyLocal = sortedKeyBuf_.Get<IDS_T>();95+ LocalTensor<CAST_T> sortedKeyLocal = sortedKeyBuf_.Get<CAST_T>();
90 LocalTensor<uint8_t> sharedTmpBuffer = sharedTmpBuf_.Get<uint8_t>();96 LocalTensor<uint8_t> sharedTmpBuffer = sharedTmpBuf_.Get<uint8_t>();
91- 
92- Duplicate(sortedKeyLocal, static_cast<IDS_T>(-1), static_cast<uint32_t>(shiftOffset_ * HELP_FRE + rows));
93- LocalTensor<IDS_T> sortedDstLocal = sortedKeyLocal[shiftOffset_];
94- static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};
95- AscendC::Sort<IDS_T, false, sortConfig>(sortedDstLocal, sortedIdxLocal, idsLocal, sharedTmpBuffer, rows);
96- 
97 LocalTensor<int32_t> noDupRes = noDupBuf_.Get<int32_t>();97 LocalTensor<int32_t> noDupRes = noDupBuf_.Get<int32_t>();
98+ 
99+ Duplicate(sortedKeyLocal, static_cast<CAST_T>(-1), static_cast<uint32_t>(shiftOffset_ * HELP_FRE + rows));
100+ LocalTensor<CAST_T> sortedDstLocal = sortedKeyLocal[shiftOffset_];
101+ 
102+ if constexpr (CAST_MODE == CAST_0) {
103+ AscendC::Sort<IDS_T, false, sortConfig>(sortedDstLocal, sortedIdxLocal, idsLocal, sharedTmpBuffer, rows);
104+ } else {
105+ LocalTensor<CAST_T> idsCastLocal = idsCastQue_.AllocTensor<CAST_T>();
106+ IndicesSortCast<IDS_T, CAST_T, CAST_MODE>(idsLocal, idsCastLocal, noDupRes, rows);
107+ AscendC::Sort<CAST_T, false, sortConfig>(sortedDstLocal, sortedIdxLocal, idsCastLocal, sharedTmpBuffer, rows);
108+ idsCastQue_.FreeTensor(idsCastLocal);
109+ }
110+ 
98 Duplicate(noDupRes, 0, rows);111 Duplicate(noDupRes, 0, rows);
99- UniqueGetElm(sortedKeyLocal, noDupRes, rows, shiftOffset_, vfIndicesNum_, arNum);112+ arNum = UniqueGetElm<CAST_T>(sortedKeyLocal, noDupRes, rows);
100- arNum = ((AscendC::MicroAPI::GetSpr<AscendC::SpecialPurposeReg::AR>()) / sizeof(int32_t)) - 1;
101 UniqueStat(noDupRes, arNum);113 UniqueStat(noDupRes, arNum);
102 idsQue_.FreeTensor(idsLocal);114 idsQue_.FreeTensor(idsLocal);
103}115}
104 116 
105-template <typename X_T, typename IDS_T, uint8_t Mode>117+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
106-__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::ComputeXSum(uint32_t cols, uint32_t colsAlign, int64_t arNum)118+__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode, CAST_MODE>::ComputeXSum(uint32_t cols, uint32_t colsAlign, int64_t arNum)
107{119{
108 LocalTensor<uint32_t> sortedIdxLocal = sortedIdxBuf_.template Get<uint32_t>();120 LocalTensor<uint32_t> sortedIdxLocal = sortedIdxBuf_.template Get<uint32_t>();
109 LocalTensor<int32_t> noDupRes = noDupBuf_.template Get<int32_t>();121 LocalTensor<int32_t> noDupRes = noDupBuf_.template Get<int32_t>();
@@ -132,12 +144,12 @@ __aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::ComputeXSum(uint32_t
132 outQueueRes_.EnQue<X_T>(resLocal);144 outQueueRes_.EnQue<X_T>(resLocal);
133}145}
134 146 
135-template <typename X_T, typename IDS_T, uint8_t Mode>147+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
136-__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::CopyResToGm(148+__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode, CAST_MODE>::CopyResToGm(
137 uint32_t cols, uint32_t colsAlign, uint64_t ubOffset, int64_t& arNum)149 uint32_t cols, uint32_t colsAlign, uint64_t ubOffset, int64_t& arNum)
138{150{
139 LocalTensor<X_T> resLocal = outQueueRes_.DeQue<X_T>();151 LocalTensor<X_T> resLocal = outQueueRes_.DeQue<X_T>();
140- LocalTensor<IDS_T> sortedKeyLocal = sortedKeyBuf_.Get<IDS_T>();152+ LocalTensor<CAST_T> sortedKeyLocal = sortedKeyBuf_.Get<CAST_T>();
141 LocalTensor<int32_t> noDupRes = noDupBuf_.Get<int32_t>();153 LocalTensor<int32_t> noDupRes = noDupBuf_.Get<int32_t>();
142 154 
143 int32_t tmpIndex = shiftOffset_;155 int32_t tmpIndex = shiftOffset_;
@@ -160,8 +172,8 @@ __aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::CopyResToGm(
160 outQueueRes_.FreeTensor(resLocal);172 outQueueRes_.FreeTensor(resLocal);
161}173}
162 174 
163-template <typename X_T, typename IDS_T, uint8_t Mode>175+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
164-__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::Compute()176+__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode, CAST_MODE>::Compute()
165{177{
166 uint64_t sIdx = GetBlockIdx() / td_->aTileNum;178 uint64_t sIdx = GetBlockIdx() / td_->aTileNum;
167 uint64_t aIdx = GetBlockIdx() % td_->aTileNum;179 uint64_t aIdx = GetBlockIdx() % td_->aTileNum;
@@ -195,8 +207,8 @@ __aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::Compute()
195 }207 }
196}208}
197 209 
198-template <typename X_T, typename IDS_T, uint8_t Mode>210+template <typename X_T, typename IDS_T, uint8_t Mode, uint32_t CAST_MODE>
199-__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode>::Process()211+__aicore__ inline void KernelSimdDynSort<X_T, IDS_T, Mode, CAST_MODE>::Process()
200{212{
201 if (GetBlockIdx() >= GetBlockNum()) {213 if (GetBlockIdx() >= GetBlockNum()) {
202 return;214 return;
Mindex/unsorted_segment_common/op_kernel/arch35/unsorted_segment_struct.h+1-0
@@ -75,6 +75,7 @@ struct UnsortedSegmentSimdDynSortTilingData{
75 uint64_t sortBaseS;75 uint64_t sortBaseS;
76 uint64_t sortBaseA;76 uint64_t sortBaseA;
77 int64_t sortSharedBufSize;77 int64_t sortSharedBufSize;
78+ uint64_t idCastMode;
78};79};
79 80 
80struct UnsortedSegmentOutFlTilingData{81struct UnsortedSegmentOutFlTilingData{
Mindex/unsorted_segment_min/op_kernel/unsorted_segment_min_apt.cpp+35-5
@@ -32,6 +32,40 @@ using namespace UnsortedSegment;
32#define TEMPLATE_SORT_SIMT 410032#define TEMPLATE_SORT_SIMT 4100
33constexpr uint8_t MODE_FLAGE = 0; // 0:unsorted_segment_min;33constexpr uint8_t MODE_FLAGE = 0; // 0:unsorted_segment_min;
34 34 
35+template <typename X_T, typename SEGMENT_IDS_T, uint8_t MODE>
36+ __aicore__ inline void KernelSimdDynSortWithCast(
37+ GM_ADDR x, GM_ADDR segment_ids, GM_ADDR num_segments, GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling, TPipe &pipe)
38+{
39+ REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 7000", UnsortedSegmentSimdDynSortTilingData);
40+ GET_TILING_DATA_WITH_STRUCT(UnsortedSegmentSimdDynSortTilingData, tilingData, tiling);
41+ uint32_t cast_mode = tilingData.idCastMode;
42+ if (cast_mode == CAST_1) {
43+ UnsortedSegment::KernelSimdDynSort<X_T, SEGMENT_IDS_T, MODE, CAST_1> op(&tilingData, &pipe);
44+ op.Init(x, segment_ids, output);
45+ op.Process();
46+ } else if (cast_mode == CAST_2) {
47+ UnsortedSegment::KernelSimdDynSort<X_T, SEGMENT_IDS_T, MODE, CAST_2> op(&tilingData, &pipe);
48+ op.Init(x, segment_ids, output);
49+ op.Process();
50+ } else if (cast_mode == CAST_3) {
51+ UnsortedSegment::KernelSimdDynSort<X_T, SEGMENT_IDS_T, MODE, CAST_3> op(&tilingData, &pipe);
52+ op.Init(x, segment_ids, output);
53+ op.Process();
54+ } else if (cast_mode == CAST_4) {
55+ UnsortedSegment::KernelSimdDynSort<X_T, SEGMENT_IDS_T, MODE, CAST_4> op(&tilingData, &pipe);
56+ op.Init(x, segment_ids, output);
57+ op.Process();
58+ } else if (cast_mode == CAST_5) {
59+ UnsortedSegment::KernelSimdDynSort<X_T, SEGMENT_IDS_T, MODE, CAST_5> op(&tilingData, &pipe);
60+ op.Init(x, segment_ids, output);
61+ op.Process();
62+ } else {
63+ UnsortedSegment::KernelSimdDynSort<X_T, SEGMENT_IDS_T, MODE, CAST_0> op(&tilingData, &pipe);
64+ op.Init(x, segment_ids, output);
65+ op.Process();
66+ }
67+}
68+ 
35extern "C" __global__ __aicore__ void unsorted_segment_min(GM_ADDR x, GM_ADDR segment_ids, GM_ADDR num_segments,69extern "C" __global__ __aicore__ void unsorted_segment_min(GM_ADDR x, GM_ADDR segment_ids, GM_ADDR num_segments,
36 GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling)70 GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling)
37{71{
@@ -69,11 +103,7 @@ extern "C" __global__ __aicore__ void unsorted_segment_min(GM_ADDR x, GM_ADDR se
69 std::is_same<int64_t, DTYPE_X>::value) {103 std::is_same<int64_t, DTYPE_X>::value) {
70 return;104 return;
71 } else {105 } else {
72- REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 7000", UnsortedSegmentSimdDynSortTilingData);106+ KernelSimdDynSortWithCast<DTYPE_X, DTYPE_SEGMENT_IDS, MODE_FLAGE>(x, segment_ids, num_segments, output, workspace, tiling, pipe);
73- GET_TILING_DATA_WITH_STRUCT(UnsortedSegmentSimdDynSortTilingData, tilingData, tiling);
74- UnsortedSegment::KernelSimdDynSort<DTYPE_X, DTYPE_SEGMENT_IDS, MODE_FLAGE> op(&tilingData, &pipe);
75- op.Init(x, segment_ids, output);
76- op.Process();
77 }107 }
78 } else if (TILING_KEY_IS(TEMPLATE_ADD_TILING_KEY)) {108 } else if (TILING_KEY_IS(TEMPLATE_ADD_TILING_KEY)) {
79 REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 4000", UnsortedSegmentOutFlTilingData);109 REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 4000", UnsortedSegmentOutFlTilingData);