Star890
1800
代码
代码
Issues265
Pull Requests320
Actions
项目讨论
Wiki
分析
Star890
1800
已合并
segment_sum算子支持simd确定性计算以及性能优化 #1831
z30075199创建于 2月11日
segment_sum算子支持simd确定性计算以及性能优化 #1831
已合并
z30075199创建于 2月11日
从z30075199/ops-nn_zxzwb_simt合入到cann/ops-nnmaster
共 15 个文件变更+899-93
Mindex/gather_v2/op_api/aclnn_gather_v3.cpp+1-1
@@ -203,7 +203,7 @@ static aclnnStatus CheckParams(const aclTensor *self, int64_t dim, const aclTens
203 203 
204 auto indexSize = 1;204 auto indexSize = 1;
205 auto indexDimNum = index->GetViewShape().GetDimNum();205 auto indexDimNum = index->GetViewShape().GetDimNum();
206- for (int i = 0; i < indexDimNum ; i++) {206+ for (int i = 0; i < static_cast<int>(indexDimNum) ; i++) {
207 indexSize *= index->GetViewShape().GetDim(i);207 indexSize *= index->GetViewShape().GetDim(i);
208 }208 }
209 209 
Mindex/gather_v2/op_host/gather_v2_infershape.cpp+2-2
@@ -51,8 +51,8 @@ static std::string OtherErrMsg(const std::string& error_detail) {
51 return msg;51 return msg;
52}52}
53 53 
54-static bool CheckAndUpdateAxis(gert::InferShapeContext* context, int64_t& batch_dims, int64_t& axes_data,54+static bool CheckAndUpdateAxis(const gert::InferShapeContext* context, int64_t& batch_dims, int64_t& axes_data,
55- GatherInfo& gather_info) {55+ const GatherInfo& gather_info) {
56 int64_t x_real_dim_cnt = gather_info.x_real_dim_cnt;56 int64_t x_real_dim_cnt = gather_info.x_real_dim_cnt;
57 int64_t index_batch_dims = gather_info.index_batch_dims;57 int64_t index_batch_dims = gather_info.index_batch_dims;
58 int64_t rank_indices = gather_info.rank_indices;58 int64_t rank_indices = gather_info.rank_indices;
Mindex/gather_v2/op_kernel/arch35/gather_v2_simd.h+1-1
@@ -158,7 +158,7 @@ __aicore__ inline INDICES_T Gatherv2Simd<INDICES_T>::GetIndex(int64_t idx, int64
158 copyLen = startBIdx + maxBSize > endBIdx ? (endBIdx - startBIdx + 1) * gatherSize_ : maxBSize * gatherSize_;158 copyLen = startBIdx + maxBSize > endBIdx ? (endBIdx - startBIdx + 1) * gatherSize_ : maxBSize * gatherSize_;
159 indicesOffsetBase_ = startBIdx * gatherSize_;159 indicesOffsetBase_ = startBIdx * gatherSize_;
160 } else {160 } else {
161- copyLen = startGIdx + maxIndex_ >= gatherSize_ ? (endBIdx + 1) * gatherSize_ - idx: maxIndex_;161+ copyLen = startGIdx + maxIndex_ >= (endBIdx + 1) * gatherSize_ ? (endBIdx + 1) * gatherSize_ - idx: maxIndex_;
162 indicesOffsetBase_ = idx;162 indicesOffsetBase_ = idx;
163 }163 }
164 curIndexSize_ = copyLen;164 curIndexSize_ = copyLen;
Mindex/segment_sum/op_host/arch35/segment_sum_simd_tiling.cpp+20-15
@@ -19,7 +19,7 @@ namespace optiling {
19 19 
20const static uint64_t SIMD_ATOMIC_SUPPORT_KEY = 2000;20const static uint64_t SIMD_ATOMIC_SUPPORT_KEY = 2000;
21const static uint64_t SIMD_DETERM_KEY = 2002;21const static uint64_t SIMD_DETERM_KEY = 2002;
22-static constexpr int64_t SIMD_INNER_THRES = 32;22+static constexpr int64_t SIMD_INNER_THRES = 1024;
23static constexpr int64_t BLOCK_TILING_THRES = 512;23static constexpr int64_t BLOCK_TILING_THRES = 512;
24static constexpr int64_t INNER_ADD_NUM = 128;24static constexpr int64_t INNER_ADD_NUM = 128;
25static constexpr int64_t BASE_BLOCK_ALIGN = 512;25static constexpr int64_t BASE_BLOCK_ALIGN = 512;
@@ -41,11 +41,9 @@ bool SegmentSumSimdTiling::IsAtomicSupport()
41 41 
42bool SegmentSumSimdTiling::IsCapable() 42bool SegmentSumSimdTiling::IsCapable()
43{43{
44- bool isFloat = (dataType_ == ge::DT_FLOAT || dataType_ == ge::DT_FLOAT16 || dataType_ == ge::DT_BF16);
45 isAtomicSupport_ = IsAtomicSupport();44 isAtomicSupport_ = IsAtomicSupport();
46- isDeterministic_ = context_->GetDeterministic() == 1 && isFloat;45+ bool isInnerSimd = innerDim_ > SIMD_INNER_THRES;
47- bool isInnerSimd = innerDim_ * valueTypeBytes_ > SIMD_INNER_THRES;46+ return isInnerSimd;
48- return isInnerSimd && isAtomicSupport_;
49}47}
50 48 
51 49 
@@ -152,7 +150,7 @@ void SegmentSumSimdTiling::DoSplitColUBTiling(int64_t availableUbsize)
152 int64_t tmpColSize = BLOCK_TILING_THRES + idTypeBytes_;150 int64_t tmpColSize = BLOCK_TILING_THRES + idTypeBytes_;
153 int64_t tmpRowNum = availableUbsize / tmpColSize;151 int64_t tmpRowNum = availableUbsize / tmpColSize;
154 while (tmpRowNum - 1 > normalCoreOutterNum_) {152 while (tmpRowNum - 1 > normalCoreOutterNum_) {
155- if (tmpColSize >= innerSizeAlign + idTypeBytes_) {153+ if (tmpColSize >= static_cast<int64_t>(innerSizeAlign + idTypeBytes_)) {
156 break;154 break;
157 }155 }
158 tmpColSize += INNER_ADD_NUM;156 tmpColSize += INNER_ADD_NUM;
@@ -187,6 +185,14 @@ void SegmentSumSimdTiling::DoUBTiling()
187{185{
188 int64_t availableUbsize = ubSize_;186 int64_t availableUbsize = ubSize_;
189 187 
188+ bool isFloat = (dataType_ == ge::DT_FLOAT || dataType_ == ge::DT_FLOAT16 || dataType_ == ge::DT_BF16);
189+ isDeterministic_ = context_->GetDeterministic() == 1 && isFloat && blockNumInRow_ != 1;
190+ 
191+ // 确定性或者atomicAdd不支持的类型，需要一块32B buffer放头尾id
192+ if (!isAtomicSupport_ || isDeterministic_) {
193+ availableUbsize -= ubBlockSize_;
194+ }
195+ 
190 availableUbsize -= ubBlockSize_ - idTypeBytes_; // 预留给segmentId对齐196 availableUbsize -= ubBlockSize_ - idTypeBytes_; // 预留给segmentId对齐
191 // 如果inner小于512B，分核及切UB均不使用inner197 // 如果inner小于512B，分核及切UB均不使用inner
192 if (normalCoreInnerNum_ * valueTypeBytes_ <= BLOCK_TILING_THRES) {198 if (normalCoreInnerNum_ * valueTypeBytes_ <= BLOCK_TILING_THRES) {
@@ -231,6 +237,8 @@ void SegmentSumSimdTiling::DoMultCoreAddTiling()
231 tailCoreMultAddInners_ = innerDim_ - (usedCoreNumForMultAdd_ - 1) * normalCoreMultAddInners_;237 tailCoreMultAddInners_ = innerDim_ - (usedCoreNumForMultAdd_ - 1) * normalCoreMultAddInners_;
232 }238 }
233 int64_t mulAddUbsize = ubSize_;239 int64_t mulAddUbsize = ubSize_;
240+ multAddIdsBufferSize_ = Ops::Base::CeilAlign(NUM_TWO * blockNumInRow_ * idTypeBytes_, ubBlockSize_);
241+ mulAddUbsize -= multAddIdsBufferSize_;
234 mulAddUbsize /= blockNumInRow_ * NUM_TWO + 1;242 mulAddUbsize /= blockNumInRow_ * NUM_TWO + 1;
235 int64_t availableInnerUb = Ops::Base::FloorAlign(mulAddUbsize, static_cast<int64_t>(ubBlockSize_));243 int64_t availableInnerUb = Ops::Base::FloorAlign(mulAddUbsize, static_cast<int64_t>(ubBlockSize_));
236 244 
@@ -266,7 +274,6 @@ void SegmentSumSimdTiling::SetTilingData()
266{274{
267 tilingData_ = context_->GetTilingData<SegmentSumSimdTilingData>();275 tilingData_ = context_->GetTilingData<SegmentSumSimdTilingData>();
268 tilingData_->needCoreNum = needCoreNum_;276 tilingData_->needCoreNum = needCoreNum_;
269- tilingData_->outerDim = outerDim_;
270 tilingData_->innerDim = innerDim_;277 tilingData_->innerDim = innerDim_;
271 278 
272 tilingData_->xBufferSize = xBufferSize_;279 tilingData_->xBufferSize = xBufferSize_;
@@ -281,9 +288,7 @@ void SegmentSumSimdTiling::SetTilingData()
281 tilingData_->blockNumInCol = blockNumInCol_;288 tilingData_->blockNumInCol = blockNumInCol_;
282 289 
283 tilingData_->normalCoreInnerNum = normalCoreInnerNum_;290 tilingData_->normalCoreInnerNum = normalCoreInnerNum_;
284- tilingData_->tailCoreInnerNum = tailCoreInnerNum_;
285 tilingData_->normalCoreOutterNum = normalCoreOutterNum_;291 tilingData_->normalCoreOutterNum = normalCoreOutterNum_;
286- tilingData_->tailCoreOutterNum = tailCoreOutterNum_;
287 292 
288 tilingData_->normalCoreRowUbLoop = normalCoreRowUbLoop_;293 tilingData_->normalCoreRowUbLoop = normalCoreRowUbLoop_;
289 tilingData_->normalCoreNormalLoopOutters = normalCoreNormalLoopOutters_;294 tilingData_->normalCoreNormalLoopOutters = normalCoreNormalLoopOutters_;
@@ -301,7 +306,6 @@ void SegmentSumSimdTiling::SetTilingData()
301 306 
302 tilingData_->usedCoreNumForMultAdd = usedCoreNumForMultAdd_;307 tilingData_->usedCoreNumForMultAdd = usedCoreNumForMultAdd_;
303 tilingData_->normalCoreMultAddInners = normalCoreMultAddInners_;308 tilingData_->normalCoreMultAddInners = normalCoreMultAddInners_;
304- tilingData_->tailCoreMultAddInners = tailCoreMultAddInners_;
305 309 
306 tilingData_->normalCoreMultAddInnerLoop = normalCoreMultAddInnerLoop_;310 tilingData_->normalCoreMultAddInnerLoop = normalCoreMultAddInnerLoop_;
307 tilingData_->normalCoreMultAddNormalLoopInners = normalCoreMultAddNormalLoopInners_;311 tilingData_->normalCoreMultAddNormalLoopInners = normalCoreMultAddNormalLoopInners_;
@@ -311,18 +315,22 @@ void SegmentSumSimdTiling::SetTilingData()
311 tilingData_->tailCoreMultAddTailLoopInners = tailCoreMultAddTailLoopInners_;315 tilingData_->tailCoreMultAddTailLoopInners = tailCoreMultAddTailLoopInners_;
312 316 
313 tilingData_->multAddXBufferSize = multAddXBufferSize_;317 tilingData_->multAddXBufferSize = multAddXBufferSize_;
318+ tilingData_->multAddIdsBufferSize = multAddIdsBufferSize_;
314 tilingData_->multAddYBufferSize = multAddYBufferSize_;319 tilingData_->multAddYBufferSize = multAddYBufferSize_;
315 return;320 return;
316}321}
317 322 
318uint64_t SegmentSumSimdTiling::GetTilingKey() const323uint64_t SegmentSumSimdTiling::GetTilingKey() const
319{324{
325+ if (isDeterministic_ || !isAtomicSupport_) {
326+ return SIMD_DETERM_KEY;
327+ }
320 return SIMD_ATOMIC_SUPPORT_KEY;328 return SIMD_ATOMIC_SUPPORT_KEY;
321}329}
322 330 
323ge::graphStatus SegmentSumSimdTiling::GetWorkspaceSize()331ge::graphStatus SegmentSumSimdTiling::GetWorkspaceSize()
324{332{
325- auto useWorkspace = WS_SYS_SIZE; // 可以不用初值333+ size_t useWorkspace = WS_SYS_SIZE; // 可以不用初值
326 if (isDeterministic_ || !isAtomicSupport_) {334 if (isDeterministic_ || !isAtomicSupport_) {
327 useWorkspace += blockNumInRow_ * NUM_TWO * (innerDim_ * valueTypeBytes_ + idTypeBytes_) + idTypeBytes_; // 对齐idTypeBytes_ 头尾id最好需要间隔 cache line335 useWorkspace += blockNumInRow_ * NUM_TWO * (innerDim_ * valueTypeBytes_ + idTypeBytes_) + idTypeBytes_; // 对齐idTypeBytes_ 头尾id最好需要间隔 cache line
328 }336 }
@@ -346,7 +354,6 @@ void SegmentSumSimdTiling::DumpTilingInfo()
346 info << "tilingKey: " << GetTilingKey();354 info << "tilingKey: " << GetTilingKey();
347 info << ", UB Size: " << ubSize_;355 info << ", UB Size: " << ubSize_;
348 info << ", needCoreNum: " << tilingData_->needCoreNum;356 info << ", needCoreNum: " << tilingData_->needCoreNum;
349- info << ", outerDim: " << tilingData_->outerDim;
350 info << ", innerDim: " << tilingData_->innerDim;357 info << ", innerDim: " << tilingData_->innerDim;
351 358 
352 info << ", xBufferSize: " << tilingData_->xBufferSize;359 info << ", xBufferSize: " << tilingData_->xBufferSize;
@@ -361,9 +368,7 @@ void SegmentSumSimdTiling::DumpTilingInfo()
361 info << ", blockNumInCol: " << tilingData_->blockNumInCol;368 info << ", blockNumInCol: " << tilingData_->blockNumInCol;
362 369 
363 info << ", normalCoreInnerNum: " << tilingData_->normalCoreInnerNum;370 info << ", normalCoreInnerNum: " << tilingData_->normalCoreInnerNum;
364- info << ", tailCoreInnerNum: " << tilingData_->tailCoreInnerNum;
365 info << ", normalCoreOutterNum: " << tilingData_->normalCoreOutterNum;371 info << ", normalCoreOutterNum: " << tilingData_->normalCoreOutterNum;
366- info << ", tailCoreOutterNum: " << tilingData_->tailCoreOutterNum;
367 372
368 info << ", normalCoreRowUbLoop: " << tilingData_->normalCoreRowUbLoop;373 info << ", normalCoreRowUbLoop: " << tilingData_->normalCoreRowUbLoop;
369 info << ", normalCoreNormalLoopOutters: " << tilingData_->normalCoreNormalLoopOutters;374 info << ", normalCoreNormalLoopOutters: " << tilingData_->normalCoreNormalLoopOutters;
@@ -381,7 +386,6 @@ void SegmentSumSimdTiling::DumpTilingInfo()
381 386 
382 info << ", usedCoreNumForMultAdd: " << tilingData_->usedCoreNumForMultAdd;387 info << ", usedCoreNumForMultAdd: " << tilingData_->usedCoreNumForMultAdd;
383 info << ", normalCoreMultAddInners: " << tilingData_->normalCoreMultAddInners;388 info << ", normalCoreMultAddInners: " << tilingData_->normalCoreMultAddInners;
384- info << ", tailCoreMultAddInners: " << tilingData_->tailCoreMultAddInners;
385 389 
386 info << ", normalCoreMultAddInnerLoop: " << tilingData_->normalCoreMultAddInnerLoop;390 info << ", normalCoreMultAddInnerLoop: " << tilingData_->normalCoreMultAddInnerLoop;
387 info << ", normalCoreMultAddNormalLoopInners: " << tilingData_->normalCoreMultAddNormalLoopInners;391 info << ", normalCoreMultAddNormalLoopInners: " << tilingData_->normalCoreMultAddNormalLoopInners;
@@ -391,6 +395,7 @@ void SegmentSumSimdTiling::DumpTilingInfo()
391 info << ", tailCoreMultAddTailLoopInners: " << tilingData_->tailCoreMultAddTailLoopInners;395 info << ", tailCoreMultAddTailLoopInners: " << tilingData_->tailCoreMultAddTailLoopInners;
392 396 
393 info << ", multAddXBufferSize: " << tilingData_->multAddXBufferSize;397 info << ", multAddXBufferSize: " << tilingData_->multAddXBufferSize;
398+ info << ", multAddIdsBufferSize: " << tilingData_->multAddIdsBufferSize;
394 info << ", multAddYBufferSize: " << tilingData_->multAddYBufferSize;399 info << ", multAddYBufferSize: " << tilingData_->multAddYBufferSize;
395 400
396 OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());401 OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());
Mindex/segment_sum/op_host/arch35/segment_sum_simd_tiling.h+4-3
@@ -57,8 +57,8 @@ private:
57 int64_t blockNumInCol_ = 0;57 int64_t blockNumInCol_ = 0;
58 int64_t normalCoreInnerNum_ = 0;58 int64_t normalCoreInnerNum_ = 0;
59 int64_t normalCoreOutterNum_ = 0;59 int64_t normalCoreOutterNum_ = 0;
60- int64_t tailCoreInnerNum_ = 0;60+ int64_t tailCoreInnerNum_ = 0; // 列尾核列上处理的inner数
61- int64_t tailCoreOutterNum_ = 0;61+ int64_t tailCoreOutterNum_ = 0; // 行尾核行上处理的行数
CANN-robot
CANN-robotCANN-robot2月11日

代码结构与可维护性: 变量名 'tailCoreOutterNum_' 存在拼写错误。'Outter' 应为 'Outer'。该拼写错误与同一类中的其他变量名（如 'normalCoreOutterNum_'）不一致，且与常见的 'inner'/'outer' 命名惯例不符。这种不一致性会降低代码的可读性和可维护性，可能导致开发者在理解变量含义时产生困惑。

问题类型: 代码结构与可维护性 文件路径: index/segment_sum/op_host/arch35/segment_sum_simd_tiling.h 行号: 61 问题代码:

int64_t tailCoreOutterNum_ = 0; // 行尾核行上处理的行数

修改建议:

将变量名 'tailCoreOutterNum_' 和 'normalCoreOutterNum_' 统一更正为 'tailCoreOuterNum_' 和 'normalCoreOuterNum_'，以遵循正确的英文拼写和一致的命名规范。

此评论由代码审查工具自动生成

likedislike
62 62 
63 int64_t normalCoreRowUbLoop_ = 0;63 int64_t normalCoreRowUbLoop_ = 0;
64 int64_t normalCoreNormalLoopOutters_ = 0;64 int64_t normalCoreNormalLoopOutters_ = 0;
@@ -80,7 +80,7 @@ private:
80 80 
81 int64_t usedCoreNumForMultAdd_ = 0;81 int64_t usedCoreNumForMultAdd_ = 0;
82 int64_t normalCoreMultAddInners_ = 0;82 int64_t normalCoreMultAddInners_ = 0;
83- int64_t tailCoreMultAddInners_ = 0;83+ int64_t tailCoreMultAddInners_ = 0; // 多核累加尾核处理的inner数
84 84 
85 int64_t normalCoreMultAddInnerLoop_ = 0;85 int64_t normalCoreMultAddInnerLoop_ = 0;
86 int64_t normalCoreMultAddNormalLoopInners_ = 0;86 int64_t normalCoreMultAddNormalLoopInners_ = 0;
@@ -90,6 +90,7 @@ private:
90 int64_t tailCoreMultAddTailLoopInners_ = 0;90 int64_t tailCoreMultAddTailLoopInners_ = 0;
91 91 
92 int64_t multAddXBufferSize_ = 0;92 int64_t multAddXBufferSize_ = 0;
93+ int64_t multAddIdsBufferSize_ = 0;
93 int64_t multAddYBufferSize_ = 0;94 int64_t multAddYBufferSize_ = 0;
94 95 
95 bool isDeterministic_ = false;96 bool isDeterministic_ = false;
Mindex/segment_sum/op_host/arch35/segment_sum_simt_tiling.cpp+43-4
@@ -19,6 +19,10 @@ namespace optiling {
19 19 
20static constexpr uint64_t SIMT_DCACHE_SIZE = static_cast<uint64_t>(32 * 1024);20static constexpr uint64_t SIMT_DCACHE_SIZE = static_cast<uint64_t>(32 * 1024);
21static constexpr uint64_t TEMPLATE_SIMT = 1000;21static constexpr uint64_t TEMPLATE_SIMT = 1000;
22+static constexpr uint32_t DOUBLE = 2;
23+static constexpr uint32_t ROWS_IN_WORKSPACE = 128;
24+static constexpr uint32_t RESERVED_WS_SIZE = 16 * 1024 * 1024;
25+static const std::set<ge::DataType> deterministicType = {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16};
22 26 
23bool SegmentSumSimtTiling::IsCapable()27bool SegmentSumSimtTiling::IsCapable()
24{28{
@@ -28,9 +32,23 @@ bool SegmentSumSimtTiling::IsCapable()
28ge::graphStatus SegmentSumSimtTiling::DoOpTiling()32ge::graphStatus SegmentSumSimtTiling::DoOpTiling()
29{33{
30 ubSize_ -= SIMT_DCACHE_SIZE;34 ubSize_ -= SIMT_DCACHE_SIZE;
35+ if (context_->GetDeterministic() && deterministicType.find(dataType_) != deterministicType.end()) {
36+ isDeterministic_ = 1;
37+ }
31 uint64_t outputSize = segmentNum_ * innerDim_;38 uint64_t outputSize = segmentNum_ * innerDim_;
32 initNumPerCore_ = outputSize / totalCoreNum_;39 initNumPerCore_ = outputSize / totalCoreNum_;
33 initNumTailCore_ = outputSize - (totalCoreNum_ - 1) * initNumPerCore_;40 initNumTailCore_ = outputSize - (totalCoreNum_ - 1) * initNumPerCore_;
41+
42+ segIdsPerCore_ = outerDim_ / totalCoreNum_;
43+ segIdsTailCore_ = outerDim_ - segIdsPerCore_ * (totalCoreNum_-1);
44+ maxSegIdsInUb = (ubSize_ / DOUBLE - ubBlockSize_) / idTypeBytes_;
45+ segIdsPerLoop_ = maxSegIdsInUb > segIdsPerCore_ ? segIdsPerCore_ : maxSegIdsInUb;
46+ segIdsPerLoopTailCore_ = maxSegIdsInUb > segIdsTailCore_ ? segIdsTailCore_ : maxSegIdsInUb;
47+ loopTimes_ = Ops::Base::CeilDiv(segIdsPerCore_, static_cast<uint64_t>(segIdsPerLoop_));
48+ loopTimesTailCore_ = Ops::Base::CeilDiv(segIdsTailCore_, static_cast<uint64_t>(segIdsPerLoopTailCore_));
49+ segIdsTailLoop_ = segIdsPerCore_ - (loopTimes_ - 1) * segIdsPerLoop_;
50+ segIdsTailLoopTailCore_ = segIdsTailCore_ - (loopTimesTailCore_ - 1) * segIdsPerLoopTailCore_;
51+ 
34 SetTilingData();52 SetTilingData();
35 return ge::GRAPH_SUCCESS;53 return ge::GRAPH_SUCCESS;
36}54}
@@ -42,18 +60,30 @@ void SegmentSumSimtTiling::SetTilingData()
42 tilingData_->innerDim = innerDim_;60 tilingData_->innerDim = innerDim_;
43 tilingData_->initNumPerCore = initNumPerCore_;61 tilingData_->initNumPerCore = initNumPerCore_;
44 tilingData_->initNumTailCore = initNumTailCore_;62 tilingData_->initNumTailCore = initNumTailCore_;
63+ tilingData_->isDeterministic = isDeterministic_;
64+ tilingData_->maxSegIdsInUb = maxSegIdsInUb;
65+ tilingData_->loopTimes = loopTimes_;
66+ tilingData_->loopTimesTailCore = loopTimesTailCore_;
67+ tilingData_->segIdsPerLoop = segIdsPerLoop_;
68+ tilingData_->segIdsPerLoopTailCore = segIdsPerLoopTailCore_;
69+ tilingData_->segIdsTailLoop = segIdsTailLoop_;
70+ tilingData_->segIdsTailLoopTailCore = segIdsTailLoopTailCore_;
45}71}
46 72 
47uint64_t SegmentSumSimtTiling::GetTilingKey() const73uint64_t SegmentSumSimtTiling::GetTilingKey() const
48{74{
49- uint64_t tilingKey = TEMPLATE_SIMT;75+ return TEMPLATE_SIMT;
50- return tilingKey;
51}76}
52 77 
53ge::graphStatus SegmentSumSimtTiling::GetWorkspaceSize()78ge::graphStatus SegmentSumSimtTiling::GetWorkspaceSize()
54{79{
55 auto currentWorkspace = context_->GetWorkspaceSizes(1);80 auto currentWorkspace = context_->GetWorkspaceSizes(1);
56- currentWorkspace[0] = 0;81+ currentWorkspace[0] = RESERVED_WS_SIZE;
82+ if (isDeterministic_ == 1) {
83+ uint64_t ws_size = ROWS_IN_WORKSPACE * innerDim_ * valueTypeBytes_ +
84+ ROWS_IN_WORKSPACE * idTypeBytes_;
85+ currentWorkspace[0] += ws_size ;
86+ }
57 return ge::GRAPH_SUCCESS;87 return ge::GRAPH_SUCCESS;
58}88}
59 89 
@@ -74,7 +104,16 @@ void SegmentSumSimtTiling::DumpTilingInfo()
74 info << ", outerDim: " << outerDim_;104 info << ", outerDim: " << outerDim_;
75 info << ", innerDim: " << innerDim_;105 info << ", innerDim: " << innerDim_;
76 info << ", initNumPerCore: " << initNumPerCore_;106 info << ", initNumPerCore: " << initNumPerCore_;
77- info << ", initNumTailCore: " <<initNumTailCore_;107+ info << ", initNumTailCore: " << initNumTailCore_;
108+ info << ", isDeterministic: " << isDeterministic_;
109+ info << ", segIdsPerCore: " << segIdsPerCore_;
110+ info << ", segIdsTailCore: " << segIdsTailCore_;
111+ info << ", segIdsPerLoop: " << segIdsPerLoop_;
112+ info << ", segIdsPerLoopTailCore: " << segIdsPerLoopTailCore_;
113+ info << ", segIdsTailLoop: " << segIdsTailLoop_;
114+ info << ", segIdsTailLoopTailCore: " << segIdsTailLoopTailCore_;
115+ info << ", loopTimes: " << loopTimes_;
116+ info << ", loopTimesTailCore: " << loopTimesTailCore_;
78 OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());117 OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());
79}118}
80 119 
Mindex/segment_sum/op_host/arch35/segment_sum_simt_tiling.h+10-0
@@ -44,6 +44,16 @@ private:
44 SegmentSumSimtTilingData* tilingData_;44 SegmentSumSimtTilingData* tilingData_;
45 uint64_t initNumPerCore_{0};45 uint64_t initNumPerCore_{0};
46 uint64_t initNumTailCore_{0};46 uint64_t initNumTailCore_{0};
47+ uint32_t isDeterministic_{0};
48+ uint32_t maxSegIdsInUb{0};
CANN-robot
CANN-robotCANN-robot2月11日

命名规范与代码可读性: 新增的成员变量 maxSegIdsInUb 命名风格与其他变量不一致。该类中其他成员变量均使用后缀下划线 _ 来标识成员变量（如 initNumPerCore_、initNumTailCore_、isDeterministic_ 等），但 maxSegIdsInUb 缺少此后缀。这种不一致性会破坏代码风格的统一性，可能影响团队协作时的代码可读性和维护性。统一的命名约定有助于快速区分局部变量、参数和成员变量。

问题类型: 命名规范与代码可读性 文件路径: index/segment_sum/op_host/arch35/segment_sum_simt_tiling.h 行号: 48 问题代码:

    uint32_t maxSegIdsInUb{0};

修改建议:

将变量名修改为与其他成员变量一致的风格，添加下划线后缀，即改为 `maxSegIdsInUb_`。

此评论由代码审查工具自动生成

likedislike
49+ uint64_t segIdsPerCore_{0};
50+ uint64_t segIdsTailCore_{0};
51+ uint32_t segIdsPerLoop_{0};
52+ uint32_t segIdsPerLoopTailCore_{0};
53+ uint32_t segIdsTailLoop_{0};
54+ uint32_t segIdsTailLoopTailCore_{0};
55+ uint32_t loopTimes_{0};
56+ uint32_t loopTimesTailCore_{0};
47};57};
48} // namespace optiling58} // namespace optiling
49#endif // SEGMENT_SUM_SIMT_TILING_H59#endif // SEGMENT_SUM_SIMT_TILING_H
Mindex/segment_sum/op_host/arch35/segment_sum_tiling_base.cpp+1-1
@@ -95,7 +95,7 @@ ge::graphStatus SegmentSumBaseTiling::GetShapeAttrsInfo()
95 outerDim_ = dataShape.GetDim(0);95 outerDim_ = dataShape.GetDim(0);
96 96
97 OP_CHECK_IF(97 OP_CHECK_IF(
98- outerDim_ != segmentIdsShape.GetDim(0),98+ static_cast<int64_t>(outerDim_) != segmentIdsShape.GetDim(0),
99 OP_LOGE(99 OP_LOGE(
100 context_->GetNodeName(), "the dimension 0 of data shape should be same with segment_ids."),100 context_->GetNodeName(), "the dimension 0 of data shape should be same with segment_ids."),
101 return ge::GRAPH_FAILED);101 return ge::GRAPH_FAILED);
Mindex/segment_sum/op_kernel/arch35/segment_sum_simd.h+18-15
@@ -27,7 +27,7 @@ class SegmentSumSimd
27{27{
28public:28public:
29 __aicore__ inline SegmentSumSimd(void){};29 __aicore__ inline SegmentSumSimd(void){};
30- __aicore__ inline void Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR y, GM_ADDR workspace, TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData);30+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR y, TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData);
31 __aicore__ inline void Process();31 __aicore__ inline void Process();
32 __aicore__ inline void CopyInX(int32_t copyCount, int32_t burstLen, int64_t xGmOffset);32 __aicore__ inline void CopyInX(int32_t copyCount, int32_t burstLen, int64_t xGmOffset);
33 __aicore__ inline void CopyInSegmentIds(int32_t burstLen, int64_t segmentIdsGmOffset);33 __aicore__ inline void CopyInSegmentIds(int32_t burstLen, int64_t segmentIdsGmOffset);
@@ -72,7 +72,7 @@ private:
72 72 
73template <typename T1, typename T2>73template <typename T1, typename T2>
74__aicore__ inline void SegmentSumSimd<T1, T2>::Init(74__aicore__ inline void SegmentSumSimd<T1, T2>::Init(
75- GM_ADDR x, GM_ADDR segmentIds, GM_ADDR y, GM_ADDR workspace, AscendC::TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData)75+ GM_ADDR x, GM_ADDR segmentIds, GM_ADDR y, AscendC::TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData)
76{76{
77 tilingData_ = tilingData;77 tilingData_ = tilingData;
78 blockIdx_ = GetBlockIdx();78 blockIdx_ = GetBlockIdx();
@@ -173,9 +173,8 @@ __aicore__ inline void SegmentSumSimd<T1, T2>::ComputeSumAndCopyOut(LocalTensor<
173 if (curId == preId_) {173 if (curId == preId_) {
174 Add(yLocal, yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);174 Add(yLocal, yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
175 } else if (curId != preId_ && preId_ == -1) {175 } else if (curId != preId_ && preId_ == -1) {
176- DataCopy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInnersAlign); // 偏移使用 xLocal[i * curLoopInners]也行，DataCopy搬运数量必须向上对齐176+ Copy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
177 preId_ = curId;177 preId_ = curId;
178- // Add(yLocal, yLocal, xLocal[i * curLoopInnersAlign], curLoopInners); // 也可以yBuf初始清零然后Add
179 } else { // curId != preId_178 } else { // curId != preId_
180 event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));179 event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
181 SetFlag<HardEvent::V_MTE3>(eventId);180 SetFlag<HardEvent::V_MTE3>(eventId);
@@ -193,7 +192,7 @@ __aicore__ inline void SegmentSumSimd<T1, T2>::ComputeSumAndCopyOut(LocalTensor<
193 WaitFlag<HardEvent::MTE3_V>(eventId1);192 WaitFlag<HardEvent::MTE3_V>(eventId1);
194 193 
195 preId_ = curId;194 preId_ = curId;
196- DataCopy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInnersAlign); // 上对齐32B195+ Copy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
197 }196 }
198 }197 }
199 198
@@ -210,21 +209,25 @@ __aicore__ inline void SegmentSumSimd<T1, T2>::Process()
210 }209 }
211 210 
212 LocalTensor<T1> yLocal = yBuf_.Get<T1>();211 LocalTensor<T1> yLocal = yBuf_.Get<T1>();
213- // Duplicate<T1>(yLocal, 0, normalLoopInners_);212+ 
213+ T2 curId;
214+ int64_t curLoopInners;
215+ int64_t curLoopInnersAlign;
216+ int64_t colOffset;
217+ int64_t curLoopOutters;
218+ int64_t xGmOffset;
219+ int64_t segmentIdsGmOffset;
214 220 
215 for (int64_t col = 0; col < colUbLoop_; col++) {221 for (int64_t col = 0; col < colUbLoop_; col++) {
216 isFirstId_ = true;222 isFirstId_ = true;
217 preId_ = -1;223 preId_ = -1;
218- T2 curId;224+ curLoopInners = col == colUbLoop_ - 1 ? tailLoopInners_ : normalLoopInners_;
219- int64_t curLoopInners;225+ curLoopInnersAlign = (curLoopInners + blockNumT1_ - 1) / blockNumT1_ * blockNumT1_;
220- int64_t colOffset = col * normalLoopInners_;226+ colOffset = col * normalLoopInners_;
221 for (int64_t row = 0; row < rowUbLoop_; row++) {227 for (int64_t row = 0; row < rowUbLoop_; row++) {
222- int64_t curLoopOutters = row == rowUbLoop_ - 1 ? tailLoopOutters_ : normalLoopOutters_;228+ curLoopOutters = row == rowUbLoop_ - 1 ? tailLoopOutters_ : normalLoopOutters_;
223- curLoopInners = col == colUbLoop_ - 1 ? tailLoopInners_ : normalLoopInners_;229+ xGmOffset = row * normalLoopOutters_ * tilingData_->innerDim + col * normalLoopInners_;
224- int64_t curLoopInnersAlign = (curLoopInners + blockNumT1_ - 1) / blockNumT1_ * blockNumT1_;230+ segmentIdsGmOffset = row * normalLoopOutters_;
225- 
226- int64_t xGmOffset = row * normalLoopOutters_ * tilingData_->innerDim + col * normalLoopInners_;
227- int64_t segmentIdsGmOffset = row * normalLoopOutters_;
228 231 
229 CopyInX(curLoopOutters, curLoopInners, xGmOffset);232 CopyInX(curLoopOutters, curLoopInners, xGmOffset);
230 CopyInSegmentIds(curLoopOutters, segmentIdsGmOffset);233 CopyInSegmentIds(curLoopOutters, segmentIdsGmOffset);
Aindex/segment_sum/op_kernel/arch35/segment_sum_simd_determ.h+290-0
@@ -0,0 +1,290 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef SEGMENT_SUM_SIMD_DETERM_H
12+#define SEGMENT_SUM_SIMD_DETERM_H
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "../inc/platform.h"
17+#include "segment_sum_struct.h"
18+ 
19+namespace SegmentSum {
20+using namespace AscendC;
21+ 
22+constexpr uint32_t DOUBLE = 2;
23+ 
24+ 
25+template <typename T1, typename T2>
26+class SegmentSumSimdDeterm
27+{
28+public:
29+ __aicore__ inline SegmentSumSimdDeterm(void){};
30+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR y, GM_ADDR workspace, TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData);
31+ __aicore__ inline void Process();
32+ __aicore__ inline void CopyInX(int32_t copyCount, int32_t burstLen, int64_t xGmOffset);
33+ __aicore__ inline void CopyInSegmentIds(int32_t burstLen, int64_t segmentIdsGmOffset);
34+ __aicore__ inline void CopyOutY(LocalTensor<T1>& yLocal, int32_t burstLen, T2 id, int64_t colOffset);
35+ __aicore__ inline void ComputeSumAndCopyOut(LocalTensor<T1>& yLocal, int32_t curLoopOutters, int32_t curLoopInners, int32_t curLoopInnersAlign, int64_t colOffset, T2& curId);
36+ __aicore__ inline void CopyOutSegIdWorkspace(LocalTensor<T2>& tmpLocal);
37+ __aicore__ inline void CopyOutSumWorkspace(LocalTensor<T1>& yLocal, int32_t burstLen, int64_t colOffset, int32_t writePostion);
CANN-robot
CANN-robotCANN-robot2月11日

拼写错误: 函数参数名'writePostion'拼写错误，应为'writePosition'。

问题类型: 拼写错误 文件路径: index/segment_sum/op_kernel/arch35/segment_sum_simd_determ.h 行号: 37 问题代码:

__aicore__ inline void CopyOutSumWorkspace(LocalTensor<T1>& yLocal, int32_t burstLen, int64_t colOffset, int32_t writePostion)

修改建议:

将参数名更正为'writePosition'。

此评论由代码审查工具自动生成

likedislike
38+ 
39+ 
40+private:
41+ GlobalTensor<T1> xGm_;
42+ GlobalTensor<T2> segmentIdsGm_;
43+ GlobalTensor<T1> yGm_;
44+ GlobalTensor<T1> sumWorkspace_;
45+ GlobalTensor<T2> segIdWorkspace_;
46+
47+ TQue<QuePosition::VECIN, BUFFER_NUM> xQue_;
48+ TQue<QuePosition::VECIN, BUFFER_NUM> segmentIdsQue_;
49+ TBuf<QuePosition::VECCALC> yBuf_;
50+ TBuf<QuePosition::VECCALC> tmpBuf_;
51+ 
52+ const SegmentSumSimdTilingData* tilingData_;
53+ 
54+ uint32_t blockIdx_ = 0;
55+ uint32_t rowCoreIdx_ = 0; // 行核idx
56+ uint32_t colCoreIdx_ = 0; // 列核idx
57+ 
58+ int64_t rowGmOffset_ = 0; // 当前核处理的数据块在GM行上的偏移
59+ int64_t colGmOffset_ = 0; // 当前核处理的数据块在GM列上的偏移
60+ 
61+ int64_t rowUbLoop_ = 0; // 当前核的ub在行上的循环次数
62+ int64_t colUbLoop_ = 0; // 当前核的ub在列上的循环次数
63+ 
64+ int64_t normalLoopOutters_ = 0; // 当前核ub正常循环一次处理的行数
65+ int64_t tailLoopOutters_ = 0; // 当前核ub尾循环一次处理的行数
66+ int64_t normalLoopInners_ = 0; // 当前核ub正常循环一次处理的列数
67+ int64_t tailLoopInners_ = 0; // 当前核ub尾循环一次处理的列数
68+ 
69+ T2 preId_ = -1;
70+ T2 position0_ = -1;
71+ T2 position1_ = -1;
72+ bool isStartRowCore_ = false;
73+ bool isEndRowCore_ = false;
74+ bool isFirstId_ = true;
75+ constexpr static int32_t blockNumT1_ = platform::GetUbBlockSize() / sizeof(T1);
76+ // constexpr static int32_t BLOCK_SIZE = platform::GetUbBlockSize();
77+ 
78+};
79+ 
80+template <typename T1, typename T2>
81+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::Init(
82+ GM_ADDR x, GM_ADDR segmentIds, GM_ADDR y, GM_ADDR workspace, AscendC::TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData)
83+{
84+ tilingData_ = tilingData;
85+ blockIdx_ = GetBlockIdx();
86+ 
87+ if (blockIdx_ >= tilingData_->needCoreNum) {
88+ return;
89+ }
90+ 
91+ rowCoreIdx_ = blockIdx_ / tilingData_->blockNumInCol;
92+ colCoreIdx_ = blockIdx_ % tilingData_->blockNumInCol;
93+ isStartRowCore_ = rowCoreIdx_ == 0; // 首行核
94+ isEndRowCore_ = rowCoreIdx_ == tilingData_->blockNumInRow - 1; // 尾行核
95+ 
96+ rowGmOffset_ = rowCoreIdx_ * tilingData_->normalCoreOutterNum;
97+ colGmOffset_ = colCoreIdx_ * tilingData_->normalCoreInnerNum;
98+ 
99+ rowUbLoop_ = rowCoreIdx_ == tilingData_->blockNumInRow - 1 ? tilingData_->tailCoreRowUbLoop : tilingData_->normalCoreRowUbLoop;
100+ colUbLoop_ = colCoreIdx_ == tilingData_->blockNumInCol - 1 ? tilingData_->tailCoreColUbLoop : tilingData_->normalCoreColUbLoop;
101+ 
102+ normalLoopOutters_ = rowCoreIdx_ == tilingData_->blockNumInRow - 1 ? tilingData_->tailCoreNormalLoopOutters : tilingData_->normalCoreNormalLoopOutters;
103+ tailLoopOutters_ = rowCoreIdx_ == tilingData_->blockNumInRow - 1 ? tilingData_->tailCoreTailLoopOutters : tilingData_->normalCoreTailLoopOutters;
104+ normalLoopInners_ = colCoreIdx_ == tilingData_->blockNumInCol - 1 ? tilingData_->tailCoreNormalLoopInners : tilingData_->normalCoreNormalLoopInners;
105+ tailLoopInners_ = colCoreIdx_ == tilingData_->blockNumInCol - 1 ? tilingData_->tailCoreTailLoopInners : tilingData_->normalCoreTailLoopInners;
106+ uint32_t segIdAddrOffset = (tilingData_->blockNumInRow * DOUBLE * tilingData_->innerDim * sizeof(T1) + sizeof(T2) - 1) / sizeof(T2);
107+
108+ xGm_.SetGlobalBuffer((__gm__ T1*)x + rowGmOffset_ * tilingData_->innerDim + colGmOffset_);
109+ segmentIdsGm_.SetGlobalBuffer((__gm__ T2*)segmentIds + rowGmOffset_);
110+ yGm_.SetGlobalBuffer((__gm__ T1*)y + colGmOffset_);
111+ sumWorkspace_.SetGlobalBuffer((__gm__ T1*)workspace + rowCoreIdx_ * DOUBLE * tilingData_->innerDim + colGmOffset_);
112+ segIdWorkspace_.SetGlobalBuffer((__gm__ T2*)workspace + segIdAddrOffset + rowCoreIdx_ * DOUBLE);
113+ 
114+ pipeIn.InitBuffer(xQue_, BUFFER_NUM, tilingData_->xBufferSize);
115+ pipeIn.InitBuffer(segmentIdsQue_, BUFFER_NUM, tilingData_->segmentIdBufferSize);
116+ pipeIn.InitBuffer(yBuf_, tilingData_->yBufferSize);
117+ pipeIn.InitBuffer(tmpBuf_, platform::GetUbBlockSize()); // 放头尾id
118+ 
119+}
120+ 
121+template <typename T1, typename T2>
122+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::CopyInX(int32_t copyCount, int32_t burstLen, int64_t xGmOffset)
123+{
124+ int64_t gmStride = tilingData_->innerDim - burstLen;
125+ LocalTensor<T1> xLocal = xQue_.AllocTensor<T1>();
126+ DataCopyPadExtParams<T1> dataCopyPadExtParams;
127+ dataCopyPadExtParams.isPad = false;
128+ dataCopyPadExtParams.leftPadding = 0;
129+ dataCopyPadExtParams.rightPadding = 0;
130+ dataCopyPadExtParams.paddingValue = 0;
131+ 
132+ DataCopyExtParams dataCoptExtParams;
133+ dataCoptExtParams.blockCount = copyCount;
134+ dataCoptExtParams.blockLen = burstLen * sizeof(T1);
135+ dataCoptExtParams.srcStride = gmStride * sizeof(T1);
136+ dataCoptExtParams.dstStride = 0;
137+ DataCopyPad(xLocal, xGm_[xGmOffset], dataCoptExtParams, dataCopyPadExtParams);
138+ xQue_.EnQue(xLocal);
139+}
140+ 
141+template <typename T1, typename T2>
142+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::CopyInSegmentIds(int32_t burstLen, int64_t segmentIdsGmOffset)
143+{
144+ LocalTensor<T2> segmentIdsLocal = segmentIdsQue_.AllocTensor<T2>();
145+ DataCopyPadExtParams<T2> dataCopyPadExtParams;
146+ dataCopyPadExtParams.isPad = false;
147+ dataCopyPadExtParams.leftPadding = 0;
148+ dataCopyPadExtParams.rightPadding = 0;
149+ dataCopyPadExtParams.paddingValue = 0;
150+ 
151+ DataCopyExtParams dataCoptExtParams;
152+ dataCoptExtParams.blockCount = 1;
153+ dataCoptExtParams.blockLen = burstLen * sizeof(T2);
154+ dataCoptExtParams.srcStride = 0;
155+ dataCoptExtParams.dstStride = 0;
156+ DataCopyPad(segmentIdsLocal, segmentIdsGm_[segmentIdsGmOffset], dataCoptExtParams, dataCopyPadExtParams);
157+ segmentIdsQue_.EnQue(segmentIdsLocal);
158+}
159+ 
160+template <typename T1, typename T2>
161+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::CopyOutY(LocalTensor<T1>& yLocal, int32_t burstLen, T2 id, int64_t colOffset)
162+{
163+ DataCopyExtParams dataCoptExtParams;
164+ dataCoptExtParams.blockCount = 1;
165+ dataCoptExtParams.blockLen = burstLen * sizeof(T1);
166+ dataCoptExtParams.srcStride = 0;
167+ dataCoptExtParams.dstStride = 0;
168+ DataCopyPad(yGm_[id * tilingData_->innerDim + colOffset], yLocal, dataCoptExtParams);
169+}
170+ 
171+template <typename T1, typename T2>
172+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::CopyOutSegIdWorkspace(LocalTensor<T2>& tmpLocal)
173+{
174+ DataCopyExtParams dataCoptExtParams;
175+ dataCoptExtParams.blockCount = 1;
176+ dataCoptExtParams.blockLen = DOUBLE * sizeof(T2);
177+ dataCoptExtParams.srcStride = 0;
178+ dataCoptExtParams.dstStride = 0;
179+ DataCopyPad(segIdWorkspace_, tmpLocal, dataCoptExtParams);
180+}
181+ 
182+template <typename T1, typename T2>
183+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::CopyOutSumWorkspace(LocalTensor<T1>& yLocal, int32_t burstLen, int64_t colOffset, int32_t writePostion)
184+{
185+ DataCopyExtParams dataCoptExtParams;
186+ dataCoptExtParams.blockCount = 1;
187+ dataCoptExtParams.blockLen = burstLen * sizeof(T1);
188+ dataCoptExtParams.srcStride = 0;
189+ dataCoptExtParams.dstStride = 0;
190+ DataCopyPad(sumWorkspace_[writePostion * tilingData_->innerDim + colOffset], yLocal, dataCoptExtParams);
191+}
192+ 
193+template <typename T1, typename T2>
194+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::ComputeSumAndCopyOut(LocalTensor<T1>& yLocal, int32_t curLoopOutters, int32_t curLoopInners, int32_t curLoopInnersAlign, int64_t colOffset, T2& curId)
195+{
196+ LocalTensor<T1> xLocal = xQue_.DeQue<T1>();
197+ LocalTensor<T2> segmentIdsLocal = segmentIdsQue_.DeQue<T2>();
198+ 
199+ // GetValue获取id前需要插同步
200+ event_t eventId3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
201+ SetFlag<HardEvent::MTE2_S>(eventId3);
202+ WaitFlag<HardEvent::MTE2_S>(eventId3);
203+ 
204+ for (int32_t i = 0; i < curLoopOutters; i++) {
205+ curId = segmentIdsLocal.GetValue(i);
206+ if (curId == preId_) {
207+ Add(yLocal, yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
208+ } else if (curId != preId_ && preId_ == -1) {
209+ Copy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
210+ preId_ = curId;
211+ } else { // curId != preId_
212+ event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
213+ SetFlag<HardEvent::V_MTE3>(eventId);
214+ WaitFlag<HardEvent::V_MTE3>(eventId);
215+ if (isFirstId_ && !isStartRowCore_) {
216+ CopyOutSumWorkspace(yLocal, curLoopInners, colOffset, 0);
217+ position0_ = preId_;
218+ isFirstId_ = false;
219+ } else {
220+ CopyOutY(yLocal, curLoopInners, preId_, colOffset);
221+ }
222+ event_t eventId1 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
223+ SetFlag<HardEvent::MTE3_V>(eventId1);
224+ WaitFlag<HardEvent::MTE3_V>(eventId1);
225+ 
226+ preId_ = curId;
227+ Copy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
228+ }
229+ }
230+ xQue_.FreeTensor(xLocal);
231+ segmentIdsQue_.FreeTensor(segmentIdsLocal);
232+}
233+ 
234+ 
235+template <typename T1, typename T2>
236+__aicore__ inline void SegmentSumSimdDeterm<T1, T2>::Process()
237+{
238+ if (blockIdx_ >= tilingData_->needCoreNum) {
239+ return;
240+ }
241+ 
242+ LocalTensor<T1> yLocal = yBuf_.Get<T1>();
243+ 
244+ int64_t curLoopInners;
245+ int64_t curLoopInnersAlign;
246+ int64_t colOffset;
247+ int64_t xGmOffset;
248+ int64_t segmentIdsGmOffset;
249+ int64_t curLoopOutters;
250+ T2 curId;
251+ 
252+ for (int64_t col = 0; col < colUbLoop_; col++) {
253+ isFirstId_ = true;
254+ preId_ = -1;
255+ curLoopInners = col == colUbLoop_ - 1 ? tailLoopInners_ : normalLoopInners_;
256+ curLoopInnersAlign = (curLoopInners + blockNumT1_ - 1) / blockNumT1_ * blockNumT1_;
257+ colOffset = col * normalLoopInners_;
258+ for (int64_t row = 0; row < rowUbLoop_; row++) {
259+ curLoopOutters = row == rowUbLoop_ - 1 ? tailLoopOutters_ : normalLoopOutters_;
260+ xGmOffset = row * normalLoopOutters_ * tilingData_->innerDim + col * normalLoopInners_;
261+ segmentIdsGmOffset = row * normalLoopOutters_;
262+ 
263+ CopyInX(curLoopOutters, curLoopInners, xGmOffset);
264+ CopyInSegmentIds(curLoopOutters, segmentIdsGmOffset);
265+ 
266+ ComputeSumAndCopyOut(yLocal, curLoopOutters, curLoopInners, curLoopInnersAlign, colOffset, curId);
267+ }
268+ event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
269+ SetFlag<HardEvent::V_MTE3>(eventId);
270+ WaitFlag<HardEvent::V_MTE3>(eventId);
271+ // PipeBarrier<PIPE_ALL>();
272+ if (isEndRowCore_ && !isFirstId_) {
273+ CopyOutY(yLocal, curLoopInners, curId, colOffset);
274+ } else {
275+ CopyOutSumWorkspace(yLocal, curLoopInners, colOffset, 1);
276+ position1_ = curId;
277+ }
278+ }
279+ LocalTensor<T2> tmpLocal = tmpBuf_.Get<T2>();
280+ tmpLocal.SetValue(0, position0_);
281+ tmpLocal.SetValue(1, position1_);
282+ event_t eventIdSToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
283+ SetFlag<HardEvent::S_MTE3>(eventIdSToMte3);
284+ WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3);
285+ CopyOutSegIdWorkspace(tmpLocal);
286+}
287+ 
288+ 
289+}
290+#endif
Aindex/segment_sum/op_kernel/arch35/segment_sum_simd_mult_core_add.h+206-0
@@ -0,0 +1,206 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef SEGMENT_SUM_SIMD_MULT_CORE_ADD_H
12+#define SEGMENT_SUM_SIMD_MULT_CORE_ADD_H
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "../inc/platform.h"
17+#include "segment_sum_struct.h"
18+ 
19+namespace SegmentSum {
20+using namespace AscendC;
21+ 
22+template <typename T1, typename T2>
23+class SegmentSumMultiCoreAdd
24+{
25+public:
26+ __aicore__ inline SegmentSumMultiCoreAdd(void){};
27+ __aicore__ inline void Init(GM_ADDR y, GM_ADDR workspace, TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData);
28+ __aicore__ inline void Process();
29+ __aicore__ inline void CopyInIds(LocalTensor<T2>& segmentIdsLocal);
30+ __aicore__ inline void CopyInSum(int32_t burstLen, int64_t colOffset);
31+ __aicore__ inline void CopyOutY(LocalTensor<T1>& yLocal, int32_t burstLen, T2 id, int64_t colOffset);
32+ __aicore__ inline void ComputeAndCopyOut(LocalTensor<T2> segmentIdsLocal, LocalTensor<T1>& yLocal, int32_t curLoopInners, int32_t curLoopInnersAlign, int64_t colOffset);
33+ 
34+private:
35+ GlobalTensor<T1> yGm_;
36+ GlobalTensor<T1> sumWorkspace_;
37+ GlobalTensor<T2> segIdWorkspace_;
38+ 
39+ TQue<QuePosition::VECIN, BUFFER_NUM> xQue_;
40+ TBuf<QuePosition::VECCALC> segmentIdsBuf_;
41+ TBuf<QuePosition::VECCALC> yBuf_;
42+ 
43+ const SegmentSumSimdTilingData* tilingData_;
44+ 
45+ uint32_t blockIdx_ = 0;
46+ 
47+ int64_t colGmOffset_ = 0; // 当前核处理的数据块在GM列上的偏移
48+ 
49+ int64_t rowUbLoop_ = 0; // 当前核的ub在行上的循环次数
50+ int64_t colUbLoop_ = 0; // 当前核的ub在列上的循环次数
51+ 
52+ int64_t normalLoopInners_ = 0; // 当前核ub正常循环一次处理的列数
53+ int64_t tailLoopInners_ = 0; // 当前核ub尾循环一次处理的列数
54+ 
55+ constexpr static int32_t blockNumT1_ = platform::GetUbBlockSize() / sizeof(T1);
56+};
57+ 
58+template <typename T1, typename T2>
59+__aicore__ inline void SegmentSumMultiCoreAdd<T1, T2>::Init(GM_ADDR y, GM_ADDR workspace, AscendC::TPipe& pipeIn, const SegmentSumSimdTilingData* tilingData)
60+{
61+ tilingData_ = tilingData;
62+ blockIdx_ = GetBlockIdx();
63+ 
64+ if (blockIdx_ >= tilingData_->usedCoreNumForMultAdd) {
65+ return;
66+ }
67+ 
68+ colGmOffset_ = blockIdx_ * tilingData_->normalCoreMultAddInners;
69+ rowUbLoop_ = DOUBLE * tilingData_->blockNumInRow; // no need
70+ colUbLoop_ = blockIdx_ == tilingData_->usedCoreNumForMultAdd - 1 ? tilingData_->tailCoreMultAddInnerLoop : tilingData_->normalCoreMultAddInnerLoop;
71+ 
72+ normalLoopInners_ = blockIdx_ == tilingData_->usedCoreNumForMultAdd - 1 ? tilingData_->tailCoreMultAddNormalLoopInners : tilingData_->normalCoreMultAddNormalLoopInners;
73+ tailLoopInners_ = blockIdx_ == tilingData_->usedCoreNumForMultAdd - 1 ? tilingData_->tailCoreMultAddTailLoopInners : tilingData_->normalCoreMultAddTailLoopInners;
74+ uint32_t segIdAddrOffset = (tilingData_->blockNumInRow * DOUBLE * tilingData_->innerDim * sizeof(T1) + sizeof(T2) - 1) / sizeof(T2);
75+ 
76+ yGm_.SetGlobalBuffer((__gm__ T1*)y + colGmOffset_);
77+ sumWorkspace_.SetGlobalBuffer((__gm__ T1*)workspace + colGmOffset_);
78+ segIdWorkspace_.SetGlobalBuffer((__gm__ T2*)workspace + segIdAddrOffset);
79+ 
80+ pipeIn.InitBuffer(xQue_, BUFFER_NUM, tilingData_->multAddXBufferSize);
81+ pipeIn.InitBuffer(segmentIdsBuf_, tilingData_->multAddIdsBufferSize);
82+ pipeIn.InitBuffer(yBuf_, tilingData_->multAddYBufferSize);
83+}
84+ 
85+ 
86+template <typename T1, typename T2>
87+__aicore__ inline void SegmentSumMultiCoreAdd<T1, T2>::CopyInIds(LocalTensor<T2>& segmentIdsLocal)
88+{
89+ DataCopyPadExtParams<T2> dataCopyPadExtParams;
90+ dataCopyPadExtParams.isPad = false;
91+ dataCopyPadExtParams.leftPadding = 0;
92+ dataCopyPadExtParams.rightPadding = 0;
93+ dataCopyPadExtParams.paddingValue = 0;
94+ 
95+ DataCopyExtParams dataCoptExtParams;
96+ dataCoptExtParams.blockCount = 1;
97+ dataCoptExtParams.blockLen = tilingData_->blockNumInRow * DOUBLE * sizeof(T2);
98+ dataCoptExtParams.srcStride = 0;
99+ dataCoptExtParams.dstStride = 0;
100+ DataCopyPad(segmentIdsLocal, segIdWorkspace_, dataCoptExtParams, dataCopyPadExtParams);
101+}
102+ 
103+ 
104+template <typename T1, typename T2>
105+__aicore__ inline void SegmentSumMultiCoreAdd<T1, T2>::CopyInSum(int32_t burstLen, int64_t colOffset)
106+{
107+ int64_t gmStride = tilingData_->innerDim - burstLen;
108+ LocalTensor<T1> xLocal = xQue_.AllocTensor<T1>();
109+ DataCopyPadExtParams<T1> dataCopyPadExtParams;
110+ dataCopyPadExtParams.isPad = false;
111+ dataCopyPadExtParams.leftPadding = 0;
112+ dataCopyPadExtParams.rightPadding = 0;
113+ dataCopyPadExtParams.paddingValue = 0;
114+ 
115+ DataCopyExtParams dataCoptExtParams;
116+ dataCoptExtParams.blockCount = tilingData_->blockNumInRow * DOUBLE;
117+ dataCoptExtParams.blockLen = burstLen * sizeof(T1);
118+ dataCoptExtParams.srcStride = gmStride * sizeof(T1);
119+ dataCoptExtParams.dstStride = 0;
120+ DataCopyPad(xLocal, sumWorkspace_[colOffset], dataCoptExtParams, dataCopyPadExtParams);
121+ xQue_.EnQue(xLocal);
122+}
123+ 
124+ 
125+template <typename T1, typename T2>
126+__aicore__ inline void SegmentSumMultiCoreAdd<T1, T2>::CopyOutY(LocalTensor<T1>& yLocal, int32_t burstLen, T2 id, int64_t colOffset)
127+{
128+ DataCopyExtParams dataCoptExtParams;
129+ dataCoptExtParams.blockCount = 1;
130+ dataCoptExtParams.blockLen = burstLen * sizeof(T1);
131+ dataCoptExtParams.srcStride = 0;
132+ dataCoptExtParams.dstStride = 0;
133+ DataCopyPad(yGm_[id * tilingData_->innerDim + colOffset], yLocal, dataCoptExtParams);
134+}
135+ 
136+ 
137+template <typename T1, typename T2>
138+__aicore__ inline void SegmentSumMultiCoreAdd<T1, T2>::ComputeAndCopyOut(LocalTensor<T2> segmentIdsLocal, LocalTensor<T1>& yLocal, int32_t curLoopInners, int32_t curLoopInnersAlign, int64_t colOffset)
139+{
140+ LocalTensor<T1> xLocal = xQue_.DeQue<T1>();
141+ Copy(yLocal, xLocal[curLoopInnersAlign], curLoopInners);
142+
143+ // GetValue获取id前需要插同步
144+ event_t eventId1 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
145+ SetFlag<HardEvent::MTE2_S>(eventId1);
146+ WaitFlag<HardEvent::MTE2_S>(eventId1);
147+ 
148+ T2 preId = segmentIdsLocal.GetValue(1);
149+ 
150+ for (int32_t i = 2; i < tilingData_->blockNumInRow * DOUBLE; i++) {
151+ T2 curId = segmentIdsLocal.GetValue(i);
152+ if (curId == preId) {
153+ Add(yLocal, yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
154+ } else if (curId == -1) {
155+ continue;
156+ } else { // curId != preId
157+ event_t eventId2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
158+ SetFlag<HardEvent::V_MTE3>(eventId2);
159+ WaitFlag<HardEvent::V_MTE3>(eventId2);
160+ 
161+ CopyOutY(yLocal, curLoopInners, preId, colOffset);
162+ preId = curId;
163+ 
164+ event_t eventId3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
165+ SetFlag<HardEvent::MTE3_V>(eventId3);
166+ WaitFlag<HardEvent::MTE3_V>(eventId3);
167+ 
168+ Copy(yLocal, xLocal[i * curLoopInnersAlign], curLoopInners);
169+ }
170+ }
171+ event_t eventId4 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
172+ SetFlag<HardEvent::V_MTE3>(eventId4);
173+ WaitFlag<HardEvent::V_MTE3>(eventId4);
174+
175+ CopyOutY(yLocal, curLoopInners, preId, colOffset);
176+ xQue_.FreeTensor(xLocal);
177+}
178+ 
179+ 
180+template <typename T1, typename T2>
181+__aicore__ inline void SegmentSumMultiCoreAdd<T1, T2>::Process()
182+{
183+ if (blockIdx_ >= tilingData_->usedCoreNumForMultAdd) {
184+ return;
185+ }
186+ 
187+ LocalTensor<T2> segmentIdsLocal = segmentIdsBuf_.Get<T2>();
188+ LocalTensor<T1> yLocal = yBuf_.Get<T1>();
189+ CopyInIds(segmentIdsLocal);
190+ 
191+ int64_t curLoopInners;
192+ int64_t curLoopInnersAlign;
193+ int64_t colOffset;
194+ 
195+ for (int64_t col = 0; col < colUbLoop_; col++) {
196+ curLoopInners = col == colUbLoop_ - 1 ? tailLoopInners_ : normalLoopInners_;
197+ curLoopInnersAlign = (curLoopInners + blockNumT1_ - 1) / blockNumT1_ * blockNumT1_;
198+ colOffset = col * normalLoopInners_;
199+ 
200+ CopyInSum(curLoopInners, colOffset);
201+ ComputeAndCopyOut(segmentIdsLocal, yLocal, curLoopInners, curLoopInnersAlign, colOffset);
202+ }
203+}
204+ 
205+}
206+#endif
Mindex/segment_sum/op_kernel/arch35/segment_sum_simt.h+243-41
@@ -12,17 +12,20 @@
12#define SEGMENT_SUM_SIMT_H12#define SEGMENT_SUM_SIMT_H
13 13 
14#include "kernel_operator.h"14#include "kernel_operator.h"
15-#include "segment_sum_struct.h"15+#include "../inc/platform.h"
16#include "simt_api/common_functions.h"16#include "simt_api/common_functions.h"
17#include "simt_api/asc_simt.h"17#include "simt_api/asc_simt.h"
18#include "simt_api/asc_fp16.h"18#include "simt_api/asc_fp16.h"
19#include "simt_api/asc_bf16.h"19#include "simt_api/asc_bf16.h"
20+#include "segment_sum_struct.h"
20 21 
21 22 
22namespace SegmentSum23namespace SegmentSum
23{24{
24-constexpr int64_t DOUBLE = 2;25+constexpr uint32_t TMP_ROWS_PER_CORE = 2;
26+constexpr uint32_t TMP_ROWS_TOTAL = 128;
25constexpr uint32_t MAX_THREAD_NUM = 2048;27constexpr uint32_t MAX_THREAD_NUM = 2048;
28+constexpr uint32_t UINT32_MAX_VALUE = 4294967295;
26using namespace AscendC;29using namespace AscendC;
27 30 
28template <typename TX, typename Index>31template <typename TX, typename Index>
@@ -31,76 +34,275 @@ class SegmentSumSimt
31public:34public:
32 __aicore__ inline SegmentSumSimt(const SegmentSumSimtTilingData* __restrict tilingData, TPipe* pipe)35 __aicore__ inline SegmentSumSimt(const SegmentSumSimtTilingData* __restrict tilingData, TPipe* pipe)
33 : tilingData_(tilingData), pipe_(pipe){};36 : tilingData_(tilingData), pipe_(pipe){};
34- __aicore__ inline void Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR output);37+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR output, GM_ADDR workspace);
35 __aicore__ inline void Process();38 __aicore__ inline void Process();
39+ __aicore__ inline void CopySegmentIdsIn(LocalTensor<Index> segmentIdsLocal, int64_t offset, int32_t length);
40+ __aicore__ inline int32_t GetUniqueSegIdCount(uint32_t segmentIdsNum);
41+ __aicore__ inline void ProcessEachLoop(uint32_t segmentIdsNum, uint64_t baseOffset);
42+ __aicore__ inline void ProcessEachLoopForDeterminstic(uint32_t segmentIdsNum, uint64_t baseOffset);
36 43 
37private:44private:
38- AscendC::GlobalTensor<TX> xGm, outputGm, outputGmInit;45+ GlobalTensor<TX> xGm, outputGm, outputGmInit, tmpRowWs;
39- AscendC::GlobalTensor<Index> segmentIdsGm;46+ GlobalTensor<Index> segmentIdsGm, tmpIdWs;
47+ TBuf<TPosition::VECCALC> uniqueIdPosBuf_;
48+ TBuf<TPosition::VECCALC> segmentIdsBuf_;
40 TPipe* pipe_;49 TPipe* pipe_;
41 const SegmentSumSimtTilingData* tilingData_;50 const SegmentSumSimtTilingData* tilingData_;
51+ int32_t loopTimes;
52+ uint32_t segIds, tailSegIds; // 整/尾循环处理的id个数
53+ uint32_t threadNum, threadBlock;
54+ LocalTensor<int32_t> uniqueIdPosLocal;
55+ LocalTensor<Index> segmentIdsLocalShift, segmentIdsLocal;
56+ int32_t coreId;
57+ int32_t coreNum;
42};58};
43 59 
44template <typename TX, typename Index>60template <typename TX, typename Index>
45-__aicore__ inline void SegmentSumSimt<TX, Index>::Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR output)61+__aicore__ inline void SegmentSumSimt<TX, Index>::Init(GM_ADDR x, GM_ADDR segmentIds, GM_ADDR output, GM_ADDR workspace)
46{62{
63+ // common
64+ coreId = GetBlockIdx();
65+ coreNum = GetBlockNum();
47 xGm.SetGlobalBuffer((__gm__ TX*)(x));66 xGm.SetGlobalBuffer((__gm__ TX*)(x));
48 segmentIdsGm.SetGlobalBuffer((__gm__ Index*)(segmentIds));67 segmentIdsGm.SetGlobalBuffer((__gm__ Index*)(segmentIds));
49 outputGm.SetGlobalBuffer((__gm__ TX*)(output));68 outputGm.SetGlobalBuffer((__gm__ TX*)(output));
50- 69+ pipe_->InitBuffer(uniqueIdPosBuf_, tilingData_->maxSegIdsInUb * sizeof(Index));
51- uint64_t initCoreReal = GetBlockIdx() == (GetBlockNum() - 1) ? 70+ pipe_->InitBuffer(segmentIdsBuf_, tilingData_->maxSegIdsInUb * sizeof(Index) + platform::GetUbBlockSize());
71+ // determinstic
72+ if (tilingData_->isDeterministic == 1) {
73+ tmpIdWs.SetGlobalBuffer((__gm__ Index*)(workspace));
74+ tmpRowWs.SetGlobalBuffer((__gm__ TX*)((__gm__ Index*)workspace + TMP_ROWS_TOTAL));
75+ }
76+ // clear
77+ uint64_t initCoreReal = coreId == (coreNum - 1) ?
52 tilingData_->initNumTailCore : tilingData_->initNumPerCore;78 tilingData_->initNumTailCore : tilingData_->initNumPerCore;
53- uint64_t outputGmOffset = GetBlockIdx() * tilingData_->initNumPerCore;79+ uint64_t outputGmOffset = coreId * tilingData_->initNumPerCore;
54 outputGmInit.SetGlobalBuffer((__gm__ TX *)(output) + outputGmOffset);80 outputGmInit.SetGlobalBuffer((__gm__ TX *)(output) + outputGmOffset);
55 InitGlobalMemory(outputGmInit, initCoreReal, static_cast<TX>(0));81 InitGlobalMemory(outputGmInit, initCoreReal, static_cast<TX>(0));
56 SyncAll();82 SyncAll();
83+}
84+ 
85+template <typename TX, typename Index>
86+__aicore__ inline void SegmentSumSimt<TX, Index>::CopySegmentIdsIn(
87+ LocalTensor<Index> segmentIdsLocal, int64_t offset, int32_t length)
88+{
89+ DataCopyPadExtParams<Index> dataCopyPadExtParams = {false, 0, 0, 0};
90+ DataCopyExtParams dataCoptExtParams;
CANN-robotCANN-robot
CANN-robotCANN-robot2月11日

代码结构与可维护性: CopySegmentIdsIn函数中，变量名'dataCoptExtParams'拼写错误，应为'dataCopyExtParams'。这种拼写错误虽然不影响功能，但影响代码可读性和维护性。

问题类型: 代码结构与可维护性 文件路径: index/segment_sum/op_kernel/arch35/segment_sum_simt.h 行号: 90 问题代码:

DataCopyExtParams dataCoptExtParams;

修改建议:

将变量名更正为'dataCopyExtParams'以保持一致性。

此评论由代码审查工具自动生成

likedislike
CANN-robotCANN-robot2月11日

代码可读性/维护性: 变量名 'dataCoptExtParams' 存在拼写错误，应为 'dataCopyExtParams'。

问题类型: 代码可读性/维护性 文件路径: index/segment_sum/op_kernel/arch35/segment_sum_simt.h 行号: 90 问题代码:

    DataCopyExtParams dataCoptExtParams;

修改建议:

将变量名更正为 `dataCopyExtParams`。

此评论由代码审查工具自动生成

likedislike
91+ dataCoptExtParams.blockCount = 1;
92+ dataCoptExtParams.blockLen = length * sizeof(Index);
93+ dataCoptExtParams.srcStride = 0;
94+ dataCoptExtParams.dstStride = 0;
95+
96+ DataCopyPad(segmentIdsLocal, segmentIdsGm[offset], dataCoptExtParams, dataCopyPadExtParams);
57 return;97 return;
58}98}
59 99 
60template <typename TX, typename Index>100template <typename TX, typename Index>
61-__simt_vf__ __launch_bounds__(MAX_THREAD_NUM) inline void Compute(101+__aicore__ inline int32_t SegmentSumSimt<TX, Index>::GetUniqueSegIdCount(uint32_t segmentIdsNum)
62- __gm__ TX* xAddr, __gm__ Index* segmentIdsAddr, __gm__ TX* outputAddr, uint64_t colSize, uint64_t outerDim)
63{102{
64- int16_t blockNum = blockDim.y;103+ __ubuf__ Index* segmentIdsAddr = (__ubuf__ Index*)segmentIdsLocalShift.GetPhyAddr();
65- int16_t blockIdx = threadIdx.y;104+ __ubuf__ int32_t* segIdsPosAddr = (__ubuf__ int32_t*)uniqueIdPosLocal.GetPhyAddr();
66- for (Index i = get_block_idx() * blockNum + blockIdx; i < outerDim; i += get_block_num() * blockNum) {105+ uint32_t vl = platform::GetVRegSize() / sizeof(Index);
67- if (blockIdx > 0 && segmentIdsAddr[i] == segmentIdsAddr[i - 1]) {106+ uint16_t loopCnt = (uint16_t)(ops::CeilDiv(static_cast<uint32_t>(segmentIdsNum), vl));
68- continue;107+ uint32_t maskCount = segmentIdsNum;
69- }108+ uint32_t offset = platform::GetUbBlockSize() / sizeof(Index);
70- for (uint64_t j = threadIdx.x; j < colSize; j += blockDim.x) {109+ __VEC_SCOPE__
71- TX res = xAddr[i * colSize + j];110+ {
72- int16_t endRow = blockNum - 1;111+ AscendC::MicroAPI::RegTensor<int32_t> orderReg;
73- for (Index k = blockIdx + 1; k < blockNum; k++) {112+ AscendC::MicroAPI::RegTensor<int32_t> selReg;
74- if (segmentIdsAddr[i] != segmentIdsAddr[i + k - blockIdx]) {113+ AscendC::MicroAPI::RegTensor<Index> indicesReg;
75- endRow = k - 1;114+ AscendC::MicroAPI::RegTensor<Index> indicesShiftOneReg;
76- break;115+ AscendC::MicroAPI::MaskReg cmpMask;
77- }116+ AscendC::MicroAPI::MaskReg maskRegUpdate;
78- res += xAddr[(i + k - blockIdx) * colSize + j];117+ AscendC::MicroAPI::UnalignReg u0;
79- }118+ MicroAPI::UnalignReg ureg;
80- uint64_t outOffset = segmentIdsAddr[i] * colSize + j;119+ AscendC::MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
81- if (blockIdx != 0 && endRow != blockNum - 1) {120+ int32_t vciStart = 0;
82- outputAddr[outOffset] = res;121+ for (uint16_t i = 0; i < loopCnt; ++i) {
122+ vciStart = i * vl;
123+ auto segIdsOffset = segmentIdsAddr + offset + i * vl;
124+ AscendC::MicroAPI::Arange(orderReg, vciStart);
125+ maskRegUpdate = AscendC::MicroAPI::UpdateMask<Index>(maskCount);
126+ AscendC::MicroAPI::DataCopy(indicesReg, segIdsOffset);
127+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, segIdsOffset - 1);
128+ AscendC::MicroAPI::DataCopyUnAlign<Index>(indicesShiftOneReg, u0, segIdsOffset - 1);
129+ AscendC::MicroAPI::Compare<Index, CMPMODE::NE>(cmpMask, indicesReg, indicesShiftOneReg, maskRegUpdate);
130+ 
131+ if constexpr (IsSameType<Index, int64_t>::value) {
132+ AscendC::MicroAPI::MaskReg maskHalf;
133+ AscendC::MicroAPI::MaskPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskHalf, cmpMask);
134+ // vSQZ
135+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(
136+ selReg, orderReg, maskHalf);
83 } else {137 } else {
84- asc_atomic_add(outputAddr + outOffset, res);138+ // vSQZ
85- } 139+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(
140+ selReg, orderReg, cmpMask);
141+ }
142+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
143+ segIdsPosAddr, selReg, ureg);
144+ }
145+ AscendC::MicroAPI::DataCopyUnAlignPost(segIdsPosAddr, ureg);
146+ }
147+ return ((AscendC::MicroAPI::GetSpr<AscendC::SpecialPurposeReg::AR>()) / sizeof(int32_t));
148+}
149+ 
150+template <typename TX, typename Index>
151+__simt_vf__ __launch_bounds__(MAX_THREAD_NUM) inline void Compute(
152+ __gm__ TX* xAddr, __gm__ Index* segmentIdsAddr, __gm__ TX* outputAddr, __local_mem__ uint32_t* uniqueIdPosAddr,
153+ uint32_t colSize, uint32_t uniqueIdNum, uint32_t baseOffset, uint32_t idsNum)
154+{
155+ for (uint32_t i = threadIdx.y; i < uniqueIdNum; i += blockDim.y) {
156+ TX res = 0;
157+ uint32_t nextIdPos = (i != uniqueIdNum - 1) ? uniqueIdPosAddr[i+1] : idsNum;
158+ for (int32_t j = uniqueIdPosAddr[i]; j < nextIdPos; j++) {
159+ int32_t offset = (baseOffset + j) * colSize + threadIdx.x;
160+ res += xAddr[offset];
161+ }
162+ int32_t outputOffset = segmentIdsAddr[baseOffset + uniqueIdPosAddr[i]] * colSize + threadIdx.x;
163+ if (i == 0 || i == uniqueIdNum - 1) {
164+ asc_atomic_add(outputAddr + outputOffset, res);
165+ } else {
166+ outputAddr[outputOffset] = res;
86 }167 }
87-
88 }168 }
89 return;169 return;
90}170}
91 171 
172+template <typename TX, typename Index>
173+__simt_vf__ __launch_bounds__(MAX_THREAD_NUM) inline void ComputeForDeterminstic(
174+ __gm__ TX* xAddr, __gm__ Index* segmentIdsAddr, __gm__ TX* outputAddr, __local_mem__ uint32_t* uniqueIdPosAddr,
175+ __gm__ TX* tmpRowWs, __gm__ Index* tmpIdWs, uint32_t colSize, uint32_t uniqueIdNum, uint32_t baseOffset, uint32_t idsNum)
176+{
177+ if (threadIdx.y == 0) { // use thread(0/1, 0) clear segids workspace
178+ for (uint32_t i = threadIdx.x; i < TMP_ROWS_PER_CORE; i++) {
179+ tmpIdWs[get_block_idx() * TMP_ROWS_PER_CORE + i] = static_cast<Index>(-1);
180+ }
181+ }
182+ for (uint32_t i = threadIdx.y; i < uniqueIdNum; i += blockDim.y) {
183+ TX res = 0;
184+ uint32_t nextIdPos = (i != uniqueIdNum - 1) ? uniqueIdPosAddr[i+1] : idsNum;
185+ for (int32_t j = uniqueIdPosAddr[i]; j < nextIdPos; j++) {
186+ int32_t offset = (baseOffset + j) * colSize + threadIdx.x;
187+ res += xAddr[offset];
188+ }
189+ if (i == 0 || i == uniqueIdNum - 1) { // 如果uniqueid个数是1，那么只搬首行，那么尾行在ws就是随机值+（-1）id
CANN-robot
CANN-robotCANN-robot2月11日

代码可读性/维护性: 代码注释中存在拼写错误 'determinstic'，应为 'deterministic'。这个错误与第74行的错误相同，出现在函数名和注释中。

问题类型: 代码可读性/维护性 文件路径: index/segment_sum/op_kernel/arch35/segment_sum_simt.h 行号: 189 问题代码:

        if (i == 0 || i == uniqueIdNum - 1) {   // 如果uniqueid个数是1，那么只搬首行，那么尾行在ws就是随机值+（-1）id

修改建议:

将函数名 `ComputeForDeterminstic` 和所有相关注释中的 'determinstic' 更正为 'deterministic'。

此评论由代码审查工具自动生成

likedislike
190+ int32_t bias = (i == 0) ? 0 : 1;
191+ int32_t offsetInWs = get_block_idx() * TMP_ROWS_PER_CORE + bias;
192+ tmpRowWs[offsetInWs * colSize + threadIdx.x] = res;
193+ if (threadIdx.x == 0) {
194+ tmpIdWs[offsetInWs] = segmentIdsAddr[baseOffset + uniqueIdPosAddr[i]];
195+ }
196+ } else {
197+ outputAddr[segmentIdsAddr[baseOffset + uniqueIdPosAddr[i]] * colSize + threadIdx.x] = res;
198+ }
199+ }
200+ return;
201+}
202+ 
203+template <typename TX, typename Index>
204+__simt_vf__ __launch_bounds__(MAX_THREAD_NUM) inline void ComputeInWs(__gm__ TX* tmpRowWs, __gm__ Index* segmentIdsWs,
205+ __gm__ TX* outputAddr, uint32_t colSize)
206+{
207+ for (uint32_t i = threadIdx.y; i < TMP_ROWS_TOTAL; i += blockDim.y) {
208+ if (segmentIdsWs[i] == -1) {
209+ continue;
210+ }
211+ if (threadIdx.y > 0) {
212+ int32_t prevId = (segmentIdsWs[i - 1] == -1) ? (i - 2) : (i - 1);
213+ if (segmentIdsWs[i] == segmentIdsWs[prevId]) { // 如果当前线程组对应id非0且跟前一行id相同，则跳过
214+ continue;
215+ }
216+ }
217+ TX res = tmpRowWs[i * colSize + threadIdx.x]; // 走到这里要么是第0组的线程，要么是当前线程组对应的id跟上一个非-1 id不同，就是一个新id的起始
218+ for (Index j = i + 1; j < TMP_ROWS_TOTAL; j++) { // 从当前组线程开始往后加，知道id跟自己对应的id不同
CANN-robotCANN-robot
CANN-robotCANN-robot2月11日

代码结构与可维护性: ComputeInWs函数中，第218行注释：'从当前组线程开始往后加，知道id跟自己对应的id不同'，其中'知道'应为'直到'。

问题类型: 代码结构与可维护性 文件路径: index/segment_sum/op_kernel/arch35/segment_sum_simt.h 行号: 218 问题代码:

for (Index j = i + 1; j < TMP_ROWS_TOTAL; j++) {  // 从当前组线程开始往后加，知道id跟自己对应的id不同

修改建议:

将注释中的'知道'修改为'直到'。

此评论由代码审查工具自动生成

likedislike
CANN-robotCANN-robot2月11日

代码可读性/维护性: 代码注释中存在错别字 '知道'，应为 '直到'。

问题类型: 代码可读性/维护性 文件路径: index/segment_sum/op_kernel/arch35/segment_sum_simt.h 行号: 218 问题代码:

        for (Index j = i + 1; j < TMP_ROWS_TOTAL; j++) {  // 从当前组线程开始往后加，知道id跟自己对应的id不同

修改建议:

将注释更正为：`// 从当前组线程开始往后加，直到id跟自己对应的id不同`。

此评论由代码审查工具自动生成

likedislike
219+ if (segmentIdsWs[j] == -1) {
220+ continue;
221+ }
222+ if (segmentIdsWs[i] != segmentIdsWs[j]) {
223+ break;
224+ }
225+ res += tmpRowWs[j * colSize + threadIdx.x];
226+ }
227+ asc_atomic_add(outputAddr + segmentIdsWs[i] * colSize + threadIdx.x, res);
228+ }
229+}
230+ 
231+ 
232+template <typename TX, typename Index>
233+__aicore__ inline void SegmentSumSimt<TX, Index>::ProcessEachLoop(uint32_t segmentIdsNum, uint64_t baseOffset)
234+{
235+ auto vectorWaitMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
236+ SetFlag<HardEvent::MTE2_V>(vectorWaitMTE2);
237+ WaitFlag<HardEvent::MTE2_V>(vectorWaitMTE2);
238+ Duplicate(segmentIdsLocalShift, static_cast<Index>(-1), platform::GetUbBlockSize() / sizeof(Index));
239+ uint32_t uniqueIdNum = GetUniqueSegIdCount(segmentIdsNum);
240+ asc_vf_call<Compute<TX, Index>>(dim3{threadNum, threadBlock},
241+ (__gm__ TX*)xGm.GetPhyAddr(),
242+ (__gm__ Index*)segmentIdsGm.GetPhyAddr(),
243+ (__gm__ TX*)outputGm.GetPhyAddr(),
244+ (__local_mem__ uint32_t*)(uniqueIdPosLocal.GetPhyAddr()),
245+ tilingData_->innerDim, uniqueIdNum, baseOffset, segmentIdsNum);
246+}
247+ 
248+template <typename TX, typename Index>
249+__aicore__ inline void SegmentSumSimt<TX, Index>::ProcessEachLoopForDeterminstic(uint32_t segmentIdsNum, uint64_t baseOffset)
250+{
251+ auto vectorWaitMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
252+ SetFlag<HardEvent::MTE2_V>(vectorWaitMTE2);
253+ WaitFlag<HardEvent::MTE2_V>(vectorWaitMTE2);
254+ Duplicate(segmentIdsLocalShift, static_cast<Index>(-1), platform::GetUbBlockSize() / sizeof(Index));
255+ uint32_t uniqueIdNum = GetUniqueSegIdCount(segmentIdsNum);
256+ asc_vf_call<ComputeForDeterminstic<TX, Index>>(dim3{threadNum, threadBlock},
257+ (__gm__ TX*)xGm.GetPhyAddr(),
258+ (__gm__ Index*)segmentIdsGm.GetPhyAddr(),
259+ (__gm__ TX*)outputGm.GetPhyAddr(),
260+ (__local_mem__ uint32_t*)(uniqueIdPosLocal.GetPhyAddr()),
261+ (__gm__ TX*)tmpRowWs.GetPhyAddr(),
262+ (__gm__ Index*)tmpIdWs.GetPhyAddr(),
263+ tilingData_->innerDim, uniqueIdNum, baseOffset, segmentIdsNum);
264+ SyncAll(); // wait for all core write ws
265+ if (coreId > 0) {
266+ return;
267+ }
268+ asc_vf_call<ComputeInWs<TX, Index>>(dim3{threadNum, threadBlock},
269+ (__gm__ TX*)tmpRowWs.GetPhyAddr(),
270+ (__gm__ Index*)tmpIdWs.GetPhyAddr(),
271+ (__gm__ TX*)outputGm.GetPhyAddr(),
272+ tilingData_->innerDim);
273+}
274+ 
92template <typename TX, typename Index>275template <typename TX, typename Index>
93__aicore__ inline void SegmentSumSimt<TX, Index>::Process()276__aicore__ inline void SegmentSumSimt<TX, Index>::Process()
94{277{
95- __gm__ TX* input = (__gm__ TX*)xGm.GetPhyAddr();278+ uniqueIdPosLocal = uniqueIdPosBuf_.Get<int32_t>();
96- __gm__ Index* segmentIds = (__gm__ Index*)segmentIdsGm.GetPhyAddr();279+ segmentIdsLocalShift = segmentIdsBuf_.Get<Index>();
97- __gm__ TX* output = (__gm__ TX*)outputGm.GetPhyAddr();280+ segmentIdsLocal = segmentIdsLocalShift[platform::GetUbBlockSize() / sizeof(Index)];
98 281 
99- uint32_t threadNum = MAX_THREAD_NUM > tilingData_->innerDim ? tilingData_->innerDim : MAX_THREAD_NUM;282+ threadNum = MAX_THREAD_NUM > tilingData_->innerDim ? tilingData_->innerDim : MAX_THREAD_NUM;
100- uint32_t threadBlock = MAX_THREAD_NUM / threadNum;283+ threadBlock = MAX_THREAD_NUM / threadNum;
101- asc_vf_call<Compute<TX, Index>>(dim3{threadNum, threadBlock}, // block dim {x, y}284+ 
102- input, segmentIds, output, tilingData_->innerDim, tilingData_->outerDim);285+ loopTimes = coreId == (coreNum - 1) ? tilingData_->loopTimesTailCore : tilingData_->loopTimes;
103- return;286+ segIds = coreId == (coreNum - 1) ? tilingData_->segIdsPerLoopTailCore : tilingData_->segIdsPerLoop;
287+ tailSegIds = coreId == (coreNum - 1) ? tilingData_->segIdsTailLoopTailCore : tilingData_->segIdsTailLoop;
288+
289+ uint64_t segIdsPerCore = tilingData_->outerDim / coreNum;
290+ uint64_t baseOffset = coreId * segIdsPerCore;
291+ if (tilingData_->isDeterministic == 1) {
292+ for (int32_t i = 0; i < loopTimes - 1; i++){
293+ CopySegmentIdsIn(segmentIdsLocal, baseOffset + i * segIds, segIds);
294+ ProcessEachLoopForDeterminstic(segIds, baseOffset + i * segIds);
295+ }
296+ CopySegmentIdsIn(segmentIdsLocal, baseOffset + (loopTimes - 1) * segIds, tailSegIds);
297+ ProcessEachLoopForDeterminstic(tailSegIds, baseOffset + (loopTimes - 1) * segIds);
298+ } else {
299+ for (int32_t i = 0; i < loopTimes - 1; i++){
300+ CopySegmentIdsIn(segmentIdsLocal, baseOffset + i * segIds, segIds);
301+ ProcessEachLoop(segIds, baseOffset + i * segIds);
302+ }
303+ CopySegmentIdsIn(segmentIdsLocal, baseOffset + (loopTimes - 1) * segIds, tailSegIds);
304+ ProcessEachLoop(tailSegIds, baseOffset + (loopTimes - 1) * segIds);
305+ }
104}306}
105}307}
106#endif308#endif
Mindex/segment_sum/op_kernel/arch35/segment_sum_struct.h+9-4
@@ -17,12 +17,19 @@ struct SegmentSumSimtTilingData
17 uint64_t innerDim{0};17 uint64_t innerDim{0};
18 uint64_t initNumPerCore{0};18 uint64_t initNumPerCore{0};
19 uint64_t initNumTailCore{0};19 uint64_t initNumTailCore{0};
20+ uint32_t isDeterministic{0};
21+ uint32_t maxSegIdsInUb{0};
22+ int32_t loopTimes{0}; // 整核循环次数
23+ int32_t loopTimesTailCore{0}; // 尾核循环次数
24+ uint32_t segIdsPerLoop{0}; //整核整循环处理多少id
25+ uint32_t segIdsPerLoopTailCore{0}; //尾核整循环处理多少id
26+ uint32_t segIdsTailLoop{0}; // 整核尾循环处理多少id
27+ uint32_t segIdsTailLoopTailCore{0}; //尾核尾循环处理多少id
20};28};
21 29 
22struct SegmentSumSimdTilingData30struct SegmentSumSimdTilingData
23{31{
24 int64_t needCoreNum{0};32 int64_t needCoreNum{0};
25- int64_t outerDim{0}; // kernel没用
26 int64_t innerDim{0};33 int64_t innerDim{0};
27 34 
28 int64_t xBufferSize{0};35 int64_t xBufferSize{0};
@@ -37,9 +44,7 @@ struct SegmentSumSimdTilingData
37 int64_t blockNumInCol{0}; // 列切分的核数44 int64_t blockNumInCol{0}; // 列切分的核数
38 45 
39 int64_t normalCoreInnerNum{0}; // 正常列核列上处理的inner数46 int64_t normalCoreInnerNum{0}; // 正常列核列上处理的inner数
40- int64_t tailCoreInnerNum{0}; // 列尾核列上处理的inner数
41 int64_t normalCoreOutterNum{0}; // 正常行核行上处理的行数47 int64_t normalCoreOutterNum{0}; // 正常行核行上处理的行数
42- int64_t tailCoreOutterNum{0}; // 行尾核行上处理的行数
43 48 
44 int64_t normalCoreRowUbLoop{0}; // 正常行核ub在行上的循环次数49 int64_t normalCoreRowUbLoop{0}; // 正常行核ub在行上的循环次数
45 int64_t normalCoreNormalLoopOutters{0}; // 正常行核ub正常循环一次处理的行数50 int64_t normalCoreNormalLoopOutters{0}; // 正常行核ub正常循环一次处理的行数
@@ -57,7 +62,6 @@ struct SegmentSumSimdTilingData
57 62 
58 int64_t usedCoreNumForMultAdd{0}; // 多核累加使用的核数63 int64_t usedCoreNumForMultAdd{0}; // 多核累加使用的核数
59 int64_t normalCoreMultAddInners{0}; // 多核累加正常核处理的inner数64 int64_t normalCoreMultAddInners{0}; // 多核累加正常核处理的inner数
60- int64_t tailCoreMultAddInners{0}; // 多核累加尾核处理的inner数
61 65 
62 int64_t normalCoreMultAddInnerLoop{0}; // 多核累加正常核列循环次数66 int64_t normalCoreMultAddInnerLoop{0}; // 多核累加正常核列循环次数
63 int64_t normalCoreMultAddNormalLoopInners{0}; // 多核累加正常核正常循环处理的inner数67 int64_t normalCoreMultAddNormalLoopInners{0}; // 多核累加正常核正常循环处理的inner数
@@ -67,6 +71,7 @@ struct SegmentSumSimdTilingData
67 int64_t tailCoreMultAddTailLoopInners{0}; // 多核累加尾核尾循环处理的inner数71 int64_t tailCoreMultAddTailLoopInners{0}; // 多核累加尾核尾循环处理的inner数
68 72 
69 int64_t multAddXBufferSize{0};73 int64_t multAddXBufferSize{0};
74+ int64_t multAddIdsBufferSize{0};
70 int64_t multAddYBufferSize{0};75 int64_t multAddYBufferSize{0};
71};76};
72 77 
Mindex/segment_sum/op_kernel/segment_sum_apt.cpp+32-2
@@ -16,6 +16,8 @@
16#include "arch35/segment_sum_simt.h"16#include "arch35/segment_sum_simt.h"
17#include "arch35/segment_sum_simd.h"17#include "arch35/segment_sum_simd.h"
18#include "arch35/clear_output.h"18#include "arch35/clear_output.h"
19+#include "arch35/segment_sum_simd_determ.h"
20+#include "arch35/segment_sum_simd_mult_core_add.h"
19#include "arch35/segment_sum_struct.h"21#include "arch35/segment_sum_struct.h"
20 22 
21using namespace AscendC;23using namespace AscendC;
@@ -23,6 +25,7 @@ using namespace SegmentSum;
23 25 
24#define TEMPLATE_SIMT_TILING_KEY 100026#define TEMPLATE_SIMT_TILING_KEY 1000
25#define SIMD_ATOMIC_SUPPORT_TILING_KEY 200027#define SIMD_ATOMIC_SUPPORT_TILING_KEY 2000
28+#define SIMD_DETERM_TILING_KEY 2002
26 29 
27template <typename T>30template <typename T>
28__aicore__ inline void invokeTemplateAllClear(31__aicore__ inline void invokeTemplateAllClear(
@@ -35,12 +38,34 @@ __aicore__ inline void invokeTemplateAllClear(
35 pipeIn.Reset();38 pipeIn.Reset();
36}39}
37 40 
41+template <typename T1, typename T2>
42+__aicore__ inline void invokeTemplateSimdDeterm(
43+ GM_ADDR x, GM_ADDR segment_ids, GM_ADDR y, GM_ADDR workspace, const SegmentSumSimdTilingData* tilingData, AscendC::TPipe& pipeIn)
44+{
45+ SegmentSumSimdDeterm<T1, T2> op;
46+ op.Init(x, segment_ids, y, workspace, pipeIn, tilingData);
47+ op.Process();
48+ SyncAll();
49+ pipeIn.Reset();
50+}
51+ 
52+template <typename T1, typename T2>
53+__aicore__ inline void invokeTemplateMultiCoreAdd(
54+ GM_ADDR y, GM_ADDR workspace, const SegmentSumSimdTilingData* tilingData, AscendC::TPipe& pipeIn)
55+{
56+ SegmentSumMultiCoreAdd<T1, T2> op;
57+ op.Init(y, workspace, pipeIn, tilingData);
58+ op.Process();
59+}
60+ 
61+ 
38extern "C" __global__ __aicore__ void segment_sum(GM_ADDR x, GM_ADDR segment_ids,62extern "C" __global__ __aicore__ void segment_sum(GM_ADDR x, GM_ADDR segment_ids,
39 GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling)63 GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling)
40{64{
41 REGISTER_TILING_DEFAULT(SegmentSumTilingData);65 REGISTER_TILING_DEFAULT(SegmentSumTilingData);
42 REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 1000", SegmentSumSimtTilingData);66 REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 1000", SegmentSumSimtTilingData);
43 REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 2000", SegmentSumSimdTilingData);67 REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 2000", SegmentSumSimdTilingData);
68+ REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 2002", SegmentSumSimdTilingData);
44 69 
45 TPipe pipe;70 TPipe pipe;
46 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);71 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
@@ -48,15 +73,20 @@ extern "C" __global__ __aicore__ void segment_sum(GM_ADDR x, GM_ADDR segment_ids
48 GET_TILING_DATA_WITH_STRUCT(SegmentSumSimtTilingData, simtTilingData, tiling);73 GET_TILING_DATA_WITH_STRUCT(SegmentSumSimtTilingData, simtTilingData, tiling);
49 const SegmentSumSimtTilingData* __restrict tilingData = &simtTilingData; 74 const SegmentSumSimtTilingData* __restrict tilingData = &simtTilingData;
50 SegmentSumSimt<DTYPE_X, DTYPE_SEGMENT_IDS> op(tilingData, &pipe);75 SegmentSumSimt<DTYPE_X, DTYPE_SEGMENT_IDS> op(tilingData, &pipe);
51- op.Init(x, segment_ids, output);76+ op.Init(x, segment_ids, output, workspace);
52 op.Process();77 op.Process();
53 } else if (TILING_KEY_IS(SIMD_ATOMIC_SUPPORT_TILING_KEY)) {78 } else if (TILING_KEY_IS(SIMD_ATOMIC_SUPPORT_TILING_KEY)) {
54 GET_TILING_DATA_WITH_STRUCT(SegmentSumSimdTilingData, tilingData, tiling);79 GET_TILING_DATA_WITH_STRUCT(SegmentSumSimdTilingData, tilingData, tiling);
55 if constexpr (!(std::is_same_v<DTYPE_X, uint32_t> || std::is_same_v<DTYPE_X, uint64_t> || std::is_same_v<DTYPE_X, int64_t>)) {80 if constexpr (!(std::is_same_v<DTYPE_X, uint32_t> || std::is_same_v<DTYPE_X, uint64_t> || std::is_same_v<DTYPE_X, int64_t>)) {
56 invokeTemplateAllClear<DTYPE_X>(output, &tilingData, pipe);81 invokeTemplateAllClear<DTYPE_X>(output, &tilingData, pipe);
57 SegmentSumSimd<DTYPE_X, DTYPE_SEGMENT_IDS> op;82 SegmentSumSimd<DTYPE_X, DTYPE_SEGMENT_IDS> op;
58- op.Init(x, segment_ids, output, workspace, pipe, &tilingData);83+ op.Init(x, segment_ids, output, pipe, &tilingData);
59 op.Process();84 op.Process();
60 }85 }
86+ } else if (TILING_KEY_IS(SIMD_DETERM_TILING_KEY)) {
87+ GET_TILING_DATA_WITH_STRUCT(SegmentSumSimdTilingData, tilingData, tiling);
88+ invokeTemplateAllClear<DTYPE_X>(output, &tilingData, pipe);
89+ invokeTemplateSimdDeterm<DTYPE_X, DTYPE_SEGMENT_IDS>(x, segment_ids, output, workspace, &tilingData, pipe);
90+ invokeTemplateMultiCoreAdd<DTYPE_X, DTYPE_SEGMENT_IDS>(output, workspace, &tilingData, pipe);
61 }91 }
62}92}
Mindex/segment_sum/tests/ut/op_host/arch35/test_segment_sum_tiling.cpp+19-4文件内容审核中，请稍后刷新重试