Star889
1799
代码
代码
Issues259
Pull Requests294
Actions
项目讨论
Wiki
分析
Star889
1799
已合并
A5增加scatter_add_with_sorted算子 #4175
wkkk0528创建于 4月24日
A5增加scatter_add_with_sorted算子 #4175
已合并
wkkk0528创建于 4月24日
从wkkk0528/ops-nnscatter_add_with_sorted_backup_0423合入到cann/ops-nnmaster
共 29 个文件变更+4708-63
Mindex/scatter_add_with_sorted/README.md+1-1
@@ -4,7 +4,7 @@
4 4 
5| 产品 | 是否支持 |5| 产品 | 是否支持 |
6| :----------------------------------------------------------- | :------: |6| :----------------------------------------------------------- | :------: |
7-| <term>Ascend 950PR/Ascend 950DT</term> | × |7+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
9| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |9| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |
10| <term>Atlas 200I/500 A2 推理产品</term> | × |10| <term>Atlas 200I/500 A2 推理产品</term> | × |
Aindex/scatter_add_with_sorted/op_host/config/ascend950/scatter_add_with_sorted_binary.json+517-0
@@ -0,0 +1,517 @@
1+{
2+ "op_type": "ScatterAddWithSorted",
C
Cchenqi3174月28日

ops_list 未加算子列表

likedislike
W
wkkk0528
4月28日 评论：
C
chenqi317
4月29日 评论：
W
wkkk0528
4月29日 评论：
3+ "op_list": [
4+ {
5+ "bin_filename": "ScatterAddWithSorted_22ad8a8eb086e1d6db857bcf4228e9c4",
6+ "inputs": [
7+ {
8+ "name": "var",
9+ "index": 0,
10+ "dtype": "bfloat16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "value",
19+ "index": 1,
20+ "dtype": "bfloat16",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "sorted_index",
29+ "index": 2,
30+ "dtype": "int32",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ },
37+ {
38+ "name": "pos",
39+ "index": 3,
40+ "dtype": "int32",
41+ "format": "ND",
42+ "paramType": "optional",
43+ "shape": [
44+ -2
45+ ]
46+ }
47+ ],
48+ "outputs": [
49+ {
50+ "name": "var",
51+ "index": 0,
52+ "dtype": "bfloat16",
53+ "format": "ND",
54+ "paramType": "required",
55+ "shape": [
56+ -2
57+ ]
58+ }
59+ ],
60+ "attrs": [
61+ {
62+ "name": "reduction",
63+ "dtype": "string",
64+ "value": "add"
65+ }
66+ ]
67+ },
68+ {
69+ "bin_filename": "ScatterAddWithSorted_cc00eba9a84e1dea2f67cca1c6ee2717",
70+ "inputs": [
71+ {
72+ "name": "var",
73+ "index": 0,
74+ "dtype": "bfloat16",
75+ "format": "ND",
76+ "paramType": "required",
77+ "shape": [
78+ -2
79+ ]
80+ },
81+ {
82+ "name": "value",
83+ "index": 1,
84+ "dtype": "bfloat16",
85+ "format": "ND",
86+ "paramType": "required",
87+ "shape": [
88+ -2
89+ ]
90+ },
91+ {
92+ "name": "sorted_index",
93+ "index": 2,
94+ "dtype": "int64",
95+ "format": "ND",
96+ "paramType": "required",
97+ "shape": [
98+ -2
99+ ]
100+ },
101+ {
102+ "name": "pos",
103+ "index": 3,
104+ "dtype": "int64",
105+ "format": "ND",
106+ "paramType": "optional",
107+ "shape": [
108+ -2
109+ ]
110+ }
111+ ],
112+ "outputs": [
113+ {
114+ "name": "var",
115+ "index": 0,
116+ "dtype": "bfloat16",
117+ "format": "ND",
118+ "paramType": "required",
119+ "shape": [
120+ -2
121+ ]
122+ }
123+ ],
124+ "attrs": [
125+ {
126+ "name": "reduction",
127+ "dtype": "string",
128+ "value": "add"
129+ }
130+ ]
131+ },
132+ {
133+ "bin_filename": "ScatterAddWithSorted_c1a7e7a7f195f1d1f7edde59cbd3416c",
134+ "inputs": [
135+ {
136+ "name": "var",
137+ "index": 0,
138+ "dtype": "float16",
139+ "format": "ND",
140+ "paramType": "required",
141+ "shape": [
142+ -2
143+ ]
144+ },
145+ {
146+ "name": "value",
147+ "index": 1,
148+ "dtype": "float16",
149+ "format": "ND",
150+ "paramType": "required",
151+ "shape": [
152+ -2
153+ ]
154+ },
155+ {
156+ "name": "sorted_index",
157+ "index": 2,
158+ "dtype": "int32",
159+ "format": "ND",
160+ "paramType": "required",
161+ "shape": [
162+ -2
163+ ]
164+ },
165+ {
166+ "name": "pos",
167+ "index": 3,
168+ "dtype": "int32",
169+ "format": "ND",
170+ "paramType": "optional",
171+ "shape": [
172+ -2
173+ ]
174+ }
175+ ],
176+ "outputs": [
177+ {
178+ "name": "var",
179+ "index": 0,
180+ "dtype": "float16",
181+ "format": "ND",
182+ "paramType": "required",
183+ "shape": [
184+ -2
185+ ]
186+ }
187+ ],
188+ "attrs": [
189+ {
190+ "name": "reduction",
191+ "dtype": "string",
192+ "value": "add"
193+ }
194+ ]
195+ },
196+ {
197+ "bin_filename": "ScatterAddWithSorted_3ff6028e4f20ba619f7ba596d551290f",
198+ "inputs": [
199+ {
200+ "name": "var",
201+ "index": 0,
202+ "dtype": "float16",
203+ "format": "ND",
204+ "paramType": "required",
205+ "shape": [
206+ -2
207+ ]
208+ },
209+ {
210+ "name": "value",
211+ "index": 1,
212+ "dtype": "float16",
213+ "format": "ND",
214+ "paramType": "required",
215+ "shape": [
216+ -2
217+ ]
218+ },
219+ {
220+ "name": "sorted_index",
221+ "index": 2,
222+ "dtype": "int64",
223+ "format": "ND",
224+ "paramType": "required",
225+ "shape": [
226+ -2
227+ ]
228+ },
229+ {
230+ "name": "pos",
231+ "index": 3,
232+ "dtype": "int64",
233+ "format": "ND",
234+ "paramType": "optional",
235+ "shape": [
236+ -2
237+ ]
238+ }
239+ ],
240+ "outputs": [
241+ {
242+ "name": "var",
243+ "index": 0,
244+ "dtype": "float16",
245+ "format": "ND",
246+ "paramType": "required",
247+ "shape": [
248+ -2
249+ ]
250+ }
251+ ],
252+ "attrs": [
253+ {
254+ "name": "reduction",
255+ "dtype": "string",
256+ "value": "add"
257+ }
258+ ]
259+ },
260+ {
261+ "bin_filename": "ScatterAddWithSorted_56789086143b2e7a1658bb5105f29618",
262+ "inputs": [
263+ {
264+ "name": "var",
265+ "index": 0,
266+ "dtype": "float32",
267+ "format": "ND",
268+ "paramType": "required",
269+ "shape": [
270+ -2
271+ ]
272+ },
273+ {
274+ "name": "value",
275+ "index": 1,
276+ "dtype": "float32",
277+ "format": "ND",
278+ "paramType": "required",
279+ "shape": [
280+ -2
281+ ]
282+ },
283+ {
284+ "name": "sorted_index",
285+ "index": 2,
286+ "dtype": "int32",
287+ "format": "ND",
288+ "paramType": "required",
289+ "shape": [
290+ -2
291+ ]
292+ },
293+ {
294+ "name": "pos",
295+ "index": 3,
296+ "dtype": "int32",
297+ "format": "ND",
298+ "paramType": "optional",
299+ "shape": [
300+ -2
301+ ]
302+ }
303+ ],
304+ "outputs": [
305+ {
306+ "name": "var",
307+ "index": 0,
308+ "dtype": "float32",
309+ "format": "ND",
310+ "paramType": "required",
311+ "shape": [
312+ -2
313+ ]
314+ }
315+ ],
316+ "attrs": [
317+ {
318+ "name": "reduction",
319+ "dtype": "string",
320+ "value": "add"
321+ }
322+ ]
323+ },
324+ {
325+ "bin_filename": "ScatterAddWithSorted_0af26dd59815502caf51679adc2ca37b",
326+ "inputs": [
327+ {
328+ "name": "var",
329+ "index": 0,
330+ "dtype": "float32",
331+ "format": "ND",
332+ "paramType": "required",
333+ "shape": [
334+ -2
335+ ]
336+ },
337+ {
338+ "name": "value",
339+ "index": 1,
340+ "dtype": "float32",
341+ "format": "ND",
342+ "paramType": "required",
343+ "shape": [
344+ -2
345+ ]
346+ },
347+ {
348+ "name": "sorted_index",
349+ "index": 2,
350+ "dtype": "int64",
351+ "format": "ND",
352+ "paramType": "required",
353+ "shape": [
354+ -2
355+ ]
356+ },
357+ {
358+ "name": "pos",
359+ "index": 3,
360+ "dtype": "int64",
361+ "format": "ND",
362+ "paramType": "optional",
363+ "shape": [
364+ -2
365+ ]
366+ }
367+ ],
368+ "outputs": [
369+ {
370+ "name": "var",
371+ "index": 0,
372+ "dtype": "float32",
373+ "format": "ND",
374+ "paramType": "required",
375+ "shape": [
376+ -2
377+ ]
378+ }
379+ ],
380+ "attrs": [
381+ {
382+ "name": "reduction",
383+ "dtype": "string",
384+ "value": "add"
385+ }
386+ ]
387+ },
388+ {
389+ "bin_filename": "ScatterAddWithSorted_4b15a6a1be01b2275383765f63aef2a4",
390+ "inputs": [
391+ {
392+ "name": "var",
393+ "index": 0,
394+ "dtype": "int32",
395+ "format": "ND",
396+ "paramType": "required",
397+ "shape": [
398+ -2
399+ ]
400+ },
401+ {
402+ "name": "value",
403+ "index": 1,
404+ "dtype": "int32",
405+ "format": "ND",
406+ "paramType": "required",
407+ "shape": [
408+ -2
409+ ]
410+ },
411+ {
412+ "name": "sorted_index",
413+ "index": 2,
414+ "dtype": "int32",
415+ "format": "ND",
416+ "paramType": "required",
417+ "shape": [
418+ -2
419+ ]
420+ },
421+ {
422+ "name": "pos",
423+ "index": 3,
424+ "dtype": "int32",
425+ "format": "ND",
426+ "paramType": "optional",
427+ "shape": [
428+ -2
429+ ]
430+ }
431+ ],
432+ "outputs": [
433+ {
434+ "name": "var",
435+ "index": 0,
436+ "dtype": "int32",
437+ "format": "ND",
438+ "paramType": "required",
439+ "shape": [
440+ -2
441+ ]
442+ }
443+ ],
444+ "attrs": [
445+ {
446+ "name": "reduction",
447+ "dtype": "string",
448+ "value": "add"
449+ }
450+ ]
451+ },
452+ {
453+ "bin_filename": "ScatterAddWithSorted_f08901324f8af8e628268575e08b214b",
454+ "inputs": [
455+ {
456+ "name": "var",
457+ "index": 0,
458+ "dtype": "int32",
459+ "format": "ND",
460+ "paramType": "required",
461+ "shape": [
462+ -2
463+ ]
464+ },
465+ {
466+ "name": "value",
467+ "index": 1,
468+ "dtype": "int32",
469+ "format": "ND",
470+ "paramType": "required",
471+ "shape": [
472+ -2
473+ ]
474+ },
475+ {
476+ "name": "sorted_index",
477+ "index": 2,
478+ "dtype": "int64",
479+ "format": "ND",
480+ "paramType": "required",
481+ "shape": [
482+ -2
483+ ]
484+ },
485+ {
486+ "name": "pos",
487+ "index": 3,
488+ "dtype": "int64",
489+ "format": "ND",
490+ "paramType": "optional",
491+ "shape": [
492+ -2
493+ ]
494+ }
495+ ],
496+ "outputs": [
497+ {
498+ "name": "var",
499+ "index": 0,
500+ "dtype": "int32",
501+ "format": "ND",
502+ "paramType": "required",
503+ "shape": [
504+ -2
505+ ]
506+ }
507+ ],
508+ "attrs": [
509+ {
510+ "name": "reduction",
511+ "dtype": "string",
512+ "value": "add"
513+ }
514+ ]
515+ }
516+ ]
517+ }
Aindex/scatter_add_with_sorted/op_host/config/ascend950/scatter_add_with_sorted_simplified_key.ini+14-0
@@ -0,0 +1,14 @@
1+; 该文件主要影响 opc 工具 编译二进制kernel时，--simplified_key_mode 选项中填写的值，格式如下所示：
2+; [某算子]
3+; default=xx
4+; ascendxx=xx
5+; 其中，default为默认mode，ascendxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6+; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7+; 2)如果仅有default配置：各个版本按default配置
8+; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9+; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10+; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11+; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12+ 
13+[ScatterAddWithSorted]
14+default=0
Mindex/scatter_add_with_sorted/op_host/op_api/scatter_add_with_sorted.cpp+6-3
@@ -16,6 +16,7 @@
16#include "opdev/make_op_executor.h"16#include "opdev/make_op_executor.h"
17#include "opdev/op_dfx.h"17#include "opdev/op_dfx.h"
18#include "opdev/platform.h"18#include "opdev/platform.h"
19+#include "op_api/aclnn_util.h"
19 20 
20using namespace op;21using namespace op;
21 22 
@@ -27,10 +28,12 @@ const aclTensor* ScatterAddWithSorted(
27 const std::string& reduction, aclOpExecutor* executor)28 const std::string& reduction, aclOpExecutor* executor)
28{29{
29 L0_DFX(ScatterAddWithSorted, self, value, sorted_index, pos);30 L0_DFX(ScatterAddWithSorted, self, value, sorted_index, pos);
30- 31+ // ASCEND950平台支持DT_FLOAT, DT_FLOAT16, DT_INT32, DT_BF16
31 auto socVersion = GetCurrentPlatformInfo().GetSocVersion();32 auto socVersion = GetCurrentPlatformInfo().GetSocVersion();
32- if (socVersion != SocVersion::ASCEND910B && socVersion != SocVersion::ASCEND910_93) {33+ if (socVersion != SocVersion::ASCEND910B && socVersion != SocVersion::ASCEND910_93 &&
S
Ssunday4月28日

这里平台不支持时只打印了日志，没有中断返回，后续仍会继续走 ADD_TO_LAUNCHER_LIST_AICORE。在不支持平台上可能继续下发 kernel，行为不可控。建议在该分支直接返回失败（例如 return nullptr）以终止执行。

likedislike
W
wkkk0528
4月28日 评论：
33- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "[ScatterAddWithSorted] only support ASCEND910B and ASCEND910_93");34+ !Ops::NN::AclnnUtil::IsRegbase()) {
35+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "[ScatterAddWithSorted] only support ASCEND910B, ASCEND910_93 and ASCEND950");
36+ return nullptr;
34 }37 }
35 38 
36 auto selfOut = const_cast<aclTensor*>(self);39 auto selfOut = const_cast<aclTensor*>(self);
Mindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_def.cpp+53-0
@@ -56,6 +56,59 @@ public:
56 this->AICore().AddConfig("ascend910b");56 this->AICore().AddConfig("ascend910b");
57 this->AICore().AddConfig("ascend910_93");57 this->AICore().AddConfig("ascend910_93");
58 58 
59+ OpAICoreConfig config_950;
60+ config_950.Input("var")
61+ .ParamType(REQUIRED)
62+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16,
63+ ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16})
64+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
65+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
66+ .UnknownShapeFormat(
67+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
68+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
69+ config_950.Input("value")
70+ .ParamType(REQUIRED)
71+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16,
72+ ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16})
73+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
74+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
75+ .UnknownShapeFormat(
76+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
77+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
78+ config_950.Input("sorted_index")
79+ .ParamType(REQUIRED)
80+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
81+ ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
82+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
83+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
84+ .UnknownShapeFormat(
85+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
86+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
87+ config_950.Input("pos")
88+ .ParamType(OPTIONAL)
89+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
90+ ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
91+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
92+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
93+ .UnknownShapeFormat(
94+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
95+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
96+ config_950.Output("var")
97+ .ParamType(REQUIRED)
98+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16,
99+ ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16})
100+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
101+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
102+ .UnknownShapeFormat(
103+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
104+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
105+ config_950.ExtendCfgInfo("opFile.value", "scatter_add_with_sorted_apt");
106+ config_950.DynamicCompileStaticFlag(true)
107+ .DynamicFormatFlag(true)
108+ .DynamicRankSupportFlag(true)
109+ .DynamicShapeSupportFlag(true);
110+ this->AICore().AddConfig("ascend950", config_950);
111+ 
59 OpAICoreConfig config_kirin = GetKirinCoreConfig();112 OpAICoreConfig config_kirin = GetKirinCoreConfig();
60 this->AICore().AddConfig("kirinx90", config_kirin);113 this->AICore().AddConfig("kirinx90", config_kirin);
61 this->AICore().AddConfig("kirin9030", config_kirin);114 this->AICore().AddConfig("kirin9030", config_kirin);
Aindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_simd_tiling.cpp+382-0
@@ -0,0 +1,382 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_add_with_sorted_simd_tiling.cpp
13+ * \brief scatter_add_with_sorted_simd_tiling
14+ */
15+ 
16+#include "scatter_add_with_sorted_simd_tiling.h"
17+ 
18+namespace optiling {
19+ 
20+static constexpr int64_t SIMD_INNER_THRES = 128;
21+static constexpr int64_t BASE_BLOCK_ALIGN = 512;
22+static constexpr int64_t SINGLE_CORE_THRESHOLD = 4 * 1024;
23+static constexpr int64_t CACHE_ALIGN_SIZE = 128;
24+ 
25+static constexpr int64_t NUM_FOUR = 4;
26+static constexpr int64_t COL_TILING_THRES = 8 * 1024;
27+constexpr uint64_t BUFFER_NUM = 2;
28+constexpr int64_t ASCENDC_TOOLS_WORKSPACE = static_cast<int64_t>(16) * 1024 * 1024;
29+ 
30+bool ScatterAddWithSortedSimdTiling::IsCapable()
31+{
32+ bool isSimd = varShape_[1] * updatesDtypeSize_ >= SIMD_INNER_THRES;
33+ return isSimd;
34+}
35+ 
36+std::set<int64_t> FindFactor(int64_t usedCoreNum)
37+{
38+ std::set<int64_t> result;
39+ uint64_t upbound = std::ceil(std::sqrt(usedCoreNum) + 1);
40+ 
41+ for (uint64_t m = 1; m < upbound; m++) {
42+ uint64_t y = usedCoreNum / m;
43+ result.insert(m);
44+ result.insert(y);
45+ }
46+ return result;
47+}
48+ 
49+void ScatterAddWithSortedSimdTiling::AutoTilingRowCol(int64_t& rowTileNum, int64_t& colTileNum, int64_t usedCoreNum, int64_t rowTotalNum, int64_t colTotalNum)
50+{
51+ int64_t tmpEleNum = BASE_BLOCK_ALIGN / updatesDtypeSize_;
52+ int64_t colBlockTotalNum = (colTotalNum + tmpEleNum - 1) / tmpEleNum;
53+ usedCoreNum = std::min(usedCoreNum, std::max<int64_t>(1, rowTotalNum * colBlockTotalNum * tmpEleNum / (SINGLE_CORE_THRESHOLD)));
54+ 
55+ 
56+ std::set<int64_t> cutSet = FindFactor(usedCoreNum);
57+ std::vector<std::vector<int64_t>> allTiling;
58+ 
59+ for (int64_t m : cutSet) {
60+ if (m > rowTotalNum) {
61+ continue;
62+ }
63+ 
64+ int64_t n = usedCoreNum / m;
65+ n = n < 1 ? 1 : n;
66+ if (n > colBlockTotalNum) {
67+ continue;
68+ }
69+ 
70+ int64_t rowNormalBlock = Ops::Base::CeilDiv(rowTotalNum, m);
71+ int64_t mReal = Ops::Base::CeilDiv(rowTotalNum, rowNormalBlock);
72+ int64_t rowTailBlock = rowTotalNum - (mReal - 1) * rowNormalBlock;
73+ 
74+ int64_t colNormalBlock = Ops::Base::CeilDiv(colBlockTotalNum, n);
75+ int64_t nReal = Ops::Base::CeilDiv(colBlockTotalNum, colNormalBlock);
76+ int64_t colTailBlock = colBlockTotalNum - (nReal - 1) * colNormalBlock;
77+ 
78+ int64_t blockNormal = rowNormalBlock * colNormalBlock;
79+ int64_t blockTail = rowTailBlock * colTailBlock;
80+ int64_t delta = blockNormal - blockTail;
81+ allTiling.push_back({m, n, m * n, delta});
82+ }
83+ 
84+ std::sort(allTiling.begin(), allTiling.end(), [](const std::vector<int64_t>& a, const std::vector<int64_t>& b) {
85+ constexpr int NIndex = 1;
86+ constexpr int DeltaIndex = 3;
87+ return std::make_pair(a[NIndex], a[DeltaIndex]) < std::make_pair(b[NIndex], b[DeltaIndex]);
88+ });
89+ 
90+ int64_t allTilingSize = static_cast<int64_t>(allTiling.size());
91+ while (allTilingSize > 1 && static_cast<int64_t>(indicesNum_) / allTiling[0][0] < std::min<int64_t>(NUM_FOUR, indicesNum_)) {
92+ allTiling.erase(allTiling.begin());
93+ allTilingSize = static_cast<int64_t>(allTiling.size());
94+ }
95+ rowTileNum = static_cast<uint16_t>(allTiling[0][0]);
96+ colTileNum = static_cast<uint16_t>(allTiling[0][1]);
97+}
98+ 
99+void ScatterAddWithSortedSimdTiling::DoBlockTiling()
100+{
101+ ubBlock_ = static_cast<int64_t>(Ops::Base::GetUbBlockSize(context_));
102+ availableUbsize = ubSize_;
103+ availableUbsize = Ops::Base::FloorAlign(static_cast<int64_t>(availableUbsize / BUFFER_NUM), ubBlock_);
104+ FrontAndBackIndex = Ops::Base::CeilAlign(
105+ static_cast<int64_t>(2 * indicesDtypeSize_), ubBlock_);
106+ resUb = availableUbsize - 3 * static_cast<int64_t>(Ops::Base::CeilAlign(static_cast<int64_t>(varShape_[1] * updatesDtypeSize_), ubBlock_));
107+ resUb = isDeterminTemplate_ ? std::max<int64_t>(0, resUb - FrontAndBackIndex) : resUb;
108+ if (resUb >= COL_TILING_THRES) {
109+ coreNumInCol_ = 1;
110+ normalCoreColNum_ = varShape_[1];
111+ tailCoreColNum_ = varShape_[1];
112+ coreNumInRow_ = std::min(totalCoreNum_, indicesNum_);
113+ normalCoreRowNum_ = Ops::Base::CeilDiv(static_cast<int64_t>(indicesNum_), coreNumInRow_);
114+ coreNumInRow_ = Ops::Base::CeilDiv(static_cast<int64_t>(indicesNum_), normalCoreRowNum_);
115+ tailCoreRowNum_ = indicesNum_ - (coreNumInRow_ - 1) * normalCoreRowNum_;
116+ } else {
117+ int64_t rowTileNum = 0;
118+ int64_t colTileNum = 0;
119+ AutoTilingRowCol(rowTileNum, colTileNum, totalCoreNum_, indicesNum_, varShape_[1]);
120+ 
121+ normalCoreRowNum_ = Ops::Base::CeilDiv(static_cast<int64_t>(indicesNum_), rowTileNum);
122+ coreNumInRow_ = Ops::Base::CeilDiv(static_cast<int64_t>(indicesNum_), normalCoreRowNum_);
123+ tailCoreRowNum_ = indicesNum_ - (coreNumInRow_ - 1) * normalCoreRowNum_;
124+ 
125+ normalCoreColNum_ = Ops::Base::CeilDiv(static_cast<int64_t>(varShape_[1]), colTileNum);
126+ coreNumInCol_ = Ops::Base::CeilDiv(static_cast<int64_t>(varShape_[1]), normalCoreColNum_);
127+ tailCoreColNum_ = varShape_[1] - (coreNumInCol_ - 1) * normalCoreColNum_;
128+ }
129+ needCoreNum_ = coreNumInCol_ * coreNumInRow_;
130+}
131+ 
132+void ScatterAddWithSortedSimdTiling::DoUBTiling()
133+{
134+ int64_t rowNumInUb = 0;
135+ if (resUb >= COL_TILING_THRES) {
136+ int64_t colSizeAlign = Ops::Base::CeilAlign(
137+ static_cast<int64_t>(normalCoreColNum_ * updatesDtypeSize_), ubBlock_);
138+ int64_t tmpRowNum = Ops::Base::FloorAlign(resUb / 2, ubBlock_) / indicesDtypeSize_;
139+ rowNumInUb = std::min(tmpRowNum, normalCoreRowNum_);
140+ rowNumInUb = std::max<int64_t>(1, rowNumInUb);
141+ updatesBufferSize_ = colSizeAlign;
142+ outBufferSize_ = colSizeAlign;
143+ normalCoreColUbLoop_ = 1;
144+ normalCoreNormalLoopCols_ = normalCoreColNum_;
145+ normalCoreTailLoopCols_ = normalCoreColNum_;
146+ tailCoreColUbLoop_ = 1;
147+ tailCoreNormalLoopCols_ = tailCoreColNum_;
148+ tailCoreTailLoopCols_ = tailCoreColNum_;
149+ } else {
150+ int64_t tmpRowNum = COL_TILING_THRES / 2 / indicesDtypeSize_;
151+ rowNumInUb = std::min(tmpRowNum, normalCoreRowNum_);
152+ int64_t colSizeInUb = Ops::Base::FloorAlign((availableUbsize - COL_TILING_THRES) / 3, ubBlock_);
153+ updatesBufferSize_ = colSizeInUb;
154+ outBufferSize_ = colSizeInUb;
155+ int64_t colNumInUb = colSizeInUb / updatesDtypeSize_;
156+ colNumInUb = std::max<int64_t>(1, colNumInUb);
157+ normalCoreColUbLoop_ = Ops::Base::CeilDiv(normalCoreColNum_, colNumInUb);
158+ normalCoreNormalLoopCols_ = Ops::Base::CeilDiv(normalCoreColNum_, normalCoreColUbLoop_);
159+ normalCoreTailLoopCols_ = normalCoreColNum_ - (normalCoreColUbLoop_ - 1) * normalCoreNormalLoopCols_;
160+ tailCoreColUbLoop_ = Ops::Base::CeilDiv(tailCoreColNum_, colNumInUb);
161+ tailCoreNormalLoopCols_ = Ops::Base::CeilDiv(tailCoreColNum_, tailCoreColUbLoop_);
162+ tailCoreTailLoopCols_ = tailCoreColNum_ - (tailCoreColUbLoop_ - 1) * tailCoreNormalLoopCols_;
163+ }
164+ indicesBufferSize_ = Ops::Base::CeilAlign(static_cast<int64_t>(rowNumInUb * indicesDtypeSize_), ubBlock_);
165+ posBufferSize_ = Ops::Base::CeilAlign(static_cast<int64_t>(rowNumInUb * indicesDtypeSize_), ubBlock_);
166+ normalCoreRowUbLoop_ = Ops::Base::CeilDiv(normalCoreRowNum_, rowNumInUb);
167+ normalCoreNormalLoopRows_ = Ops::Base::CeilDiv(normalCoreRowNum_, normalCoreRowUbLoop_);
168+ normalCoreTailLoopRows_ = normalCoreRowNum_ - (normalCoreRowUbLoop_ - 1) * normalCoreNormalLoopRows_;
169+ tailCoreRowUbLoop_ = Ops::Base::CeilDiv(tailCoreRowNum_, rowNumInUb);
170+ tailCoreNormalLoopRows_ = Ops::Base::CeilDiv(tailCoreRowNum_, tailCoreRowUbLoop_);
171+ tailCoreTailLoopRows_ = tailCoreRowNum_ - (tailCoreRowUbLoop_ - 1) * tailCoreNormalLoopRows_;
172+}
173+ 
174+void ScatterAddWithSortedSimdTiling::DeterminTemplateUbTiling()
175+{
176+ vecAlignSize_ = Ops::Base::CeilAlign(static_cast<int64_t>(varShape_[1] * updatesDtypeSize_), BASE_BLOCK_ALIGN);
177+ int64_t resUbForDetermin = ubSize_ - 3 * static_cast<int64_t>(Ops::Base::CeilAlign(
178+ static_cast<int64_t>(varShape_[1] * updatesDtypeSize_), ubBlock_));
179+ // do UBTILING
180+ if (resUbForDetermin >=
181+ (coreNumInRow_ * CACHE_ALIGN_SIZE)) {
182+ int64_t colSizeAlignDetermin =
183+ Ops::Base::CeilAlign(static_cast<int64_t>(varShape_[1] * updatesDtypeSize_), ubBlock_);
184+ updatesDeterminBufferSize_ = colSizeAlignDetermin;
185+ outBufferDeterminSize_ = colSizeAlignDetermin;
186+ normalCoreColDetermNum_ = varShape_[1];
187+ colNumInUbDeterm = varShape_[1];
188+ colNumInUbDeterm = std::max<int64_t>(1, colNumInUbDeterm);
189+ coreNumInColDeterm_ = 1;
190+ tailCoreColNumDeterm_ = varShape_[1];
191+ 
192+ } else {
193+ int64_t resUbForUpdates = ubSize_ - Ops::Base::FloorAlign((coreNumInRow_ * CACHE_ALIGN_SIZE), ubBlock_);
194+ int64_t copyUpdatesInUb = Ops::Base::FloorAlign(resUbForUpdates / 3, ubBlock_);
195+ updatesDeterminBufferSize_ = copyUpdatesInUb;
196+ outBufferDeterminSize_ = copyUpdatesInUb;
197+ normalCoreColDetermNum_ = Ops::Base::CeilDiv(static_cast<int64_t>(varShape_[1]), coreNumInCol_);
198+ 
199+ colNumInUbDeterm = copyUpdatesInUb / updatesDtypeSize_;
200+ colNumInUbDeterm = std::max<int64_t>(1, colNumInUbDeterm);
201+ 
202+ coreNumInColDeterm_ = Ops::Base::CeilDiv(static_cast<int64_t>(varShape_[1]), normalCoreColDetermNum_);
203+ tailCoreColNumDeterm_ = varShape_[1] - (coreNumInColDeterm_ - 1) * normalCoreColDetermNum_;
204+ }
205+ 
206+ indicesWorkspaceBufferSize_ = Ops::Base::FloorAlign((coreNumInRow_ * CACHE_ALIGN_SIZE), ubBlock_);
207+ normalCoreColUbDetermLoop_ = Ops::Base::CeilDiv(normalCoreColDetermNum_, colNumInUbDeterm);
208+ normalCoreNormalLoopDetermCols_ = Ops::Base::CeilDiv(normalCoreColDetermNum_, normalCoreColUbDetermLoop_);
209+ normalCoreTailLoopDetermCols_ =
210+ normalCoreColDetermNum_ - (normalCoreColUbDetermLoop_ - 1) * normalCoreNormalLoopDetermCols_;
211+ 
212+ tailCoreColUbDetermLoop_ = Ops::Base::CeilDiv(tailCoreColNumDeterm_, colNumInUbDeterm);
213+ tailCoreNormalLoopDetermCols_ = Ops::Base::CeilDiv(tailCoreColNumDeterm_, tailCoreColUbDetermLoop_);
214+ tailCoreTailLoopDetermCols_ =
215+ tailCoreColNumDeterm_ - (tailCoreColUbDetermLoop_ - 1) * tailCoreNormalLoopDetermCols_;
216+}
217+ 
218+void ScatterAddWithSortedSimdTiling::SetTilingData()
219+{
220+ tilingData_ = context_->GetTilingData<ScatterAddWithSortedSimdTilingData>();
221+ tilingData_->needCoreNum = needCoreNum_;
222+ tilingData_->indicesNum = indicesNum_;
223+ tilingData_->updatesInner = varShape_[1];
224+ tilingData_->withPos = hasPos_;
225+ 
226+ tilingData_->updatesBufferSize = updatesBufferSize_;
227+ tilingData_->outBufferSize = outBufferSize_;
228+ tilingData_->indicesBufferSize = indicesBufferSize_;
229+ tilingData_->posBufferSize = posBufferSize_;
230+ tilingData_->FrontAndBackIndexSize = FrontAndBackIndex;
231+ 
232+ tilingData_->coreNumInRow = coreNumInRow_;
233+ tilingData_->coreNumInCol = coreNumInCol_;
234+ 
235+ tilingData_->normalCoreColNum = normalCoreColNum_;
236+ tilingData_->tailCoreColNum = tailCoreColNum_;
237+ tilingData_->normalCoreRowNum = normalCoreRowNum_;
238+ tilingData_->tailCoreRowNum = tailCoreRowNum_;
239+ 
240+ tilingData_->normalCoreRowUbLoop = normalCoreRowUbLoop_;
241+ tilingData_->normalCoreNormalLoopRows = normalCoreNormalLoopRows_;
242+ tilingData_->normalCoreTailLoopRows = normalCoreTailLoopRows_;
243+ tilingData_->tailCoreRowUbLoop = tailCoreRowUbLoop_;
244+ tilingData_->tailCoreNormalLoopRows = tailCoreNormalLoopRows_;
245+ tilingData_->tailCoreTailLoopRows = tailCoreTailLoopRows_;
246+ 
247+ tilingData_->normalCoreColUbLoop = normalCoreColUbLoop_;
248+ tilingData_->normalCoreNormalLoopCols = normalCoreNormalLoopCols_;
249+ tilingData_->normalCoreTailLoopCols = normalCoreTailLoopCols_;
250+ tilingData_->tailCoreColUbLoop = tailCoreColUbLoop_;
251+ tilingData_->tailCoreNormalLoopCols = tailCoreNormalLoopCols_;
252+ tilingData_->tailCoreTailLoopCols = tailCoreTailLoopCols_;
253+ 
254+ tilingData_->vecAlignSize = vecAlignSize_;
255+ tilingData_->indicesWorkspaceBufferSize = indicesWorkspaceBufferSize_;
256+ tilingData_->coreNumInColDeterm = coreNumInColDeterm_;
257+ tilingData_->tailCoreColUbDetermLoop = tailCoreColUbDetermLoop_;
258+ tilingData_->normalCoreColUbDetermLoop = normalCoreColUbDetermLoop_;
259+ tilingData_->tailCoreNormalLoopDetermCols = tailCoreNormalLoopDetermCols_;
260+ tilingData_->normalCoreNormalLoopDetermCols = normalCoreNormalLoopDetermCols_;
261+ tilingData_->tailCoreTailLoopDetermCols = tailCoreTailLoopDetermCols_;
262+ tilingData_->normalCoreTailLoopDetermCols = normalCoreTailLoopDetermCols_;
263+ tilingData_->updatesDeterminBufferSize = updatesDeterminBufferSize_;
264+ tilingData_->outBufferDeterminSize = outBufferDeterminSize_;
265+ tilingData_->normalCoreColDetermNum = normalCoreColDetermNum_;
266+ tilingData_->tailCoreColNumDeterm = tailCoreColNumDeterm_;
267+ tilingData_->ubBlock = ubBlock_;
268+ tilingData_->tilingKey = GetTilingKey();
269+ return;
270+}
271+ 
272+ge::graphStatus ScatterAddWithSortedSimdTiling::DoOpTiling()
273+{
274+ if (varShape_[0] * varShape_[1] * indicesNum_ == 0) {
275+ needCoreNum_ = 1;
276+ SetTilingData();
277+ return ge::GRAPH_SUCCESS;
278+ }
279+ 
280+ DoBlockTiling();
281+ DoUBTiling();
282+ 
283+ if (isDeterminTemplate_) {
284+ DeterminTemplateUbTiling();
285+ }
286+ 
287+ SetTilingData();
288+ 
289+ return ge::GRAPH_SUCCESS;
290+}
291+ 
292+uint64_t ScatterAddWithSortedSimdTiling::GetTilingKey() const
293+{
294+ if (varShape_[0] * varShape_[1] * indicesNum_ == 0) {
295+ return GET_TPL_TILING_KEY(TPL_MODE_EMPTY, TPL_SCALAR_FALSE, TPL_DETERM_FALSE, TPL_ADDR_B32);
296+ }
297+ 
298+ uint64_t isScalar = isUpdateScalar_ ? TPL_SCALAR_TRUE : TPL_SCALAR_FALSE;
299+ uint64_t isDeterm = isDeterminTemplate_ ? TPL_DETERM_TRUE : TPL_DETERM_FALSE;
300+ return GET_TPL_TILING_KEY(TPL_MODE_SIMD, isScalar, isDeterm, TPL_ADDR_B32);
301+}
302+ 
303+ge::graphStatus ScatterAddWithSortedSimdTiling::GetWorkspaceSize()
304+{
305+ size_t useWorkspace = 0;
306+ if (isDeterminTemplate_) {
307+ useWorkspace += coreNumInRow_ * 2 * vecAlignSize_ + coreNumInRow_ * CACHE_ALIGN_SIZE;
308+ }
309+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());
310+ uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
311+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
312+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
313+ currentWorkspace[0] = useWorkspace + sysWorkspaceSize;
314+ 
315+ return ge::GRAPH_SUCCESS;
316+}
317+ 
318+ge::graphStatus ScatterAddWithSortedSimdTiling::PostTiling()
319+{
320+ OP_LOGD(context_->GetNodeName(), "ScatterAddWithSortedSimdTiling simd PostTiling enter.");
321+ context_->SetBlockDim(needCoreNum_);
322+ context_->SetScheduleMode(1);
323+ return ge::GRAPH_SUCCESS;
324+}
325+ 
326+void ScatterAddWithSortedSimdTiling::DumpTilingInfo()
327+{
328+ std::ostringstream info1;
329+ info1 << "tilingKey: " << GetTilingKey();
330+ info1 << ", UB Size: " << ubSize_;
331+ info1 << ", needCoreNum: " << tilingData_->needCoreNum;
332+ info1 << ", indicesNum: " << tilingData_->indicesNum;
333+ info1 << ", updatesInner: " << tilingData_->updatesInner;
334+ 
335+ info1 << ", updatesBufferSize: " << tilingData_->updatesBufferSize;
336+ info1 << ", outBufferSize: " << tilingData_->outBufferSize;
337+ info1 << ", indicesBufferSize: " << tilingData_->indicesBufferSize;
338+ info1 << ", posBufferSize: " << tilingData_->posBufferSize;
339+ 
340+ info1 << ", coreNumInRow: " << tilingData_->coreNumInRow;
341+ info1 << ", coreNumInCol: " << tilingData_->coreNumInCol;
342+ 
343+ info1 << ", normalCoreColNum: " << tilingData_->normalCoreColNum;
344+ info1 << ", normalCoreRowNum: " << tilingData_->normalCoreRowNum;
345+ info1 << ", tailCoreRowNum: " << tilingData_->tailCoreRowNum;
346+ 
347+ info1 << ", normalCoreRowUbLoop: " << tilingData_->normalCoreRowUbLoop;
348+ info1 << ", tailCoreColNum: " << tilingData_->tailCoreColNum;
349+ info1 << ", normalCoreTailLoopRows: " << tilingData_->normalCoreTailLoopRows;
350+ info1 << ", tailCoreRowUbLoop: " << tilingData_->tailCoreRowUbLoop;
351+ info1 << ", tailCoreNormalLoopRows: " << tilingData_->tailCoreNormalLoopRows;
352+ info1 << ", tailCoreTailLoopRows: " << tilingData_->tailCoreTailLoopRows;
353+ OP_LOGI(context_->GetNodeName(), "%s", info1.str().c_str());
354+ 
355+ std::ostringstream info2;
356+ 
357+ info2 << ", normalCoreColUbLoop: " << tilingData_->normalCoreColUbLoop;
358+ info2 << ", normalCoreNormalLoopCols: " << tilingData_->normalCoreNormalLoopCols;
359+ info2 << ", normalCoreTailLoopCols: " << tilingData_->normalCoreTailLoopCols;
360+ info2 << ", tailCoreColUbLoop: " << tilingData_->tailCoreColUbLoop;
361+ info2 << ", tailCoreNormalLoopCols: " << tilingData_->tailCoreNormalLoopCols;
362+ info2 << ", tailCoreTailLoopCols: " << tilingData_->tailCoreTailLoopCols;
363+ 
364+ info2 << ", vecAlignSize: " << tilingData_->vecAlignSize;
365+ info2 << ", indicesWorkspaceBufferSize: " << tilingData_->indicesWorkspaceBufferSize;
366+ info2 << ", coreNumInColDeterm: " << tilingData_->coreNumInColDeterm;
367+ info2 << ", tailCoreColUbDetermLoop: " << tilingData_->tailCoreColUbDetermLoop;
368+ info2 << ", normalCoreColUbDetermLoop: " << tilingData_->normalCoreColUbDetermLoop;
369+ info2 << ", tailCoreNormalLoopDetermCols: " << tilingData_->tailCoreNormalLoopDetermCols;
370+ info2 << ", normalCoreNormalLoopDetermCols: " << tilingData_->normalCoreNormalLoopDetermCols;
371+ info2 << ", tailCoreTailLoopDetermCols: " << tilingData_->tailCoreTailLoopDetermCols;
372+ info2 << ", normalCoreTailLoopDetermCols: " << tilingData_->normalCoreTailLoopDetermCols;
373+ info2 << ", updatesDeterminBufferSize: " << tilingData_->updatesDeterminBufferSize;
374+ info2 << ", outBufferDeterminSize: " << tilingData_->outBufferDeterminSize;
375+ info2 << ", normalCoreColDetermNum: " << tilingData_->normalCoreColDetermNum;
376+ info2 << ", tailCoreColNumDeterm: " << tilingData_->tailCoreColNumDeterm;
377+ 
378+ OP_LOGI(context_->GetNodeName(), "%s", info2.str().c_str());
379+}
380+ 
381+REGISTER_TILING_TEMPLATE("ScatterAddWithSorted", ScatterAddWithSortedSimdTiling, 0);
382+} // namespace optiling
Aindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_simd_tiling.h+104-0
@@ -0,0 +1,104 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_add_with_sorted_simd_tiling.h
13+ * \brief scatter_add_with_sorted_simd_tiling
14+ */
15+#ifndef SCATTER_ADD_WITH_SORTED_SIMD_TILING_H
16+#define SCATTER_ADD_WITH_SORTED_SIMD_TILING_H
17+ 
18+#include "scatter_add_with_sorted_tiling_base.h"
19+#include "op_api/op_util.h"
20+#include "op_host/tiling_templates_registry.h"
21+#include "util/math_util.h"
22+ 
23+namespace optiling {
24+ 
25+class ScatterAddWithSortedSimdTiling : public ScatterAddWithSortedBaseTiling
26+{
27+public:
28+ explicit ScatterAddWithSortedSimdTiling(gert::TilingContext* context) : ScatterAddWithSortedBaseTiling(context)
29+ {}
30+ ~ScatterAddWithSortedSimdTiling() override
31+ {}
32+ 
33+private:
34+ bool IsCapable() override;
35+ ge::graphStatus DoOpTiling() override;
36+ ge::graphStatus PostTiling() override;
37+ ge::graphStatus GetWorkspaceSize() override;
38+ uint64_t GetTilingKey() const override;
39+ void DumpTilingInfo() override;
40+ void SetTilingData() override;
41+ void DoBlockTiling();
42+ void DoUBTiling();
43+ void AutoTilingRowCol(
44+ int64_t& rowTileNum, int64_t& colTileNum, int64_t usedCoreNum, int64_t rowTotalNum, int64_t colTotalNum);
45+ void DeterminTemplateUbTiling();
46+ 
47+private:
48+ ScatterAddWithSortedSimdTilingData* tilingData_;
49+ int64_t normalCoreColNum_ = 0;
50+ int64_t coreNumInCol_ = 0;
51+ int64_t tailCoreColNum_ = 0;
52+ int64_t coreNumInRow_ = 0;
53+ int64_t normalCoreRowNum_ = 0;
54+ int64_t tailCoreRowNum_ = 0;
55+ int64_t needCoreNum_ = 0;
56+ int64_t availableUbsize = 0;
57+ int64_t updatesBufferSize_ = 0;
58+ int64_t indicesBufferSize_ = 0;
59+ int64_t FrontAndBackIndex = 0;
60+ int64_t posBufferSize_ = 0;
61+ int64_t outBufferSize_ = 0;
62+ int64_t normalCoreRowUbLoop_ = 0;
63+ int64_t normalCoreNormalLoopRows_ = 0;
64+ int64_t normalCoreTailLoopRows_ = 0;
65+ int64_t tailCoreRowUbLoop_ = 0;
66+ int64_t tailCoreNormalLoopRows_ = 0;
67+ int64_t tailCoreTailLoopRows_ = 0;
68+ int64_t normalCoreColUbLoop_ = 0;
69+ int64_t normalCoreTailLoopCols_ = 0;
70+ int64_t normalCoreNormalLoopCols_ = 0;
71+ int64_t tailCoreColUbLoop_ = 0;
72+ int64_t tailCoreNormalLoopCols_ = 0;
73+ int64_t tailCoreTailLoopCols_ = 0;
74+ int64_t vecAlignSize_ = 0;
75+ int64_t colNumInUbDeterm = 0;
76+ 
77+ int64_t resUb = 0;
78+ int64_t usedCoreNumForDetermin = 0;
79+ int64_t normalCoreDeterminCols_ = 0;
80+ int64_t tailCoreDeterminCols_ = 0;
81+ int64_t normalCoreDeterminColsUbLoop_ = 0;
82+ int64_t updatesDeterminBufferSize_ = 0;
83+ int64_t outBufferDeterminSize_ = 0;
84+ int64_t indicesBufferDeterminSize_ = 0;
85+ 
86+ int64_t normalCoreDeterminNormalLoopCols_ = 0;
87+ int64_t normalCoreDeterminTailLoopCols_ = 0;
88+ int64_t tailCoreDeterminColsUbLoop_ = 0;
89+ int64_t tailCoreDeterminNormalLoopCols_ = 0;
90+ int64_t tailCoreDeterminTailLoopCols_ = 0;
91+ int64_t indicesWorkspaceBufferSize_ = 0;
92+ int64_t normalCoreColUbDetermLoop_ = 0;
93+ int64_t normalCoreNormalLoopDetermCols_ = 0;
94+ int64_t normalCoreTailLoopDetermCols_ = 0;
95+ int64_t normalCoreColDetermNum_ = 0;
96+ int64_t coreNumInColDeterm_ = 0;
97+ int64_t tailCoreColNumDeterm_ = 0;
98+ int64_t tailCoreColUbDetermLoop_ = 0;
99+ int64_t tailCoreNormalLoopDetermCols_ = 0;
100+ int64_t tailCoreTailLoopDetermCols_ = 0;
101+ int64_t ubBlock_ = 0;
102+};
103+} // namespace optiling
104+#endif // SCATTER_ADD_WITH_SORTED_SIMD_TILING_H
Aindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_simt_tiling.cpp+118-0
@@ -0,0 +1,118 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_add_with_sorted_simt_tiling.cpp
13+ * \brief scatter_add_with_sorted_simt_tiling
14+ */
15+ 
16+#include "scatter_add_with_sorted_simt_tiling.h"
17+ 
18+namespace optiling {
19+ 
20+static constexpr int64_t SIMT_INNER_THRES = 512;
21+static constexpr int64_t DOUBLE = 2;
22+ 
23+bool ScatterAddWithSortedSimtTiling::IsCapable()
24+{
25+ bool isFloat = (varDtype_ == ge::DT_FLOAT || varDtype_ == ge::DT_FLOAT16 || varDtype_ == ge::DT_BF16);
26+ if (context_->GetDeterministic() && !isUpdateScalar_ && isFloat) {
27+ isDeterminTemplate_ = 1;
28+ }
29+ bool isSimt = varShape_[1] * updatesDtypeSize_ < SIMT_INNER_THRES;
30+ return isSimt;
31+}
32+ 
33+void ScatterAddWithSortedSimtTiling::SetTilingData()
34+{
35+ tilingData_ = context_->GetTilingData<ScatterAddWithSortedSimtTilingData>();
36+ tilingData_->varShape[0] = varShape_[0];
37+ tilingData_->varShape[1] = varShape_[1];
38+ tilingData_->indicesNum = indicesNum_;
39+ tilingData_->normBlockIndices = normBlockIndices_;
40+ tilingData_->tailBlockIndices = tailBlockIndices_;
41+ tilingData_->usedCoreNum = usedCoreNum_;
42+ tilingData_->withPos = hasPos_;
43+ tilingData_->tilingKey = GetTilingKey();
44+ 
45+ return;
46+}
47+ 
48+ge::graphStatus ScatterAddWithSortedSimtTiling::DoOpTiling()
49+{
50+ if (varShape_[0] * varShape_[1] * indicesNum_ == 0) {
51+ usedCoreNum_ = 1;
52+ SetTilingData();
53+ return ge::GRAPH_SUCCESS;
54+ }
55+ normBlockIndices_ = Ops::Base::CeilDiv(indicesNum_, totalCoreNum_);
56+ usedCoreNum_ = Ops::Base::CeilDiv(static_cast<int64_t>(indicesNum_), normBlockIndices_);
57+ tailBlockIndices_ = indicesNum_ - (usedCoreNum_ - 1) * normBlockIndices_;
58+ SetTilingData();
59+ 
60+ return ge::GRAPH_SUCCESS;
61+}
62+ 
63+uint64_t ScatterAddWithSortedSimtTiling::GetTilingKey() const
64+{
65+ if (varShape_[0] * varShape_[1] * indicesNum_ == 0) {
66+ return GET_TPL_TILING_KEY(TPL_MODE_EMPTY, TPL_SCALAR_FALSE, TPL_DETERM_FALSE, TPL_ADDR_B32);
67+ }
68+ 
69+ uint64_t isScalar = isUpdateScalar_ ? TPL_SCALAR_TRUE : TPL_SCALAR_FALSE;
70+ uint64_t isDeterm = isDeterminTemplate_ ? TPL_DETERM_TRUE : TPL_DETERM_FALSE;
71+ uint64_t addrType =
72+ ((varShape_[1] * indicesNum_ > INT32_MAX) || (varSize_ > INT32_MAX)) ? TPL_ADDR_B64 : TPL_ADDR_B32;
73+ return GET_TPL_TILING_KEY(TPL_MODE_SIMT, isScalar, isDeterm, addrType);
74+}
75+ 
76+ge::graphStatus ScatterAddWithSortedSimtTiling::GetWorkspaceSize()
77+{
78+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());
79+ workspacesSize_ = ascendcPlatform.GetLibApiWorkSpaceSize();
80+ if (isDeterminTemplate_) {
81+ uint32_t userWorkspacesSize =
82+ (usedCoreNum_ * DOUBLE * varShape_[1]) * varTypeSize_ + usedCoreNum_ * DOUBLE * indicesDtypeSize_;
83+ userWorkspacesSize = Ops::Base::CeilAlign(userWorkspacesSize, static_cast<uint32_t>(indicesDtypeSize_));
84+ workspacesSize_ += userWorkspacesSize;
85+ }
86+ 
87+ return ge::GRAPH_SUCCESS;
88+}
89+ 
90+ge::graphStatus ScatterAddWithSortedSimtTiling::PostTiling()
91+{
92+ OP_LOGD(context_->GetNodeName(), "ScatterAddWithSortedSimtTiling simt PostTiling enter.");
93+ auto workspaces = context_->GetWorkspaceSizes(1);
94+ OP_CHECK_NULL_WITH_CONTEXT(context_, workspaces);
95+ workspaces[0] = workspacesSize_;
96+ 
97+ context_->SetBlockDim(usedCoreNum_);
98+ context_->SetScheduleMode(1);
99+ 
100+ return ge::GRAPH_SUCCESS;
101+}
102+ 
103+void ScatterAddWithSortedSimtTiling::DumpTilingInfo()
104+{
105+ std::ostringstream info;
106+ info << "tilingKey: " << GetTilingKey();
107+ info << ", UB Size: " << ubSize_;
108+ info << ", usedCoreNum: " << tilingData_->usedCoreNum;
109+ info << ", varShape[0]: " << tilingData_->varShape[0];
110+ info << ", varShape[1]: " << tilingData_->varShape[1];
111+ info << ", normBlockIndices: " << tilingData_->normBlockIndices;
112+ info << ", tailBlockIndices: " << tilingData_->tailBlockIndices;
113+ 
114+ OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str());
115+}
116+ 
117+REGISTER_TILING_TEMPLATE("ScatterAddWithSorted", ScatterAddWithSortedSimtTiling, 1);
118+} // namespace optiling
Aindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_simt_tiling.h+51-0
@@ -0,0 +1,51 @@
1+/**
2+* Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+* CANN Open Software License Agreement Version 2.0 (the "License")
5+* Please refer to the License for details. You may not use this file except in compliance with the License.
6+* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+* See LICENSE in the root of the software repository for the full text of the License.
9+*/
10+ 
11+/*!
12+* \file scatter_add_with_sorted_simt_tiling.h
13+* \brief scatter_add_with_sorted_simt_tiling
14+*/
15+#ifndef SCATTER_ADD_WITH_SORTED_SIMT_TILING_H
16+#define SCATTER_ADD_WITH_SORTED_SIMT_TILING_H
17+ 
18+#include "scatter_add_with_sorted_tiling_base.h"
19+#include "op_api/op_util.h"
20+#include "op_host/tiling_templates_registry.h"
21+#include "util/math_util.h"
22+ 
23+namespace optiling {
24+ 
25+class ScatterAddWithSortedSimtTiling : public ScatterAddWithSortedBaseTiling
26+{
27+public:
28+ explicit ScatterAddWithSortedSimtTiling(gert::TilingContext* context) : ScatterAddWithSortedBaseTiling(context)
29+ {}
30+ ~ScatterAddWithSortedSimtTiling() override
31+ {}
32+ 
33+private:
34+ bool IsCapable() override;
35+ ge::graphStatus DoOpTiling() override;
36+ ge::graphStatus PostTiling() override;
37+ ge::graphStatus GetWorkspaceSize() override;
38+ uint64_t GetTilingKey() const override;
39+ void DumpTilingInfo() override;
40+ void SetTilingData() override;
41+ 
42+private:
43+ ScatterAddWithSortedSimtTilingData* tilingData_;
44+ int64_t normBlockIndices_ = 0;
45+ int64_t tailBlockIndices_ = 0;
46+ int64_t usedCoreNum_ = 0;
47+ int64_t workspacesSize_ = 0;
48+
49+};
50+}
51+#endif // SCATTER_ADD_WITH_SORTED_SIMT_TILING_H
Mindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_tiling.cpp+14-0
@@ -19,8 +19,12 @@
19#include "platform/platform_info.h"19#include "platform/platform_info.h"
20#include "tiling/platform/platform_ascendc.h"20#include "tiling/platform/platform_ascendc.h"
21#include "scatter_add_with_sorted_tiling.h"21#include "scatter_add_with_sorted_tiling.h"
22+#include "op_host/tiling_util.h"
23+#include "index/scatter_add_with_sorted/op_host/scatter_add_with_sorted_tiling_base.h"
22 24 
23using namespace std;25using namespace std;
26+using Ops::NN::Optiling::TilingRegistry;
27+using namespace AscendC;
24 28 
25namespace {29namespace {
26const int DT_FLOAT32_TYPE = 1;30const int DT_FLOAT32_TYPE = 1;
@@ -317,8 +321,18 @@ void ScatterAddWithSortedTiling::TilingDataPrint() const
317 OP_LOGD(tilingContext, "max_ub: %lu.", max_ub);321 OP_LOGD(tilingContext, "max_ub: %lu.", max_ub);
318}322}
319 323 
324+ge::graphStatus ScatterAddWithSortedTilingForAscendC(gert::TilingContext* context)
325+{
326+ return Ops::NN::Optiling::TilingRegistry::GetInstance().DoTilingImpl(context);
327+}
328+ 
320ge::graphStatus TilingScatterAddWithSorted(gert::TilingContext* context)329ge::graphStatus TilingScatterAddWithSorted(gert::TilingContext* context)
321{330{
331+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context)) {
332+ OP_LOGD(context->GetNodeName(), "ScatterAddWithSorted is ascendc. runing ascendc tiling.");
333+ return ScatterAddWithSortedTilingForAscendC(context);
334+ }
335+
322 ScatterAddWithSortedTiling tilingObject(context);336 ScatterAddWithSortedTiling tilingObject(context);
323 if (tilingObject.Init() != ge::GRAPH_SUCCESS) {337 if (tilingObject.Init() != ge::GRAPH_SUCCESS) {
324 return ge::GRAPH_FAILED;338 return ge::GRAPH_FAILED;
Aindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_tiling_base.cpp+231-0
@@ -0,0 +1,231 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_add_with_sorted_tiling_base.cpp
13+ * \brief scatter_add_with_sorted_tiling_base
14+ */
15+ 
16+#include "scatter_add_with_sorted_tiling_base.h"
17+ 
18+using namespace AscendC;
19+using namespace ge;
20+ 
21+namespace optiling {
22+ 
23+constexpr int64_t VAR_IDX = 0;
24+constexpr int64_t SORTED_INDEX_IDX = 2;
25+constexpr int64_t UPDATES_IDX = 1;
26+constexpr int64_t POS_IDX = 3;
27+ 
28+static constexpr int64_t BASE_BLOCK_ALIGN = 512;
29+static constexpr int64_t SINGLE_CORE_THRESHOLD = 4 * 1024;
30+static constexpr int64_t BLOCK_TILING_THRES = 512;
31+static constexpr int64_t INNER_ADD_NUM = 128;
32+ 
33+static const std::set<ge::DataType> INDICES_DTYPE_SET = {ge::DT_INT32, ge::DT_INT64};
34+static const std::set<ge::DataType> VAR_DTYPE_SET = {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_BF16};
35+ 
36+static std::string ToString(std::set<ge::DataType> supportDtypes)
37+{
38+ std::stringstream ss;
39+ for (const auto& element : supportDtypes) {
40+ ss << element << " ";
41+ }
42+ 
43+ return ss.str();
44+}
45+ 
46+template <typename T>
47+static std::string ToString(const T* value, size_t size)
48+{
49+ std::string r = "[";
50+ for (size_t i = 0; i < size; i++) {
51+ r = r + std::to_string(value[i]) + ", ";
52+ }
53+ r = r + "]";
54+ return r;
55+}
56+ 
57+ge::graphStatus ScatterAddWithSortedBaseTiling::GetPlatformInfo()
58+{
59+ auto platformInfo = context_->GetPlatformInfo();
60+ OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(opName, "fail to get platform info"), return ge::GRAPH_FAILED);
61+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
62+ auto aivNum = ascendcPlatform.GetCoreNumAiv();
63+ OP_CHECK_IF((aivNum <= 0), OP_LOGE(opName, "fail to get coreNum."), return ge::GRAPH_FAILED);
64+ totalCoreNum_ = aivNum;
65+ uint64_t ubSizePlatForm = 0;
66+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
67+ ubSize_ = ubSizePlatForm;
68+ return ge::GRAPH_SUCCESS;
69+}
70+ 
71+ge::graphStatus ScatterAddWithSortedBaseTiling::GetShapeAttrsInfo()
72+{
73+ auto var = context_->GetInputShape(VAR_IDX);
74+ OP_CHECK_NULL_WITH_CONTEXT(context_, var);
75+ auto varShape = var->GetStorageShape();
76+ varSize_ = varShape.GetShapeSize();
C
Cchenqi3174月28日

varShape_[1] = ( varShape_[0] != 0) ? varSize_ / varShape_[0] : 0

likedislike
W
wkkk0528
4月28日 评论：
77+ varShape_[0] = varShape.GetDim(0);
78+ varShape_[1] = (varShape_[0] != 0) ? varSize_ / varShape_[0] : 0;
79+ auto indices = context_->GetInputShape(SORTED_INDEX_IDX);
80+ OP_CHECK_NULL_WITH_CONTEXT(context_, indices);
81+ auto indiceShape = indices->GetStorageShape();
82+ indicesNum_ = indiceShape.GetShapeSize();
83+ 
84+ auto updates = context_->GetInputShape(UPDATES_IDX);
85+ OP_CHECK_NULL_WITH_CONTEXT(context_, updates);
86+ auto updateShape = updates->GetStorageShape();
87+ updatesSize_ = updateShape.GetShapeSize();
88+ uint64_t updatesDims = updateShape.GetDimNum();
89+ 
90+ if (updatesDims == 0 || (updatesDims == 1 && updatesSize_ == 1)) {
91+ isUpdateScalar_ = 1;
92+ } else {
93+ OP_CHECK_IF(
94+ CheckUpdatesShape(varShape, indiceShape, updateShape) != ge::GRAPH_SUCCESS,
95+ OP_LOGE(opName, "update shape check failed."), return ge::GRAPH_FAILED);
96+ }
97+ 
98+ auto pos = context_->GetInputShape(POS_IDX);
99+ if (pos != nullptr) {
100+ auto posShape = pos->GetStorageShape();
101+ OP_CHECK_IF(
102+ static_cast<uint64_t>(posShape.GetShapeSize()) != indicesNum_,
103+ OP_LOGE(opName, "pos shape must be equal to indices shape."), return ge::GRAPH_FAILED);
104+ hasPos_ = true;
105+ }
106+ 
C
Cchenqi3174月28日

格式

likedislike
W
wkkk0528
4月28日 评论：
107+ OP_CHECK_IF(
108+ CheckInputDtype() != ge::GRAPH_SUCCESS, OP_LOGE(opName, "input dtype check failed."), return ge::GRAPH_FAILED);
109+ 
110+ bool isFloat = (varDtype_ == ge::DT_FLOAT || varDtype_ == ge::DT_FLOAT16 || varDtype_ == ge::DT_BF16);
111+ if (context_->GetDeterministic() && !isUpdateScalar_ && isFloat) {
112+ isDeterminTemplate_ = 1;
113+ }
114+ return ge::GRAPH_SUCCESS;
115+}
116+ 
117+ge::graphStatus ScatterAddWithSortedBaseTiling::CheckInputDtype()
118+{
119+ auto indicesPtr = context_->GetInputDesc(SORTED_INDEX_IDX);
120+ OP_CHECK_NULL_WITH_CONTEXT(context_, indicesPtr);
121+ indicesDtype_ = indicesPtr->GetDataType();
122+ OP_CHECK_IF(
123+ (INDICES_DTYPE_SET.find(indicesDtype_) == INDICES_DTYPE_SET.end()),
124+ OP_LOGE(
125+ opName, "indices data dtype only support %s, currently, please check.",
126+ ToString(INDICES_DTYPE_SET).c_str()),
127+ return ge::GRAPH_FAILED);
128+ indicesDtypeSize_ = ge::GetSizeByDataType(indicesDtype_);
129+ OP_CHECK_IF(indicesDtypeSize_ <= 0, OP_LOGE(opName, "get indicesDtype size fail."), return ge::GRAPH_FAILED);
130+ 
131+ if (hasPos_) {
132+ auto posPtr = context_->GetInputDesc(POS_IDX);
133+ OP_CHECK_NULL_WITH_CONTEXT(context_, posPtr);
134+ posDtype_ = posPtr->GetDataType();
135+ OP_CHECK_IF(
136+ (INDICES_DTYPE_SET.find(posDtype_) == INDICES_DTYPE_SET.end()),
137+ OP_LOGE(
138+ opName, "pos data dtype only support %s, currently, please check.",
139+ ToString(INDICES_DTYPE_SET).c_str()),
140+ return ge::GRAPH_FAILED);
141+ posDtypeSize_ = ge::GetSizeByDataType(posDtype_);
142+ OP_CHECK_IF(posDtypeSize_ <= 0, OP_LOGE(opName, "get posDtype size fail."), return ge::GRAPH_FAILED);
143+ }
144+ 
145+ auto dataPtr = context_->GetInputDesc(VAR_IDX);
146+ OP_CHECK_NULL_WITH_CONTEXT(context_, dataPtr);
147+ varDtype_ = dataPtr->GetDataType();
148+ OP_CHECK_IF(
149+ (VAR_DTYPE_SET.find(varDtype_) == VAR_DTYPE_SET.end()),
150+ OP_LOGE(opName, "var data dtype only support %s, please check.", ToString(VAR_DTYPE_SET).c_str()),
151+ return ge::GRAPH_FAILED);
152+ varTypeSize_ = ge::GetSizeByDataType(varDtype_);
153+ OP_CHECK_IF(varTypeSize_ <= 0, OP_LOGE(opName, "get dataType size fail."), return ge::GRAPH_FAILED);
154+ auto updatePtr = context_->GetInputDesc(UPDATES_IDX);
155+ OP_CHECK_NULL_WITH_CONTEXT(context_, updatePtr);
156+ auto updatesType = updatePtr->GetDataType();
157+ OP_CHECK_IF(
158+ (VAR_DTYPE_SET.find(updatesType) == VAR_DTYPE_SET.end()),
159+ OP_LOGE(opName, "updates data dtype only support %s currently, please check.", ToString(VAR_DTYPE_SET).c_str()),
160+ return ge::GRAPH_FAILED);
161+ updatesDtypeSize_ = ge::GetSizeByDataType(updatesType);
162+ OP_CHECK_IF(
163+ (updatesType != varDtype_), OP_LOGE(opName, "expected updates dtype to be equal to var dtype, please check."),
164+ return ge::GRAPH_FAILED);
165+ return ge::GRAPH_SUCCESS;
166+}
167+ 
168+ge::graphStatus ScatterAddWithSortedBaseTiling::CheckUpdatesShape(
169+ const gert::Shape& varShape, const gert::Shape& indicesShape, const gert::Shape& updatesShape)
170+{
171+ uint64_t varDimNum = static_cast<uint64_t>(varShape.GetDimNum());
172+ uint64_t indicesDimNum = static_cast<uint64_t>(indicesShape.GetDimNum());
173+ uint64_t updatesDimNum = static_cast<uint64_t>(updatesShape.GetDimNum());
174+ OP_CHECK_IF(
175+ (updatesDimNum != indicesDimNum + varDimNum - 1),
176+ OP_LOGE(opName, "updatesDimNum must have the same number of indicesDimNum add varDimNum - 1, please check."),
177+ return ge::GRAPH_FAILED);
178+ for (uint64_t i = 0; i < indicesDimNum; i++) {
179+ OP_CHECK_IF(
180+ (static_cast<uint32_t>(updatesShape.GetDim(i)) != static_cast<uint32_t>(indicesShape.GetDim(i))),
181+ OP_LOGE(
182+ opName,
183+ "updatesShape should be equal to the shape of 'indices' concats the shape of 'var' except for the "
184+ "first dimension."),
185+ return ge::GRAPH_FAILED);
186+ }
187+ 
188+ for (uint64_t i = 1; i < varDimNum; i++) {
189+ OP_CHECK_IF(
190+ (static_cast<uint32_t>(updatesShape.GetDim(i + indicesDimNum - 1)) !=
191+ static_cast<uint32_t>(varShape.GetDim(i))),
192+ OP_LOGE(
193+ opName,
194+ "updatesShape should be equal to the shape of 'indices' concats the shape of 'var' except for the "
195+ "first dimension."),
196+ return ge::GRAPH_FAILED);
197+ }
198+ return ge::GRAPH_SUCCESS;
199+}
200+ 
201+bool ScatterAddWithSortedBaseTiling::IsCapable()
202+{
203+ return true;
204+}
205+ 
206+ge::graphStatus ScatterAddWithSortedBaseTiling::DoOpTiling()
207+{
208+ return ge::GRAPH_SUCCESS;
209+}
210+ 
211+ge::graphStatus ScatterAddWithSortedBaseTiling::DoLibApiTiling()
212+{
213+ return ge::GRAPH_SUCCESS;
214+}
215+ 
216+uint64_t ScatterAddWithSortedBaseTiling::GetTilingKey() const
217+{
218+ return 0;
219+}
220+ 
221+ge::graphStatus ScatterAddWithSortedBaseTiling::GetWorkspaceSize()
222+{
223+ return ge::GRAPH_SUCCESS;
224+}
225+ 
226+ge::graphStatus ScatterAddWithSortedBaseTiling::PostTiling()
227+{
228+ return ge::GRAPH_SUCCESS;
229+}
230+ 
231+} // namespace optiling
Aindex/scatter_add_with_sorted/op_host/scatter_add_with_sorted_tiling_base.h+84-0
@@ -0,0 +1,84 @@
1+/**
2+* Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+* CANN Open Software License Agreement Version 2.0 (the "License")
5+* Please refer to the License for details. You may not use this file except in compliance with the License.
6+* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+* See LICENSE in the root of the software repository for the full text of the License.
9+*/
10+ 
11+/*!
12+* \file scatter_add_with_sorted_tiling_base.h
13+* \brief scatter_add_with_sorted_tiling_base
14+*/
15+#ifndef SCATTER_ADD_WITH_SORTED_TILING_BASE_H
16+#define SCATTER_ADD_WITH_SORTED_TILING_BASE_H
17+ 
18+#include <cstdint>
19+
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "register/op_impl_registry.h"
22+#include "util/math_util.h"
23+#include "register/tilingdata_base.h"
24+#include "tiling/tiling_api.h"
25+#include "op_host/tiling_base.h"
26+#include "exe_graph/runtime/shape.h"
27+#include "op_host/tiling_templates_registry.h"
28+#include "log/log.h"
29+#include "error_util.h"
30+#include "op_common/op_host/util/platform_util.h"
31+#include "tiling/platform/platform_ascendc.h"
32+#include "op_common/atvoss/broadcast/broadcast_tiling.h"
33+#include "index/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_struct.h"
34+ 
35+namespace optiling {
36+ 
37+class ScatterAddWithSortedBaseTiling : public Ops::NN::Optiling::TilingBaseClass
38+{
39+public:
40+ explicit ScatterAddWithSortedBaseTiling(gert::TilingContext* context) : TilingBaseClass(context)
41+ {}
42+ ~ScatterAddWithSortedBaseTiling() override
43+ {}
44+ 
45+protected:
46+ bool IsCapable() override;
47+ ge::graphStatus GetShapeAttrsInfo() override;
48+ ge::graphStatus GetPlatformInfo() override;
49+ ge::graphStatus DoOpTiling() override;
50+ ge::graphStatus DoLibApiTiling() override;
51+ uint64_t GetTilingKey() const override;
52+ ge::graphStatus GetWorkspaceSize() override;
53+ ge::graphStatus PostTiling() override;
54+ void DumpTilingInfo() override
55+ {}
56+ virtual void SetTilingData() = 0;
57+ ge::graphStatus CheckInputDtype();
58+ ge::graphStatus CheckUpdatesShape(
59+ const gert::Shape& varShape, const gert::Shape& indicesShape, const gert::Shape& updatesShape);
60+ 
61+public:
62+ uint64_t varShape_[2] = {0, 0};
63+ uint64_t totalCoreNum_ = 1;
64+ uint64_t ubSize_ = 0;
65+ uint64_t varSize_ = 0;
66+ uint64_t indicesNum_ = 0;
67+ uint64_t updatesSize_ = 0;
68+ uint32_t isUpdateScalar_ = 0;
69+ uint64_t indicesDtypeSize_ = 0;
70+ uint64_t posDtypeSize_ = 0;
71+ uint64_t varTypeSize_ = 0;
72+ uint64_t updatesDtypeSize_ = 0;
73+ int64_t isDeterminTemplate_ = 0;
74+ 
75+ bool hasPos_ = false;
76+ 
77+ ge::DataType indicesDtype_ = ge::DT_UNDEFINED;
78+ ge::DataType varDtype_ = ge::DT_UNDEFINED;
79+ ge::DataType posDtype_ = ge::DT_UNDEFINED;
80+ 
81+ const char* opName = "ScatterAddWithSorted";
82+};
83+} // namespace optiling
84+#endif // SCATTER_ADD_WITH_SORTED_TILING_BASE_H
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_determ_workspace.h+254-0
@@ -0,0 +1,254 @@
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+#ifndef SCATTER_ADD_WITH_SORTED_SIMD_DETERM_WORKSPACE_H
11+#define SCATTER_ADD_WITH_SORTED_SIMD_DETERM_WORKSPACE_H
12+ 
13+#include "kernel_operator.h"
14+#include "kernel_tiling/kernel_tiling.h"
15+#include "../inc/platform.h"
16+#include "scatter_add_with_sorted_struct.h"
17+#include "scatter_add_with_sorted_simd_common.h"
18+ 
19+namespace ScatterAddWithSorted {
20+using namespace AscendC;
21+ 
22+template <typename T, typename U>
23+class ScatterAddWithSortedSimdDtermWorkspace {
24+public:
25+ __aicore__ inline ScatterAddWithSortedSimdDtermWorkspace(void)
26+ {}
27+ 
28+ __aicore__ inline void Init(
29+ GM_ADDR var, GM_ADDR varRef, GM_ADDR workspace, TPipe& pipeIn,
30+ const ScatterAddWithSortedSimdTilingData* tilingData);
31+ 
32+ __aicore__ inline void CopyInIndicesWorkspace(LocalTensor<U>& indicesWorkspaceLocal);
33+ 
34+ __aicore__ inline void CopyIn(
35+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
36+ uint32_t srcStride = 0, uint32_t dstStride = 0);
37+ 
38+ __aicore__ inline void CopyOut(
39+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
40+ uint32_t srcStride = 0, uint32_t dstStride = 0);
41+ 
42+ __aicore__ inline void Process();
43+ 
44+private:
45+ GlobalTensor<U> indexWorkspace_;
46+ GlobalTensor<T> sumWorkspace_;
47+ GlobalTensor<T> varRefGm_;
48+ 
49+ TQue<QuePosition::VECIN, kBufferNum> updatesQueue_;
50+ TBuf<QuePosition::VECCALC> indicesWorkspaceBuf_;
51+ TBuf<QuePosition::VECCALC> yBuf_;
52+ 
53+ const ScatterAddWithSortedSimdTilingData* tilingData_ = nullptr;
54+ 
55+ uint32_t blockIdx_ = 0;
56+ uint32_t rowCoreIdx_ = 0;
57+ uint32_t colCoreIdx_ = 0;
58+ 
59+ int64_t colUbLoop_ = 0;
60+ int64_t normalLoopCols_ = 0;
61+ int64_t tailLoopCols_ = 0;
62+ int64_t colGmOffset_ = 0;
63+ 
64+ constexpr static int32_t blockNumT_ = platform::GetUbBlockSize() / sizeof(T);
65+ constexpr static int32_t ubStride_ = platform::GetUbBlockSize() / sizeof(U); // 32B / sizeof(U)
66+};
67+ 
68+template <typename T, typename U>
69+__aicore__ inline void ScatterAddWithSortedSimdDtermWorkspace<T, U>::Init(
70+ GM_ADDR var, GM_ADDR varRef, GM_ADDR workspace, TPipe& pipeIn, const ScatterAddWithSortedSimdTilingData* tilingData)
71+{
72+ (void)var;
73+ 
74+ tilingData_ = tilingData;
75+ blockIdx_ = GetBlockIdx();
76+ 
77+ if (blockIdx_ >= tilingData_->needCoreNum) {
78+ return;
79+ }
80+ 
81+ rowCoreIdx_ = blockIdx_ / tilingData_->coreNumInCol;
82+ colCoreIdx_ = blockIdx_ % tilingData_->coreNumInCol;
83+ 
84+ colGmOffset_ = static_cast<int64_t>(colCoreIdx_) * tilingData_->normalCoreColDetermNum;
85+ 
86+ colUbLoop_ = (colCoreIdx_ == tilingData_->coreNumInColDeterm - 1) ? tilingData_->tailCoreColUbDetermLoop :
87+ tilingData_->normalCoreColUbDetermLoop;
88+ 
89+ normalLoopCols_ = (colCoreIdx_ == tilingData_->coreNumInColDeterm - 1) ?
90+ tilingData_->tailCoreNormalLoopDetermCols :
91+ tilingData_->normalCoreNormalLoopDetermCols;
92+ 
93+ tailLoopCols_ = (colCoreIdx_ == tilingData_->coreNumInColDeterm - 1) ? tilingData_->tailCoreTailLoopDetermCols :
94+ tilingData_->normalCoreTailLoopDetermCols;
95+ 
96+ int64_t sumAreaBytes = static_cast<int64_t>(tilingData_->coreNumInRow) * kDouble * tilingData_->vecAlignSize;
97+ int64_t indexOffsetU = (sumAreaBytes + sizeof(U) - 1) / sizeof(U);
98+ 
99+ indexWorkspace_.SetGlobalBuffer((__gm__ U*)workspace + indexOffsetU);
100+ sumWorkspace_.SetGlobalBuffer((__gm__ T*)workspace);
101+ varRefGm_.SetGlobalBuffer((__gm__ T*)varRef);
102+ 
103+ pipeIn.InitBuffer(updatesQueue_, kBufferNum, tilingData_->updatesDeterminBufferSize);
104+ pipeIn.InitBuffer(indicesWorkspaceBuf_, tilingData_->indicesWorkspaceBufferSize);
105+ pipeIn.InitBuffer(yBuf_, tilingData_->outBufferDeterminSize);
106+}
107+ 
108+template <typename T, typename U>
109+__aicore__ inline void ScatterAddWithSortedSimdDtermWorkspace<T, U>::CopyInIndicesWorkspace(
110+ LocalTensor<U>& indicesWorkspaceLocal)
111+{
112+ DataCopyPadExtParams<U> pad;
113+ pad.isPad = false;
114+ pad.leftPadding = 0;
115+ pad.rightPadding = 0;
116+ pad.paddingValue = 0;
117+ 
118+ DataCopyExtParams ext;
119+ ext.blockCount = tilingData_->coreNumInRow;
120+ ext.blockLen = kDouble * sizeof(U);
121+ ext.srcStride = kCacheLineSize - kDouble * sizeof(U);
122+ ext.dstStride = 0;
123+ 
124+ DataCopyPad(indicesWorkspaceLocal, indexWorkspace_, ext, pad);
125+}
126+ 
127+template <typename T, typename U>
128+__aicore__ inline void ScatterAddWithSortedSimdDtermWorkspace<T, U>::CopyIn(
129+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
130+ uint32_t srcStride, uint32_t dstStride)
131+{
132+ KernelUtil::CopyIn<T>(dstLocal, srcGm, offset, nBurst, copyLen, srcStride, dstStride);
133+}
134+ 
135+template <typename T, typename U>
136+__aicore__ inline void ScatterAddWithSortedSimdDtermWorkspace<T, U>::CopyOut(
137+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
138+ uint32_t srcStride, uint32_t dstStride)
139+{
140+ KernelUtil::CopyOut<T>(dstGm, srcLocal, offset, nBurst, copyLen, srcStride, dstStride);
141+}
142+ 
143+template <typename T, typename U>
144+__aicore__ inline void ScatterAddWithSortedSimdDtermWorkspace<T, U>::Process()
145+{
146+ if (blockIdx_ >= tilingData_->needCoreNum) {
147+ return;
148+ }
149+ 
150+ LocalTensor<U> indicesWorkspaceLocal = indicesWorkspaceBuf_.Get<U>();
151+ CopyInIndicesWorkspace(indicesWorkspaceLocal);
152+ 
153+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
154+ SetFlag<HardEvent::MTE2_S>(ev);
155+ WaitFlag<HardEvent::MTE2_S>(ev);
156+ 
157+ for (int64_t col = 0; col < colUbLoop_; col++) {
158+ int64_t curLoopCols = (col == colUbLoop_ - 1) ? tailLoopCols_ : normalLoopCols_;
159+ int64_t curLoopColsAlign = (curLoopCols + blockNumT_ - 1) / blockNumT_ * blockNumT_;
160+ int64_t colOffset = colGmOffset_ + col * normalLoopCols_;
161+ 
162+ U currentHead = indicesWorkspaceLocal.GetValue(ubStride_ * rowCoreIdx_ + 0);
163+ U currentTail = indicesWorkspaceLocal.GetValue(ubStride_ * rowCoreIdx_ + 1);
164+ 
165+ U prevTail = static_cast<U>(-1);
166+ if (rowCoreIdx_ > 0) {
167+ prevTail = indicesWorkspaceLocal.GetValue(ubStride_ * (rowCoreIdx_ - 1) + 1);
168+ }
169+ 
170+ if (currentHead != static_cast<U>(-1)) {
171+ bool handledByPrev = (rowCoreIdx_ > 0 && prevTail == currentHead);
172+ if (!handledByPrev) {
173+ uint64_t headOffset =
174+ (static_cast<uint64_t>(rowCoreIdx_) * kDouble) * (tilingData_->vecAlignSize / sizeof(T)) +
175+ static_cast<uint64_t>(colOffset);
176+ LocalTensor<T> upd = updatesQueue_.AllocTensor<T>();
177+ CopyIn(upd, sumWorkspace_, headOffset, 1, static_cast<uint32_t>(curLoopCols));
178+ updatesQueue_.EnQue(upd);
179+ upd = updatesQueue_.DeQue<T>();
180+ 
181+ LocalTensor<T> yLocal = yBuf_.Get<T>();
182+ DataCopy(yLocal, upd, curLoopColsAlign);
183+ updatesQueue_.FreeTensor(upd);
184+ 
185+ KernelUtil::WaitVToMte3();
186+ SetAtomicAdd<T>();
187+ uint64_t dstOffset = static_cast<uint64_t>(currentHead) * tilingData_->updatesInner + colOffset;
188+ CopyOut(varRefGm_, yLocal, dstOffset, 1, static_cast<uint32_t>(curLoopCols));
189+ SetAtomicNone();
190+ KernelUtil::WaitMte3ToV();
191+ }
192+ }
193+ 
194+ if (currentTail != static_cast<U>(-1)) {
195+ bool handledByPrev = (rowCoreIdx_ > 0 && prevTail == currentTail);
196+ if (!handledByPrev) {
197+ uint64_t tailOffset =
198+ (static_cast<uint64_t>(rowCoreIdx_) * kDouble + 1) * (tilingData_->vecAlignSize / sizeof(T)) +
199+ static_cast<uint64_t>(colOffset);
200+ 
201+ LocalTensor<T> upd = updatesQueue_.AllocTensor<T>();
202+ CopyIn(upd, sumWorkspace_, tailOffset, 1, static_cast<uint32_t>(curLoopCols));
203+ updatesQueue_.EnQue(upd);
204+ upd = updatesQueue_.DeQue<T>();
205+ 
206+ LocalTensor<T> yLocal = yBuf_.Get<T>();
207+ DataCopy(yLocal, upd, curLoopColsAlign);
208+ updatesQueue_.FreeTensor(upd);
209+ 
210+ int32_t nextCore = static_cast<int32_t>(rowCoreIdx_) + 1;
211+ U tailIdx = currentTail;
212+ 
213+ while (nextCore < static_cast<int32_t>(tilingData_->coreNumInRow)) {
214+ U nh = indicesWorkspaceLocal.GetValue(ubStride_ * nextCore + 0);
215+ U nt = indicesWorkspaceLocal.GetValue(ubStride_ * nextCore + 1);
216+ 
217+ bool merged = false;
218+ uint64_t off = 0;
219+ if ((nh == static_cast<U>(-1) && nt == tailIdx) || (nh == tailIdx && nh != static_cast<U>(-1))) {
220+ uint64_t coreMul = (nh == static_cast<U>(-1)) ?
221+ (static_cast<uint64_t>(nextCore) * kDouble + 1) :
222+ (static_cast<uint64_t>(nextCore) * kDouble);
223+ off = coreMul * (tilingData_->vecAlignSize / sizeof(T)) + static_cast<uint64_t>(colOffset);
224+ merged = true;
225+ }
226+ if (merged) {
227+ LocalTensor<T> tmp = updatesQueue_.AllocTensor<T>();
C
Cchenqi3174月28日

提取公共代码

likedislike
W
wkkk0528
4月28日 评论：
228+ CopyIn(tmp, sumWorkspace_, off, 1, static_cast<uint32_t>(curLoopCols));
229+ updatesQueue_.EnQue(tmp);
230+ tmp = updatesQueue_.DeQue<T>();
231+ 
232+ Add(yLocal, yLocal, tmp, curLoopCols);
233+ updatesQueue_.FreeTensor(tmp);
234+ }
235+ 
236+ if (!merged || nt == static_cast<U>(-1)) {
237+ break;
238+ }
239+ nextCore++;
240+ }
241+ 
242+ KernelUtil::WaitVToMte3();
243+ SetAtomicAdd<T>();
244+ uint64_t dstOffset = static_cast<uint64_t>(tailIdx) * tilingData_->updatesInner + colOffset;
245+ CopyOut(varRefGm_, yLocal, dstOffset, 1, static_cast<uint32_t>(curLoopCols));
246+ SetAtomicNone();
247+ KernelUtil::WaitMte3ToV();
248+ }
249+ }
250+ }
251+}
252+ 
253+} // namespace ScatterAddWithSorted
254+#endif // SCATTER_ADD_WITH_SORTED_SIMD_DETERM_WORKSPACE_H
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_simd.h+244-0
@@ -0,0 +1,244 @@
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+#ifndef SCATTER_ADD_WITH_SORTED_SIMD_H
11+#define SCATTER_ADD_WITH_SORTED_SIMD_H
12+ 
13+#include "kernel_operator.h"
14+#include "kernel_tiling/kernel_tiling.h"
15+#include "../inc/platform.h"
16+#include "scatter_add_with_sorted_struct.h"
17+#include "scatter_add_with_sorted_simd_common.h"
18+ 
19+namespace ScatterAddWithSorted {
20+using namespace AscendC;
21+ 
22+template <typename T, typename U, bool updatesIsScalar, bool withPos>
23+class ScatterAddWithSortedSIMD {
24+public:
25+ __aicore__ inline ScatterAddWithSortedSIMD(void)
26+ {}
27+ __aicore__ inline void Init(
28+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR varRef, GM_ADDR workspace, TPipe& pipeIn,
29+ const ScatterAddWithSortedSimdTilingData* tilingData);
30+ 
31+ __aicore__ inline void CopyIn(
32+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
33+ uint32_t srcStride = 0, uint32_t dstStride = 0);
34+ 
35+ __aicore__ inline void CopyInIndices(
36+ LocalTensor<U> dstLocal, GlobalTensor<U> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
37+ uint32_t srcStride = 0, uint32_t dstStride = 0);
38+ 
39+ __aicore__ inline void CopyOut(
40+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
41+ uint32_t srcStride = 0, uint32_t dstStride = 0);
42+ 
43+ __aicore__ inline void BroadcastUpdatesScalar(
44+ LocalTensor<T> updatesLocal, GlobalTensor<T> updatesGm, int32_t count);
45+ 
46+ __aicore__ inline void ComputeSumAndCopyOut(
47+ LocalTensor<T>& yLocal, int32_t curLoopRows, int32_t curLoopCols, int32_t curLoopColsAlign, int64_t colOffset,
48+ U& curId, int64_t indicesGmOffset);
49+ 
50+ __aicore__ inline void Process();
51+ 
52+private:
53+ GlobalTensor<T> varGm_;
54+ GlobalTensor<U> indicesGm_;
55+ GlobalTensor<U> posGm_;
56+ GlobalTensor<T> updatesGm_;
57+ GlobalTensor<T> varRefGm_;
58+ 
59+ TQue<QuePosition::VECIN, kBufferNum> updatesQueue_;
60+ TQue<QuePosition::VECIN, kBufferNum> indicesQueue_;
61+ TQue<QuePosition::VECIN, kBufferNum> posQueue_;
62+ TQue<QuePosition::VECOUT, TMP_BUFFER_NUM> tmpQue_;
63+ TBuf<QuePosition::VECCALC> yBuf_;
64+ 
65+ const ScatterAddWithSortedSimdTilingData* tilingData_ = nullptr;
66+ 
67+ uint32_t blockIdx_ = 0;
68+ uint32_t rowCoreIdx_ = 0;
69+ uint32_t colCoreIdx_ = 0;
70+ 
71+ int64_t rowGmOffset_ = 0;
72+ int64_t colGmOffset_ = 0;
73+ 
74+ int64_t rowUbLoop_ = 0;
75+ int64_t colUbLoop_ = 0;
76+ 
77+ int64_t normalLoopRows_ = 0;
78+ int64_t tailLoopRows_ = 0;
79+ int64_t normalLoopCols_ = 0;
80+ int64_t tailLoopCols_ = 0;
81+ 
82+ U preId_ = static_cast<U>(-1);
83+ 
84+ bool isStartRowCore_ = false;
85+ bool isEndRowCore_ = false;
86+ 
87+ constexpr static int32_t blockNumT_ = platform::GetUbBlockSize() / sizeof(T);
88+};
89+ 
90+template <typename T, typename U, bool updatesIsScalar, bool withPos>
91+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::Init(
92+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR varRef, GM_ADDR workspace, TPipe& pipeIn,
93+ const ScatterAddWithSortedSimdTilingData* tilingData)
94+{
95+ tilingData_ = tilingData;
96+ blockIdx_ = GetBlockIdx();
97+ if (!KernelUtil::InitSimdBase(
98+ tilingData_, blockIdx_, rowCoreIdx_, colCoreIdx_, isStartRowCore_, isEndRowCore_, rowGmOffset_,
99+ colGmOffset_, rowUbLoop_, colUbLoop_, normalLoopRows_, tailLoopRows_, normalLoopCols_, tailLoopCols_,
100+ updatesGm_, indicesGm_, updates, indices)) {
101+ return;
102+ }
103+ 
104+ if constexpr (withPos) {
105+ posGm_.SetGlobalBuffer((__gm__ U*)(pos));
106+ }
107+ varRefGm_.SetGlobalBuffer((__gm__ T*)(varRef));
108+ 
109+ pipeIn.InitBuffer(updatesQueue_, kBufferNum, tilingData_->updatesBufferSize);
110+ pipeIn.InitBuffer(indicesQueue_, kBufferNum, tilingData_->indicesBufferSize);
111+ if constexpr (withPos) {
112+ pipeIn.InitBuffer(posQueue_, kBufferNum, tilingData_->posBufferSize);
113+ }
114+ pipeIn.InitBuffer(yBuf_, tilingData_->outBufferSize);
115+ pipeIn.InitBuffer(tmpQue_, TMP_BUFFER_NUM, tilingData_->outBufferSize);
116+}
117+ 
118+template <typename T, typename U, bool updatesIsScalar, bool withPos>
119+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::CopyIn(
120+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
121+ uint32_t srcStride, uint32_t dstStride)
122+{
123+ KernelUtil::CopyIn<T>(dstLocal, srcGm, offset, nBurst, copyLen, srcStride, dstStride);
124+}
125+ 
126+template <typename T, typename U, bool updatesIsScalar, bool withPos>
127+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::CopyInIndices(
128+ LocalTensor<U> dstLocal, GlobalTensor<U> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
129+ uint32_t srcStride, uint32_t dstStride)
130+{
131+ KernelUtil::CopyInIndices<U>(dstLocal, srcGm, offset, nBurst, copyLen, srcStride, dstStride);
132+}
133+ 
134+template <typename T, typename U, bool updatesIsScalar, bool withPos>
135+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::CopyOut(
136+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
137+ uint32_t srcStride, uint32_t dstStride)
138+{
139+ KernelUtil::CopyOut<T>(dstGm, srcLocal, offset, nBurst, copyLen, srcStride, dstStride);
140+}
141+ 
142+template <typename T, typename U, bool updatesIsScalar, bool withPos>
143+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::BroadcastUpdatesScalar(
144+ LocalTensor<T> updatesLocal, GlobalTensor<T> updatesGm, int32_t count)
145+{
146+ KernelUtil::BroadcastScalar<T>(updatesLocal, updatesGm, static_cast<uint32_t>(count));
147+}
148+ 
149+template <typename T, typename U, bool updatesIsScalar, bool withPos>
150+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::ComputeSumAndCopyOut(
151+ LocalTensor<T>& yLocal, int32_t curLoopRows, int32_t curLoopCols, int32_t curLoopColsAlign, int64_t colOffset,
152+ U& curId, int64_t indicesGmOffset)
153+{
154+ LocalTensor<U> indicesLocal = indicesQueue_.DeQue<U>();
155+ LocalTensor<U> posLocal;
156+ if constexpr (withPos) {
157+ posLocal = posQueue_.DeQue<U>();
158+ }
159+ event_t evMte2S = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
160+ SetFlag<HardEvent::MTE2_S>(evMte2S);
161+ WaitFlag<HardEvent::MTE2_S>(evMte2S);
162+ 
163+ for (int32_t i = 0; i < curLoopRows; i++) {
164+ curId = indicesLocal.GetValue(i);
165+ if constexpr (updatesIsScalar) {
166+ LocalTensor<T> upd = updatesQueue_.AllocTensor<T>();
167+ BroadcastUpdatesScalar(upd, updatesGm_, curLoopColsAlign);
168+ updatesQueue_.EnQue(upd);
169+ } else {
170+ LocalTensor<T> upd = updatesQueue_.AllocTensor<T>();
171+ uint64_t updatesOffset = 0;
172+ if constexpr (withPos) {
173+ U posId = posLocal.GetValue(i);
174+ updatesOffset = static_cast<uint64_t>(posId) * tilingData_->updatesInner + colOffset;
175+ } else {
176+ updatesOffset = static_cast<uint64_t>(indicesGmOffset + i) * tilingData_->updatesInner + colOffset;
177+ }
178+ CopyIn(upd, updatesGm_, updatesOffset, 1, static_cast<uint32_t>(curLoopCols));
179+ updatesQueue_.EnQue(upd);
180+ }
181+ 
182+ LocalTensor<T> updatesLocal = updatesQueue_.DeQue<T>();
183+ 
184+ if (KernelUtil::AccumulateOrInit(yLocal, updatesLocal, preId_, curId, curLoopCols, curLoopColsAlign)) {
185+ LocalTensor<T> tmpLocal = tmpQue_.AllocTensor<T>();
186+ DataCopy(tmpLocal, yLocal, curLoopColsAlign);
187+ tmpQue_.EnQue(tmpLocal);
188+ LocalTensor<T> outLocal = tmpQue_.DeQue<T>();
189+ 
190+ uint64_t dstOffset = static_cast<uint64_t>(preId_) * tilingData_->updatesInner + colOffset;
191+ KernelUtil::WaitVToMte3();
192+ SetAtomicAdd<T>();
193+ CopyOut(varRefGm_, outLocal, dstOffset, 1, static_cast<uint32_t>(curLoopCols));
194+ SetAtomicNone();
195+ 
196+ tmpQue_.FreeTensor(outLocal);
197+ preId_ = curId;
198+ DataCopy(yLocal, updatesLocal, curLoopColsAlign);
199+ }
200+ updatesQueue_.FreeTensor(updatesLocal);
201+ }
202+ indicesQueue_.FreeTensor(indicesLocal);
203+ if constexpr (withPos) {
204+ posQueue_.FreeTensor(posLocal);
205+ }
206+}
207+ 
208+template <typename T, typename U, bool updatesIsScalar, bool withPos>
209+__aicore__ inline void ScatterAddWithSortedSIMD<T, U, updatesIsScalar, withPos>::Process()
210+{
211+ if (blockIdx_ >= tilingData_->needCoreNum) {
212+ return;
213+ }
214+ LocalTensor<T> yLocal = yBuf_.Get<T>();
215+ for (int64_t col = 0; col < colUbLoop_; col++) {
216+ preId_ = static_cast<U>(-1);
217+ U curId = static_cast<U>(-1);
218+ int64_t curLoopCols = (col == colUbLoop_ - 1) ? tailLoopCols_ : normalLoopCols_;
219+ int64_t curLoopColsAlign = (curLoopCols + blockNumT_ - 1) / blockNumT_ * blockNumT_;
220+ int64_t colOffset = colGmOffset_ + col * normalLoopCols_;
221+ 
222+ for (int64_t row = 0; row < rowUbLoop_; row++) {
223+ int64_t curLoopRows = (row == rowUbLoop_ - 1) ? tailLoopRows_ : normalLoopRows_;
224+ int64_t indicesGmOffset = rowGmOffset_ + row * normalLoopRows_;
225+ 
226+ KernelUtil::CopyInRowIndices<U, withPos>(
227+ indicesQueue_, posQueue_, indicesGm_, posGm_, indicesGmOffset, static_cast<uint32_t>(curLoopRows));
228+ ComputeSumAndCopyOut(
229+ yLocal, static_cast<int32_t>(curLoopRows), static_cast<int32_t>(curLoopCols),
230+ static_cast<int32_t>(curLoopColsAlign), colOffset, curId, indicesGmOffset);
231+ }
232+ if (preId_ != static_cast<U>(-1)) {
233+ KernelUtil::WaitVToMte3();
234+ SetAtomicAdd<T>();
235+ uint64_t dstOffset = static_cast<uint64_t>(preId_) * tilingData_->updatesInner + colOffset;
236+ CopyOut(varRefGm_, yLocal, dstOffset, 1, static_cast<uint32_t>(curLoopCols));
237+ SetAtomicNone();
238+ KernelUtil::WaitMte3ToV();
239+ }
240+ }
241+}
242+ 
243+} // namespace ScatterAddWithSorted
244+#endif // SCATTER_ADD_WITH_SORTED_SIMD_H
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_simd_common.h+188-0
@@ -0,0 +1,188 @@
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 SCATTER_ADD_WITH_SORTED_SIMD_COMMON_H
12+#define SCATTER_ADD_WITH_SORTED_SIMD_COMMON_H
13+ 
14+#include "kernel_operator.h"
15+#include "../inc/platform.h"
16+#include "scatter_add_with_sorted_struct.h"
17+ 
18+namespace ScatterAddWithSorted {
19+using namespace AscendC;
20+ 
21+constexpr uint32_t kBufferNum = 2; // double buffer
22+constexpr uint32_t kCacheLineSize = 128; // workspace index区按cacheline对齐
23+constexpr uint32_t kDouble = 2; // workspace: head/tail 两行
24+constexpr uint32_t TMP_BUFFER_NUM = 2;
25+namespace KernelUtil {
26+ 
27+__aicore__ inline void InitRowColTiling(
28+ const ScatterAddWithSortedSimdTilingData* tiling, uint32_t blockIdx, uint32_t& rowCoreIdx, uint32_t& colCoreIdx,
29+ bool& isStartRowCore, bool& isEndRowCore, int64_t& rowGmOffset, int64_t& colGmOffset, int64_t& rowUbLoop,
30+ int64_t& colUbLoop, int64_t& normalLoopRows, int64_t& tailLoopRows, int64_t& normalLoopCols, int64_t& tailLoopCols)
31+{
32+ rowCoreIdx = blockIdx / tiling->coreNumInCol;
33+ colCoreIdx = blockIdx % tiling->coreNumInCol;
34+ 
35+ isStartRowCore = (rowCoreIdx == 0);
36+ isEndRowCore = (rowCoreIdx == tiling->coreNumInRow - 1);
37+ 
38+ rowGmOffset = static_cast<int64_t>(rowCoreIdx) * tiling->normalCoreRowNum;
39+ colGmOffset = static_cast<int64_t>(colCoreIdx) * tiling->normalCoreColNum;
40+ 
41+ rowUbLoop = (rowCoreIdx == tiling->coreNumInRow - 1) ? tiling->tailCoreRowUbLoop : tiling->normalCoreRowUbLoop;
42+ 
43+ colUbLoop = (colCoreIdx == tiling->coreNumInCol - 1) ? tiling->tailCoreColUbLoop : tiling->normalCoreColUbLoop;
44+ 
45+ normalLoopRows =
46+ (rowCoreIdx == tiling->coreNumInRow - 1) ? tiling->tailCoreNormalLoopRows : tiling->normalCoreNormalLoopRows;
47+ 
48+ tailLoopRows =
49+ (rowCoreIdx == tiling->coreNumInRow - 1) ? tiling->tailCoreTailLoopRows : tiling->normalCoreTailLoopRows;
50+ 
51+ normalLoopCols =
52+ (colCoreIdx == tiling->coreNumInCol - 1) ? tiling->tailCoreNormalLoopCols : tiling->normalCoreNormalLoopCols;
53+ 
54+ tailLoopCols =
55+ (colCoreIdx == tiling->coreNumInCol - 1) ? tiling->tailCoreTailLoopCols : tiling->normalCoreTailLoopCols;
56+}
57+ 
58+template <typename T>
59+__aicore__ inline void CopyIn(
60+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
61+ uint32_t srcStride = 0, uint32_t dstStride = 0)
62+{
63+ DataCopyPadExtParams<T> pad;
64+ pad.isPad = false;
65+ pad.leftPadding = 0;
66+ pad.rightPadding = 0;
67+ pad.paddingValue = 0;
68+ 
69+ DataCopyExtParams ext;
70+ ext.blockCount = nBurst;
71+ ext.blockLen = copyLen * sizeof(T);
72+ ext.srcStride = srcStride * sizeof(T);
73+ ext.dstStride = dstStride * sizeof(T);
74+ 
75+ DataCopyPad(dstLocal, srcGm[offset], ext, pad);
76+}
77+ 
78+template <typename T>
79+__aicore__ inline void CopyOut(
80+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
81+ uint32_t srcStride = 0, uint32_t dstStride = 0)
82+{
83+ DataCopyExtParams ext;
84+ ext.blockCount = nBurst;
85+ ext.blockLen = copyLen * sizeof(T);
86+ ext.srcStride = srcStride * sizeof(T);
87+ ext.dstStride = dstStride * sizeof(T);
88+ 
89+ DataCopyPad(dstGm[offset], srcLocal, ext);
90+}
91+ 
92+template <typename U>
93+__aicore__ inline void CopyInIndices(
94+ LocalTensor<U> dstLocal, GlobalTensor<U> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
95+ uint32_t srcStride = 0, uint32_t dstStride = 0)
96+{
97+ CopyIn<U>(dstLocal, srcGm, offset, nBurst, copyLen, srcStride, dstStride);
98+}
99+ 
100+template <typename T>
101+__aicore__ inline void BroadcastScalar(LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint32_t count)
102+{
103+ T val = srcGm.GetValue(0);
104+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
105+ SetFlag<HardEvent::S_V>(ev);
106+ WaitFlag<HardEvent::S_V>(ev);
107+ Duplicate(dstLocal, val, count);
108+}
109+ 
110+__aicore__ inline void WaitVToMte3()
111+{
112+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
113+ SetFlag<HardEvent::V_MTE3>(ev);
114+ WaitFlag<HardEvent::V_MTE3>(ev);
115+}
116+ 
117+__aicore__ inline void WaitMte3ToV()
118+{
119+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
120+ SetFlag<HardEvent::MTE3_V>(ev);
121+ WaitFlag<HardEvent::MTE3_V>(ev);
122+}
123+ 
124+template <typename T>
125+__aicore__ inline void AtomicAddCopyOutSync(
126+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t dstOffset, uint32_t copyLen)
127+{
128+ WaitVToMte3();
129+ SetAtomicAdd<T>();
130+ CopyOut<T>(dstGm, srcLocal, dstOffset, 1, copyLen);
131+ SetAtomicNone();
132+ WaitMte3ToV();
133+}
134+ 
135+template <typename T, typename U>
136+__aicore__ inline bool InitSimdBase(
137+ const ScatterAddWithSortedSimdTilingData* tilingData, uint32_t blockIdx, uint32_t& rowCoreIdx, uint32_t& colCoreIdx,
138+ bool& isStartRowCore, bool& isEndRowCore, int64_t& rowGmOffset, int64_t& colGmOffset, int64_t& rowUbLoop,
139+ int64_t& colUbLoop, int64_t& normalLoopRows, int64_t& tailLoopRows, int64_t& normalLoopCols, int64_t& tailLoopCols,
140+ GlobalTensor<T>& updatesGm, GlobalTensor<U>& indicesGm, GM_ADDR updates, GM_ADDR indices)
141+{
142+ if (blockIdx >= tilingData->needCoreNum) {
143+ return false;
144+ }
145+ InitRowColTiling(
146+ tilingData, blockIdx, rowCoreIdx, colCoreIdx, isStartRowCore, isEndRowCore, rowGmOffset, colGmOffset, rowUbLoop,
147+ colUbLoop, normalLoopRows, tailLoopRows, normalLoopCols, tailLoopCols);
148+ updatesGm.SetGlobalBuffer((__gm__ T*)(updates));
149+ indicesGm.SetGlobalBuffer((__gm__ U*)(indices));
150+ return true;
151+}
152+ 
153+template <typename T, typename U>
154+__aicore__ inline bool AccumulateOrInit(
155+ LocalTensor<T>& yLocal, LocalTensor<T>& updatesLocal, U& preId, U curId, int32_t curLoopCols,
156+ int32_t curLoopColsAlign)
157+{
158+ if (curId == preId) {
159+ Add(yLocal, yLocal, updatesLocal, curLoopCols);
160+ return false;
161+ }
162+ if (preId == static_cast<U>(-1)) {
163+ DataCopy(yLocal, updatesLocal, curLoopColsAlign);
164+ preId = curId;
165+ return false;
166+ }
167+ return true;
168+}
169+ 
170+template <typename U, bool withPos>
171+__aicore__ inline void CopyInRowIndices(
172+ TQue<QuePosition::VECIN, kBufferNum>& indicesQueue, TQue<QuePosition::VECIN, kBufferNum>& posQueue,
173+ GlobalTensor<U>& indicesGm, GlobalTensor<U>& posGm, int64_t indicesGmOffset, uint32_t curLoopRows)
174+{
175+ LocalTensor<U> indicesLocal = indicesQueue.AllocTensor<U>();
176+ CopyInIndices<U>(indicesLocal, indicesGm, indicesGmOffset, 1, curLoopRows);
177+ indicesQueue.EnQue(indicesLocal);
178+ if constexpr (withPos) {
179+ LocalTensor<U> posLocal = posQueue.AllocTensor<U>();
180+ CopyInIndices<U>(posLocal, posGm, indicesGmOffset, 1, curLoopRows);
181+ posQueue.EnQue(posLocal);
182+ }
183+}
184+ 
185+} // namespace KernelUtil
186+} // namespace ScatterAddWithSorted
187+ 
188+#endif // SCATTER_ADD_WITH_SORTED_SIMD_COMMON_H
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_simd_determ.h+370-0
@@ -0,0 +1,370 @@
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 SCATTER_ADD_WITH_SORTED_SIMD_DETERM_H
12+#define SCATTER_ADD_WITH_SORTED_SIMD_DETERM_H
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "../inc/platform.h"
17+#include "scatter_add_with_sorted_struct.h"
18+#include "scatter_add_with_sorted_simd_common.h"
19+ 
20+namespace ScatterAddWithSorted {
21+using namespace AscendC;
22+ 
23+template <typename T, typename U, bool withPos>
24+class ScatterAddWithSortedSimdDterm {
25+public:
26+ __aicore__ inline ScatterAddWithSortedSimdDterm(void)
27+ {}
28+ 
29+ __aicore__ inline void CopyIn(
30+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
31+ uint32_t srcStride = 0, uint32_t dstStride = 0);
32+ __aicore__ inline void Init(
33+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR varRef, GM_ADDR workspace, TPipe& pipeIn,
34+ const ScatterAddWithSortedSimdTilingData* tilingData);
35+ __aicore__ inline void CopyOut(
36+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
37+ uint32_t srcStride = 0, uint32_t dstStride = 0);
38+ __aicore__ inline void CopyInIndices(
39+ LocalTensor<U> dstLocal, GlobalTensor<U> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
40+ uint32_t srcStride = 0, uint32_t dstStride = 0);
41+ __aicore__ inline void CopyOutIndexToWorkspace(LocalTensor<U>& tmpLocal);
42+ __aicore__ inline void CopyOutToWorkspace(
43+ LocalTensor<T>& dataLocal, int32_t burstLen, int64_t colOffset, int32_t writePosition);
44+ __aicore__ inline void Process();
45+ __aicore__ inline void ComputeSumAndCopyOut(
46+ LocalTensor<T>& yLocal, int32_t curLoopRows, int32_t curLoopCols, int32_t curLoopColsAlign, int64_t colOffset,
47+ U& curId, int64_t row, int64_t indicesGmOffset);
48+ 
49+private:
50+ GlobalTensor<T> varGm_;
51+ GlobalTensor<U> indicesGm_;
52+ GlobalTensor<U> posGm_;
53+ GlobalTensor<T> updatesGm_;
54+ GlobalTensor<T> varRefGm_;
55+ GlobalTensor<T> sumWorkspace_;
56+ GlobalTensor<U> indexWorkspace_;
57+ 
58+ TQue<QuePosition::VECIN, kBufferNum> updatesQueue_;
59+ TQue<QuePosition::VECIN, kBufferNum> indicesQueue_;
60+ TQue<QuePosition::VECIN, kBufferNum> posQueue_;
61+ TQue<QuePosition::VECIN, 1> lastCoreLastIndexQueue_;
62+ TQue<QuePosition::VECIN, 1> nextCoreFirstIndexQueue_;
63+ TQue<QuePosition::VECOUT, TMP_BUFFER_NUM> tmpQue_;
64+ TBuf<QuePosition::VECCALC> yBuf_;
65+ TBuf<QuePosition::VECCALC> tmpBuf_;
66+ 
67+ const ScatterAddWithSortedSimdTilingData* tilingData_ = nullptr;
68+ 
69+ uint32_t rowCoreIdx_ = 0;
70+ uint32_t blockIdx_ = 0;
71+ uint32_t colCoreIdx_ = 0;
72+ int64_t rowGmOffset_ = 0;
73+ int64_t rowUbLoop_ = 0;
74+ int64_t colGmOffset_ = 0;
75+ int64_t colUbLoop_ = 0;
76+ int64_t tailLoopCols_ = 0;
77+ int64_t tailLoopRows_ = 0;
78+ int64_t normalLoopRows_ = 0;
79+ int64_t normalLoopCols_ = 0;
80+ 
81+ U preId_ = static_cast<U>(-1);
82+ U prevCoreLastIndex_ = static_cast<U>(-2);
83+ U nextCoreFirstIndex_ = static_cast<U>(-2);
84+ 
85+ U headIndex_ = static_cast<U>(-1);
86+ U tailIndex_ = static_cast<U>(-1);
87+ 
88+ bool isStartRowCore_ = false;
89+ bool isEndRowCore_ = false;
90+ 
91+ bool headToWorkspace_ = false;
92+ bool tailToWorkspace_ = false;
93+ 
94+ constexpr static int32_t blockNumT_ = platform::GetUbBlockSize() / sizeof(T);
95+};
96+ 
97+template <typename T, typename U, bool withPos>
98+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::Init(
99+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR varRef, GM_ADDR workspace, TPipe& pipeIn,
100+ const ScatterAddWithSortedSimdTilingData* tilingData)
101+{
102+ tilingData_ = tilingData;
103+ blockIdx_ = GetBlockIdx();
104+ if (!KernelUtil::InitSimdBase(
105+ tilingData_, blockIdx_, rowCoreIdx_, colCoreIdx_, isStartRowCore_, isEndRowCore_, rowGmOffset_,
106+ colGmOffset_, rowUbLoop_, colUbLoop_, normalLoopRows_, tailLoopRows_, normalLoopCols_, tailLoopCols_,
107+ updatesGm_, indicesGm_, updates, indices)) {
108+ return;
109+ }
110+ 
111+ int64_t sumAreaBytes = static_cast<int64_t>(tilingData_->coreNumInRow) * kDouble * tilingData_->vecAlignSize;
112+ int64_t indexAreaOffsetBytes = sumAreaBytes + static_cast<int64_t>(rowCoreIdx_) * kCacheLineSize;
113+ 
114+ int64_t sumOffsetT = static_cast<int64_t>(rowCoreIdx_) * kDouble * tilingData_->vecAlignSize / sizeof(T);
115+ int64_t indexOffsetU = (indexAreaOffsetBytes + sizeof(U) - 1) / sizeof(U);
116+ 
117+ if constexpr (withPos) {
118+ posGm_.SetGlobalBuffer((__gm__ U*)(pos));
119+ }
120+ varRefGm_.SetGlobalBuffer((__gm__ T*)(varRef));
121+ 
122+ sumWorkspace_.SetGlobalBuffer((__gm__ T*)workspace + sumOffsetT);
123+ indexWorkspace_.SetGlobalBuffer((__gm__ U*)workspace + indexOffsetU);
124+ 
125+ pipeIn.InitBuffer(updatesQueue_, kBufferNum, tilingData_->updatesBufferSize);
126+ pipeIn.InitBuffer(indicesQueue_, kBufferNum, tilingData_->indicesBufferSize);
127+ if constexpr (withPos) {
128+ pipeIn.InitBuffer(posQueue_, kBufferNum, tilingData_->posBufferSize);
129+ }
130+ pipeIn.InitBuffer(lastCoreLastIndexQueue_, 1, tilingData_->FrontAndBackIndexSize / 2);
131+ pipeIn.InitBuffer(nextCoreFirstIndexQueue_, 1, tilingData_->FrontAndBackIndexSize / 2);
132+ 
133+ pipeIn.InitBuffer(yBuf_, tilingData_->outBufferSize);
134+ pipeIn.InitBuffer(tmpQue_, TMP_BUFFER_NUM, tilingData_->outBufferSize);
135+ pipeIn.InitBuffer(tmpBuf_, platform::GetUbBlockSize());
136+}
137+ 
138+template <typename T, typename U, bool withPos>
139+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::CopyIn(
140+ LocalTensor<T> dstLocal, GlobalTensor<T> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
141+ uint32_t srcStride, uint32_t dstStride)
142+{
143+ KernelUtil::CopyIn<T>(dstLocal, srcGm, offset, nBurst, copyLen, srcStride, dstStride);
144+}
145+ 
146+template <typename T, typename U, bool withPos>
147+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::CopyInIndices(
148+ LocalTensor<U> dstLocal, GlobalTensor<U> srcGm, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
149+ uint32_t srcStride, uint32_t dstStride)
150+{
151+ KernelUtil::CopyInIndices<U>(dstLocal, srcGm, offset, nBurst, copyLen, srcStride, dstStride);
152+}
153+ 
154+template <typename T, typename U, bool withPos>
155+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::CopyOut(
156+ GlobalTensor<T> dstGm, LocalTensor<T> srcLocal, uint64_t offset, uint32_t nBurst, uint32_t copyLen,
157+ uint32_t srcStride, uint32_t dstStride)
158+{
159+ KernelUtil::CopyOut<T>(dstGm, srcLocal, offset, nBurst, copyLen, srcStride, dstStride);
160+}
161+ 
162+template <typename T, typename U, bool withPos>
163+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::CopyOutToWorkspace(
164+ LocalTensor<T>& dataLocal, int32_t burstLen, int64_t colOffset, int32_t writePosition)
165+{
166+ DataCopyExtParams ext;
167+ ext.blockCount = 1;
168+ ext.blockLen = static_cast<uint32_t>(burstLen) * sizeof(T);
169+ ext.srcStride = 0;
170+ ext.dstStride = 0;
171+ 
172+ uint64_t workspaceOffset = static_cast<uint64_t>(writePosition) * (tilingData_->vecAlignSize / sizeof(T)) +
173+ static_cast<uint64_t>(colOffset);
174+ DataCopyPad(sumWorkspace_[workspaceOffset], dataLocal, ext);
175+}
176+ 
177+template <typename T, typename U, bool withPos>
178+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::CopyOutIndexToWorkspace(LocalTensor<U>& tmpLocal)
179+{
180+ DataCopyExtParams ext;
181+ ext.blockCount = 1;
182+ ext.blockLen = kDouble * sizeof(U);
183+ ext.srcStride = 0;
184+ ext.dstStride = 0;
185+ DataCopyPad(indexWorkspace_, tmpLocal, ext);
186+}
187+ 
188+template <typename T, typename U, bool withPos>
189+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::ComputeSumAndCopyOut(
190+ LocalTensor<T>& yLocal, int32_t curLoopRows, int32_t curLoopCols, int32_t curLoopColsAlign, int64_t colOffset,
191+ U& curId, int64_t row, int64_t indicesGmOffset)
192+{
193+ LocalTensor<U> indicesLocal = indicesQueue_.DeQue<U>();
194+ LocalTensor<U> posLocal;
195+ if constexpr (withPos) {
196+ posLocal = posQueue_.DeQue<U>();
197+ }
198+ event_t evMte2S = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
199+ SetFlag<HardEvent::MTE2_S>(evMte2S);
200+ WaitFlag<HardEvent::MTE2_S>(evMte2S);
201+ 
202+ if (row == 0) {
203+ headIndex_ = indicesLocal.GetValue(0);
204+ if (!isStartRowCore_ && headIndex_ == prevCoreLastIndex_) {
205+ headToWorkspace_ = true;
206+ }
207+ }
208+ 
209+ if (row == rowUbLoop_ - 1) {
210+ tailIndex_ = indicesLocal.GetValue(curLoopRows - 1);
211+ if (!isEndRowCore_ && tailIndex_ == nextCoreFirstIndex_) {
212+ tailToWorkspace_ = true;
213+ }
214+ }
215+ 
216+ for (int32_t i = 0; i < curLoopRows; i++) {
217+ curId = indicesLocal.GetValue(i);
218+ uint64_t updatesOffset = 0;
219+ if constexpr (withPos) {
220+ U posId = posLocal.GetValue(i);
221+ updatesOffset = static_cast<uint64_t>(posId) * tilingData_->updatesInner + colOffset;
222+ } else {
223+ updatesOffset = static_cast<uint64_t>(indicesGmOffset + i) * tilingData_->updatesInner + colOffset;
224+ }
225+ LocalTensor<T> updatesLocal = updatesQueue_.AllocTensor<T>();
226+ CopyIn(updatesLocal, updatesGm_, updatesOffset, 1, static_cast<uint32_t>(curLoopCols));
227+ updatesQueue_.EnQue(updatesLocal);
228+ updatesLocal = updatesQueue_.DeQue<T>();
229+ 
230+ if (KernelUtil::AccumulateOrInit(yLocal, updatesLocal, preId_, curId, curLoopCols, curLoopColsAlign)) {
231+ LocalTensor<T> tmpLocal = tmpQue_.AllocTensor<T>();
232+ DataCopy(tmpLocal, yLocal, curLoopColsAlign);
233+ tmpQue_.EnQue(tmpLocal);
234+ LocalTensor<T> outLocal = tmpQue_.DeQue<T>();
235+ 
236+ if (preId_ == headIndex_ && headToWorkspace_) {
237+ KernelUtil::WaitVToMte3();
238+ CopyOutToWorkspace(outLocal, curLoopCols, colOffset, 0);
239+ } else if (preId_ == tailIndex_ && tailToWorkspace_) {
240+ KernelUtil::WaitVToMte3();
241+ CopyOutToWorkspace(outLocal, curLoopCols, colOffset, 1);
242+ } else {
243+ uint64_t dstOffset = static_cast<uint64_t>(preId_) * tilingData_->updatesInner + colOffset;
244+ KernelUtil::WaitVToMte3();
245+ SetAtomicAdd<T>();
246+ CopyOut(varRefGm_, outLocal, dstOffset, 1, static_cast<uint32_t>(curLoopCols));
247+ SetAtomicNone();
248+ }
249+ 
250+ tmpQue_.FreeTensor(outLocal);
251+ preId_ = curId;
252+ DataCopy(yLocal, updatesLocal, curLoopColsAlign);
253+ }
254+ 
255+ updatesQueue_.FreeTensor(updatesLocal);
256+ }
257+ 
258+ indicesQueue_.FreeTensor(indicesLocal);
259+ if constexpr (withPos) {
260+ posQueue_.FreeTensor(posLocal);
261+ }
262+}
263+ 
264+template <typename T, typename U, bool withPos>
265+__aicore__ inline void ScatterAddWithSortedSimdDterm<T, U, withPos>::Process()
266+{
267+ if (blockIdx_ >= tilingData_->needCoreNum) {
268+ return;
269+ }
270+ 
271+ if (!isStartRowCore_) {
272+ LocalTensor<U> tmp = lastCoreLastIndexQueue_.AllocTensor<U>();
273+ CopyInIndices(tmp, indicesGm_, static_cast<uint64_t>(rowCoreIdx_) * tilingData_->normalCoreRowNum - 1, 1, 1);
274+ 
275+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
276+ SetFlag<HardEvent::MTE2_S>(ev);
277+ WaitFlag<HardEvent::MTE2_S>(ev);
278+ 
279+ prevCoreLastIndex_ = tmp.GetValue(0);
280+ lastCoreLastIndexQueue_.FreeTensor(tmp);
281+ } else {
282+ prevCoreLastIndex_ = static_cast<U>(-1);
283+ }
284+ 
285+ if (!isEndRowCore_) {
286+ LocalTensor<U> tmp = nextCoreFirstIndexQueue_.AllocTensor<U>();
287+ CopyInIndices(tmp, indicesGm_, static_cast<uint64_t>(rowCoreIdx_ + 1) * tilingData_->normalCoreRowNum, 1, 1);
288+ 
289+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
290+ SetFlag<HardEvent::MTE2_S>(ev);
291+ WaitFlag<HardEvent::MTE2_S>(ev);
292+ 
293+ nextCoreFirstIndex_ = tmp.GetValue(0);
294+ nextCoreFirstIndexQueue_.FreeTensor(tmp);
295+ } else {
296+ nextCoreFirstIndex_ = static_cast<U>(-1);
297+ }
298+ 
299+ LocalTensor<T> yLocal = yBuf_.Get<T>();
300+ 
301+ for (int64_t col = 0; col < colUbLoop_; col++) {
302+ preId_ = static_cast<U>(-1);
303+ U curId = static_cast<U>(-1);
304+ 
305+ headIndex_ = static_cast<U>(-1);
306+ tailIndex_ = static_cast<U>(-1);
307+ headToWorkspace_ = false;
308+ tailToWorkspace_ = false;
309+ 
310+ int64_t curLoopCols = (col == colUbLoop_ - 1) ? tailLoopCols_ : normalLoopCols_;
311+ int64_t colOffset = colGmOffset_ + col * normalLoopCols_;
312+ int64_t curLoopColsAlign = (curLoopCols + blockNumT_ - 1) / blockNumT_ * blockNumT_;
313+ 
314+ for (int64_t row = 0; row < rowUbLoop_; row++) {
315+ int64_t indicesGmOffset = rowGmOffset_ + row * normalLoopRows_;
316+ int64_t curLoopRows = (row == rowUbLoop_ - 1) ? tailLoopRows_ : normalLoopRows_;
317+ KernelUtil::CopyInRowIndices<U, withPos>(
318+ indicesQueue_, posQueue_, indicesGm_, posGm_, indicesGmOffset, static_cast<uint32_t>(curLoopRows));
319+ ComputeSumAndCopyOut(
320+ yLocal, static_cast<int32_t>(curLoopRows), static_cast<int32_t>(curLoopCols),
321+ static_cast<int32_t>(curLoopColsAlign), colOffset, curId, row, indicesGmOffset);
322+ }
323+ if (preId_ != static_cast<U>(-1)) {
324+ bool isHeadBoundary = (preId_ == headIndex_ && headToWorkspace_);
325+ bool isTailBoundary = (preId_ == tailIndex_ && tailToWorkspace_);
326+ 
327+ LocalTensor<T> tmpLocal = tmpQue_.AllocTensor<T>();
328+ DataCopy(tmpLocal, yLocal, curLoopColsAlign);
329+ tmpQue_.EnQue(tmpLocal);
330+ LocalTensor<T> outLocal = tmpQue_.DeQue<T>();
331+ 
332+ if (isHeadBoundary && isTailBoundary) {
333+ KernelUtil::WaitVToMte3();
334+ CopyOutToWorkspace(outLocal, curLoopCols, colOffset, 1);
335+ } else if (isTailBoundary) {
336+ KernelUtil::WaitVToMte3();
337+ CopyOutToWorkspace(outLocal, curLoopCols, colOffset, 1);
338+ } else if (isHeadBoundary) {
339+ KernelUtil::WaitVToMte3();
340+ CopyOutToWorkspace(outLocal, curLoopCols, colOffset, 0);
341+ } else {
342+ uint64_t dstOffset = static_cast<uint64_t>(preId_) * tilingData_->updatesInner + colOffset;
343+ KernelUtil::WaitVToMte3();
344+ SetAtomicAdd<T>();
345+ CopyOut(varRefGm_, outLocal, dstOffset, 1, static_cast<uint32_t>(curLoopCols));
346+ SetAtomicNone();
347+ }
348+ tmpQue_.FreeTensor(outLocal);
349+ }
350+ }
351+ LocalTensor<U> tmpIndicesLocal = tmpBuf_.Get<U>();
352+ U headFinalIndex = headToWorkspace_ ? headIndex_ : static_cast<U>(-1);
353+ U tailFinalIndex = tailToWorkspace_ ? tailIndex_ : static_cast<U>(-1);
354+ 
355+ if (headToWorkspace_ && tailToWorkspace_ && headIndex_ == tailIndex_) {
356+ headFinalIndex = static_cast<U>(-1);
357+ }
358+ 
359+ tmpIndicesLocal.SetValue(0, headFinalIndex);
360+ tmpIndicesLocal.SetValue(1, tailFinalIndex);
361+ 
362+ event_t ev = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
363+ SetFlag<HardEvent::S_MTE3>(ev);
364+ WaitFlag<HardEvent::S_MTE3>(ev);
365+ 
366+ CopyOutIndexToWorkspace(tmpIndicesLocal);
367+}
368+ 
369+} // namespace ScatterAddWithSorted
370+#endif // SCATTER_ADD_WITH_SORTED_SIMD_DETERM_H
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_simt.h+205-0
@@ -0,0 +1,205 @@
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 SCATTER_ADD_WITH_SORTED_SIMT_H
12+#define SCATTER_ADD_WITH_SORTED_SIMT_H
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "../inc/platform.h"
17+#include "scatter_add_with_sorted_struct.h"
18+#include "simt_api/common_functions.h"
19+#include "simt_api/asc_simt.h"
20+#include "simt_api/asc_fp16.h"
21+#include "simt_api/asc_bf16.h"
22+ 
23+namespace ScatterAddWithSorted {
24+using namespace AscendC;
25+ 
26+#ifdef __DAV_FPGA__
27+constexpr uint32_t THREAD_NUM = 256;
28+#else
29+constexpr uint32_t THREAD_NUM = 1024;
30+#endif
31+constexpr uint32_t THREAD_MIN_INDICES_NUM = 10;
32+constexpr uint32_t TEN = 10;
33+ 
34+template <typename T, typename U, typename ADDR_T, bool isUpdateScalar, bool withPos>
35+__simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM) inline void ScatterAddWithSortedSimtCompute(
36+ ADDR_T totalCol, ADDR_T indicesBlockOffset, ADDR_T indicesBlockNum, ADDR_T varFirstDimSize, ADDR_T magic,
37+ ADDR_T shift, T updateScalarValue, __gm__ T* var, __gm__ U* sortedIndices, __gm__ U* pos, __gm__ T* updates)
38+{
39+ ADDR_T startIndices = indicesBlockOffset;
40+ ADDR_T endIndices = startIndices + indicesBlockNum;
41+ ADDR_T updatesBlockOffset = startIndices * totalCol;
42+ ADDR_T updatesBlockNum = indicesBlockNum * totalCol;
43+ ADDR_T updatesBlockEnd = updatesBlockOffset + updatesBlockNum;
44+ for (ADDR_T i = updatesBlockOffset + threadIdx.x; i < updatesBlockEnd; i += blockDim.x) {
45+ ADDR_T indiceRow = Simt::UintDiv(i, magic, shift);
46+ U varRow = sortedIndices[indiceRow];
47+ if (varRow < 0 || varRow >= varFirstDimSize) {
48+ continue;
49+ }
50+ if (indiceRow != startIndices) {
51+ if (varRow == sortedIndices[indiceRow - 1]) { // 当前线程对应索引与前一个索引相同，由前一个线程处理
52+ continue;
53+ }
54+ }
55+ 
56+ ADDR_T colIdx = i - indiceRow * totalCol;
57+ T sumRes = 0;
58+ for (ADDR_T j = indiceRow; j < endIndices; j++) {
59+ if (sortedIndices[j] != varRow) {
60+ break;
61+ }
62+ if constexpr (isUpdateScalar) {
63+ sumRes += updateScalarValue;
64+ } else {
65+ if constexpr (withPos) {
66+ sumRes += updates[pos[j] * totalCol + colIdx];
67+ } else {
68+ sumRes += updates[j * totalCol + colIdx];
69+ }
70+ }
71+ }
72+ 
73+ ADDR_T varGmOffset = varRow * totalCol + colIdx;
74+ asc_atomic_add(var + varGmOffset, sumRes);
75+ }
76+}
77+ 
78+template <typename T, typename U, typename ADDR_T, bool isUpdateScalar, bool withPos>
79+__simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM) inline void ScatterAddWithSortedSimtComputeCutIndices(
80+ ADDR_T threadIndicesNum, ADDR_T tailThreadRowIndicesNum, ADDR_T indicesBlockOffset, ADDR_T varFirstDimSize,
81+ T updateScalarValue, __gm__ T* var, __gm__ U* sortedIndices, __gm__ U* pos, __gm__ T* updates)
82+{
83+ uint32_t threadIdxRow = threadIdx.y;
84+ uint32_t threadNumRow = blockDim.y;
85+ uint32_t colIdx = threadIdx.x;
86+ uint32_t totalCol = blockDim.x;
87+ 
88+ ADDR_T startIndices = indicesBlockOffset + threadIdxRow * threadIndicesNum;
89+ ADDR_T endIndices = startIndices + threadIndicesNum;
90+ if (threadIdxRow == threadNumRow - 1) {
91+ endIndices = startIndices + tailThreadRowIndicesNum;
92+ }
93+ 
94+ T sumRes = 0;
95+ U preIndices = sortedIndices[startIndices];
96+ U varRow = preIndices;
97+ for (ADDR_T i = startIndices; i < endIndices; i += 1) {
98+ varRow = sortedIndices[i];
99+ 
100+ if (preIndices != varRow && preIndices >= 0 && preIndices < varFirstDimSize) {
101+ ADDR_T varGmOffsetLast = preIndices * totalCol + colIdx;
102+ asc_atomic_add(var + varGmOffsetLast, sumRes);
103+ sumRes = 0;
104+ preIndices = varRow;
105+ }
106+ if constexpr (isUpdateScalar) {
107+ sumRes += updateScalarValue;
108+ } else {
109+ if constexpr (withPos) {
110+ sumRes += updates[pos[i] * totalCol + colIdx];
111+ } else {
112+ sumRes += updates[i * totalCol + colIdx];
113+ }
114+ }
115+ }
116+ if (varRow >= 0 && varRow < varFirstDimSize) {
117+ ADDR_T varGmOffset = varRow * totalCol + colIdx;
118+ asc_atomic_add(var + varGmOffset, sumRes);
119+ }
120+}
121+ 
122+template<typename T, typename U, typename ADDR_T, bool isUpdateScalar, bool withPos>
123+class ScatterAddWithSortedSIMT {
124+public:
125+ __aicore__ inline ScatterAddWithSortedSIMT(const ScatterAddWithSortedSimtTilingData& tilingData)
126+ : tilingData_(tilingData){};
127+ 
128+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos);
129+ __aicore__ inline void Process();
130+ 
131+private:
132+ AscendC::GlobalTensor<T> varGm_;
133+ AscendC::GlobalTensor<U> indicesGm_;
134+ AscendC::GlobalTensor<U> posGm_;
135+ AscendC::GlobalTensor<T> updatesGm_;
136+ const ScatterAddWithSortedSimtTilingData& tilingData_;
137+ 
138+ uint32_t blockIdx_ = 0;
139+ ADDR_T indicesBlockNum_ = 0;
140+};
141+ 
142+template<typename T, typename U, typename ADDR_T, bool isUpdateScalar, bool withPos>
143+__aicore__ inline void ScatterAddWithSortedSIMT<T, U, ADDR_T, isUpdateScalar, withPos>::Init(
144+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos)
145+{
146+ blockIdx_ = GetBlockIdx();
147+ if (blockIdx_ >= tilingData_.usedCoreNum) {
148+ return;
149+ }
150+
151+ varGm_.SetGlobalBuffer((__gm__ T *)(var));
152+ updatesGm_.SetGlobalBuffer((__gm__ T *)(updates));
153+ indicesGm_.SetGlobalBuffer((__gm__ U*)indices);
154+ posGm_.SetGlobalBuffer((__gm__ U*)pos);
155+ 
156+ indicesBlockNum_ = tilingData_.normBlockIndices;
157+ if (blockIdx_ == tilingData_.usedCoreNum - 1) {
158+ indicesBlockNum_ = tilingData_.tailBlockIndices;
159+ }
160+}
161+ 
162+template<typename T, typename U, typename ADDR_T, bool isUpdateScalar, bool withPos>
163+__aicore__ inline void ScatterAddWithSortedSIMT<T, U, ADDR_T, isUpdateScalar, withPos>::Process()
164+{
165+ if (blockIdx_ >= tilingData_.usedCoreNum) {
166+ return;
167+ }
168+ 
169+ ADDR_T totalCol = static_cast<uint32_t>(tilingData_.varShape[1]);
170+ ADDR_T varFirstDimSize = static_cast<ADDR_T>(tilingData_.varShape[0]);
171+ 
172+ uint32_t blockIdx = blockIdx_;
173+ T updateScalarValue = ((__gm__ T*)(updatesGm_.GetPhyAddr()))[0];
174+ ADDR_T indicesBlockNum = indicesBlockNum_;
175+ ADDR_T indicesBlockOffset = blockIdx * tilingData_.normBlockIndices;
176+ ADDR_T updatesBlockNum = indicesBlockNum_ * totalCol;
177+ 
178+ bool isNoCutIndices = tilingData_.indicesNum <= tilingData_.varShape[0] * TWO;
179+ isNoCutIndices = isNoCutIndices || (tilingData_.indicesNum <= tilingData_.varShape[0] * TEN && tilingData_.varShape[1] <= TWO);
180+ if (isNoCutIndices) {
181+ ADDR_T magic = 0;
182+ ADDR_T shift = 0;
183+ GetUintDivMagicAndShift(magic, shift, totalCol);
184+ asc_vf_call<ScatterAddWithSortedSimtCompute<T, U, ADDR_T, isUpdateScalar, withPos>>(dim3(THREAD_NUM),
185+ totalCol, indicesBlockOffset, indicesBlockNum, varFirstDimSize, magic, shift, updateScalarValue,
186+ (__gm__ T*)(varGm_.GetPhyAddr()), (__gm__ U*)(indicesGm_.GetPhyAddr()), (__gm__ U*)(posGm_.GetPhyAddr()),
187+ (__gm__ T*)(updatesGm_.GetPhyAddr()));
188+ } else {
189+ uint32_t currentMaxThread = updatesBlockNum >= THREAD_NUM ? THREAD_NUM : updatesBlockNum;
190+ uint32_t threadNumRow = currentMaxThread / totalCol;
191+ threadNumRow = AscendC::Std::min(static_cast<ADDR_T>(threadNumRow),
192+ ops::CeilDiv(indicesBlockNum, static_cast<ADDR_T>(THREAD_MIN_INDICES_NUM)));
193+ ADDR_T threadIndicesNum = ops::CeilDiv(indicesBlockNum, static_cast<ADDR_T>(threadNumRow));
194+ threadNumRow = ops::CeilDiv(indicesBlockNum, threadIndicesNum);
195+ ADDR_T tailThreadRowIndicesNum = indicesBlockNum - threadIndicesNum * (threadNumRow - 1);
196+
197+ asc_vf_call<ScatterAddWithSortedSimtComputeCutIndices<T, U, ADDR_T, isUpdateScalar, withPos>>(
198+ dim3({static_cast<uint32_t>(totalCol), threadNumRow}), threadIndicesNum, tailThreadRowIndicesNum,
199+ indicesBlockOffset, varFirstDimSize, updateScalarValue, (__gm__ T*)(varGm_.GetPhyAddr()),
200+ (__gm__ U*)(indicesGm_.GetPhyAddr()), (__gm__ U*)(posGm_.GetPhyAddr()), (__gm__ T*)(updatesGm_.GetPhyAddr()));
201+ }
202+}
203+ 
204+}
205+#endif
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_simt_determ.h+293-0
@@ -0,0 +1,293 @@
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 SCATTER_ADD_WITH_SORTED_DETERM_SIMT_H
12+#define SCATTER_ADD_WITH_SORTED_DETERM_SIMT_H
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "../inc/platform.h"
17+#include "scatter_add_with_sorted_struct.h"
18+#include "simt_api/common_functions.h"
19+#include "simt_api/asc_simt.h"
20+#include "simt_api/asc_fp16.h"
21+#include "simt_api/asc_bf16.h"
22+ 
23+namespace ScatterAddWithSorted {
24+using namespace AscendC;
25+ 
26+#ifdef __DAV_FPGA__
27+constexpr uint32_t THREAD_NUM_DETERM = 256;
28+#else
29+constexpr uint32_t THREAD_NUM_DETERM = 1024;
30+#endif
31+constexpr uint32_t TWO = 2;
32+ 
33+ 
34+template <typename T, typename U, typename ADDR_T, bool isStartRowCore, bool isEndRowCore, bool withPos>
35+__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM_DETERM) inline void ScatterAddWithSortedSimtDetermStep1(
36+ ADDR_T totalCol, ADDR_T indicesBlockOffset, ADDR_T indicesBlockNum, ADDR_T varFirstDimSize, ADDR_T magic,
37+ ADDR_T shift, uint32_t blockIdx, __gm__ T* var, __gm__ U* sortedIndices, __gm__ U* pos, __gm__ T* updates,
38+ __gm__ T* sumWorkspace, __gm__ U* indexWorkspace)
39+{
40+ ADDR_T startIndices = indicesBlockOffset;
41+ ADDR_T endIndices = startIndices + indicesBlockNum;
42+ ADDR_T updatesBlockOffset = startIndices * totalCol;
43+ ADDR_T updatesBlockNum = indicesBlockNum * totalCol;
44+ ADDR_T updatesBlockEnd = updatesBlockOffset + updatesBlockNum;
45+ 
46+ for (ADDR_T i = updatesBlockOffset + threadIdx.x; i < updatesBlockEnd; i += blockDim.x) {
47+ uint32_t flag1 = 0;
48+ uint32_t flag2 = 0;
49+ ADDR_T indiceRow = Simt::UintDiv(i, magic, shift);
50+ U varRow = sortedIndices[indiceRow];
51+ if (varRow < 0 || varRow >= varFirstDimSize) { // 跳过无效索引
52+ continue;
53+ }
54+ 
55+ bool isCoreFirstIndices = indiceRow == startIndices;
56+ if (!isCoreFirstIndices) {
57+ if (varRow == sortedIndices[indiceRow - 1]) { // 当前线程对应索引与前一个索引相同，由前一个线程处理
58+ continue;
59+ }
60+ }
61+ if constexpr (!isStartRowCore) {
62+ if (isCoreFirstIndices && varRow == sortedIndices[startIndices - 1]) { // 判断头索引与前一个核的尾索引是否相同
63+ flag1 = 1;
64+ }
65+ }
66+ 
67+ T sumRes = 0;
68+ ADDR_T currentRow = 0;
69+ ADDR_T colIdx = i - indiceRow * totalCol;
70+ for (ADDR_T j = indiceRow; j < endIndices; j++) {
71+ if (sortedIndices[j] != varRow) {
72+ break;
73+ }
74+ currentRow = j;
75+ if constexpr (withPos) {
76+ sumRes += updates[pos[j] * totalCol + colIdx];
77+ } else {
78+ sumRes += updates[j * totalCol + colIdx];
79+ }
80+ }
81+ 
82+ bool isCoreLastIndices = currentRow == endIndices - 1;
83+ if constexpr (!isEndRowCore) {
84+ if (isCoreLastIndices && varRow == sortedIndices[endIndices]) { // 判断尾索引与后一个核的头索引是否相同
85+ flag2 = 1;
86+ }
87+ }
88+ 
89+ ADDR_T varGmOffset = varRow * totalCol + colIdx;
90+ U indexWorkspace1 = 0;
91+ U indexWorkspace2 = 0;
92+ if (isCoreFirstIndices && isCoreLastIndices) {
93+ if (flag1 && flag2) { // 特殊场景，整个核处理的都是相同的索引，且与前后核相邻索引相同
94+ sumWorkspace[totalCol * TWO * blockIdx + colIdx] = 0;
95+ sumWorkspace[totalCol * (TWO * blockIdx + 1) + colIdx] = sumRes;
96+ indexWorkspace1 = varRow;
97+ indexWorkspace2 = varRow;
98+ } else if (flag1) {
99+ sumWorkspace[totalCol * TWO * blockIdx + colIdx] = sumRes;
100+ indexWorkspace1 = varRow;
101+ indexWorkspace2 = -1;
102+ } else if (flag2) {
103+ sumWorkspace[totalCol * (TWO * blockIdx + 1) + colIdx] = sumRes;
104+ indexWorkspace1 = -1;
105+ indexWorkspace2 = varRow;
106+ } else {
107+ asc_atomic_add(var + varGmOffset, sumRes);
108+ indexWorkspace1 = -1;
109+ indexWorkspace2 = -1;
110+ }
111+ if (colIdx == 0) {
112+ indexWorkspace[TWO * blockIdx] = indexWorkspace1;
113+ indexWorkspace[TWO * blockIdx + 1] = indexWorkspace2;
114+ }
115+ } else if (isCoreFirstIndices) {
116+ if (flag1) {
117+ sumWorkspace[totalCol * TWO * blockIdx + colIdx] = sumRes;
118+ indexWorkspace1 = varRow;
119+ } else {
120+ asc_atomic_add(var + varGmOffset, sumRes);
121+ indexWorkspace1 = -1;
122+ }
123+ if (colIdx == 0) {
124+ indexWorkspace[TWO * blockIdx] = indexWorkspace1;
125+ }
126+ } else if (isCoreLastIndices) {
127+ if (flag2) {
128+ sumWorkspace[totalCol * (TWO * blockIdx + 1) + colIdx] = sumRes;
129+ indexWorkspace2 = varRow;
130+ } else {
131+ asc_atomic_add(var + varGmOffset, sumRes);
132+ indexWorkspace2 = -1;
133+ }
134+ if (colIdx == 0) {
135+ indexWorkspace[TWO * blockIdx + 1] = indexWorkspace2;
136+ }
137+ } else {
138+ asc_atomic_add(var + varGmOffset, sumRes);
139+ }
140+ }
141+}
142+ 
143+template <typename T, typename U, typename ADDR_T>
144+__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM_DETERM) inline void ScatterAddWithSortedSimtDetermStep2(
145+ uint32_t totalCol, uint32_t blockIdx, uint32_t sumWorkspaceEnd, uint32_t magic, uint32_t shift, __gm__ T* var,
146+ __gm__ T* sumWorkspace, __gm__ U* indexWorkspace)
147+{
148+ if (indexWorkspace[blockIdx * TWO + 1] == indexWorkspace[blockIdx * TWO]) { // 当前核对应的2个workspace索引相等则跳过
149+ return;
150+ }
151+ if (indexWorkspace[blockIdx * TWO + 1] == -1) {
152+ return;
153+ }
154+ 
155+ T workspaceSumRes = 0;
156+ uint32_t indicesStartOffset = blockIdx * TWO + 1; // 每个核从对应的第2行开始遍历
157+ U varRow = indexWorkspace[indicesStartOffset];
158+ 
159+ uint32_t dataWorkspaceStart = totalCol * indicesStartOffset;
160+ for (uint32_t i = dataWorkspaceStart + threadIdx.x; i < sumWorkspaceEnd; i += blockDim.x) {
161+ uint32_t indiceRow = Simt::UintDiv(i, magic, shift);
162+ uint32_t colIdx = i - indiceRow * totalCol;
163+ ADDR_T varGmOffset = varRow * totalCol + colIdx;
164+ if (indexWorkspace[indiceRow] == varRow) {
165+ workspaceSumRes += sumWorkspace[i];
166+ } else {
167+ asc_atomic_add(var + varGmOffset, workspaceSumRes);
168+ break;
169+ }
170+ }
171+}
172+ 
173+template<typename T, typename U, typename ADDR_T, bool withPos>
174+class ScatterAddWithSortedDetermSIMT {
175+public:
176+ __aicore__ inline ScatterAddWithSortedDetermSIMT(const ScatterAddWithSortedSimtTilingData& tilingData)
177+ : tilingData_(tilingData){};
178+ 
179+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR workspace);
180+ __aicore__ inline void Process();
181+ __aicore__ inline void ProcessSimtStep1();
182+ __aicore__ inline void ProcessSimtStep2();
183+ 
184+private:
185+ AscendC::GlobalTensor<T> varGm_;
186+ AscendC::GlobalTensor<U> indicesGm_;
187+ AscendC::GlobalTensor<U> posGm_;
188+ AscendC::GlobalTensor<T> updatesGm_;
189+ AscendC::GlobalTensor<T> sumWorkspace_;
190+ AscendC::GlobalTensor<U> indexWorkspace_;
191+ bool isStartRowCore_ = false;
192+ bool isEndRowCore_ = false;
193+ const ScatterAddWithSortedSimtTilingData& tilingData_;
194+ 
195+ uint32_t blockIdx_ = 0;
196+ uint32_t blockNum_ = 0;
197+ ADDR_T indicesBlockNum_ = 0;
198+ static constexpr U workspaceIndicesDefault_ = -2;
199+};
200+ 
201+template<typename T, typename U, typename ADDR_T, bool withPos>
202+__aicore__ inline void ScatterAddWithSortedDetermSIMT<T, U, ADDR_T, withPos>::Init(
203+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR workspace)
204+{
205+ blockIdx_ = GetBlockIdx();
206+ blockNum_ = GetBlockNum();
207+ if (blockIdx_ >= tilingData_.usedCoreNum) {
208+ return;
209+ }
210+ 
211+ varGm_.SetGlobalBuffer((__gm__ T *)(var));
212+ updatesGm_.SetGlobalBuffer((__gm__ T *)(updates));
213+ indicesGm_.SetGlobalBuffer((__gm__ U*)indices);
214+ posGm_.SetGlobalBuffer((__gm__ U*)pos);
215+ 
216+ isStartRowCore_ = blockIdx_ == 0; // 首核
217+ isEndRowCore_ = blockIdx_ == tilingData_.usedCoreNum - 1; // 尾核
218+ 
219+ uint32_t workspaceIndicesOffset = tilingData_.usedCoreNum * TWO * tilingData_.varShape[1] * sizeof(T);
220+ workspaceIndicesOffset = ops::CeilAlign(workspaceIndicesOffset, static_cast<uint32_t>(sizeof(U)));
221+ sumWorkspace_.SetGlobalBuffer((__gm__ T*)workspace);
222+ indexWorkspace_.SetGlobalBuffer((__gm__ U*)(workspace + workspaceIndicesOffset));
223+ 
224+ indicesBlockNum_ = tilingData_.normBlockIndices;
225+ if (blockIdx_ == tilingData_.usedCoreNum - 1) {
226+ indicesBlockNum_ = tilingData_.tailBlockIndices;
227+ }
228+}
229+ 
230+template<typename T, typename U, typename ADDR_T, bool withPos>
231+__aicore__ inline void ScatterAddWithSortedDetermSIMT<T, U, ADDR_T, withPos>::ProcessSimtStep1()
232+{
233+ uint32_t blockIdx = blockIdx_;
234+ ADDR_T totalCol = static_cast<ADDR_T>(tilingData_.varShape[1]);
235+ ADDR_T varFirstDimSize = static_cast<ADDR_T>(tilingData_.varShape[0]);
236+ ADDR_T indicesBlockNum = indicesBlockNum_;
237+ ADDR_T indicesBlockOffset = blockIdx * tilingData_.normBlockIndices;
238+ ADDR_T magic = 0;
239+ ADDR_T shift = 0;
240+ GetUintDivMagicAndShift(magic, shift, totalCol);
241+ 
242+ if (isStartRowCore_ && isEndRowCore_) {
243+ asc_vf_call<ScatterAddWithSortedSimtDetermStep1<T, U, ADDR_T, true, true, withPos>>(dim3(THREAD_NUM_DETERM),
244+ totalCol, indicesBlockOffset, indicesBlockNum, varFirstDimSize, magic, shift, blockIdx,
245+ (__gm__ T*)(varGm_.GetPhyAddr()), (__gm__ U*)(indicesGm_.GetPhyAddr()), (__gm__ U*)(posGm_.GetPhyAddr()),
246+ (__gm__ T*)(updatesGm_.GetPhyAddr()), (__gm__ T*)(sumWorkspace_.GetPhyAddr()), (__gm__ U*)(indexWorkspace_.GetPhyAddr()));
247+ } else if (isStartRowCore_ && !isEndRowCore_) {
248+ asc_vf_call<ScatterAddWithSortedSimtDetermStep1<T, U, ADDR_T, true, false, withPos>>(dim3(THREAD_NUM_DETERM),
249+ totalCol, indicesBlockOffset, indicesBlockNum, varFirstDimSize, magic, shift, blockIdx,
250+ (__gm__ T*)(varGm_.GetPhyAddr()), (__gm__ U*)(indicesGm_.GetPhyAddr()), (__gm__ U*)(posGm_.GetPhyAddr()),
251+ (__gm__ T*)(updatesGm_.GetPhyAddr()), (__gm__ T*)(sumWorkspace_.GetPhyAddr()), (__gm__ U*)(indexWorkspace_.GetPhyAddr()));
252+ } else if (!isStartRowCore_ && isEndRowCore_) {
253+ asc_vf_call<ScatterAddWithSortedSimtDetermStep1<T, U, ADDR_T, false, true, withPos>>(dim3(THREAD_NUM_DETERM),
254+ totalCol, indicesBlockOffset, indicesBlockNum, varFirstDimSize, magic, shift, blockIdx,
255+ (__gm__ T*)(varGm_.GetPhyAddr()), (__gm__ U*)(indicesGm_.GetPhyAddr()), (__gm__ U*)(posGm_.GetPhyAddr()),
256+ (__gm__ T*)(updatesGm_.GetPhyAddr()), (__gm__ T*)(sumWorkspace_.GetPhyAddr()), (__gm__ U*)(indexWorkspace_.GetPhyAddr()));
257+ } else {
258+ asc_vf_call<ScatterAddWithSortedSimtDetermStep1<T, U, ADDR_T, false, false, withPos>>(dim3(THREAD_NUM_DETERM),
259+ totalCol, indicesBlockOffset, indicesBlockNum, varFirstDimSize, magic, shift, blockIdx,
260+ (__gm__ T*)(varGm_.GetPhyAddr()), (__gm__ U*)(indicesGm_.GetPhyAddr()), (__gm__ U*)(posGm_.GetPhyAddr()),
261+ (__gm__ T*)(updatesGm_.GetPhyAddr()), (__gm__ T*)(sumWorkspace_.GetPhyAddr()), (__gm__ U*)(indexWorkspace_.GetPhyAddr()));
262+ }
263+}
264+ 
265+template<typename T, typename U, typename ADDR_T, bool withPos>
266+__aicore__ inline void ScatterAddWithSortedDetermSIMT<T, U, ADDR_T, withPos>::ProcessSimtStep2()
267+{
268+ uint32_t blockIdx = blockIdx_;
269+ uint32_t totalCol = static_cast<uint32_t>(tilingData_.varShape[1]);
270+ uint32_t sumWorkspaceEnd = blockNum_ * TWO * totalCol;
271+ uint32_t magic = 0;
272+ uint32_t shift = 0;
273+ GetUintDivMagicAndShift(magic, shift, totalCol);
274+ 
275+ asc_vf_call<ScatterAddWithSortedSimtDetermStep2<T, U, ADDR_T>>(dim3(totalCol),
276+ totalCol, blockIdx, sumWorkspaceEnd, magic, shift, (__gm__ T*)(varGm_.GetPhyAddr()),
277+ (__gm__ T*)(sumWorkspace_.GetPhyAddr()), (__gm__ U*)(indexWorkspace_.GetPhyAddr()));
278+}
279+ 
280+template<typename T, typename U, typename ADDR_T, bool withPos>
281+__aicore__ inline void ScatterAddWithSortedDetermSIMT<T, U, ADDR_T, withPos>::Process()
282+{
283+ if (blockIdx_ >= tilingData_.usedCoreNum) {
284+ return;
285+ }
286+ 
287+ ProcessSimtStep1();
288+ SyncAll();
289+ ProcessSimtStep2();
290+}
291+ 
292+}
293+#endif
Aindex/scatter_add_with_sorted/op_kernel/arch35/scatter_add_with_sorted_struct.h+151-0
@@ -0,0 +1,151 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_add_with_sorted_struct.h
13+ * \brief tiling base data
14+ */
15+ 
16+#ifndef SCATTER_ADD_WITH_SORTED_STRUCT_H
17+#define SCATTER_ADD_WITH_SORTED_STRUCT_H
18+ 
19+#include "ascendc/host_api/tiling/template_argument.h"
20+ 
21+#define TPL_MODE_SIMD 0
22+#define TPL_MODE_SIMT 1
23+#define TPL_MODE_EMPTY 2
24+ 
25+#define TPL_SCALAR_FALSE 0
26+#define TPL_SCALAR_TRUE 1
27+ 
28+#define TPL_DETERM_FALSE 0
29+#define TPL_DETERM_TRUE 1
30+ 
31+#define TPL_ADDR_B32 0
32+#define TPL_ADDR_B64 1
33+ 
34+ASCENDC_TPL_ARGS_DECL(
35+ ScatterAddWithSorted,
36+ ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, 2, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMD, TPL_MODE_SIMT, TPL_MODE_EMPTY),
37+ ASCENDC_TPL_UINT_DECL(IS_SCALAR, 1, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE, TPL_SCALAR_TRUE),
38+ ASCENDC_TPL_UINT_DECL(IS_DETERM, 1, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE, TPL_DETERM_TRUE),
39+ ASCENDC_TPL_UINT_DECL(ADDR_TYPE, 1, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32, TPL_ADDR_B64));
40+ 
41+ASCENDC_TPL_SEL(
42+ ASCENDC_TPL_ARGS_SEL(
43+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMD),
44+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
45+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
46+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)),
47+ ASCENDC_TPL_ARGS_SEL(
48+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMD),
49+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_TRUE),
50+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
51+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)),
52+ ASCENDC_TPL_ARGS_SEL(
53+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMD),
54+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
55+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_TRUE),
56+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)),
57+ ASCENDC_TPL_ARGS_SEL(
58+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMT),
59+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
60+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
61+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)),
62+ ASCENDC_TPL_ARGS_SEL(
63+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMT),
64+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_TRUE),
65+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
66+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)),
67+ ASCENDC_TPL_ARGS_SEL(
68+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMT),
69+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
70+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
71+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B64)),
72+ ASCENDC_TPL_ARGS_SEL(
73+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMT),
74+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_TRUE),
75+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
76+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B64)),
77+ ASCENDC_TPL_ARGS_SEL(
78+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMT),
79+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
80+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_TRUE),
81+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)),
82+ ASCENDC_TPL_ARGS_SEL(
83+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_SIMT),
84+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
85+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_TRUE),
86+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B64)),
87+ ASCENDC_TPL_ARGS_SEL(
88+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_EMPTY),
89+ ASCENDC_TPL_UINT_SEL(IS_SCALAR, ASCENDC_TPL_UI_LIST, TPL_SCALAR_FALSE),
90+ ASCENDC_TPL_UINT_SEL(IS_DETERM, ASCENDC_TPL_UI_LIST, TPL_DETERM_FALSE),
91+ ASCENDC_TPL_UINT_SEL(ADDR_TYPE, ASCENDC_TPL_UI_LIST, TPL_ADDR_B32)));
92+ 
93+struct ScatterAddWithSortedSimtTilingData {
94+ int64_t varShape[2];
95+ int64_t indicesNum;
96+ int64_t normBlockIndices;
97+ int64_t tailBlockIndices;
98+ int64_t usedCoreNum;
99+ bool withPos;
100+ uint64_t tilingKey{0};
101+};
102+struct ScatterAddWithSortedSimdTilingData {
103+ int64_t needCoreNum{0};
104+ int64_t indicesNum{0};
105+ int64_t updatesInner{0};
106+ 
107+ int64_t updatesBufferSize{0};
108+ int64_t outBufferSize{0};
109+ int64_t indicesBufferSize{0};
110+ int64_t posBufferSize{0};
111+ int64_t FrontAndBackIndexSize{0};
112+ 
113+ int64_t normalCoreColNum{0};
114+ int64_t tailCoreColNum{0};
115+ int64_t normalCoreRowNum{0};
116+ int64_t tailCoreRowNum{0};
117+ 
118+ int64_t normalCoreRowUbLoop{0};
119+ int64_t normalCoreNormalLoopRows{0};
120+ int64_t normalCoreTailLoopRows{0};
121+ int64_t tailCoreRowUbLoop{0};
122+ int64_t tailCoreNormalLoopRows{0};
123+ int64_t tailCoreTailLoopRows{0};
124+ 
125+ int64_t normalCoreColUbLoop{0};
126+ int64_t normalCoreNormalLoopCols{0};
127+ int64_t normalCoreTailLoopCols{0};
128+ int64_t tailCoreColUbLoop{0};
129+ int64_t tailCoreNormalLoopCols{0};
130+ int64_t tailCoreTailLoopCols{0};
131+ int64_t coreNumInRow{0};
132+ int64_t coreNumInCol{0};
133+ 
134+ int64_t vecAlignSize{0};
135+ int64_t indicesWorkspaceBufferSize{0};
136+ int64_t coreNumInColDeterm{0};
137+ int64_t tailCoreColUbDetermLoop{0};
138+ int64_t normalCoreColUbDetermLoop{0};
139+ int64_t tailCoreNormalLoopDetermCols{0};
140+ int64_t normalCoreNormalLoopDetermCols{0};
141+ int64_t tailCoreTailLoopDetermCols{0};
142+ int64_t normalCoreTailLoopDetermCols{0};
143+ int64_t updatesDeterminBufferSize{0};
144+ int64_t outBufferDeterminSize{0};
145+ int64_t normalCoreColDetermNum{0};
146+ int64_t tailCoreColNumDeterm{0};
147+ int64_t ubBlock{0};
148+ bool withPos{false};
149+ uint64_t tilingKey{0};
150+};
151+#endif
Aindex/scatter_add_with_sorted/op_kernel/scatter_add_with_sorted_apt.cpp+129-0
@@ -0,0 +1,129 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_add_with_sorted_apt.cpp
13+ * \brief scatter_add_with_sorted_apt.cpp
14+ */
15+#include "arch35/scatter_add_with_sorted_struct.h"
16+#include "arch35/scatter_add_with_sorted_simd.h"
17+#include "arch35/scatter_add_with_sorted_simd_determ.h"
18+#include "arch35/scatter_add_with_sorted_determ_workspace.h"
19+#include "arch35/scatter_add_with_sorted_simt.h"
20+#include "arch35/scatter_add_with_sorted_simt_determ.h"
21+ 
22+using namespace AscendC;
23+using namespace ScatterAddWithSorted;
24+ 
25+template <uint64_t TEMPLATE_MODE, uint64_t IS_SCALAR, uint64_t IS_DETERM, uint64_t ADDR_TYPE>
26+__global__ __aicore__ void scatter_add_with_sorted(
27+ GM_ADDR var, GM_ADDR updates, GM_ADDR indices, GM_ADDR pos, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
28+{
29+ TPipe pipe;
30+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
31+ REGISTER_TILING_DEFAULT(ScatterAddWithSortedSimdTilingData);
32+ 
33+ if constexpr (TEMPLATE_MODE == TPL_MODE_EMPTY) {
34+ return;
35+ } else if constexpr (TEMPLATE_MODE == TPL_MODE_SIMD) {
36+ REGISTER_TILING_FOR_TILINGKEY("TEMPLATE_MODE == TPL_MODE_SIMD", ScatterAddWithSortedSimdTilingData);
37+ GET_TILING_DATA_WITH_STRUCT(ScatterAddWithSortedSimdTilingData, tilingData, tiling);
38+ 
39+ if constexpr (IS_DETERM == TPL_DETERM_TRUE) {
40+ if (tilingData.withPos) {
41+ ScatterAddWithSortedSimdDterm<DTYPE_VAR, DTYPE_SORTED_INDEX, true> op;
42+ op.Init(var, updates, indices, pos, y, workspace, pipe, &tilingData);
43+ op.Process();
44+ } else {
45+ ScatterAddWithSortedSimdDterm<DTYPE_VAR, DTYPE_SORTED_INDEX, false> op;
46+ op.Init(var, updates, indices, pos, y, workspace, pipe, &tilingData);
47+ op.Process();
48+ }
49+ SyncAll();
50+ pipe.Reset();
51+ ScatterAddWithSortedSimdDtermWorkspace<DTYPE_VAR, DTYPE_SORTED_INDEX> workspaceOp;
52+ workspaceOp.Init(var, y, workspace, pipe, &tilingData);
53+ workspaceOp.Process();
54+ } else if constexpr (IS_SCALAR == TPL_SCALAR_TRUE) {
55+ ScatterAddWithSortedSIMD<DTYPE_VAR, DTYPE_SORTED_INDEX, true, false> op;
56+ op.Init(var, updates, indices, pos, y, workspace, pipe, &tilingData);
57+ op.Process();
58+ } else {
59+ if (tilingData.withPos) {
60+ ScatterAddWithSortedSIMD<DTYPE_VAR, DTYPE_SORTED_INDEX, false, true> op;
61+ op.Init(var, updates, indices, pos, y, workspace, pipe, &tilingData);
62+ op.Process();
63+ } else {
64+ ScatterAddWithSortedSIMD<DTYPE_VAR, DTYPE_SORTED_INDEX, false, false> op;
65+ op.Init(var, updates, indices, pos, y, workspace, pipe, &tilingData);
66+ op.Process();
67+ }
68+ }
69+ } else if constexpr (TEMPLATE_MODE == TPL_MODE_SIMT) {
70+ REGISTER_TILING_FOR_TILINGKEY("TEMPLATE_MODE == TPL_MODE_SIMT", ScatterAddWithSortedSimtTilingData);
71+ GET_TILING_DATA_WITH_STRUCT(ScatterAddWithSortedSimtTilingData, tilingData, tiling);
72+ 
73+ if constexpr (IS_DETERM == TPL_DETERM_TRUE) {
74+ if constexpr (ADDR_TYPE == TPL_ADDR_B32) {
75+ if (tilingData.withPos) {
C
Cchenqi3174月28日
已过期

withPos 为什么不用tilingkey

likedislike
W
wkkk0528
4月28日 评论：
W
wkkk0528
4月29日 评论：
76+ ScatterAddWithSortedDetermSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint32_t, true> op(tilingData);
77+ op.Init(var, updates, indices, pos, workspace);
78+ op.Process();
79+ } else {
80+ ScatterAddWithSortedDetermSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint32_t, false> op(tilingData);
81+ op.Init(var, updates, indices, pos, workspace);
82+ op.Process();
83+ }
84+ } else {
85+ if (tilingData.withPos) {
86+ ScatterAddWithSortedDetermSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint64_t, true> op(tilingData);
87+ op.Init(var, updates, indices, pos, workspace);
88+ op.Process();
89+ } else {
90+ ScatterAddWithSortedDetermSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint64_t, false> op(tilingData);
91+ op.Init(var, updates, indices, pos, workspace);
92+ op.Process();
93+ }
94+ }
95+ } else if constexpr (IS_SCALAR == TPL_SCALAR_TRUE) {
96+ if constexpr (ADDR_TYPE == TPL_ADDR_B32) {
97+ ScatterAddWithSortedSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint32_t, true, true> op(tilingData);
98+ op.Init(var, updates, indices, pos);
99+ op.Process();
100+ } else {
101+ ScatterAddWithSortedSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint64_t, true, true> op(tilingData);
102+ op.Init(var, updates, indices, pos);
103+ op.Process();
104+ }
105+ } else {
106+ if constexpr (ADDR_TYPE == TPL_ADDR_B32) {
107+ if (tilingData.withPos) {
108+ ScatterAddWithSortedSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint32_t, false, true> op(tilingData);
109+ op.Init(var, updates, indices, pos);
110+ op.Process();
111+ } else {
112+ ScatterAddWithSortedSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint32_t, false, false> op(tilingData);
113+ op.Init(var, updates, indices, pos);
114+ op.Process();
115+ }
116+ } else {
117+ if (tilingData.withPos) {
118+ ScatterAddWithSortedSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint64_t, false, true> op(tilingData);
119+ op.Init(var, updates, indices, pos);
120+ op.Process();
121+ } else {
122+ ScatterAddWithSortedSIMT<DTYPE_VAR, DTYPE_SORTED_INDEX, uint64_t, false, false> op(tilingData);
123+ op.Init(var, updates, indices, pos);
124+ op.Process();
125+ }
126+ }
127+ }
128+ }
129+}
Mindex/scatter_add_with_sorted/tests/ut/op_host/test_scatter_add_with_sorted_tiling.cpp+844-0
@@ -1267,3 +1267,847 @@ TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_bfloat16_add)
1267 auto tiling_key = tiling_context->GetTilingKey();1267 auto tiling_key = tiling_context->GetTilingKey();
1268 ASSERT_EQ(tiling_key, 16);1268 ASSERT_EQ(tiling_key, 16);
1269}1269}
1270+ 
1271+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simt_float32)
1272+{
1273+ std::string op_type("ScatterAddWithSorted");
1274+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1275+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1276+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1277+ 
1278+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1279+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1280+ "Intrinsic_fix_pipe_l0c2out": false,
1281+ "Intrinsic_data_move_l12ub": true,
1282+ "Intrinsic_data_move_l0c2ub": true,
1283+ "Intrinsic_data_move_out2l1_nd2nz": false,
1284+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1285+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1286+ "CORE_NUM": 48}
1287+ })";
1288+ map<string, string> soc_infos;
1289+ map<string, string> aicore_spec;
1290+ map<string, string> intrinsics;
1291+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1292+ 
1293+ fe::PlatFormInfos platform_info;
1294+ platform_info.Init();
1295+ struct ScatterAddWithSortedCompileInfo {
1296+ int32_t totalCoreNum = 30;
1297+ uint64_t ubSizePlatForm = 0;
1298+ uint64_t workspaceSize = 0;
1299+ } compile_info;
1300+ 
1301+ auto kernel_holder =
1302+ gert::KernelRunContextFaker()
1303+ .KernelIONum(2, 1)
1304+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1305+ .Outputs({&compile_info})
1306+ .Build();
1307+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1308+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1309+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1310+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1311+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1312+ "AICoreintrinsicDtypeMap", intrinsics);
1313+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1314+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1315+ 
1316+ auto param = gert::TilingData::CreateCap(4096);
1317+ ASSERT_NE(param, nullptr);
1318+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1319+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1320+ gert::StorageShape input_shape = {{10, 16}, {10, 16}};
1321+ gert::StorageShape src_shape = {{8, 16}, {8, 16}};
1322+ gert::StorageShape indic_shape = {{8}, {8}};
1323+ gert::StorageShape pos_shape = {{8}, {8}};
1324+ gert::StorageShape output_shape = {{10, 16}, {10, 16}};
1325+ auto holder = gert::TilingContextFaker()
1326+ .SetOpType("ScatterAddWithSorted")
1327+ .NodeIoNum(4, 1)
1328+ .IrInstanceNum({1, 1, 1, 1})
1329+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1330+ .OutputShapes({&output_shape})
1331+ .CompileInfo(&compile_info)
1332+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1333+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1334+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1335+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1336+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1337+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1338+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1339+ .TilingData(param.get())
1340+ .Workspace(ws_size)
1341+ .Build();
1342+ 
1343+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1344+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1345+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1346+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1347+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1348+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1349+ 
1350+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1351+ auto tiling_key = tiling_context->GetTilingKey();
1352+ ASSERT_EQ(tiling_key, 1);
1353+}
1354+ 
1355+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_float32)
1356+{
1357+ std::string op_type("ScatterAddWithSorted");
1358+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1359+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1360+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1361+ 
1362+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1363+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1364+ "Intrinsic_fix_pipe_l0c2out": false,
1365+ "Intrinsic_data_move_l12ub": true,
1366+ "Intrinsic_data_move_l0c2ub": true,
1367+ "Intrinsic_data_move_out2l1_nd2nz": false,
1368+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1369+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1370+ "CORE_NUM": 48}
1371+ })";
1372+ map<string, string> soc_infos;
1373+ map<string, string> aicore_spec;
1374+ map<string, string> intrinsics;
1375+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1376+ 
1377+ fe::PlatFormInfos platform_info;
1378+ platform_info.Init();
1379+ struct ScatterAddWithSortedCompileInfo {
1380+ int32_t totalCoreNum = 30;
1381+ uint64_t ubSizePlatForm = 0;
1382+ uint64_t workspaceSize = 0;
1383+ } compile_info;
1384+ 
1385+ auto kernel_holder =
1386+ gert::KernelRunContextFaker()
1387+ .KernelIONum(2, 1)
1388+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1389+ .Outputs({&compile_info})
1390+ .Build();
1391+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1392+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1393+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1394+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1395+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1396+ "AICoreintrinsicDtypeMap", intrinsics);
1397+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1398+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1399+ 
1400+ auto param = gert::TilingData::CreateCap(4096);
1401+ ASSERT_NE(param, nullptr);
1402+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1403+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1404+ gert::StorageShape input_shape = {{65535, 4096}, {65535, 4096}};
1405+ gert::StorageShape src_shape = {{63806, 4096}, {63806, 4096}};
1406+ gert::StorageShape indic_shape = {{63806}, {63806}};
1407+ gert::StorageShape pos_shape = {{63806}, {63806}};
1408+ gert::StorageShape output_shape = {{65535, 4096}, {65535, 4096}};
1409+ auto holder = gert::TilingContextFaker()
1410+ .SetOpType("ScatterAddWithSorted")
1411+ .NodeIoNum(4, 1)
1412+ .IrInstanceNum({1, 1, 1, 1})
1413+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1414+ .OutputShapes({&output_shape})
1415+ .CompileInfo(&compile_info)
1416+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1417+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1418+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1419+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1420+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1421+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1422+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1423+ .TilingData(param.get())
1424+ .Workspace(ws_size)
1425+ .Build();
1426+ 
1427+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1428+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1429+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1430+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1431+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1432+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1433+ 
1434+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1435+ auto tiling_key = tiling_context->GetTilingKey();
1436+ ASSERT_EQ(tiling_key, 0);
1437+}
1438+ 
1439+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_deterministic_float32)
1440+{
1441+ std::string op_type("ScatterAddWithSorted");
1442+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1443+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1444+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1445+ 
1446+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1447+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1448+ "Intrinsic_fix_pipe_l0c2out": false,
1449+ "Intrinsic_data_move_l12ub": true,
1450+ "Intrinsic_data_move_l0c2ub": true,
1451+ "Intrinsic_data_move_out2l1_nd2nz": false,
1452+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1453+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1454+ "CORE_NUM": 48}
1455+ })";
1456+ map<string, string> soc_infos;
1457+ map<string, string> aicore_spec;
1458+ map<string, string> intrinsics;
1459+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1460+ 
1461+ fe::PlatFormInfos platform_info;
1462+ platform_info.Init();
1463+ struct ScatterAddWithSortedCompileInfo {
1464+ int32_t totalCoreNum = 30;
1465+ uint64_t ubSizePlatForm = 0;
1466+ uint64_t workspaceSize = 0;
1467+ } compile_info;
1468+ 
1469+ auto kernel_holder =
1470+ gert::KernelRunContextFaker()
1471+ .KernelIONum(2, 1)
1472+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1473+ .Outputs({&compile_info})
1474+ .Build();
1475+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1476+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1477+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1478+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1479+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1480+ "AICoreintrinsicDtypeMap", intrinsics);
1481+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1482+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1483+ 
1484+ auto param = gert::TilingData::CreateCap(4096);
1485+ ASSERT_NE(param, nullptr);
1486+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1487+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1488+ gert::StorageShape input_shape = {{65535, 4096}, {65535, 4096}};
1489+ gert::StorageShape src_shape = {{63806, 4096}, {63806, 4096}};
1490+ gert::StorageShape indic_shape = {{63806}, {63806}};
1491+ gert::StorageShape pos_shape = {{63806}, {63806}};
1492+ gert::StorageShape output_shape = {{65535, 4096}, {65535, 4096}};
1493+ int32_t deterministic_info = 1;
1494+ auto holder = gert::TilingContextFaker()
1495+ .SetOpType("ScatterAddWithSorted")
1496+ .NodeIoNum(4, 1)
1497+ .IrInstanceNum({1, 1, 1, 1})
1498+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1499+ .OutputShapes({&output_shape})
1500+ .CompileInfo(&compile_info)
1501+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1502+ .DeterministicInfo(deterministic_info)
1503+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1504+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1505+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1506+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1507+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1508+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1509+ .TilingData(param.get())
1510+ .Workspace(ws_size)
1511+ .Build();
1512+ 
1513+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1514+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1515+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1516+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1517+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1518+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1519+ 
1520+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1521+ auto tiling_key = tiling_context->GetTilingKey();
1522+ ASSERT_EQ(tiling_key, 8);
1523+}
1524+ 
1525+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_deterministic_float16)
1526+{
1527+ std::string op_type("ScatterAddWithSorted");
1528+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1529+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1530+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1531+ 
1532+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1533+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1534+ "Intrinsic_fix_pipe_l0c2out": false,
1535+ "Intrinsic_data_move_l12ub": true,
1536+ "Intrinsic_data_move_l0c2ub": true,
1537+ "Intrinsic_data_move_out2l1_nd2nz": false,
1538+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1539+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1540+ "CORE_NUM": 48}
1541+ })";
1542+ map<string, string> soc_infos;
1543+ map<string, string> aicore_spec;
1544+ map<string, string> intrinsics;
1545+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1546+ 
1547+ fe::PlatFormInfos platform_info;
1548+ platform_info.Init();
1549+ struct ScatterAddWithSortedCompileInfo {
1550+ int32_t totalCoreNum = 30;
1551+ uint64_t ubSizePlatForm = 0;
1552+ uint64_t workspaceSize = 0;
1553+ } compile_info;
1554+ 
1555+ auto kernel_holder =
1556+ gert::KernelRunContextFaker()
1557+ .KernelIONum(2, 1)
1558+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1559+ .Outputs({&compile_info})
1560+ .Build();
1561+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1562+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1563+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1564+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1565+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1566+ "AICoreintrinsicDtypeMap", intrinsics);
1567+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1568+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1569+ 
1570+ auto param = gert::TilingData::CreateCap(4096);
1571+ ASSERT_NE(param, nullptr);
1572+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1573+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1574+ gert::StorageShape input_shape = {{65535, 4096}, {65535, 4096}};
1575+ gert::StorageShape src_shape = {{63806, 4096}, {63806, 4096}};
1576+ gert::StorageShape indic_shape = {{63806}, {63806}};
1577+ gert::StorageShape pos_shape = {{63806}, {63806}};
1578+ gert::StorageShape output_shape = {{65535, 4096}, {65535, 4096}};
1579+ int32_t deterministic_info = 1;
1580+ auto holder = gert::TilingContextFaker()
1581+ .SetOpType("ScatterAddWithSorted")
1582+ .NodeIoNum(4, 1)
1583+ .IrInstanceNum({1, 1, 1, 1})
1584+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1585+ .OutputShapes({&output_shape})
1586+ .CompileInfo(&compile_info)
1587+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1588+ .DeterministicInfo(deterministic_info)
1589+ .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
1590+ .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
1591+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1592+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1593+ .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
1594+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1595+ .TilingData(param.get())
1596+ .Workspace(ws_size)
1597+ .Build();
1598+ 
1599+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1600+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1601+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1602+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1603+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1604+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1605+ 
1606+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1607+ auto tiling_key = tiling_context->GetTilingKey();
1608+ ASSERT_EQ(tiling_key, 8);
1609+}
1610+ 
1611+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_deterministic_bf16)
1612+{
1613+ std::string op_type("ScatterAddWithSorted");
1614+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1615+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1616+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1617+ 
1618+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1619+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1620+ "Intrinsic_fix_pipe_l0c2out": false,
1621+ "Intrinsic_data_move_l12ub": true,
1622+ "Intrinsic_data_move_l0c2ub": true,
1623+ "Intrinsic_data_move_out2l1_nd2nz": false,
1624+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1625+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1626+ "CORE_NUM": 48}
1627+ })";
1628+ map<string, string> soc_infos;
1629+ map<string, string> aicore_spec;
1630+ map<string, string> intrinsics;
1631+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1632+ 
1633+ fe::PlatFormInfos platform_info;
1634+ platform_info.Init();
1635+ struct ScatterAddWithSortedCompileInfo {
1636+ int32_t totalCoreNum = 30;
1637+ uint64_t ubSizePlatForm = 0;
1638+ uint64_t workspaceSize = 0;
1639+ } compile_info;
1640+ 
1641+ auto kernel_holder =
1642+ gert::KernelRunContextFaker()
1643+ .KernelIONum(2, 1)
1644+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1645+ .Outputs({&compile_info})
1646+ .Build();
1647+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1648+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1649+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1650+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1651+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1652+ "AICoreintrinsicDtypeMap", intrinsics);
1653+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1654+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1655+ 
1656+ auto param = gert::TilingData::CreateCap(4096);
1657+ ASSERT_NE(param, nullptr);
1658+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1659+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1660+ gert::StorageShape input_shape = {{65535, 24096}, {65535, 24096}};
1661+ gert::StorageShape src_shape = {{63806, 24096}, {63806, 24096}};
1662+ gert::StorageShape indic_shape = {{63806}, {63806}};
1663+ gert::StorageShape pos_shape = {{63806}, {63806}};
1664+ gert::StorageShape output_shape = {{65535, 24096}, {65535, 24096}};
1665+ int32_t deterministic_info = 1;
1666+ auto holder = gert::TilingContextFaker()
1667+ .SetOpType("ScatterAddWithSorted")
1668+ .NodeIoNum(4, 1)
1669+ .IrInstanceNum({1, 1, 1, 1})
1670+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1671+ .OutputShapes({&output_shape})
1672+ .CompileInfo(&compile_info)
1673+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1674+ .DeterministicInfo(deterministic_info)
1675+ .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
1676+ .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
1677+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1678+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1679+ .NodeOutputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
1680+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1681+ .TilingData(param.get())
1682+ .Workspace(ws_size)
1683+ .Build();
1684+ 
1685+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1686+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1687+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1688+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1689+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1690+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1691+ 
1692+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1693+ auto tiling_key = tiling_context->GetTilingKey();
1694+ ASSERT_EQ(tiling_key, 8);
1695+}
1696+ 
1697+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_scalar_update)
1698+{
1699+ std::string op_type("ScatterAddWithSorted");
1700+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1701+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1702+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1703+ 
1704+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1705+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1706+ "Intrinsic_fix_pipe_l0c2out": false,
1707+ "Intrinsic_data_move_l12ub": true,
1708+ "Intrinsic_data_move_l0c2ub": true,
1709+ "Intrinsic_data_move_out2l1_nd2nz": false,
1710+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1711+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1712+ "CORE_NUM": 48}
1713+ })";
1714+ map<string, string> soc_infos;
1715+ map<string, string> aicore_spec;
1716+ map<string, string> intrinsics;
1717+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1718+ 
1719+ fe::PlatFormInfos platform_info;
1720+ platform_info.Init();
1721+ struct ScatterAddWithSortedCompileInfo {
1722+ int32_t totalCoreNum = 30;
1723+ uint64_t ubSizePlatForm = 0;
1724+ uint64_t workspaceSize = 0;
1725+ } compile_info;
1726+ 
1727+ auto kernel_holder =
1728+ gert::KernelRunContextFaker()
1729+ .KernelIONum(2, 1)
1730+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1731+ .Outputs({&compile_info})
1732+ .Build();
1733+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1734+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1735+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1736+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1737+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1738+ "AICoreintrinsicDtypeMap", intrinsics);
1739+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1740+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1741+ 
1742+ auto param = gert::TilingData::CreateCap(4096);
1743+ ASSERT_NE(param, nullptr);
1744+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1745+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1746+ gert::StorageShape input_shape = {{65535, 4096}, {65535, 4096}};
1747+ gert::StorageShape src_shape = {{1}, {1}};
1748+ gert::StorageShape indic_shape = {{63806}, {63806}};
1749+ gert::StorageShape pos_shape = {{63806}, {63806}};
1750+ gert::StorageShape output_shape = {{65535, 4096}, {65535, 4096}};
1751+ auto holder = gert::TilingContextFaker()
1752+ .SetOpType("ScatterAddWithSorted")
1753+ .NodeIoNum(4, 1)
1754+ .IrInstanceNum({1, 1, 1, 1})
1755+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1756+ .OutputShapes({&output_shape})
1757+ .CompileInfo(&compile_info)
1758+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1759+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1760+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1761+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1762+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1763+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1764+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1765+ .TilingData(param.get())
1766+ .Workspace(ws_size)
1767+ .Build();
1768+ 
1769+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1770+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1771+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1772+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1773+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1774+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1775+ 
1776+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1777+ auto tiling_key = tiling_context->GetTilingKey();
1778+ ASSERT_EQ(tiling_key, 4);
1779+}
1780+ 
1781+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_split_col_float32)
1782+{
1783+ std::string op_type("ScatterAddWithSorted");
1784+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1785+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1786+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1787+ 
1788+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1789+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1790+ "Intrinsic_fix_pipe_l0c2out": false,
1791+ "Intrinsic_data_move_l12ub": true,
1792+ "Intrinsic_data_move_l0c2ub": true,
1793+ "Intrinsic_data_move_out2l1_nd2nz": false,
1794+ "UB_SIZE": 32768, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1795+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1796+ "CORE_NUM": 48}
1797+ })";
1798+ map<string, string> soc_infos;
1799+ map<string, string> aicore_spec;
1800+ map<string, string> intrinsics;
1801+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1802+ 
1803+ fe::PlatFormInfos platform_info;
1804+ platform_info.Init();
1805+ struct ScatterAddWithSortedCompileInfo {
1806+ int32_t totalCoreNum = 30;
1807+ uint64_t ubSizePlatForm = 0;
1808+ uint64_t workspaceSize = 0;
1809+ } compile_info;
1810+ 
1811+ auto kernel_holder =
1812+ gert::KernelRunContextFaker()
1813+ .KernelIONum(2, 1)
1814+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1815+ .Outputs({&compile_info})
1816+ .Build();
1817+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1818+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1819+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1820+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1821+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1822+ "AICoreintrinsicDtypeMap", intrinsics);
1823+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1824+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1825+ 
1826+ auto param = gert::TilingData::CreateCap(4096);
1827+ ASSERT_NE(param, nullptr);
1828+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1829+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1830+ gert::StorageShape input_shape = {{1000, 65536}, {1000, 65536}};
1831+ gert::StorageShape src_shape = {{800, 65536}, {800, 65536}};
1832+ gert::StorageShape indic_shape = {{800}, {800}};
1833+ gert::StorageShape pos_shape = {{800}, {800}};
1834+ gert::StorageShape output_shape = {{1000, 65536}, {1000, 65536}};
1835+ auto holder = gert::TilingContextFaker()
1836+ .SetOpType("ScatterAddWithSorted")
1837+ .NodeIoNum(4, 1)
1838+ .IrInstanceNum({1, 1, 1, 1})
1839+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1840+ .OutputShapes({&output_shape})
1841+ .CompileInfo(&compile_info)
1842+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1843+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1844+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1845+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1846+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
1847+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1848+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1849+ .TilingData(param.get())
1850+ .Workspace(ws_size)
1851+ .Build();
1852+ 
1853+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1854+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1855+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1856+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1857+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1858+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1859+ 
1860+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1861+ auto tiling_key = tiling_context->GetTilingKey();
1862+ ASSERT_EQ(tiling_key, 0);
1863+}
1864+ 
1865+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_int64_indices)
1866+{
1867+ std::string op_type("ScatterAddWithSorted");
1868+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1869+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1870+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1871+ 
1872+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1873+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1874+ "Intrinsic_fix_pipe_l0c2out": false,
1875+ "Intrinsic_data_move_l12ub": true,
1876+ "Intrinsic_data_move_l0c2ub": true,
1877+ "Intrinsic_data_move_out2l1_nd2nz": false,
1878+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1879+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1880+ "CORE_NUM": 48}
1881+ })";
1882+ map<string, string> soc_infos;
1883+ map<string, string> aicore_spec;
1884+ map<string, string> intrinsics;
1885+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1886+ 
1887+ fe::PlatFormInfos platform_info;
1888+ platform_info.Init();
1889+ struct ScatterAddWithSortedCompileInfo {
1890+ int32_t totalCoreNum = 30;
1891+ uint64_t ubSizePlatForm = 0;
1892+ uint64_t workspaceSize = 0;
1893+ } compile_info;
1894+ 
1895+ auto kernel_holder =
1896+ gert::KernelRunContextFaker()
1897+ .KernelIONum(2, 1)
1898+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1899+ .Outputs({&compile_info})
1900+ .Build();
1901+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1902+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1903+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1904+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1905+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1906+ "AICoreintrinsicDtypeMap", intrinsics);
1907+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1908+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1909+ 
1910+ auto param = gert::TilingData::CreateCap(4096);
1911+ ASSERT_NE(param, nullptr);
1912+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1913+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1914+ gert::StorageShape input_shape = {{65535, 4096}, {65535, 4096}};
1915+ gert::StorageShape src_shape = {{63806, 4096}, {63806, 4096}};
1916+ gert::StorageShape indic_shape = {{63806}, {63806}};
1917+ gert::StorageShape pos_shape = {{63806}, {63806}};
1918+ gert::StorageShape output_shape = {{65535, 4096}, {65535, 4096}};
1919+ auto holder = gert::TilingContextFaker()
1920+ .SetOpType("ScatterAddWithSorted")
1921+ .NodeIoNum(4, 1)
1922+ .IrInstanceNum({1, 1, 1, 1})
1923+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
1924+ .OutputShapes({&output_shape})
1925+ .CompileInfo(&compile_info)
1926+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
1927+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1928+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1929+ .NodeInputTd(2, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
1930+ .NodeInputTd(3, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
1931+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
1932+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
1933+ .TilingData(param.get())
1934+ .Workspace(ws_size)
1935+ .Build();
1936+ 
1937+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
1938+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
1939+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1940+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1941+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1942+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
1943+ 
1944+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
1945+ auto tiling_key = tiling_context->GetTilingKey();
1946+ ASSERT_EQ(tiling_key, 0);
1947+}
1948+ 
1949+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_without_pos)
1950+{
1951+ std::string op_type("ScatterAddWithSorted");
1952+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
1953+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
1954+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
1955+ 
1956+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
1957+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
1958+ "Intrinsic_fix_pipe_l0c2out": false,
1959+ "Intrinsic_data_move_l12ub": true,
1960+ "Intrinsic_data_move_l0c2ub": true,
1961+ "Intrinsic_data_move_out2l1_nd2nz": false,
1962+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
1963+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
1964+ "CORE_NUM": 48}
1965+ })";
1966+ map<string, string> soc_infos;
1967+ map<string, string> aicore_spec;
1968+ map<string, string> intrinsics;
1969+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
1970+ 
1971+ fe::PlatFormInfos platform_info;
1972+ platform_info.Init();
1973+ struct ScatterAddWithSortedCompileInfo {
1974+ int32_t totalCoreNum = 30;
1975+ uint64_t ubSizePlatForm = 0;
1976+ uint64_t workspaceSize = 0;
1977+ } compile_info;
1978+ 
1979+ auto kernel_holder =
1980+ gert::KernelRunContextFaker()
1981+ .KernelIONum(2, 1)
1982+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
1983+ .Outputs({&compile_info})
1984+ .Build();
1985+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
1986+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
1987+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
1988+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
1989+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
1990+ "AICoreintrinsicDtypeMap", intrinsics);
1991+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
1992+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
1993+ 
1994+ auto param = gert::TilingData::CreateCap(4096);
1995+ ASSERT_NE(param, nullptr);
1996+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
1997+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
1998+ gert::StorageShape input_shape = {{65535, 4096}, {65535, 4096}};
1999+ gert::StorageShape src_shape = {{63806, 4096}, {63806, 4096}};
2000+ gert::StorageShape indic_shape = {{63806}, {63806}};
2001+ gert::StorageShape output_shape = {{65535, 4096}, {65535, 4096}};
2002+ auto holder = gert::TilingContextFaker()
2003+ .SetOpType("ScatterAddWithSorted")
2004+ .NodeIoNum(3, 1)
2005+ .IrInstanceNum({1, 1, 1})
2006+ .InputShapes({&input_shape, &src_shape, &indic_shape})
2007+ .OutputShapes({&output_shape})
2008+ .CompileInfo(&compile_info)
2009+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
2010+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
2011+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
2012+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
2013+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
2014+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
2015+ .TilingData(param.get())
2016+ .Workspace(ws_size)
2017+ .Build();
2018+ 
2019+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
2020+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
2021+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
2022+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
2023+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
2024+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
2025+ 
2026+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
2027+ auto tiling_key = tiling_context->GetTilingKey();
2028+ ASSERT_EQ(tiling_key, 0);
2029+}
2030+ 
2031+TEST_F(ScatterAddWithSortedTiling, test_scatter_add_with_sorted_simd_zero_shape)
2032+{
2033+ std::string op_type("ScatterAddWithSorted");
2034+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
2035+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
2036+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
2037+ 
2038+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}, {"NpuArch", "3510"}};
2039+ string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
2040+ "Intrinsic_fix_pipe_l0c2out": false,
2041+ "Intrinsic_data_move_l12ub": true,
2042+ "Intrinsic_data_move_l0c2ub": true,
2043+ "Intrinsic_data_move_out2l1_nd2nz": false,
2044+ "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
2045+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
2046+ "CORE_NUM": 48}
2047+ })";
2048+ map<string, string> soc_infos;
2049+ map<string, string> aicore_spec;
2050+ map<string, string> intrinsics;
2051+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
2052+ 
2053+ fe::PlatFormInfos platform_info;
2054+ platform_info.Init();
2055+ struct ScatterAddWithSortedCompileInfo {
2056+ int32_t totalCoreNum = 30;
2057+ uint64_t ubSizePlatForm = 0;
2058+ uint64_t workspaceSize = 0;
2059+ } compile_info;
2060+ 
2061+ auto kernel_holder =
2062+ gert::KernelRunContextFaker()
2063+ .KernelIONum(2, 1)
2064+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
2065+ .Outputs({&compile_info})
2066+ .Build();
2067+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
2068+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
2069+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
2070+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
2071+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
2072+ "AICoreintrinsicDtypeMap", intrinsics);
2073+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
2074+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
2075+ 
2076+ auto param = gert::TilingData::CreateCap(4096);
2077+ ASSERT_NE(param, nullptr);
2078+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
2079+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
2080+ gert::StorageShape input_shape = {{0, 0}, {0, 0}};
2081+ gert::StorageShape src_shape = {{0, 0}, {0, 0}};
2082+ gert::StorageShape indic_shape = {{0}, {0}};
2083+ gert::StorageShape pos_shape = {{0}, {0}};
2084+ gert::StorageShape output_shape = {{0, 0}, {0, 0}};
2085+ auto holder = gert::TilingContextFaker()
2086+ .SetOpType("ScatterAddWithSorted")
2087+ .NodeIoNum(4, 1)
2088+ .IrInstanceNum({1, 1, 1, 1})
2089+ .InputShapes({&input_shape, &src_shape, &indic_shape, &pos_shape})
2090+ .OutputShapes({&output_shape})
2091+ .CompileInfo(&compile_info)
2092+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
2093+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
2094+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
2095+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
2096+ .NodeInputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
2097+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
2098+ .NodeAttrs({{"reduction", Ops::NN::AnyValue::CreateFrom<std::string>("add")}})
2099+ .TilingData(param.get())
2100+ .Workspace(ws_size)
2101+ .Build();
2102+ 
2103+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
2104+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
2105+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
2106+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
2107+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
2108+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
2109+ 
2110+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
2111+ auto tiling_key = tiling_context->GetTilingKey();
2112+ ASSERT_EQ(tiling_key, 2);
2113+}
Mindex/scatter_add_with_sorted/tests/ut/op_kernel/CMakeLists.txt+1-0
@@ -9,4 +9,5 @@
9 9 
10if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE)10if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE)
11 AddOpTestCase(scatter_add_with_sorted "ascend910B1" "")11 AddOpTestCase(scatter_add_with_sorted "ascend910B1" "")
12+ AddOpTestCase(scatter_add_with_sorted "ascend950pr_9599" "-DDTYPE_VAR=float -DDTYPE_SORTED_INDEX=int32_t")
12endif()13endif()
Aindex/scatter_add_with_sorted/tests/ut/op_kernel/scatter_add_with_sorted_data/test_scatter_add_with_sorted_apt.cpp+216-0
@@ -0,0 +1,216 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <array>
12+#include <vector>
13+#include "gtest/gtest.h"
14+ 
15+#ifdef __CCE_KT_TEST__
16+#include "tikicpulib.h"
17+#include "data_utils.h"
18+#include "string.h"
19+#include <iostream>
20+#include <string>
21+#endif
22+ 
23+#include <cstdint>
24+ 
25+#include "arch35/scatter_add_with_sorted_struct.h"
26+ 
27+using namespace std;
28+ 
29+extern "C" __global__ __aicore__ void scatter_add_with_sorted(
30+ GM_ADDR var, GM_ADDR value, GM_ADDR sorted_index, GM_ADDR pos, GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling);
31+ 
32+class scatter_add_with_sorted_apt_test : public testing::Test {
33+protected:
34+ static void SetUpTestCase()
35+ {
36+ cout << "scatter_add_with_sorted_apt_test SetUp\n" << endl;
37+ }
38+ static void TearDownTestCase()
39+ {
40+ cout << "scatter_add_with_sorted_apt_test TearDown\n" << endl;
41+ }
42+};
43+ 
44+static void FillSimdTilingData(ScatterAddWithSortedSimdTilingData *td, int64_t indicesNum, int64_t updatesInner,
45+ bool withPos, uint64_t tilingKey)
46+{
47+ td->tilingKey = tilingKey;
48+ td->needCoreNum = 1;
49+ td->indicesNum = indicesNum;
50+ td->updatesInner = updatesInner;
51+ td->withPos = withPos;
52+ 
53+ td->updatesBufferSize = 4096;
54+ td->outBufferSize = 4096;
55+ td->indicesBufferSize = 512;
56+ td->posBufferSize = withPos ? 512 : 0;
57+ td->FrontAndBackIndexSize = 0;
58+ 
59+ td->coreNumInRow = 1;
60+ td->coreNumInCol = 1;
61+ td->normalCoreColNum = updatesInner;
62+ td->tailCoreColNum = updatesInner;
63+ td->normalCoreRowNum = indicesNum;
64+ td->tailCoreRowNum = 0;
65+ 
66+ td->normalCoreRowUbLoop = 1;
67+ td->normalCoreNormalLoopRows = indicesNum;
68+ td->normalCoreTailLoopRows = 0;
69+ td->tailCoreRowUbLoop = 1;
70+ td->tailCoreNormalLoopRows = 1;
71+ td->tailCoreTailLoopRows = 0;
72+ 
73+ td->normalCoreColUbLoop = 1;
74+ td->normalCoreNormalLoopCols = updatesInner;
75+ td->normalCoreTailLoopCols = updatesInner;
76+ td->tailCoreColUbLoop = 1;
77+ td->tailCoreNormalLoopCols = updatesInner;
78+ td->tailCoreTailLoopCols = updatesInner;
79+ 
80+ td->vecAlignSize = 0;
81+ td->indicesWorkspaceBufferSize = 0;
82+ td->coreNumInColDeterm = 0;
83+ td->tailCoreColUbDetermLoop = 0;
84+ td->normalCoreColUbDetermLoop = 0;
85+ td->tailCoreNormalLoopDetermCols = 0;
86+ td->normalCoreNormalLoopDetermCols = 0;
87+ td->tailCoreTailLoopDetermCols = 0;
88+ td->normalCoreTailLoopDetermCols = 0;
89+ td->updatesDeterminBufferSize = 0;
90+ td->outBufferDeterminSize = 0;
91+ td->normalCoreColDetermNum = 0;
92+ td->tailCoreColNumDeterm = 0;
93+ td->ubBlock = 32;
94+}
95+ 
96+TEST_F(scatter_add_with_sorted_apt_test, test_simd_float32_with_pos)
S
Ssunday4月28日

这组新增 kernel UT 当前主要验证可运行，没有对输出数值做断言。像重复索引聚合、排序稳定性、deterministic 语义这类核心行为容易回归而不被发现。建议补充至少 1-2 个带期望值比对的断言用例。

likedislike
W
wkkk0528
4月28日 评论：
97+{
98+ int64_t rows = 16;
99+ int64_t cols = 128;
100+ size_t var_size = (rows + 2) * cols * sizeof(float);
101+ size_t src_size = rows * cols * sizeof(float);
102+ size_t ind_size = rows * sizeof(int32_t);
103+ size_t pos_size = rows * sizeof(int32_t);
104+ size_t output_size = (rows + 2) * cols * sizeof(float);
105+ size_t tiling_data_size = sizeof(ScatterAddWithSortedSimdTilingData);
106+ 
107+ uint8_t *var = (uint8_t *)AscendC::GmAlloc(var_size);
108+ uint8_t *src = (uint8_t *)AscendC::GmAlloc(src_size);
109+ uint8_t *ind = (uint8_t *)AscendC::GmAlloc(ind_size);
110+ uint8_t *pos = (uint8_t *)AscendC::GmAlloc(pos_size);
111+ uint8_t *output = (uint8_t *)AscendC::GmAlloc(output_size);
112+ uint8_t *workspace = (uint8_t *)AscendC::GmAlloc(16 * 1024 * 1024);
113+ uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tiling_data_size);
114+ uint32_t blockDim = 1;
115+ 
116+ memset(var, 0, var_size);
117+ memset(src, 1, src_size);
118+ memset(ind, 0, ind_size);
119+ memset(pos, 0, pos_size);
120+ memset(output, 0, output_size);
121+ 
122+ ScatterAddWithSortedSimdTilingData *td = reinterpret_cast<ScatterAddWithSortedSimdTilingData *>(tiling);
123+ FillSimdTilingData(td, rows, cols, true, 0);
124+ 
125+ ICPU_SET_TILING_KEY(0);
126+ AscendC::SetKernelMode(KernelMode::AIV_MODE);
127+ ICPU_RUN_KF(scatter_add_with_sorted, blockDim, var, src, ind, pos, output, workspace, (uint8_t *)td);
128+ 
129+ AscendC::GmFree(var);
130+ AscendC::GmFree(src);
131+ AscendC::GmFree(ind);
132+ AscendC::GmFree(pos);
133+ AscendC::GmFree(output);
134+ AscendC::GmFree(workspace);
135+ AscendC::GmFree(tiling);
136+}
137+ 
138+TEST_F(scatter_add_with_sorted_apt_test, test_simd_float32_without_pos)
139+{
140+ int64_t rows = 16;
141+ int64_t cols = 128;
142+ size_t var_size = (rows + 2) * cols * sizeof(float);
143+ size_t src_size = rows * cols * sizeof(float);
144+ size_t ind_size = rows * sizeof(int32_t);
145+ size_t pos_size = rows * sizeof(int32_t);
146+ size_t output_size = (rows + 2) * cols * sizeof(float);
147+ size_t tiling_data_size = sizeof(ScatterAddWithSortedSimdTilingData);
148+ 
149+ uint8_t *var = (uint8_t *)AscendC::GmAlloc(var_size);
150+ uint8_t *src = (uint8_t *)AscendC::GmAlloc(src_size);
151+ uint8_t *ind = (uint8_t *)AscendC::GmAlloc(ind_size);
152+ uint8_t *pos = (uint8_t *)AscendC::GmAlloc(pos_size);
153+ uint8_t *output = (uint8_t *)AscendC::GmAlloc(output_size);
154+ uint8_t *workspace = (uint8_t *)AscendC::GmAlloc(16 * 1024 * 1024);
155+ uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tiling_data_size);
156+ uint32_t blockDim = 1;
157+ 
158+ memset(var, 0, var_size);
159+ memset(src, 1, src_size);
160+ memset(ind, 0, ind_size);
161+ memset(output, 0, output_size);
162+ 
163+ ScatterAddWithSortedSimdTilingData *td = reinterpret_cast<ScatterAddWithSortedSimdTilingData *>(tiling);
164+ FillSimdTilingData(td, rows, cols, false, 0);
165+ 
166+ ICPU_SET_TILING_KEY(0);
167+ AscendC::SetKernelMode(KernelMode::AIV_MODE);
168+ ICPU_RUN_KF(scatter_add_with_sorted, blockDim, var, src, ind, pos, output, workspace, (uint8_t *)td);
169+ 
170+ AscendC::GmFree(var);
171+ AscendC::GmFree(src);
172+ AscendC::GmFree(ind);
173+ AscendC::GmFree(pos);
174+ AscendC::GmFree(output);
175+ AscendC::GmFree(workspace);
176+ AscendC::GmFree(tiling);
177+}
178+ 
179+TEST_F(scatter_add_with_sorted_apt_test, test_empty_shape)
180+{
181+ int64_t rows = 0;
182+ int64_t cols = 128;
183+ size_t var_size = 2 * cols * sizeof(float);
184+ size_t src_size = 1 * sizeof(float);
185+ size_t ind_size = 1 * sizeof(int32_t);
186+ size_t pos_size = 1 * sizeof(int32_t);
187+ size_t output_size = 2 * cols * sizeof(float);
188+ size_t tiling_data_size = sizeof(ScatterAddWithSortedSimdTilingData);
189+ 
190+ uint8_t *var = (uint8_t *)AscendC::GmAlloc(var_size);
191+ uint8_t *src = (uint8_t *)AscendC::GmAlloc(src_size);
192+ uint8_t *ind = (uint8_t *)AscendC::GmAlloc(ind_size);
193+ uint8_t *pos = (uint8_t *)AscendC::GmAlloc(pos_size);
194+ uint8_t *output = (uint8_t *)AscendC::GmAlloc(output_size);
195+ uint8_t *workspace = (uint8_t *)AscendC::GmAlloc(16 * 1024 * 1024);
196+ uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tiling_data_size);
197+ uint32_t blockDim = 1;
198+ 
199+ memset(var, 0, var_size);
200+ memset(output, 0, output_size);
201+ 
202+ ScatterAddWithSortedSimdTilingData *td = reinterpret_cast<ScatterAddWithSortedSimdTilingData *>(tiling);
203+ FillSimdTilingData(td, rows, cols, false, 2);
204+ 
205+ ICPU_SET_TILING_KEY(2);
206+ AscendC::SetKernelMode(KernelMode::AIV_MODE);
207+ ICPU_RUN_KF(scatter_add_with_sorted, blockDim, var, src, ind, pos, output, workspace, (uint8_t *)td);
208+ 
209+ AscendC::GmFree(var);
210+ AscendC::GmFree(src);
211+ AscendC::GmFree(ind);
212+ AscendC::GmFree(pos);
213+ AscendC::GmFree(output);
214+ AscendC::GmFree(workspace);
215+ AscendC::GmFree(tiling);
216+}
Aindex/scatter_add_with_sorted/tests/ut/op_kernel/scatter_add_with_sorted_tiling_def.h+80-0
@@ -0,0 +1,80 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions 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 SCATTER_ADD_WITH_SORTED_TILING_DEF_H
12+#define SCATTER_ADD_WITH_SORTED_TILING_DEF_H
13+ 
14+#include "kernel_tiling/kernel_tiling.h"
15+ 
16+struct ScatterAddWithSortedTilingData {
17+ uint64_t usedCoreNum = 0;
18+ uint64_t extraTaskCore = 0;
19+ uint64_t eachCount = 0;
20+ uint64_t lastCount = 0;
21+ uint64_t inputCount = 0;
22+ uint64_t indicesCount = 0;
23+ uint64_t updatesCount = 0;
24+ uint64_t inputOneTime = 0;
25+ uint64_t updatesOneTime = 0;
26+ uint64_t updatesAlign = 0;
27+ uint64_t maxSize = 0;
28+ uint64_t eachNum = 0;
29+ uint64_t eachLoop = 0;
30+ uint64_t eachTail = 0;
31+ uint64_t lastNum = 0;
32+ uint64_t lastLoop = 0;
33+ uint64_t lastTail = 0;
34+ uint64_t updatesLoop = 0;
35+ uint64_t updatesEach = 0;
36+ uint64_t updatesLast = 0;
37+};
38+ 
39+#define DTYPE_X int64_t
40+ 
41+#pragma pack(1)
42+ 
43+#pragma pack()
44+ 
45+#define CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \
46+ __ubuf__ tilingStruct* tilingDataPointer = \
47+ reinterpret_cast<__ubuf__ tilingStruct*>((__ubuf__ uint8_t*)(tilingPointer));
48+ 
49+#define INIT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \
50+ CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer);
51+ 
52+#define GET_TILING_DATA(tilingData, tilingPointer) \
53+ ScatterAddWithSortedTilingData tilingData; \
54+ INIT_TILING_DATA(ScatterAddWithSortedTilingData, tilingDataPointer, tilingPointer); \
55+ (tilingData).usedCoreNum = tilingDataPointer->usedCoreNum; \
56+ (tilingData).extraTaskCore = tilingDataPointer->extraTaskCore; \
57+ (tilingData).eachCount = tilingDataPointer->eachCount; \
58+ (tilingData).lastCount = tilingDataPointer->lastCount; \
59+ (tilingData).inputCount = tilingDataPointer->inputCount; \
60+ (tilingData).indicesCount = tilingDataPointer->indicesCount; \
61+ (tilingData).updatesCount = tilingDataPointer->updatesCount; \
62+ (tilingData).inputOneTime = tilingDataPointer->inputOneTime; \
63+ (tilingData).updatesOneTime = tilingDataPointer->updatesOneTime; \
64+ (tilingData).updatesAlign = tilingDataPointer->updatesAlign; \
65+ (tilingData).maxSize = tilingDataPointer->maxSize; \
66+ (tilingData).eachNum = tilingDataPointer->eachNum; \
67+ (tilingData).eachLoop = tilingDataPointer->eachLoop; \
68+ (tilingData).eachTail = tilingDataPointer->eachTail; \
69+ (tilingData).lastNum = tilingDataPointer->lastNum; \
70+ (tilingData).lastLoop = tilingDataPointer->lastLoop; \
71+ (tilingData).lastTail = tilingDataPointer->lastTail; \
72+ (tilingData).updatesLoop = tilingDataPointer->updatesLoop; \
73+ (tilingData).updatesEach = tilingDataPointer->updatesEach; \
74+ (tilingData).updatesLast = tilingDataPointer->updatesLast;
75+ 
76+#define GET_TILING_DATA_WITH_STRUCT(tilingStruct, tilingData, tilingArg) \
77+ tilingStruct tilingData; \
78+ INIT_TILING_DATA(tilingStruct, tilingData##_ptr, tilingArg);
79+ 
80+#endif // SCATTER_ADD_WITH_SORTED_TILING_DEF_H
Mindex/scatter_add_with_sorted/tests/ut/op_kernel/test_scatter_add_with_sorted.cpp+3-3
@@ -15,7 +15,7 @@
15#include <array>15#include <array>
16#include <vector>16#include <vector>
17#include "gtest/gtest.h"17#include "gtest/gtest.h"
18-#include "test_scatter_add_with_sorted_tiling_def.h"18+#include "scatter_add_with_sorted_tiling_def.h"
19 19 
20#ifdef __CCE_KT_TEST__20#ifdef __CCE_KT_TEST__
21#include "tikicpulib.h"21#include "tikicpulib.h"
@@ -52,7 +52,7 @@ TEST_F(scatter_add_with_sorted_test, test_case_fp32)
52 size_t ind_size = 63 * sizeof(int);52 size_t ind_size = 63 * sizeof(int);
53 size_t pos_size = 63 * sizeof(int);53 size_t pos_size = 63 * sizeof(int);
54 size_t output_size = 65 * 4096 * sizeof(float);54 size_t output_size = 65 * 4096 * sizeof(float);
55- size_t tiling_data_size = sizeof(ScatterAddWithSortedTilingDataDef);55+ size_t tiling_data_size = sizeof(ScatterAddWithSortedTilingData);
56 56 
57 uint8_t* var = (uint8_t*)AscendC::GmAlloc(var_size);57 uint8_t* var = (uint8_t*)AscendC::GmAlloc(var_size);
58 uint8_t* src = (uint8_t*)AscendC::GmAlloc(src_size);58 uint8_t* src = (uint8_t*)AscendC::GmAlloc(src_size);
@@ -80,7 +80,7 @@ TEST_F(scatter_add_with_sorted_test, test_case_fp32)
80 ReadFile(path + "/scatter_add_with_sorted_data/pos.bin", pos_size, pos, pos_size);80 ReadFile(path + "/scatter_add_with_sorted_data/pos.bin", pos_size, pos, pos_size);
81 ReadFile(path + "/scatter_add_with_sorted_data/tiling.bin", tiling_data_size, tiling, tiling_data_size);81 ReadFile(path + "/scatter_add_with_sorted_data/tiling.bin", tiling_data_size, tiling, tiling_data_size);
82 82 
83- ScatterAddWithSortedTilingDataDef* tilingDatafromBin = reinterpret_cast<ScatterAddWithSortedTilingDataDef*>(tiling);83+ ScatterAddWithSortedTilingData* tilingDatafromBin = reinterpret_cast<ScatterAddWithSortedTilingData*>(tiling);
84 84 
85 ICPU_SET_TILING_KEY(11);85 ICPU_SET_TILING_KEY(11);
86 AscendC::SetKernelMode(KernelMode::AIV_MODE);86 AscendC::SetKernelMode(KernelMode::AIV_MODE);
Mindex/scatter_add_with_sorted/tests/ut/op_kernel/test_scatter_add_with_sorted_tiling_def.h+32-31
@@ -4,20 +4,16 @@
4 * CANN Open Software License Agreement Version 2.0 (the "License").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
6 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,6 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
8 * See LICENSE in the root of the software repository for the full text of the License.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
11-/*!11+#ifndef SCATTER_ADD_WITH_SORTED_TILING_DEF_H
12- * \file test_scatter_add_with_sorted_tiling.h12+#define SCATTER_ADD_WITH_SORTED_TILING_DEF_H
13- * \brief
14- */
15-#ifndef _FAST_OP_TEST_SCATTER_ADD_WITH_SORTED_TILING_H_
16-#define _FAST_OP_TEST_SCATTER_ADD_WITH_SORTED_TILING_H_
17 13 
18#include "kernel_tiling/kernel_tiling.h"14#include "kernel_tiling/kernel_tiling.h"
19 15 
20-struct ScatterAddWithSortedTilingDataDef {16+struct ScatterAddWithSortedTilingData {
21 uint64_t usedCoreNum = 0;17 uint64_t usedCoreNum = 0;
22 uint64_t extraTaskCore = 0;18 uint64_t extraTaskCore = 0;
23 uint64_t eachCount = 0;19 uint64_t eachCount = 0;
@@ -53,27 +49,32 @@ struct ScatterAddWithSortedTilingDataDef {
53#define INIT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \49#define INIT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \
54 CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer);50 CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer);
55 51 
56-#define GET_TILING_DATA(tilingData, tilingPointer) \52+#define GET_TILING_DATA(tilingData, tilingPointer) \
57- ScatterAddWithSortedTilingDataDef tilingData; \53+ ScatterAddWithSortedTilingData tilingData; \
58- INIT_TILING_DATA(ScatterAddWithSortedTilingDataDef, tilingDataPointer, tilingPointer); \54+ INIT_TILING_DATA(ScatterAddWithSortedTilingData, tilingDataPointer, tilingPointer); \
59- (tilingData).usedCoreNum = tilingDataPointer->usedCoreNum; \55+ (tilingData).usedCoreNum = tilingDataPointer->usedCoreNum; \
60- (tilingData).extraTaskCore = tilingDataPointer->extraTaskCore; \56+ (tilingData).extraTaskCore = tilingDataPointer->extraTaskCore; \
61- (tilingData).eachCount = tilingDataPointer->eachCount; \57+ (tilingData).eachCount = tilingDataPointer->eachCount; \
62- (tilingData).lastCount = tilingDataPointer->lastCount; \58+ (tilingData).lastCount = tilingDataPointer->lastCount; \
63- (tilingData).inputCount = tilingDataPointer->inputCount; \59+ (tilingData).inputCount = tilingDataPointer->inputCount; \
64- (tilingData).indicesCount = tilingDataPointer->indicesCount; \60+ (tilingData).indicesCount = tilingDataPointer->indicesCount; \
65- (tilingData).updatesCount = tilingDataPointer->updatesCount; \61+ (tilingData).updatesCount = tilingDataPointer->updatesCount; \
66- (tilingData).inputOneTime = tilingDataPointer->inputOneTime; \62+ (tilingData).inputOneTime = tilingDataPointer->inputOneTime; \
67- (tilingData).updatesOneTime = tilingDataPointer->updatesOneTime; \63+ (tilingData).updatesOneTime = tilingDataPointer->updatesOneTime; \
68- (tilingData).updatesAlign = tilingDataPointer->updatesAlign; \64+ (tilingData).updatesAlign = tilingDataPointer->updatesAlign; \
69- (tilingData).maxSize = tilingDataPointer->maxSize; \65+ (tilingData).maxSize = tilingDataPointer->maxSize; \
70- (tilingData).eachNum = tilingDataPointer->eachNum; \66+ (tilingData).eachNum = tilingDataPointer->eachNum; \
71- (tilingData).eachLoop = tilingDataPointer->eachLoop; \67+ (tilingData).eachLoop = tilingDataPointer->eachLoop; \
72- (tilingData).eachTail = tilingDataPointer->eachTail; \68+ (tilingData).eachTail = tilingDataPointer->eachTail; \
73- (tilingData).lastNum = tilingDataPointer->lastNum; \69+ (tilingData).lastNum = tilingDataPointer->lastNum; \
74- (tilingData).lastLoop = tilingDataPointer->lastLoop; \70+ (tilingData).lastLoop = tilingDataPointer->lastLoop; \
75- (tilingData).lastTail = tilingDataPointer->lastTail; \71+ (tilingData).lastTail = tilingDataPointer->lastTail; \
76- (tilingData).updatesLoop = tilingDataPointer->updatesLoop; \72+ (tilingData).updatesLoop = tilingDataPointer->updatesLoop; \
77- (tilingData).updatesEach = tilingDataPointer->updatesEach; \73+ (tilingData).updatesEach = tilingDataPointer->updatesEach; \
78 (tilingData).updatesLast = tilingDataPointer->updatesLast;74 (tilingData).updatesLast = tilingDataPointer->updatesLast;
79-#endif // _FAST_OP_TEST_SCATTER_ADD_WITH_SORTED_TILING_H_75+ 
76+#define GET_TILING_DATA_WITH_STRUCT(tilingStruct, tilingData, tilingArg) \
77+ tilingStruct tilingData; \
78+ INIT_TILING_DATA(tilingStruct, tilingData##_ptr, tilingArg);
79+ 
80+#endif // SCATTER_ADD_WITH_SORTED_TILING_DEF_H
Mindex/scatter_elements_v2/op_api/aclnn_scatter.cpp+50-9
@@ -40,6 +40,9 @@
40#include "opdev/op_executor.h"40#include "opdev/op_executor.h"
41#include "opdev/tensor_view_utils.h"41#include "opdev/tensor_view_utils.h"
42#include "op_api/aclnn_util.h"42#include "op_api/aclnn_util.h"
43+#include "runtime/context.h"
44+#include "acl/acl_rt.h"
45+ 
43 46 
44using namespace op;47using namespace op;
45#ifdef __cplusplus48#ifdef __cplusplus
@@ -71,8 +74,25 @@ static const std::initializer_list<op::DataType> DTYPE_SUPPORT_LIST_INDICES = {
71 op::DataType::DT_INT64, op::DataType::DT_INT32};74 op::DataType::DT_INT64, op::DataType::DT_INT32};
72 75 
73static const std::initializer_list<op::DataType> SCATTER_ADD_AICORE_REGBASE_DTYPE_SUPPORT_LIST = {76static const std::initializer_list<op::DataType> SCATTER_ADD_AICORE_REGBASE_DTYPE_SUPPORT_LIST = {
74- op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16, 77+ op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16,
75- op::DataType::DT_INT32, op::DataType::DT_INT8, op::DataType::DT_UINT8};78+ op::DataType::DT_INT32, op::DataType::DT_INT8, op::DataType::DT_UINT8};
79+ 
80+static const std::initializer_list<op::DataType> SCATTER_ADD_WITH_SORTED_950_DTYPE_SUPPORT_LIST = {
81+ op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16};
82+ 
83+static bool IsUseScatterAddWithSorted(const aclTensor* varRef)
84+{
85+ int64_t deterministicValue = 0;
86+ rtError_t retRts = aclrtGetSysParamOpt(ACL_OPT_DETERMINISTIC, &deterministicValue);
87+ if (retRts != RT_ERROR_NONE) {
88+ deterministicValue = 0;
89+ }
90+ bool isAscend950 = Ops::NN::AclnnUtil::IsRegbase();
91+ if (!(isAscend950 && deterministicValue != 0)) {
92+ return false;
93+ }
94+ return CheckType(varRef->GetDataType(), SCATTER_ADD_WITH_SORTED_950_DTYPE_SUPPORT_LIST);
95+}
76 96 
77static const std::initializer_list<DataType>& GetDtypeSupportList()97static const std::initializer_list<DataType>& GetDtypeSupportList()
78{98{
@@ -536,7 +556,6 @@ static aclnnStatus ExecScatterBase(
536 aclOpExecutor* executor)556 aclOpExecutor* executor)
537{557{
538 const std::string& reduction = GetReduceStr(reduce);558 const std::string& reduction = GetReduceStr(reduce);
539- 
540 auto ret = CheckParams(self, index, src, out);559 auto ret = CheckParams(self, index, src, out);
541 CHECK_RET(ret == ACLNN_SUCCESS, ret);560 CHECK_RET(ret == ACLNN_SUCCESS, ret);
542 561 
@@ -594,7 +613,7 @@ static aclnnStatus ExecScatterBase(
594 // index的步长为[1,0]或者[x,1,0] 且x不为0，是index expand场景，走scatteraddwithsorted, 超过16777216的FP32无法精准表示整数613 // index的步长为[1,0]或者[x,1,0] 且x不为0，是index expand场景，走scatteraddwithsorted, 超过16777216的FP32无法精准表示整数
595 bool expandFlag =614 bool expandFlag =
596 aicore910b &&615 aicore910b &&
597- ((selfDimNum == TWO_DIM && dimFinal != 1) ||616+ ((selfDimNum == TWO_DIM && indexShape[0] < MAX_EXACT_FLOAT && dimFinal != 1) ||
598 (selfDimNum == THREE_DIM && indexShape[0] * indexShape[1] < MAX_EXACT_FLOAT && dimFinal != TWO_DIM)) &&617 (selfDimNum == THREE_DIM && indexShape[0] * indexShape[1] < MAX_EXACT_FLOAT && dimFinal != TWO_DIM)) &&
599 strides[selfDimNum + NEG_TWO] == 1 && strides[selfDimNum + NEG_ONE] == 0 && shape.GetDimNum() == 1;618 strides[selfDimNum + NEG_TWO] == 1 && strides[selfDimNum + NEG_ONE] == 0 && shape.GetDimNum() == 1;
600 if (selfDimNum == THREE_DIM) {619 if (selfDimNum == THREE_DIM) {
@@ -612,10 +631,12 @@ static aclnnStatus ExecScatterBase(
612 CHECK_COND(scatterRes != nullptr, ACLNN_ERR_INNER_NULLPTR, "DoScatterAddWithSorted failed!");631 CHECK_COND(scatterRes != nullptr, ACLNN_ERR_INNER_NULLPTR, "DoScatterAddWithSorted failed!");
613 return ACLNN_SUCCESS;632 return ACLNN_SUCCESS;
614 }633 }
615- 634+ 
616 bool expandFlagRegbase = scatterAddRegbaseSupport && selfDimNum == TWO_DIM && dimFinal != 1 &&635 bool expandFlagRegbase = scatterAddRegbaseSupport && selfDimNum == TWO_DIM && dimFinal != 1 &&
617- strides[selfDimNum + NEG_TWO] == 1 && strides[selfDimNum + NEG_ONE] == 0 && shape.GetDimNum() == 1;636+ strides[selfDimNum + NEG_TWO] == 1 && strides[selfDimNum + NEG_ONE] == 0 &&
618- if (expandFlagRegbase && IsMeetScatterAddShape(selfContiguous->GetViewShape(), indexShape, srcContiguous->GetViewShape(), dimFinal)) {637+ shape.GetDimNum() == 1;
638+ if (expandFlagRegbase &&
639+ IsMeetScatterAddShape(selfContiguous->GetViewShape(), indexShape, srcContiguous->GetViewShape(), dimFinal)) {
619 OP_LOGD("Use AICORE for ScatterAdd.");640 OP_LOGD("Use AICORE for ScatterAdd.");
620 op::Shape newViewShape;641 op::Shape newViewShape;
621 newViewShape.SetDimNum(dimFinal + 1);642 newViewShape.SetDimNum(dimFinal + 1);
@@ -623,8 +644,28 @@ static aclnnStatus ExecScatterBase(
623 auto indexTmp = executor->CreateView(index, newViewShape, index->GetViewOffset());644 auto indexTmp = executor->CreateView(index, newViewShape, index->GetViewOffset());
624 CHECK_RET(indexTmp != nullptr, ACLNN_ERR_INNER_NULLPTR);645 CHECK_RET(indexTmp != nullptr, ACLNN_ERR_INNER_NULLPTR);
625 indexTmp->SetDataType(index->GetDataType());646 indexTmp->SetDataType(index->GetDataType());
626- 647+ 
627- const aclTensor *scatterRes = l0op::ScatterAdd(selfContiguous, indexTmp, srcContiguous, false, executor);648+ const aclTensor* scatterRes = nullptr;
649+ if (IsUseScatterAddWithSorted(selfContiguous)) {
650+ auto indexSize = static_cast<int64_t>(indexTmp->Size());
651+ if (indexSize > 1) {
652+ auto indicesType = indexTmp->GetDataType();
653+ auto sortResult = l0op::Sort(indexTmp, -1, false, true, indicesType, executor);
654+ auto sortIdxOut = std::get<0>(sortResult);
655+ auto posIdx = std::get<1>(sortResult);
656+ CHECK_RET(sortIdxOut != nullptr && posIdx != nullptr, ACLNN_ERR_INNER_NULLPTR);
657+ scatterRes =
658+ l0op::ScatterAddWithSorted(selfContiguous, srcContiguous, sortIdxOut, posIdx, "add", executor);
659+ } else {
660+ const aclTensor* posTensor =
661+ executor->ConvertToTensor(executor->AllocScalar(0), op::DataType::DT_INT32);
662+ CHECK_RET(posTensor != nullptr, ACLNN_ERR_INNER_NULLPTR);
663+ scatterRes =
664+ l0op::ScatterAddWithSorted(selfContiguous, srcContiguous, indexTmp, posTensor, "add", executor);
665+ }
666+ } else {
667+ scatterRes = l0op::ScatterAdd(selfContiguous, indexTmp, srcContiguous, false, executor);
668+ }
628 CHECK_RET(scatterRes != nullptr, ACLNN_ERR_INNER_NULLPTR);669 CHECK_RET(scatterRes != nullptr, ACLNN_ERR_INNER_NULLPTR);
629 670 
630 scatterRes = needUnsqueeze ? l0op::SqueezeNd(scatterRes, squeezeDim, executor) : scatterRes;671 scatterRes = needUnsqueeze ? l0op::SqueezeNd(scatterRes, squeezeDim, executor) : scatterRes;
Mindex/tf_scatter_add/op_host/op_api/aclnn_tf_scatter_add.cpp+72-15
@@ -26,6 +26,11 @@
26#include "opdev/op_dfx.h"26#include "opdev/op_dfx.h"
27#include "opdev/tensor_view_utils.h"27#include "opdev/tensor_view_utils.h"
28#include "opdev/op_executor.h"28#include "opdev/op_executor.h"
29+#include "index/scatter_add_with_sorted/op_host/op_api/scatter_add_with_sorted.h"
30+#include "level0/sort.h"
31+#include "runtime/context.h"
32+#include "acl/acl_rt.h"
33+#include "op_api/aclnn_util.h"
29 34 
30using namespace op;35using namespace op;
31#ifdef __cplusplus36#ifdef __cplusplus
@@ -45,6 +50,24 @@ static const std::initializer_list<op::DataType> INDEX_DTYPE_SUPPORT_LIST = {
45 50 
46static const std::initializer_list<op::DataType> NULL_SUPPORT_LIST = {};51static const std::initializer_list<op::DataType> NULL_SUPPORT_LIST = {};
47 52 
53+static const std::initializer_list<op::DataType> SCATTER_ADD_WITH_SORTED_950_DTYPE_SUPPORT_LIST = {
54+ op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16};
55+ 
56+static bool IsUseScatterAddWithSorted(const aclTensor* varRef)
57+{
58+ int64_t deterministicValue = 0;
59+ rtError_t retRts = aclrtGetSysParamOpt(ACL_OPT_DETERMINISTIC, &deterministicValue);
60+ if (retRts != RT_ERROR_NONE) {
61+ deterministicValue = 0;
62+ }
63+ bool isAscend950 = Ops::NN::AclnnUtil::IsRegbase();
64+ 
65+ if (!(isAscend950 && deterministicValue != 0)) {
66+ return false;
67+ }
68+ return CheckType(varRef->GetDataType(), SCATTER_ADD_WITH_SORTED_950_DTYPE_SUPPORT_LIST);
69+}
70+ 
48static bool CheckNotNull(aclTensor* varRef, const aclTensor* indices, const aclTensor* updates)71static bool CheckNotNull(aclTensor* varRef, const aclTensor* indices, const aclTensor* updates)
49{72{
50 OP_CHECK_NULL(varRef, return false);73 OP_CHECK_NULL(varRef, return false);
@@ -55,9 +78,11 @@ static bool CheckNotNull(aclTensor* varRef, const aclTensor* indices, const aclT
55 78 
56static const std::initializer_list<DataType>& GetDtypeSupportList()79static const std::initializer_list<DataType>& GetDtypeSupportList()
57{80{
58- if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910 ||81+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
59- (GetCurrentPlatformInfo().GetSocVersion() >= SocVersion::ASCEND910B &&82+ auto socVersion = GetCurrentPlatformInfo().GetSocVersion();
60- GetCurrentPlatformInfo().GetSocVersion() <= SocVersion::ASCEND910E)) {83+ if (Ops::NN::AclnnUtil::IsRegbase(curArch) ||
84+ socVersion == SocVersion::ASCEND910 ||
85+ (socVersion >= SocVersion::ASCEND910B && socVersion <= SocVersion::ASCEND910E)) {
61 return ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST;86 return ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST;
62 }87 }
63 return NULL_SUPPORT_LIST;88 return NULL_SUPPORT_LIST;
@@ -173,6 +198,28 @@ static aclnnStatus CheckParams(aclTensor* varRef, const aclTensor* indices, cons
173 return ACLNN_SUCCESS;198 return ACLNN_SUCCESS;
174}199}
175 200 
201+static const aclTensor* DoScatterAddWithSortedForTfScatterAdd(
202+ const aclTensor* varRef, const aclTensor* indices, const aclTensor* updates, aclOpExecutor* executor)
203+{
204+ auto indexSize = static_cast<int64_t>(indices->Size());
205+ const aclTensor* scatterAddRes = nullptr;
206+ if (indexSize > 1) {
207+ // 直接对 indices 排序，输出 indices 类型与输入相同
208+ auto indicesType = indices->GetDataType();
209+ auto sortResult = l0op::Sort(indices, -1, false, true, indicesType, executor);
210+ auto sortIdxOut = std::get<0>(sortResult);
211+ auto posIdx = std::get<1>(sortResult);
212+ CHECK_RET(sortIdxOut != nullptr && posIdx != nullptr, nullptr);
213+ scatterAddRes = l0op::ScatterAddWithSorted(varRef, updates, sortIdxOut, posIdx, "add", executor);
214+ } else {
215+ // indexSize == 1 时，不需要 Sort，直接使用原始 indices
216+ const aclTensor* posTensor = executor->ConvertToTensor(executor->AllocScalar(0), op::DataType::DT_INT32);
217+ CHECK_RET(posTensor != nullptr, nullptr);
218+ scatterAddRes = l0op::ScatterAddWithSorted(varRef, updates, indices, posTensor, "add", executor);
219+ }
220+ return scatterAddRes;
221+}
222+ 
176aclnnStatus aclnnTfScatterAddGetWorkspaceSize(223aclnnStatus aclnnTfScatterAddGetWorkspaceSize(
177 aclTensor* varRef, const aclTensor* indices, const aclTensor* updates, uint64_t* workspaceSize,224 aclTensor* varRef, const aclTensor* indices, const aclTensor* updates, uint64_t* workspaceSize,
178 aclOpExecutor** executor)225 aclOpExecutor** executor)
@@ -215,23 +262,33 @@ aclnnStatus aclnnTfScatterAddGetWorkspaceSize(
215 auto updatesContiguousFloat = l0op::Cast(updatesContiguous, op::DataType::DT_FLOAT, uniqueExecutor.get());262 auto updatesContiguousFloat = l0op::Cast(updatesContiguous, op::DataType::DT_FLOAT, uniqueExecutor.get());
216 CHECK_RET(updatesContiguousFloat != nullptr, ACLNN_ERR_INNER_NULLPTR);263 CHECK_RET(updatesContiguousFloat != nullptr, ACLNN_ERR_INNER_NULLPTR);
217 264 
218- auto scatterAddResFloat =265+ const aclTensor* scatterAddResFloat = nullptr;
219- useScatterNd ?266+ if (!useScatterNd && IsUseScatterAddWithSorted(varRefContiguousFloat)) {
220- l0op::ScatterNdAdd(267+ scatterAddResFloat = DoScatterAddWithSortedForTfScatterAdd(
221- varRefContiguousFloat, indicesContiguous, updatesContiguousFloat, false, uniqueExecutor.get()) :268+ varRefContiguousFloat, indicesContiguous, updatesContiguousFloat, uniqueExecutor.get());
222- l0op::ScatterAdd(269+ } else {
223- varRefContiguousFloat, indicesContiguous, updatesContiguousFloat, false, uniqueExecutor.get());270+ scatterAddResFloat =
271+ useScatterNd ?
272+ l0op::ScatterNdAdd(
273+ varRefContiguousFloat, indicesContiguous, updatesContiguousFloat, false, uniqueExecutor.get()) :
274+ l0op::ScatterAdd(
275+ varRefContiguousFloat, indicesContiguous, updatesContiguousFloat, false, uniqueExecutor.get());
276+ }
224 CHECK_RET(scatterAddResFloat != nullptr, ACLNN_ERR_INNER_NULLPTR);277 CHECK_RET(scatterAddResFloat != nullptr, ACLNN_ERR_INNER_NULLPTR);
225 278 
226 scatterAddRes = l0op::Cast(scatterAddResFloat, op::DataType::DT_BF16, uniqueExecutor.get());279 scatterAddRes = l0op::Cast(scatterAddResFloat, op::DataType::DT_BF16, uniqueExecutor.get());
227 CHECK_RET(scatterAddRes != nullptr, ACLNN_ERR_INNER_NULLPTR);280 CHECK_RET(scatterAddRes != nullptr, ACLNN_ERR_INNER_NULLPTR);
228 } else {281 } else {
229- // 执行L0算子282+ if (!useScatterNd && IsUseScatterAddWithSorted(varRefContiguous)) {
230- scatterAddRes =283+ scatterAddRes = DoScatterAddWithSortedForTfScatterAdd(
231- useScatterNd ?284+ varRefContiguous, indicesContiguous, updatesContiguous, uniqueExecutor.get());
232- l0op::ScatterNdAdd(285+ } else {
233- varRefContiguous, indicesContiguous, updatesContiguous, false, uniqueExecutor.get()) :286+ scatterAddRes =
234- l0op::ScatterAdd(varRefContiguous, indicesContiguous, updatesContiguous, false, uniqueExecutor.get());287+ useScatterNd ?
288+ l0op::ScatterNdAdd(
289+ varRefContiguous, indicesContiguous, updatesContiguous, false, uniqueExecutor.get()) :
290+ l0op::ScatterAdd(varRefContiguous, indicesContiguous, updatesContiguous, false, uniqueExecutor.get());
291+ }
235 CHECK_RET(scatterAddRes != nullptr, ACLNN_ERR_INNER_NULLPTR);292 CHECK_RET(scatterAddRes != nullptr, ACLNN_ERR_INNER_NULLPTR);
236 }293 }
237 294 
Mscripts/kernel/binary_config/ascendc_config.json+1-1
@@ -543,7 +543,7 @@
543 {"name":"ScatterAdd", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "high_precision", "compile_options": {"ascend950": ["--cce-no-dcache-flush"]}},543 {"name":"ScatterAdd", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "high_precision", "compile_options": {"ascend950": ["--cce-no-dcache-flush"]}},
544 {"name":"ScatterNd", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "high_precision", "compile_options": {"ascend950": ["--cce-no-dcache-flush"]}},544 {"name":"ScatterNd", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "high_precision", "compile_options": {"ascend950": ["--cce-no-dcache-flush"]}},
545 {"name": "ScatterNdUpdate", "compute_units": ["ascend950", "ascend910b", "ascend910_93"], "auto_sync": true, "impl_mode": "", "compile_options": {"ascend950": ["--cce-no-dcache-flush"]}},545 {"name": "ScatterNdUpdate", "compute_units": ["ascend950", "ascend910b", "ascend910_93"], "auto_sync": true, "impl_mode": "", "compile_options": {"ascend950": ["--cce-no-dcache-flush"]}},
546- {"name":"ScatterAddWithSorted", "compute_units": ["ascend910b", "ascend910_93", "kirinx90", "kirin9030"], "auto_sync" : false},546+ {"name":"ScatterAddWithSorted", "compute_units": ["ascend910b", "ascend910_93", "kirinx90", "kirin9030", "ascend950"], "auto_sync" : false, "impl_mode" : ""},
547 {"name":"Scatter", "compute_units": ["ascend950"], "auto_sync" : false, "impl_mode" : "" },547 {"name":"Scatter", "compute_units": ["ascend950"], "auto_sync" : false, "impl_mode" : "" },
548 {"name":"BucketizeV2", "compute_units": ["ascend950"], "auto_sync" : false, "impl_mode" : "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},548 {"name":"BucketizeV2", "compute_units": ["ascend950"], "auto_sync" : false, "impl_mode" : "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},
549 {"name":"InplaceIndexAddWithSorted", "compute_units": ["ascend910b", "ascend910_93", "kirinx90", "kirin9030"], "auto_sync" : false},549 {"name":"InplaceIndexAddWithSorted", "compute_units": ["ascend910b", "ascend910_93", "kirinx90", "kirin9030"], "auto_sync" : false},