已合并
新增算子ScatterSub #1164
gengxueru创建于 1月29日
新增算子ScatterSub #1164
已合并
gengxueru创建于 1月29日
27 个文件变更+2606-597
Mindex/scatter_add/op_host/arch35/scatter_add_tiling_base.cpp+24-24文件内容审核中,请稍后刷新重试
Mindex/scatter_add/op_host/arch35/scatter_add_tiling_base.h+2-1
@@ -58,7 +58,6 @@ TILING_DATA_FIELD_DEF(uint64_t, isDeterministic);
58TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleCol);58TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleCol);
59TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsLoopSize);59TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsLoopSize);
60TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsTailUbFactor);60TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsTailUbFactor);
61-TILING_DATA_FIELD_DEF(uint64_t, ubSize);
62TILING_DATA_FIELD_DEF(uint64_t, perCoreHandleRows); //反量化按var[0]分核或indices分核61TILING_DATA_FIELD_DEF(uint64_t, perCoreHandleRows); //反量化按var[0]分核或indices分核
63TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleRows);62TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleRows);
64TILING_DATA_FIELD_DEF(uint64_t, rowsInUb);63TILING_DATA_FIELD_DEF(uint64_t, rowsInUb);
@@ -83,9 +82,11 @@ TILING_DATA_FIELD_DEF(uint64_t, normBlockCol); // 整核分块列数
83TILING_DATA_FIELD_DEF(uint64_t, tailBlockCol); // 列尾核分块列数82TILING_DATA_FIELD_DEF(uint64_t, tailBlockCol); // 列尾核分块列数
84TILING_DATA_FIELD_DEF(uint64_t, ubFactorRow); // UB每次循环搬运的行数83TILING_DATA_FIELD_DEF(uint64_t, ubFactorRow); // UB每次循环搬运的行数
85TILING_DATA_FIELD_DEF(uint64_t, ubFactorCol); // UB每次循环搬运的列数84TILING_DATA_FIELD_DEF(uint64_t, ubFactorCol); // UB每次循环搬运的列数
85+TILING_DATA_FIELD_DEF(uint64_t, indicesCastMode);
86END_TILING_DATA_DEF;86END_TILING_DATA_DEF;
87 87 
88REGISTER_TILING_DATA_CLASS(ScatterAdd, ScatterAddTilingData)88REGISTER_TILING_DATA_CLASS(ScatterAdd, ScatterAddTilingData)
89+REGISTER_TILING_DATA_CLASS(ScatterSub, ScatterAddTilingData)
89 90 
90ge::graphStatus ScatterAddTilingForAscendC(gert::TilingContext* context);91ge::graphStatus ScatterAddTilingForAscendC(gert::TilingContext* context);
91 92 
Aindex/scatter_add/op_host/op_api/aclnn_scatter_add.h+61-0
@@ -0,0 +1,61 @@
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+#ifndef OP_API_INC_SCATTER_OUT_H_
11+#define OP_API_INC_SCATTER_OUT_H_
12+ 
13+#include "aclnn/aclnn_base.h"
14+#include "aclnn_util.h"
15+ 
16+#ifdef __cplusplus
17+extern "C" {
18+#endif
19+ 
20+/**
21+ * @brief aclnnScatterAdd的第一段接口,根据具体的计算流程,计算workspace大小。
22+ * @domain aclnn_ops_infer
23+ *
24+ * 算子功能: 将源tensor中的值按指定的轴方向和index tensor中的位置关系逐个填入输出tensor中,
25+ * 若有多于一个src值被填入到self的同一位置,那么这些值将会在这一位置上进行累加
26+ * @param [in] self: npu device侧的aclTensor, 数据类型支持FLOAT16, FLOAT32, INT32, INT8, UINT8,
27+ * 支持非连续的Tensor,数据格式支持ND,
28+ * @param [in] dim: host侧的num, 数据类型支持INT64。
29+ * @param [in] index: npu device侧的aclTensor,数据类型支持INT32, int64类型,dim反向的维度数量需要与src相同。
30+ * 支持非连续的Tensor,数据格式支持ND。
31+ * @param [in] src: npu device侧的aclTensor,数据类型支持FLOAT16, FLOAT32, INT32, INT8,
32+ * UINT8类型,dim反向的维度数量需要与src相同。 支持非连续的Tensor,数据格式支持ND,且数据类型与self保持一致。
33+ * @param [in] out: npu device侧的aclTensor, 数据类型支持FLOAT16, FLOAT32, INT32, INT8, UINT8,
34+ * 数据类型,数据格式,tensor shape需要与self保持一致
35+ * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
36+ * @param [out] executor: 返回op执行器,包含算子计算流程。
37+ * @return aclnnStatus: 返回状态码
38+ */
39+ACLNN_API aclnnStatus aclnnScatterAddGetWorkspaceSize(const aclTensor* self, int64_t dim, const aclTensor* index,
40+ const aclTensor* src, aclTensor* out, uint64_t* workspaceSize,
41+ aclOpExecutor** executor);
42+ 
43+/**
44+ * @brief: aclnnScatterAdd的第二段接口,用于执行计算
45+ *
46+ * 算子功能: 将源tensor中的值按指定的轴方向和index tensor中的位置关系逐个填入输出tensor中,
47+ * 若有多于一个src值被填入到self的同一位置,那么这些值将会在这一位置上进行累加
48+ * @param [in] workspace: 在npu device侧申请的workspace内存起址。
49+ * @param [in] workspaceSize: 在npu device侧申请的workspace大小,由第一段接口aclnnScatterAddGetWorkspaceSize获取。
50+ * @param [in] stream: acl stream流。
51+ * @param [in] executor: op执行器,包含了算子计算流程。
52+ * @return aclnnStatus: 返回状态码。
53+ */
54+ACLNN_API aclnnStatus aclnnScatterAdd(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
55+ const aclrtStream stream);
56+ 
57+#ifdef __cplusplus
58+}
59+#endif
60+ 
61+#endif // OP_API_INC_SCATTER_OUT_H_
Aindex/scatter_add/op_host/op_api/scatter_add.h+22-0
@@ -0,0 +1,22 @@
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 PTA_NPU_OP_API_INC_LEVEL0_OP_SCATTER_ADD_H_
12+#define PTA_NPU_OP_API_INC_LEVEL0_OP_SCATTER_ADD_H_
13+ 
14+#include "opdev/op_executor.h"
15+ 
16+namespace l0op {
17+const aclTensor *ScatterAddWithAxis(const aclTensor *self, int64_t dim, const aclTensor *index,
18+ const aclTensor *src, aclOpExecutor *executor);
19+ 
20+}
21+ 
22+#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_SCATTER_ADD_H_
Mindex/scatter_add/op_kernel/arch35/scatter_add_common.h+208-52
@@ -28,10 +28,23 @@ constexpr uint64_t UB_AGLIN_VALUE = 32;
28constexpr uint64_t SORT_PAD_NUM = 2;28constexpr uint64_t SORT_PAD_NUM = 2;
29constexpr uint64_t HASH_BUCKER_BUFFER_SIZE = 128 * sizeof(float);29constexpr uint64_t HASH_BUCKER_BUFFER_SIZE = 128 * sizeof(float);
30constexpr int64_t DOUBLE_BUFFER = 2;30constexpr int64_t DOUBLE_BUFFER = 2;
31+constexpr uint32_t TWO = 2;
32+constexpr uint32_t THREE = 3;
33+constexpr uint32_t FOUR = 4;
31constexpr uint32_t CAST_0 = 0;34constexpr uint32_t CAST_0 = 0;
32constexpr uint32_t CAST_1 = 1;35constexpr uint32_t CAST_1 = 1;
33constexpr uint32_t CAST_2 = 2;36constexpr uint32_t CAST_2 = 2;
34constexpr uint32_t CAST_3 = 3;37constexpr uint32_t CAST_3 = 3;
38+constexpr uint32_t CAST_4 = 4;
39+constexpr uint32_t CAST_5 = 5;
40+constexpr int64_t VFLEN_INT64 = platform::GetVRegSize() / sizeof(int64_t);
41+constexpr int64_t VFLEN_INT32 = platform::GetVRegSize() / sizeof(int32_t);
42+constexpr int64_t VFLEN_INT16 = platform::GetVRegSize() / sizeof(int16_t);
43+constexpr int64_t VFLEN_INT16HALF = platform::GetVRegSize() / sizeof(int16_t) / TWO;
44+constexpr int64_t VFLEN_UINT8 = platform::GetVRegSize() / sizeof(uint8_t);
45+constexpr int64_t VFLEN_UINT8HALFHALF = platform::GetVRegSize() / sizeof(uint8_t) / FOUR;
46+constexpr uint32_t ADD = 0;
47+constexpr uint32_t SUB = 1;
35 48 
36constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};49constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};
37static constexpr MicroAPI::CastTrait castTraitU82Int32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,50static constexpr MicroAPI::CastTrait castTraitU82Int32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
@@ -71,6 +84,51 @@ __aicore__ inline void CastToInt32(LocalTensor<int32_t>& dstLocal, LocalTensor<T
71 }84 }
72}85}
73 86 
87+template<typename T>
88+__aicore__ inline void NegateUpdate(LocalTensor<T>& updatesLocal, uint32_t dataLen)
89+{
90+ if constexpr (IsSameType<T, uint8_t>::value) {
91+ return ;
92+ }
93+ 
94+ __local_mem__ T* updatesAddr = (__local_mem__ T*)updatesLocal.GetPhyAddr();
95+ uint32_t loopSize = platform::GetVRegSize() / sizeof(T);
96+ uint16_t loopTimes = ops::CeilDiv(dataLen, loopSize);
97+
98+ if constexpr (IsSameType<T, bfloat16_t>::value) {
99+ __VEC_SCOPE__
100+ {
101+ MicroAPI::RegTensor<bfloat16_t> updatesValue;
102+ MicroAPI::RegTensor<bfloat16_t> scalarReg;
103+ MicroAPI::RegTensor<bfloat16_t> dstReg;
104+ MicroAPI::MaskReg maskReg;
105+ uint32_t count = dataLen;
106+ bfloat16_t scalarValue = -1;
107+ MicroAPI::Duplicate(scalarReg, scalarValue);
108+ for (uint16_t j = 0; j < loopTimes; j++) {
109+ maskReg = MicroAPI::UpdateMask<T>(count);
110+ MicroAPI::DataCopy(updatesValue, updatesAddr + loopSize * j);
111+ MicroAPI::Mul(dstReg, updatesValue, scalarReg, maskReg);
112+ MicroAPI::DataCopy(updatesAddr + loopSize * j, dstReg, maskReg);
113+ }
114+ }
115+ } else {
116+ __VEC_SCOPE__
117+ {
118+ MicroAPI::RegTensor<T> updatesValue;
119+ MicroAPI::RegTensor<T> negValue;
120+ MicroAPI::MaskReg maskReg;
121+ uint32_t count = dataLen;
122+ for (uint16_t j = 0; j < loopTimes; j++) {
123+ maskReg = MicroAPI::UpdateMask<T>(count);
124+ MicroAPI::DataCopy(updatesValue, updatesAddr + loopSize * j);
125+ MicroAPI::Neg(negValue, updatesValue, maskReg);
126+ MicroAPI::DataCopy(updatesAddr + loopSize * j, negValue, maskReg);
127+ }
128+ }
129+ }
130+}
131+ 
74template <typename T>132template <typename T>
75__aicore__ inline void CastToOrigin(LocalTensor<T>& dstLocal, LocalTensor<int32_t>& srcLocal, uint32_t dataLen)133__aicore__ inline void CastToOrigin(LocalTensor<T>& dstLocal, LocalTensor<int32_t>& srcLocal, uint32_t dataLen)
76{134{
@@ -95,21 +153,152 @@ __aicore__ inline void CastToOrigin(LocalTensor<T>& dstLocal, LocalTensor<int32_
95 }153 }
96}154}
97 155 
98-template <typename T, uint64_t bufferNum, bool updatesIsScalar>156+template <typename IDX_T, typename CAST_T, uint32_t castType>
99-__aicore__ inline void BroadcastUpdatesScalar(157+__aicore__ inline void IndicesSortCast(LocalTensor<IDX_T> indicesLocal, LocalTensor<CAST_T> indicesCastLocal,
100- TQue<QuePosition::VECIN, bufferNum> updatesQueue, GlobalTensor<T> updatesGm, int32_t count)158+ LocalTensor<int32_t> indicesCastTmpLocal, uint32_t indicesCount)
101{159{
102- if constexpr (updatesIsScalar) {160+ if constexpr (castType == CAST_4) { // int32 Cast uint8
103- T updatesValue = updatesGm.GetValue(0);161+ CompareScalar(indicesCastLocal, indicesLocal, static_cast<IDX_T>(0), CMPMODE::GE, indicesCount);
104- auto vWaitSEventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));162+ Select(indicesLocal, indicesCastLocal, indicesLocal, static_cast<IDX_T>(255), SELMODE::VSEL_TENSOR_SCALAR_MODE, indicesCount);
105- SetFlag<HardEvent::S_V>(vWaitSEventID);163+ Cast<CAST_T, IDX_T>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
106- WaitFlag<HardEvent::S_V>(vWaitSEventID);164+ } else if constexpr (castType == CAST_3) { // int64 Cast int16
107- LocalTensor<T> updatesLocal = updatesQueue.template AllocTensor<T>();165+ Cast<int32_t, IDX_T>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
108- Duplicate(updatesLocal, updatesValue, count);166+ Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
109- updatesQueue.EnQue(updatesLocal);167+ } else if constexpr (castType == CAST_5) { // int64 Cast uint8
168+ CompareScalar(indicesCastLocal, indicesLocal, static_cast<IDX_T>(0), CMPMODE::GE, indicesCount);
169+ Select(indicesLocal, indicesCastLocal, indicesLocal, static_cast<IDX_T>(255), SELMODE::VSEL_TENSOR_SCALAR_MODE, indicesCount);
170+ Cast<int32_t, IDX_T>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
171+ Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
172+ } else { // CAST_1 + CAST_2, int32 Cast int16 + int64 Cast int32
173+ Cast<CAST_T, IDX_T>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
110 }174 }
111}175}
112 176 
177+template <typename T, uint32_t scatterOp>
178+__aicore__ inline void BroadcastUpdatesScalar(LocalTensor<T> updatesLocal, GlobalTensor<T> updatesGm, int32_t count)
179+{
180+ T updatesValue = updatesGm.GetValue(0);
181+ if constexpr (scatterOp == SUB) {
182+ updatesValue = -updatesValue;
183+ }
184+ auto vWaitSEventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
185+ SetFlag<HardEvent::S_V>(vWaitSEventID);
186+ WaitFlag<HardEvent::S_V>(vWaitSEventID);
187+ Duplicate(updatesLocal, updatesValue, count);
188+}
189+ 
190+template<typename IDX_T>
191+__aicore__ void ComputeUniqueIdNumInt64(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
192+{
193+ uint32_t counter = dataLen + 1;
194+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, selReg;
195+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
196+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskHalf;
197+ AscendC::MicroAPI::UnalignReg u0, uOut;
198+ for (uint16_t i = 0; i < loopCnt; ++i) {
199+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT64);
200+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
201+ auto startAddr = indicesAddr + i * VFLEN_INT64;
202+ DataCopy(sortedIdxReg, startAddr);
203+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
204+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
205+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
206+ AscendC::MicroAPI::MaskPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskHalf, cmpMask);
207+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskHalf);
208+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
209+ uniqueIdCountsAddr, selReg, uOut);
210+ }
211+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
212+}
213+ 
214+template<typename IDX_T>
215+__aicore__ void ComputeUniqueIdNumInt32(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
216+{
217+ uint32_t counter = dataLen + 1;
218+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, selReg;
219+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
220+ AscendC::MicroAPI::MaskReg cmpMask, maskReg;
221+ AscendC::MicroAPI::UnalignReg u0, uOut;
222+ for (uint16_t i = 0; i < loopCnt; ++i) {
223+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT32);
224+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
225+ auto startAddr = indicesAddr + i * VFLEN_INT32;
226+ DataCopy(sortedIdxReg, startAddr);
227+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
228+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
229+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
230+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, cmpMask);
231+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
232+ uniqueIdCountsAddr, selReg, uOut);
233+ }
234+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
235+}
236+ 
237+template<typename IDX_T>
238+__aicore__ void ComputeUniqueIdNumInt16(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
239+{
240+ uint32_t counter = dataLen + 1;
241+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, orderReg2, selReg, selReg2;
242+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
243+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskDouble1, maskDouble2;
244+ AscendC::MicroAPI::UnalignReg u0, uOut;
245+ for (uint16_t i = 0; i < loopCnt; ++i) {
246+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT16);
247+ AscendC::MicroAPI::Arange(orderReg2, i * VFLEN_INT16 + VFLEN_INT16HALF);
248+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
249+ auto startAddr = indicesAddr + i * VFLEN_INT16;
250+ DataCopy(sortedIdxReg, startAddr);
251+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
252+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
253+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
254+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskDouble1, cmpMask);
255+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskDouble2, cmpMask);
256+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskDouble1);
257+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskDouble2);
258+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg, uOut);
259+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg2, uOut);
260+ }
261+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
262+}
263+ 
264+template<typename IDX_T>
265+__aicore__ void ComputeUniqueIdNumUint8(__local_mem__ IDX_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
266+{
267+ uint32_t counter = dataLen + 1;
268+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, orderReg2, orderReg3, orderReg4;
269+ AscendC::MicroAPI::RegTensor<int32_t> selReg, selReg2, selReg3, selReg4;
270+ AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg, sortedIdxShiftOneReg;
271+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskFour1, maskFour2, maskFour3, maskFour4;
272+ AscendC::MicroAPI::UnalignReg u0, uOut;
273+ for (uint16_t i = 0; i < loopCnt; ++i) {
274+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_UINT8);
275+ AscendC::MicroAPI::Arange(orderReg2, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF);
276+ AscendC::MicroAPI::Arange(orderReg3, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF * TWO);
277+ AscendC::MicroAPI::Arange(orderReg4, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF * THREE);
278+ maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);
279+ auto startAddr = indicesAddr + i * VFLEN_UINT8;
280+ DataCopy(sortedIdxReg, startAddr);
281+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
282+ AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
283+ AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
284+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour3, cmpMask);
285+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour4, cmpMask);
286+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour1, maskFour3);
287+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour2, maskFour3);
288+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour3, maskFour4);
289+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour4, maskFour4);
290+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskFour1);
291+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskFour2);
292+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg3, orderReg3, maskFour3);
293+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg4, orderReg4, maskFour4);
294+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg, uOut);
295+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg2, uOut);
296+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg3, uOut);
297+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg4, uOut);
298+ }
299+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
300+}
301+ 
113template<typename IDX_T>302template<typename IDX_T>
114__aicore__ uint32_t ComputeUniqueIdNum(LocalTensor<IDX_T> indicesLocal, LocalTensor<int32_t> uniqueIdCountLocal, int64_t dataLen)303__aicore__ uint32_t ComputeUniqueIdNum(LocalTensor<IDX_T> indicesLocal, LocalTensor<int32_t> uniqueIdCountLocal, int64_t dataLen)
115{304{
@@ -118,52 +307,19 @@ __aicore__ uint32_t ComputeUniqueIdNum(LocalTensor<IDX_T> indicesLocal, LocalTen
118 307 
119 constexpr int64_t vfLen = platform::GetVRegSize() / sizeof(IDX_T);308 constexpr int64_t vfLen = platform::GetVRegSize() / sizeof(IDX_T);
120 uint16_t loopCnt = ops::CeilDiv(dataLen + 1, vfLen);309 uint16_t loopCnt = ops::CeilDiv(dataLen + 1, vfLen);
121- uint32_t counter = dataLen + 1;
122 __VEC_SCOPE__310 __VEC_SCOPE__
123 {311 {
124- AscendC::MicroAPI::RegTensor<int32_t> orderReg;
125- AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxReg;
126- AscendC::MicroAPI::RegTensor<IDX_T> sortedIdxShiftOneReg;
127- AscendC::MicroAPI::RegTensor<int32_t> selReg;
128- AscendC::MicroAPI::RegTensor<int32_t> orderReg2;
129- AscendC::MicroAPI::RegTensor<int32_t> selReg2;
130- AscendC::MicroAPI::MaskReg cmpMask;
131- AscendC::MicroAPI::MaskReg maskReg;
132- AscendC::MicroAPI::UnalignReg u0;
133- AscendC::MicroAPI::UnalignReg uOut;
134 AscendC::MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();312 AscendC::MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
135 313 
136- for (uint16_t i = 0; i < loopCnt; ++i) {314+ if constexpr (std::is_same<int64_t, IDX_T>::value) {
137- AscendC::MicroAPI::Arange(orderReg, i * vfLen);315+ ComputeUniqueIdNumInt64<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
138- maskReg = AscendC::MicroAPI::UpdateMask<IDX_T>(counter);316+ } else if constexpr (std::is_same<int32_t, IDX_T>::value) {
139- auto startAddr = indicesAddr + i * vfLen;317+ ComputeUniqueIdNumInt32<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
140- DataCopy(sortedIdxReg, startAddr);318+ } else if constexpr (std::is_same<int16_t, IDX_T>::value) {
141- AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);319+ ComputeUniqueIdNumInt16<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
142- AscendC::MicroAPI::DataCopyUnAlign<IDX_T>(sortedIdxShiftOneReg, u0, startAddr - 1);320+ } else { // uint8
143- AscendC::MicroAPI::Compare<IDX_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);321+ ComputeUniqueIdNumUint8<IDX_T>(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
144- if constexpr (std::is_same<int64_t, IDX_T>::value) {
145- AscendC::MicroAPI::MaskReg maskHalf;
146- AscendC::MicroAPI::MaskPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskHalf, cmpMask);
147- AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskHalf);
148- } else if constexpr (std::is_same<int32_t, IDX_T>::value) {
149- AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, cmpMask);
150- } else { // int16
151- AscendC::MicroAPI::Arange(orderReg2, i * vfLen + vfLen / 2);
152- AscendC::MicroAPI::MaskReg maskDouble1;
153- AscendC::MicroAPI::MaskReg maskDouble2;
154- AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskDouble1, cmpMask);
155- AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskDouble2, cmpMask);
156- AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskDouble1);
157- AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskDouble2);
158- }
159- AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
160- uniqueIdCountsAddr, selReg, uOut);
161- if constexpr (std::is_same<int16_t, IDX_T>::value) {
162- AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
163- uniqueIdCountsAddr, selReg2, uOut);
164- }
165 }322 }
166- AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
167 }323 }
168 uint32_t uniqueIdNum = ((AscendC::MicroAPI::GetSpr<AscendC::SpecialPurposeReg::AR>()) / sizeof(int32_t)) - 1;324 uint32_t uniqueIdNum = ((AscendC::MicroAPI::GetSpr<AscendC::SpecialPurposeReg::AR>()) / sizeof(int32_t)) - 1;
169 return uniqueIdNum;325 return uniqueIdNum;
Mindex/scatter_add/op_kernel/arch35/scatter_add_deterministic.h+49-47
@@ -19,9 +19,11 @@
19#include "kernel_operator.h"19#include "kernel_operator.h"
20#include "../inc/platform.h"20#include "../inc/platform.h"
21#include "../inc/kernel_utils.h"21#include "../inc/kernel_utils.h"
22+#include "scatter_add_common.h"
22 23 
23namespace ScatterAdd {24namespace ScatterAdd {
24using namespace AscendC;25using namespace AscendC;
26+using namespace ScatterAddCommon;
25 27 
26constexpr uint64_t DOUBLE_BUF = 2;28constexpr uint64_t DOUBLE_BUF = 2;
27constexpr uint64_t BLOCK_SIZE = 32;29constexpr uint64_t BLOCK_SIZE = 32;
@@ -48,7 +50,7 @@ AscendC::RoundMode::UNKNOWN,
48 50 
49static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};51static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};
50 52 
51-template<typename T, typename U>53+template<typename T, typename U, uint32_t scatterOp>
52class ScatterAddDeterministicImpl {54class ScatterAddDeterministicImpl {
53public:55public:
54 __aicore__ inline ScatterAddDeterministicImpl(const ScatterAddTilingData& tilingData, TPipe& pipe) :56 __aicore__ inline ScatterAddDeterministicImpl(const ScatterAddTilingData& tilingData, TPipe& pipe) :
@@ -115,8 +117,8 @@ private:
115 uint64_t postVarAlignSizeFp32_{0};117 uint64_t postVarAlignSizeFp32_{0};
116};118};
117 119 
118-template<typename T, typename U>120+template<typename T, typename U, uint32_t scatterOp>
119-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,121+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
120 GM_ADDR varRef, GM_ADDR workspace)122 GM_ADDR varRef, GM_ADDR workspace)
121{123{
122 isDeterministic = tilingData_.isDeterministic == 1 ? true : false;124 isDeterministic = tilingData_.isDeterministic == 1 ? true : false;
@@ -162,8 +164,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::Init(GM_ADDR var, GM_
162 }164 }
163}165}
164 166 
165-template<typename T, typename U>167+template<typename T, typename U, uint32_t scatterOp>
166-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::InitWspZero()168+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::InitWspZero()
167{169{
168 InitGlobalMemory(workspaceCoreLogicCoreSumId_, tilingData_.perCoreHandleIndices, (U)(-1));170 InitGlobalMemory(workspaceCoreLogicCoreSumId_, tilingData_.perCoreHandleIndices, (U)(-1));
169 InitGlobalMemory(workspaceLogicCoreSumValue_, tilingData_.perCoreHandleIndices * postVarAlignSizeFp32_, (float)(0));171 InitGlobalMemory(workspaceLogicCoreSumValue_, tilingData_.perCoreHandleIndices * postVarAlignSizeFp32_, (float)(0));
@@ -172,8 +174,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::InitWspZero()
172 InitGlobalMemory(workspaceInt32Res_, tilingData_.varShape[1] * tilingData_.varShape[0], (int)(0));174 InitGlobalMemory(workspaceInt32Res_, tilingData_.varShape[1] * tilingData_.varShape[0], (int)(0));
173}175}
174 176 
175-template<typename T, typename U>177+template<typename T, typename U, uint32_t scatterOp>
176-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyInUpdates(int64_t offset, int64_t dataLen)178+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopyInUpdates(int64_t offset, int64_t dataLen)
177{179{
178 DataCopyExtParams inParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };180 DataCopyExtParams inParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };
179 DataCopyPadExtParams<T> padParams = { false, 0, 0, 0 };181 DataCopyPadExtParams<T> padParams = { false, 0, 0, 0 };
@@ -182,8 +184,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyInUpdates(int64_t
182 dataQueue_.EnQue(updatesLocal);184 dataQueue_.EnQue(updatesLocal);
183}185}
184 186 
185-template<typename T, typename U>187+template<typename T, typename U, uint32_t scatterOp>
186-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyOutUpdates(int64_t offset, int64_t dataLen)188+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopyOutUpdates(int64_t offset, int64_t dataLen)
187{189{
188 DataCopyExtParams outParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };190 DataCopyExtParams outParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };
189 LocalTensor<T> updatesLocal = dataQueue_.DeQue<T>();191 LocalTensor<T> updatesLocal = dataQueue_.DeQue<T>();
@@ -193,8 +195,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyOutUpdates(int64_
193 dataQueue_.FreeTensor(updatesLocal);195 dataQueue_.FreeTensor(updatesLocal);
194}196}
195 197 
196-template<typename T, typename U>198+template<typename T, typename U, uint32_t scatterOp>
197-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ProcessSingleLoopIndices(int64_t indicesOffset,199+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ProcessSingleLoopIndices(int64_t indicesOffset,
198 int64_t indicesLen)200 int64_t indicesLen)
199{201{
200 DataCopyExtParams inParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(indicesLen * sizeof(U)), 0, 0, 0 };202 DataCopyExtParams inParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(indicesLen * sizeof(U)), 0, 0, 0 };
@@ -228,8 +230,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ProcessSingleLoopIndi
228 }230 }
229}231}
230 232 
231-template<typename T, typename U>233+template<typename T, typename U, uint32_t scatterOp>
232-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ProcessAtomicAdd()234+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ProcessAtomicAdd()
233{235{
234 if (GetBlockIdx() >= tilingData_.logicCoreNum) {236 if (GetBlockIdx() >= tilingData_.logicCoreNum) {
235 return;237 return;
@@ -246,8 +248,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ProcessAtomicAdd()
246 ProcessSingleLoopIndices(indicesOffset, indicesTailUbFactor);248 ProcessSingleLoopIndices(indicesOffset, indicesTailUbFactor);
247}249}
248 250 
249-template<typename T, typename U>251+template<typename T, typename U, uint32_t scatterOp>
250-__aicore__ inline void ScatterAddDeterministicImpl<T, U>:: CopyInIndicesAndUpdates(uint64_t loopIdNum, uint64_t rowNums)252+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>:: CopyInIndicesAndUpdates(uint64_t loopIdNum, uint64_t rowNums)
251{253{
252 uint64_t updatesOffset = loopIdNum * tilingData_.indicesUbFactor * tilingData_.varShape[1];254 uint64_t updatesOffset = loopIdNum * tilingData_.indicesUbFactor * tilingData_.varShape[1];
253 LocalTensor<T> updatesLocal = updatesQueue_.template AllocTensor<T>();255 LocalTensor<T> updatesLocal = updatesQueue_.template AllocTensor<T>();
@@ -265,8 +267,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>:: CopyInIndicesAndUpda
265 indicesQue_.EnQue(indicesLocal);267 indicesQue_.EnQue(indicesLocal);
266}268}
267 269 
268-template<typename T, typename U>270+template<typename T, typename U, uint32_t scatterOp>
269-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeUniqueIdNum(LocalTensor<U>& sortedIndicesLocal, 271+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ComputeUniqueIdNum(LocalTensor<U>& sortedIndicesLocal,
270 LocalTensor<U>& uniqueIdCountLocalU, uint32_t dataLen)272 LocalTensor<U>& uniqueIdCountLocalU, uint32_t dataLen)
271{273{
272 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.AllocTensor<U>();274 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.AllocTensor<U>();
@@ -330,8 +332,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeUniqueIdNum(Lo
330 updateSumIdxQueue_.EnQue(updateSumIdxLocal);332 updateSumIdxQueue_.EnQue(updateSumIdxLocal);
331}333}
332 334 
333-template<typename T, typename U>335+template<typename T, typename U, uint32_t scatterOp>
334-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeUinqueIdTimes(LocalTensor<U>& uniqueIdCountLocalU, uint32_t uniqueIdNum)336+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ComputeUinqueIdTimes(LocalTensor<U>& uniqueIdCountLocalU, uint32_t uniqueIdNum)
335{337{
336 LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountLocalU.template ReinterpretCast<int32_t>();338 LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountLocalU.template ReinterpretCast<int32_t>();
337 __local_mem__ int32_t* uniqueIdCountsAddr = (__local_mem__ int32_t*)uniqueIdCountLocal.GetPhyAddr();339 __local_mem__ int32_t* uniqueIdCountsAddr = (__local_mem__ int32_t*)uniqueIdCountLocal.GetPhyAddr();
@@ -361,8 +363,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeUinqueIdTimes(
361 }363 }
362}364}
363 365 
364-template<typename T, typename U>366+template<typename T, typename U, uint32_t scatterOp>
365-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeSum(LocalTensor<U>& uniqueIdCountLocalU, LocalTensor<uint32_t>& updatesOriginIdexLocal,367+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ComputeSum(LocalTensor<U>& uniqueIdCountLocalU, LocalTensor<uint32_t>& updatesOriginIdexLocal,
366 uint32_t uniqueIdNum, uint32_t dataLen)368 uint32_t uniqueIdNum, uint32_t dataLen)
367{369{
368 LocalTensor<uint32_t> uniqueIdCountLocal = uniqueIdCountLocalU.template ReinterpretCast<uint32_t>();370 LocalTensor<uint32_t> uniqueIdCountLocal = uniqueIdCountLocalU.template ReinterpretCast<uint32_t>();
@@ -423,8 +425,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeSum(LocalTenso
423 updatesQueue_.FreeTensor(updatesLocal);425 updatesQueue_.FreeTensor(updatesLocal);
424}426}
425 427 
426-template<typename T, typename U>428+template<typename T, typename U, uint32_t scatterOp>
427-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ProcessSortAndSum(uint64_t loopIdx, uint32_t dataLen)429+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ProcessSortAndSum(uint64_t loopIdx, uint32_t dataLen)
428{430{
429 LocalTensor<U> indicesLocal = indicesQue_.DeQue<U>();431 LocalTensor<U> indicesLocal = indicesQue_.DeQue<U>();
430 LocalTensor<U> sortedIndicesLocal = sortedIndicesQue_.Get<U>(); 432 LocalTensor<U> sortedIndicesLocal = sortedIndicesQue_.Get<U>();
@@ -446,8 +448,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ProcessSortAndSum(uin
446 indicesQue_.FreeTensor(indicesLocal);448 indicesQue_.FreeTensor(indicesLocal);
447}449}
448 450 
449-template<typename T, typename U>451+template<typename T, typename U, uint32_t scatterOp>
450-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopySumOutToWs(uint64_t loopIdx, uint64_t dataLen)452+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopySumOutToWs(uint64_t loopIdx, uint64_t dataLen)
451{453{
452 // updatesSum和updatesSumId整块搬出去454 // updatesSum和updatesSumId整块搬出去
453 uint32_t updatesWspOffset = loopIdx * tilingData_.indicesUbFactor * postVarAlignSizeFp32_;455 uint32_t updatesWspOffset = loopIdx * tilingData_.indicesUbFactor * postVarAlignSizeFp32_;
@@ -481,8 +483,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopySumOutToWs(uint64
481 updateSumIdxQueue_.FreeTensor(updateSumIdxLocal);483 updateSumIdxQueue_.FreeTensor(updateSumIdxLocal);
482}484}
483 485 
484-template<typename T, typename U>486+template<typename T, typename U, uint32_t scatterOp>
485-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopySumAndIdxIn(uint64_t loopIdx, uint64_t dataCount)487+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopySumAndIdxIn(uint64_t loopIdx, uint64_t dataCount)
486{488{
487 LocalTensor<float> updateSumLocal = updateSumInQue_.AllocTensor<float>();489 LocalTensor<float> updateSumLocal = updateSumInQue_.AllocTensor<float>();
488 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.AllocTensor<U>();490 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.AllocTensor<U>();
@@ -502,8 +504,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopySumAndIdxIn(uint6
502 updateSumIdxQueue_.EnQue(updateSumIdxLocal);504 updateSumIdxQueue_.EnQue(updateSumIdxLocal);
503}505}
504 506 
505-template<typename T, typename U>507+template<typename T, typename U, uint32_t scatterOp>
506-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeRValueAndQuantize(uint64_t loopIdx, uint64_t dataCount) 508+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ComputeRValueAndQuantize(uint64_t loopIdx, uint64_t dataCount)
507{509{
508 LocalTensor<float> updateSumLocal = updateSumInQue_.DeQue<float>();510 LocalTensor<float> updateSumLocal = updateSumInQue_.DeQue<float>();
509 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.DeQue<U>();511 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.DeQue<U>();
@@ -537,8 +539,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeRValueAndQuant
537 updateSumIdxQueue_.FreeTensor(updateSumIdxLocal);539 updateSumIdxQueue_.FreeTensor(updateSumIdxLocal);
538}540}
539 541 
540-template<typename T, typename U>542+template<typename T, typename U, uint32_t scatterOp>
541-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::QuantizeForSum(uint64_t curRowIdx, uint32_t RCountsValue,543+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::QuantizeForSum(uint64_t curRowIdx, uint32_t RCountsValue,
542 __local_mem__ float* updateSumAddr)544 __local_mem__ float* updateSumAddr)
543{545{
544 LocalTensor<float> updateRValueLocal = RValueQue_.DeQue<float>();546 LocalTensor<float> updateRValueLocal = RValueQue_.DeQue<float>();
@@ -581,8 +583,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::QuantizeForSum(uint64
581 updateSumQue_.EnQue(updateSumIntLocal);583 updateSumQue_.EnQue(updateSumIntLocal);
582}584}
583 585 
584-template<typename T, typename U>586+template<typename T, typename U, uint32_t scatterOp>
585-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyQuantizedSumOutToIntWs(uint64_t RValueOffset)587+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopyQuantizedSumOutToIntWs(uint64_t RValueOffset)
586{588{
587 LocalTensor<int> updateSumIntLocal = updateSumQue_.DeQue<int>();589 LocalTensor<int> updateSumIntLocal = updateSumQue_.DeQue<int>();
588 DataCopyExtParams outParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(tilingData_.varShape[1] * sizeof(int)), 0, 0, 0 };590 DataCopyExtParams outParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(tilingData_.varShape[1] * sizeof(int)), 0, 0, 0 };
@@ -592,8 +594,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyQuantizedSumOutTo
592 updateSumQue_.FreeTensor(updateSumIntLocal);594 updateSumQue_.FreeTensor(updateSumIntLocal);
593}595}
594 596 
595-template<typename T, typename U>597+template<typename T, typename U, uint32_t scatterOp>
596-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyIntSumIn(uint64_t loopIdx, uint64_t dataCount)598+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopyIntSumIn(uint64_t loopIdx, uint64_t dataCount)
597{599{
598 LocalTensor<int> updateIntSumLocal = updateSumInQue_.AllocTensor<int>();600 LocalTensor<int> updateIntSumLocal = updateSumInQue_.AllocTensor<int>();
599 601 
@@ -605,8 +607,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyIntSumIn(uint64_t
605 updateSumInQue_.EnQue(updateIntSumLocal);607 updateSumInQue_.EnQue(updateIntSumLocal);
606}608}
607 609 
608-template<typename T, typename U>610+template<typename T, typename U, uint32_t scatterOp>
609-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeRValueAndDeQuantize(uint64_t loopIdx, uint64_t dataLen)611+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ComputeRValueAndDeQuantize(uint64_t loopIdx, uint64_t dataLen)
610{612{
611 LocalTensor<int> updateSumIntLocal = updateSumInQue_.DeQue<int>();613 LocalTensor<int> updateSumIntLocal = updateSumInQue_.DeQue<int>();
612 __local_mem__ int32_t * updateSumIntAddr = (__local_mem__ int32_t*)updateSumIntLocal.GetPhyAddr();614 __local_mem__ int32_t * updateSumIntAddr = (__local_mem__ int32_t*)updateSumIntLocal.GetPhyAddr();
@@ -631,8 +633,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeRValueAndDeQua
631 updateSumInQue_.FreeTensor(updateSumIntLocal);633 updateSumInQue_.FreeTensor(updateSumIntLocal);
632}634}
633 635 
634-template<typename T, typename U>636+template<typename T, typename U, uint32_t scatterOp>
635-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::DeQuantizeForSum(uint64_t curRowIdx, uint32_t RCountsValue, __local_mem__ int32_t* updateSumIntAddr)637+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::DeQuantizeForSum(uint64_t curRowIdx, uint32_t RCountsValue, __local_mem__ int32_t* updateSumIntAddr)
636{638{
637 LocalTensor<float> updateRValueLocal = RValueQue_.DeQue<float>();639 LocalTensor<float> updateRValueLocal = RValueQue_.DeQue<float>();
638 LocalTensor<T> updateSumLocal = updateSumQue_.AllocTensor<T>();640 LocalTensor<T> updateSumLocal = updateSumQue_.AllocTensor<T>();
@@ -685,8 +687,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::DeQuantizeForSum(uint
685 RValueQue_.FreeTensor(updateRValueLocal);687 RValueQue_.FreeTensor(updateRValueLocal);
686}688}
687 689 
688-template<typename T, typename U>690+template<typename T, typename U, uint32_t scatterOp>
689-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyOutDeQuantizedSum(uint64_t varOffset)691+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopyOutDeQuantizedSum(uint64_t varOffset)
690{692{
691 LocalTensor<T> invQuantDataLoacl = updateSumQue_.DeQue<T>();693 LocalTensor<T> invQuantDataLoacl = updateSumQue_.DeQue<T>();
692 DataCopyExtParams dataTCopyParam{static_cast<uint16_t>(1), static_cast<uint32_t>(tilingData_.varShape[1] * sizeof(T)), 0, 0, 0};694 DataCopyExtParams dataTCopyParam{static_cast<uint16_t>(1), static_cast<uint32_t>(tilingData_.varShape[1] * sizeof(T)), 0, 0, 0};
@@ -696,8 +698,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyOutDeQuantizedSum
696 updateSumQue_.FreeTensor(invQuantDataLoacl);698 updateSumQue_.FreeTensor(invQuantDataLoacl);
697}699}
698 700 
699-template<typename T, typename U>701+template<typename T, typename U, uint32_t scatterOp>
700-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyIdxIn(uint64_t loopIdx, uint64_t dataCount)702+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::CopyIdxIn(uint64_t loopIdx, uint64_t dataCount)
701{703{
702 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.AllocTensor<U>();704 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.AllocTensor<U>();
703 705
@@ -708,8 +710,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::CopyIdxIn(uint64_t lo
708 updateSumIdxQueue_.EnQue(updateSumIdxLocal);710 updateSumIdxQueue_.EnQue(updateSumIdxLocal);
709}711}
710 712 
711-template<typename T, typename U>713+template<typename T, typename U, uint32_t scatterOp>
712-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeRValueAndDeQuantizePro(uint64_t loopIdx, uint64_t dataLen)714+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::ComputeRValueAndDeQuantizePro(uint64_t loopIdx, uint64_t dataLen)
713{715{
714 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.DeQue<U>();716 LocalTensor<U> updateSumIdxLocal = updateSumIdxQueue_.DeQue<U>();
715 DataCopyExtParams dataCopyParam1{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(float) * tilingData_.varShape[1]), 0, 0, 0};717 DataCopyExtParams dataCopyParam1{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(float) * tilingData_.varShape[1]), 0, 0, 0};
@@ -753,8 +755,8 @@ __aicore__ inline void ScatterAddDeterministicImpl<T, U>::ComputeRValueAndDeQua
753 updateSumIdxQueue_.FreeTensor(updateSumIdxLocal);755 updateSumIdxQueue_.FreeTensor(updateSumIdxLocal);
754}756}
755 757 
756-template<typename T, typename U>758+template<typename T, typename U, uint32_t scatterOp>
757-__aicore__ inline void ScatterAddDeterministicImpl<T, U>::Process()759+__aicore__ inline void ScatterAddDeterministicImpl<T, U, scatterOp>::Process()
758{760{
759 if ((GetBlockIdx() >= tilingData_.logicCoreNum) && !isDeterministic) {761 if ((GetBlockIdx() >= tilingData_.logicCoreNum) && !isDeterministic) {
760 return;762 return;
Mindex/scatter_add/op_kernel/arch35/scatter_add_simd.h+47-31
@@ -34,7 +34,7 @@ struct UpdatesDstPosSelector { constexpr static TPosition pos = TPosition::VECIN
34template<> struct UpdatesDstPosSelector<true, false> { constexpr static TPosition pos = TPosition::VECOUT; }; // 支持atomic_add且updates不是标量,直接搬入搬出34template<> struct UpdatesDstPosSelector<true, false> { constexpr static TPosition pos = TPosition::VECOUT; }; // 支持atomic_add且updates不是标量,直接搬入搬出
35template<> struct UpdatesDstPosSelector<true, true> { constexpr static TPosition pos = TPosition::GM; }; // 支持atomic_add且updates是标量,duplicate后搬出35template<> struct UpdatesDstPosSelector<true, true> { constexpr static TPosition pos = TPosition::GM; }; // 支持atomic_add且updates是标量,duplicate后搬出
36 36 
37-template<typename T, typename U, bool updatesIsScalar>37+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
38class ScatterAddSIMDImpl {38class ScatterAddSIMDImpl {
39public:39public:
40 __aicore__ inline ScatterAddSIMDImpl(const ScatterAddTilingData& tilingData, TPipe& pipe) :40 __aicore__ inline ScatterAddSIMDImpl(const ScatterAddTilingData& tilingData, TPipe& pipe) :
@@ -74,8 +74,8 @@ private:
74 const ScatterAddTilingData& tilingData_;74 const ScatterAddTilingData& tilingData_;
75};75};
76 76 
77-template<typename T, typename U, bool updatesIsScalar>77+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
78-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,78+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
79 GM_ADDR varRef, GM_ADDR workspace)79 GM_ADDR varRef, GM_ADDR workspace)
80{80{
81 if (GetBlockIdx() >= GetBlockNum()) {81 if (GetBlockIdx() >= GetBlockNum()) {
@@ -106,8 +106,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::Init(GM_ADDR v
106 }106 }
107}107}
108 108 
109-template<typename T, typename U, bool updatesIsScalar>109+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
110-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyInUpdates(int64_t offset, int64_t dataLen)110+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyInUpdates(int64_t offset, int64_t dataLen)
111{111{
112 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };112 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };
113 DataCopyPadExtParams<T> padParams = { false, 0, 0, 0 };113 DataCopyPadExtParams<T> padParams = { false, 0, 0, 0 };
@@ -116,8 +116,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyInUpdates(
116 updatesQueue_.EnQue(updatesLocal);116 updatesQueue_.EnQue(updatesLocal);
117}117}
118 118 
119-template<typename T, typename U, bool updatesIsScalar>119+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
120-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutUpdates(int64_t offset, int64_t dataLen)120+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyOutUpdates(int64_t offset, int64_t dataLen)
121{121{
122 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };122 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };
123 LocalTensor<T> updatesLocal = updatesQueue_.template DeQue<T>();123 LocalTensor<T> updatesLocal = updatesQueue_.template DeQue<T>();
@@ -131,8 +131,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutUpdates
131 }131 }
132}132}
133 133 
134-template<typename T, typename U, bool updatesIsScalar>134+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
135-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyInVar(int64_t offset, int64_t dataLen)135+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyInVar(int64_t offset, int64_t dataLen)
136{136{
137 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };137 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };
138 DataCopyPadExtParams<T> padParams = { false, 0, 0, 0 };138 DataCopyPadExtParams<T> padParams = { false, 0, 0, 0 };
@@ -141,8 +141,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyInVar(int6
141 varInQueue_.EnQue(varLocal);141 varInQueue_.EnQue(varLocal);
142}142}
143 143 
144-template<typename T, typename U, bool updatesIsScalar>144+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
145-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutVarToWS(int64_t offset, int64_t dataLen)145+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyOutVarToWS(int64_t offset, int64_t dataLen)
146{146{
147 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(int32_t)), 0, 0, 0 };147 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(int32_t)), 0, 0, 0 };
148 LocalTensor<int32_t> varCastLocal = varCastOutQueue_.DeQue<int32_t>();148 LocalTensor<int32_t> varCastLocal = varCastOutQueue_.DeQue<int32_t>();
@@ -150,11 +150,17 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutVarToWS
150 varCastOutQueue_.FreeTensor(varCastLocal);150 varCastOutQueue_.FreeTensor(varCastLocal);
151}151}
152 152 
153-template<typename T, typename U, bool updatesIsScalar>153+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
154-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutUpdatesToWS(int64_t offset, int64_t dataLen)154+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyOutUpdatesToWS(int64_t offset, int64_t dataLen)
155{155{
156 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(int32_t)), 0, 0, 0 };156 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(int32_t)), 0, 0, 0 };
157 LocalTensor<int32_t> updatesCastLocal = updatesCastQueue_.DeQue<int32_t>();157 LocalTensor<int32_t> updatesCastLocal = updatesCastQueue_.DeQue<int32_t>();
158+ if constexpr (scatterOp == SUB && !updatesIsScalar) {
159+ NegateUpdate<int32_t>(updatesCastLocal, static_cast<uint32_t>(dataLen));
160+ auto MTE3WaitVEventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
161+ SetFlag<HardEvent::V_MTE3>(MTE3WaitVEventID);
162+ WaitFlag<HardEvent::V_MTE3>(MTE3WaitVEventID);
163+ }
158 SetAtomicAdd<int32_t>();164 SetAtomicAdd<int32_t>();
159 DataCopyPad(varCastAtomicAddGm_[offset], updatesCastLocal, outParams);165 DataCopyPad(varCastAtomicAddGm_[offset], updatesCastLocal, outParams);
160 SetAtomicNone();166 SetAtomicNone();
@@ -165,8 +171,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutUpdates
165 }171 }
166}172}
167 173 
168-template<typename T, typename U, bool updatesIsScalar>174+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
169-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyInVarFromWS(int64_t offset, int64_t dataLen)175+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyInVarFromWS(int64_t offset, int64_t dataLen)
170{176{
171 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(dataLen * sizeof(int32_t)), 0, 0, 0 };177 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(dataLen * sizeof(int32_t)), 0, 0, 0 };
172 DataCopyPadExtParams<int32_t> padParams = { false, 0, 0, 0 };178 DataCopyPadExtParams<int32_t> padParams = { false, 0, 0, 0 };
@@ -175,8 +181,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyInVarFromW
175 varCastInQueue_.EnQue(varCastLocal);181 varCastInQueue_.EnQue(varCastLocal);
176}182}
177 183 
178-template<typename T, typename U, bool updatesIsScalar>184+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
179-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutVar(int64_t offset, int64_t dataLen)185+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::CopyOutVar(int64_t offset, int64_t dataLen)
180{186{
181 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };187 DataCopyExtParams outParams = { 1, static_cast<uint32_t>(dataLen * sizeof(T)), 0, 0, 0 };
182 LocalTensor<T> varLocal = varOutQueue_.DeQue<T>();188 LocalTensor<T> varLocal = varOutQueue_.DeQue<T>();
@@ -184,8 +190,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::CopyOutVar(int
184 varOutQueue_.FreeTensor(varLocal);190 varOutQueue_.FreeTensor(varLocal);
185}191}
186 192 
187-template<typename T, typename U, bool updatesIsScalar>193+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
188-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ExcuteAtomicAdd(int64_t varRefOffset,194+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::ExcuteAtomicAdd(int64_t varRefOffset,
189 int64_t dataLen)195 int64_t dataLen)
190{196{
191 if constexpr (platform::IsSupportAtomicAddTypeSIMD<T>()) {197 if constexpr (platform::IsSupportAtomicAddTypeSIMD<T>()) {
@@ -202,8 +208,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ExcuteAtomicAd
202 }208 }
203}209}
204 210 
205-template<typename T, typename U, bool updatesIsScalar>211+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
206-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessSingleLoopIndices(int64_t indicesOffset,212+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::ProcessSingleLoopIndices(int64_t indicesOffset,
207 int64_t indicesLen)213 int64_t indicesLen)
208{214{
209 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(indicesLen * sizeof(U)), 0, 0, 0 };215 DataCopyExtParams inParams = { 1, static_cast<uint32_t>(indicesLen * sizeof(U)), 0, 0, 0 };
@@ -235,8 +241,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessSingleL
235 }241 }
236}242}
237 243 
238-template<typename T, typename U, bool updatesIsScalar>244+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
239-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessAtomicAdd()245+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::ProcessAtomicAdd()
240{246{
241 if (GetBlockIdx() >= tilingData_.atomicAddCoreNum) {247 if (GetBlockIdx() >= tilingData_.atomicAddCoreNum) {
242 return;248 return;
@@ -248,11 +254,21 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessAtomicA
248 WaitFlag<HardEvent::S_V>(vWaitSEventID);254 WaitFlag<HardEvent::S_V>(vWaitSEventID);
249 if constexpr (platform::IsSupportAtomicAddTypeSIMD<T>()) {255 if constexpr (platform::IsSupportAtomicAddTypeSIMD<T>()) {
250 LocalTensor<T> updatesLocal = updatesQueue_.template AllocTensor<T>();256 LocalTensor<T> updatesLocal = updatesQueue_.template AllocTensor<T>();
251- Duplicate(updatesLocal, updatesValue, tilingData_.updatesUbFactor);257+ if constexpr (scatterOp == ADD) {
258+ Duplicate(updatesLocal, updatesValue, tilingData_.updatesUbFactor);
259+ } else if constexpr (scatterOp == SUB) {
260+ updatesValue = -updatesValue;
261+ Duplicate(updatesLocal, updatesValue, tilingData_.updatesUbFactor);
262+ }
252 updatesQueue_.EnQue(updatesLocal);263 updatesQueue_.EnQue(updatesLocal);
253 } else {264 } else {
254 LocalTensor<int32_t> updatesCastLocal = updatesCastQueue_.AllocTensor<int32_t>();265 LocalTensor<int32_t> updatesCastLocal = updatesCastQueue_.AllocTensor<int32_t>();
255- Duplicate(updatesCastLocal, static_cast<int32_t>(updatesValue), tilingData_.updatesUbFactor);266+ if constexpr (scatterOp == ADD) {
267+ Duplicate(updatesCastLocal, static_cast<int32_t>(updatesValue), tilingData_.updatesUbFactor);
268+ } else if constexpr (scatterOp == SUB) {
269+ int32_t updatesValueInt32 = -static_cast<int32_t>(updatesValue);
270+ Duplicate(updatesCastLocal, updatesValueInt32, tilingData_.updatesUbFactor);
271+ }
256 updatesCastQueue_.EnQue(updatesCastLocal);272 updatesCastQueue_.EnQue(updatesCastLocal);
257 }273 }
258 }274 }
@@ -280,8 +296,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessAtomicA
280 }296 }
281}297}
282 298 
283-template<typename T, typename U, bool updatesIsScalar>299+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
284-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessVarToWS()300+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::ProcessVarToWS()
285{301{
286 if (GetBlockIdx() >= tilingData_.copyCoreNum) {302 if (GetBlockIdx() >= tilingData_.copyCoreNum) {
287 return;303 return;
@@ -313,8 +329,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessVarToWS
313 CopyOutVarToWS(varOffset, tailUbFactor);329 CopyOutVarToWS(varOffset, tailUbFactor);
314}330}
315 331 
316-template<typename T, typename U, bool updatesIsScalar>332+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
317-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessVarFromWS()333+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::ProcessVarFromWS()
318{334{
319 if (GetBlockIdx() >= tilingData_.copyCoreNum) {335 if (GetBlockIdx() >= tilingData_.copyCoreNum) {
320 return;336 return;
@@ -346,8 +362,8 @@ __aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::ProcessVarFrom
346 CopyOutVar(varOffset, tailUbFactor);362 CopyOutVar(varOffset, tailUbFactor);
347}363}
348 364 
349-template<typename T, typename U, bool updatesIsScalar>365+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
350-__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar>::Process()366+__aicore__ inline void ScatterAddSIMDImpl<T, U, updatesIsScalar, scatterOp>::Process()
351{367{
352 if (GetBlockIdx() >= GetBlockNum()) {368 if (GetBlockIdx() >= GetBlockNum()) {
353 return;369 return;
Mindex/scatter_add/op_kernel/arch35/scatter_add_simd_sort_support_atomicadd.h+37-28
@@ -24,7 +24,7 @@ using namespace ScatterAddCommon;
24 24 
25constexpr uint32_t SORT_PADDING = 64;25constexpr uint32_t SORT_PADDING = 64;
26 26 
27-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>27+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
28class ScatterAddSIMDSortSupportAtomicAdd {28class ScatterAddSIMDSortSupportAtomicAdd {
29public:29public:
30 __aicore__ inline ScatterAddSIMDSortSupportAtomicAdd(const ScatterAddTilingData& tilingData, TPipe& pipe) :30 __aicore__ inline ScatterAddSIMDSortSupportAtomicAdd(const ScatterAddTilingData& tilingData, TPipe& pipe) :
@@ -46,7 +46,6 @@ private:
46 TQue<QuePosition::VECIN, 1> indicesQueue_;46 TQue<QuePosition::VECIN, 1> indicesQueue_;
47 TQue<QuePosition::VECOUT, 1> outQueueRes_;47 TQue<QuePosition::VECOUT, 1> outQueueRes_;
48 TBuf<QuePosition::VECCALC> castIndicesQue_;48 TBuf<QuePosition::VECCALC> castIndicesQue_;
49- TBuf<QuePosition::VECCALC> castTmpIndicesQue_;
50 TBuf<QuePosition::VECCALC> sortIndicesQue_;49 TBuf<QuePosition::VECCALC> sortIndicesQue_;
51 TBuf<QuePosition::VECCALC> updatesOriginIdexQue_;50 TBuf<QuePosition::VECCALC> updatesOriginIdexQue_;
52 TBuf<QuePosition::VECCALC> uniqueIdCountQue_;51 TBuf<QuePosition::VECCALC> uniqueIdCountQue_;
@@ -57,8 +56,8 @@ private:
57 static constexpr uint32_t shiftOffset_ = platform::GetUbBlockSize() / sizeof(CAST_T);56 static constexpr uint32_t shiftOffset_ = platform::GetUbBlockSize() / sizeof(CAST_T);
58};57};
59 58 
60-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>59+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
61-__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType>::Init(60+__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType, scatterOp>::Init(
62 GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR varRef, GM_ADDR workspace)61 GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR varRef, GM_ADDR workspace)
63{62{
64 varRefGm_.SetGlobalBuffer((__gm__ T*)(varRef));63 varRefGm_.SetGlobalBuffer((__gm__ T*)(varRef));
@@ -72,18 +71,14 @@ __aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesI
72 pipe_.InitBuffer(uniqueIdCountQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(int32_t), UB_AGLIN_VALUE) + SORT_PADDING);71 pipe_.InitBuffer(uniqueIdCountQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(int32_t), UB_AGLIN_VALUE) + SORT_PADDING);
73 if constexpr (castType == CAST_0) {72 if constexpr (castType == CAST_0) {
74 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(U), UB_AGLIN_VALUE) + SORT_PADDING);73 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(U), UB_AGLIN_VALUE) + SORT_PADDING);
75- } else if constexpr (castType == CAST_3) {
76- pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(CAST_T), UB_AGLIN_VALUE) + SORT_PADDING);
77- pipe_.InitBuffer(castIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(CAST_T), UB_AGLIN_VALUE));
78- pipe_.InitBuffer(castTmpIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(int32_t), UB_AGLIN_VALUE));
79 } else {74 } else {
80 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(CAST_T), UB_AGLIN_VALUE) + SORT_PADDING);75 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(CAST_T), UB_AGLIN_VALUE) + SORT_PADDING);
81 pipe_.InitBuffer(castIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(CAST_T), UB_AGLIN_VALUE));76 pipe_.InitBuffer(castIndicesQue_, ops::CeilAlign(tilingData_.ubFactorRow * sizeof(CAST_T), UB_AGLIN_VALUE));
82 }77 }
83}78}
84 79 
85-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>80+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
86-__aicore__ inline uint32_t ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType>::ProcessIndices(81+__aicore__ inline uint32_t ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType, scatterOp>::ProcessIndices(
87 uint64_t blockOffsetindices, uint64_t rowLoop, uint32_t rows)82 uint64_t blockOffsetindices, uint64_t rowLoop, uint32_t rows)
88{83{
89 LocalTensor<U> indicesLocal = indicesQueue_.AllocTensor<U>();84 LocalTensor<U> indicesLocal = indicesQueue_.AllocTensor<U>();
@@ -98,16 +93,11 @@ __aicore__ inline uint32_t ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, upda
98 uint32_t uniqueIdNum = 0;93 uint32_t uniqueIdNum = 0;
99 if constexpr (castType == CAST_0) {94 if constexpr (castType == CAST_0) {
100 uniqueIdNum = SortAndComputeUniqueIdx<U>(rows, indicesLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);95 uniqueIdNum = SortAndComputeUniqueIdx<U>(rows, indicesLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);
101- } else if constexpr (castType == CAST_3) {
102- LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();
103- LocalTensor<int32_t> indicesCastTmpLocal = castTmpIndicesQue_.Get<int32_t>();
104- Cast<int32_t, U>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, rows);
105- Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, rows);
106- uniqueIdNum = SortAndComputeUniqueIdx<CAST_T>(rows, indicesCastLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);
107 } else {96 } else {
108 LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();97 LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();
109- Cast<CAST_T, U>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, rows);98+ IndicesSortCast<U, CAST_T, castType>(indicesLocal, indicesCastLocal, uniqueIdCountLocal, rows);
110- uniqueIdNum = SortAndComputeUniqueIdx<CAST_T>(rows, indicesCastLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);99+ uniqueIdNum = SortAndComputeUniqueIdx<CAST_T>(
100+ rows, indicesCastLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);
111 }101 }
112 102
113 ComputeUniqueIdTimes(uniqueIdCountLocal, uniqueIdNum); // 计算每个indices的重复度存放于uniqueIdCountLocal103 ComputeUniqueIdTimes(uniqueIdCountLocal, uniqueIdNum); // 计算每个indices的重复度存放于uniqueIdCountLocal
@@ -116,8 +106,8 @@ __aicore__ inline uint32_t ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, upda
116 return uniqueIdNum;106 return uniqueIdNum;
117}107}
118 108 
119-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>109+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
120-__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType>::ComputeUpdatesSum(110+__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType, scatterOp>::ComputeUpdatesSum(
121 uint64_t cols, uint64_t colsAlign, uint32_t uniqueIdNum)111 uint64_t cols, uint64_t colsAlign, uint32_t uniqueIdNum)
122{112{
123 LocalTensor<uint32_t> updatesOriginIdxLocal = updatesOriginIdexQue_.Get<uint32_t>();113 LocalTensor<uint32_t> updatesOriginIdxLocal = updatesOriginIdexQue_.Get<uint32_t>();
@@ -142,17 +132,27 @@ __aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesI
142 indicesOffset += uniqueTimes;132 indicesOffset += uniqueTimes;
143 }133 }
144 134 
145- updatesQueue_.FreeTensor<T>(updatesLocal);
146 outQueueRes_.EnQue<T>(resLocal);135 outQueueRes_.EnQue<T>(resLocal);
136+ if constexpr (updatesIsScalar) {
137+ updatesQueue_.EnQue<T>(updatesLocal);
138+ } else {
139+ updatesQueue_.FreeTensor<T>(updatesLocal);
140+ }
147}141}
148 142 
149-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>143+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
150-__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType>::CopyResToGm(144+__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType, scatterOp>::CopyResToGm(
151 uint32_t cols, uint32_t colsAlign, uint64_t ubOffset, uint32_t& uniqueIdNum)145 uint32_t cols, uint32_t colsAlign, uint64_t ubOffset, uint32_t& uniqueIdNum)
152{146{
153 LocalTensor<T> resLocal = outQueueRes_.DeQue<T>();147 LocalTensor<T> resLocal = outQueueRes_.DeQue<T>();
154 LocalTensor<CAST_T> indicesSortedLocal = sortIndicesQue_.Get<CAST_T>();148 LocalTensor<CAST_T> indicesSortedLocal = sortIndicesQue_.Get<CAST_T>();
155 LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountQue_.Get<int32_t>();149 LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountQue_.Get<int32_t>();
150+ if constexpr (scatterOp == SUB && !updatesIsScalar) {
151+ NegateUpdate<T>(resLocal, static_cast<uint32_t>(uniqueIdNum * colsAlign));
152+ auto MTE3WaitVEventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
153+ SetFlag<HardEvent::V_MTE3>(MTE3WaitVEventID);
154+ WaitFlag<HardEvent::V_MTE3>(MTE3WaitVEventID);
155+ }
156 156 
157 int32_t tmpIndex = shiftOffset_;157 int32_t tmpIndex = shiftOffset_;
158 158 
@@ -172,14 +172,18 @@ __aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesI
172 outQueueRes_.FreeTensor(resLocal);172 outQueueRes_.FreeTensor(resLocal);
173}173}
174 174 
175-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>175+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
176-__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType>::Compute()176+__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType, scatterOp>::Compute()
177{177{
178 if (GetBlockIdx() >= tilingData_.atomicAddCoreNum) {178 if (GetBlockIdx() >= tilingData_.atomicAddCoreNum) {
179 return;179 return;
180 }180 }
181 181 
182- BroadcastUpdatesScalar<T, static_cast<uint64_t>(1), updatesIsScalar>(updatesQueue_, updatesGm_, static_cast<int32_t>(tilingData_.ubFactorCol));182+ if constexpr (updatesIsScalar) {
183+ LocalTensor<T> updatesLocal = updatesQueue_.AllocTensor<T>();
184+ BroadcastUpdatesScalar<T, scatterOp>(updatesLocal, updatesGm_, static_cast<int32_t>(tilingData_.ubFactorCol));
185+ updatesQueue_.EnQue<T>(updatesLocal);
186+ }
183 uint64_t rowIdx = GetBlockIdx() / tilingData_.colTileNum; // 当前block在第几个分块行187 uint64_t rowIdx = GetBlockIdx() / tilingData_.colTileNum; // 当前block在第几个分块行
184 uint64_t colIdx = GetBlockIdx() % tilingData_.colTileNum; // 当前block在第几个分块列188 uint64_t colIdx = GetBlockIdx() % tilingData_.colTileNum; // 当前block在第几个分块列
185 uint64_t curCoreRows = rowIdx != (tilingData_.rowTileNum - 1) ? tilingData_.normBlockRow : tilingData_.tailBlockRow; // 当前分块行数189 uint64_t curCoreRows = rowIdx != (tilingData_.rowTileNum - 1) ? tilingData_.normBlockRow : tilingData_.tailBlockRow; // 当前分块行数
@@ -209,10 +213,15 @@ __aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesI
209 CopyResToGm(cols, colsAlign, ubOffset, uniqueIdNum);213 CopyResToGm(cols, colsAlign, ubOffset, uniqueIdNum);
210 }214 }
211 }215 }
216+ 
217+ if constexpr (updatesIsScalar) {
218+ LocalTensor<T> updatesLocal = updatesQueue_.DeQue<T>();
219+ updatesQueue_.FreeTensor<T>(updatesLocal);
220+ }
212}221}
213 222 
214-template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType>223+template<typename T, typename U, typename CAST_T, bool updatesIsScalar, uint32_t castType, uint32_t scatterOp>
215-__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType>::Process()224+__aicore__ inline void ScatterAddSIMDSortSupportAtomicAdd<T, U, CAST_T, updatesIsScalar, castType, scatterOp>::Process()
216{225{
217 if (GetBlockIdx() >= GetBlockNum()) {226 if (GetBlockIdx() >= GetBlockNum()) {
218 return;227 return;
Mindex/scatter_add/op_kernel/arch35/scatter_add_simd_support_atomicadd.h+25-12
@@ -22,7 +22,7 @@ namespace ScatterAdd {
22using namespace AscendC;22using namespace AscendC;
23using namespace ScatterAddCommon;23using namespace ScatterAddCommon;
24 24 
25-template<typename T, typename U, bool updatesIsScalar>25+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
26class ScatterAddSIMDSupportAtomicAdd {26class ScatterAddSIMDSupportAtomicAdd {
27public:27public:
28 __aicore__ inline ScatterAddSIMDSupportAtomicAdd(const ScatterAddTilingData& tilingData, TPipe& pipe) :28 __aicore__ inline ScatterAddSIMDSupportAtomicAdd(const ScatterAddTilingData& tilingData, TPipe& pipe) :
@@ -43,8 +43,8 @@ private:
43 const ScatterAddTilingData& tilingData_;43 const ScatterAddTilingData& tilingData_;
44};44};
45 45 
46-template<typename T, typename U, bool updatesIsScalar>46+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
47-__aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,47+__aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar, scatterOp>::Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,
48 GM_ADDR varRef, GM_ADDR workspace)48 GM_ADDR varRef, GM_ADDR workspace)
49{49{
50 if (GetBlockIdx() >= GetBlockNum()) {50 if (GetBlockIdx() >= GetBlockNum()) {
@@ -61,14 +61,15 @@ __aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::In
61 }61 }
62}62}
63 63 
64-template<typename T, typename U, bool updatesIsScalar>64+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
65-__aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::ProcessAtomicAdd()65+__aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar, scatterOp>::ProcessAtomicAdd()
66{66{
67- if (GetBlockIdx() >= tilingData_.atomicAddCoreNum) {67+ if constexpr (updatesIsScalar) {
68- return;68+ LocalTensor<T> updatesLocal = updatesQueue_.AllocTensor<T>();
69+ BroadcastUpdatesScalar<T, scatterOp>(updatesLocal, updatesGm_, static_cast<int32_t>(tilingData_.ubFactorCol));
70+ updatesQueue_.EnQue(updatesLocal);
69 }71 }
70 72 
71- BroadcastUpdatesScalar<T, DOUBLE_BUFFER, updatesIsScalar>(updatesQueue_, updatesGm_, static_cast<int32_t>(tilingData_.ubFactorCol));
72 uint64_t rowIdx = GetBlockIdx() / tilingData_.colTileNum; // 当前block在第几个分块行73 uint64_t rowIdx = GetBlockIdx() / tilingData_.colTileNum; // 当前block在第几个分块行
73 uint64_t colIdx = GetBlockIdx() % tilingData_.colTileNum; // 当前block在第几个分块列74 uint64_t colIdx = GetBlockIdx() % tilingData_.colTileNum; // 当前block在第几个分块列
74 uint64_t curCoreRows = rowIdx != (tilingData_.rowTileNum - 1) ? tilingData_.normBlockRow : tilingData_.tailBlockRow; // 当前分块行数75 uint64_t curCoreRows = rowIdx != (tilingData_.rowTileNum - 1) ? tilingData_.normBlockRow : tilingData_.tailBlockRow; // 当前分块行数
@@ -103,7 +104,7 @@ __aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::Pr
103 CopyOut(varRefGm_, updatesLocal, dstOffset, 1, cols);104 CopyOut(varRefGm_, updatesLocal, dstOffset, 1, cols);
104 }105 }
105 SetAtomicNone();106 SetAtomicNone();
106- updatesQueue_.FreeTensor(updatesLocal);107+ updatesQueue_.EnQue<T>(updatesLocal);
107 } else {108 } else {
108 LocalTensor<T> updatesLocal = updatesQueue_.AllocTensor<T>();109 LocalTensor<T> updatesLocal = updatesQueue_.AllocTensor<T>();
109 uint64_t offset = blockOffsetUpdate + rowLoop * tilingData_.ubFactorRow * tilingData_.varShape[1] + colLoop * tilingData_.ubFactorCol;110 uint64_t offset = blockOffsetUpdate + rowLoop * tilingData_.ubFactorRow * tilingData_.varShape[1] + colLoop * tilingData_.ubFactorCol;
@@ -115,6 +116,13 @@ __aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::Pr
115 SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);116 SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);
116 WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);117 WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);
117 118 
119+ if constexpr (scatterOp == SUB) {
120+ NegateUpdate<T>(updatesLocal, static_cast<uint32_t>(rows * colsAlign));
121+ auto MTE3WaitVEventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
122+ SetFlag<HardEvent::V_MTE3>(MTE3WaitVEventID);
123+ WaitFlag<HardEvent::V_MTE3>(MTE3WaitVEventID);
124+ }
125+ 
118 SetAtomicAdd<T>();126 SetAtomicAdd<T>();
119 for (uint64_t i = 0; i < rows; i++) {127 for (uint64_t i = 0; i < rows; i++) {
120 uint64_t dstIdx = indicesLocal.GetValue(i); // 找到当次循环对应的indices内的值,即var的索引128 uint64_t dstIdx = indicesLocal.GetValue(i); // 找到当次循环对应的indices内的值,即var的索引
@@ -130,12 +138,17 @@ __aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::Pr
130 }138 }
131 indicesQueue_.FreeTensor(indicesLocal);139 indicesQueue_.FreeTensor(indicesLocal);
132 }140 }
141+ 
142+ if constexpr (updatesIsScalar) {
143+ LocalTensor<T> updatesLocal = updatesQueue_.DeQue<T>();
144+ updatesQueue_.FreeTensor<T>(updatesLocal);
145+ }
133}146}
134 147 
135-template<typename T, typename U, bool updatesIsScalar>148+template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp>
136-__aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar>::Process()149+__aicore__ inline void ScatterAddSIMDSupportAtomicAdd<T, U, updatesIsScalar, scatterOp>::Process()
137{150{
138- if (GetBlockIdx() >= GetBlockNum()) {151+ if (GetBlockIdx() >= GetBlockNum() || GetBlockIdx() >= tilingData_.atomicAddCoreNum) {
139 return;152 return;
140 }153 }
141 154 
Mindex/scatter_add/op_kernel/arch35/scatter_add_simt.h+45-27
@@ -18,10 +18,12 @@
18#include "kernel_operator.h"18#include "kernel_operator.h"
19#include "../inc/platform.h"19#include "../inc/platform.h"
20#include "../inc/kernel_utils.h"20#include "../inc/kernel_utils.h"
21+#include "scatter_add_common.h"
21 22 
22namespace ScatterAdd23namespace ScatterAdd
23{24{
24using namespace AscendC;25using namespace AscendC;
26+using namespace ScatterAddCommon;
25 27 
26constexpr uint32_t VECTOR_LENGTH = platform::GetVRegSize();28constexpr uint32_t VECTOR_LENGTH = platform::GetVRegSize();
27constexpr uint32_t VL_B32 = VECTOR_LENGTH / sizeof(uint32_t);29constexpr uint32_t VL_B32 = VECTOR_LENGTH / sizeof(uint32_t);
@@ -38,7 +40,7 @@ constexpr uint32_t THREAD_NUM_LAUNCH_BOUND = 1024;
38static constexpr MicroAPI::CastTrait castTraitB8B162B32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,40static constexpr MicroAPI::CastTrait castTraitB8B162B32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
39 MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};41 MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
40 42 
41-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>43+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
42class ScatterAddSimt44class ScatterAddSimt
43{45{
44public:46public:
@@ -72,8 +74,8 @@ private:
72 uint64_t tailLoopLength_{0};74 uint64_t tailLoopLength_{0};
73};75};
74 76 
75-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>77+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
76-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::Init(GM_ADDR var, GM_ADDR indices,78+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::Init(GM_ADDR var, GM_ADDR indices,
77 GM_ADDR updates,79 GM_ADDR updates,
78 GM_ADDR workspace)80 GM_ADDR workspace)
79{81{
@@ -98,8 +100,8 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
98 }100 }
99}101}
100 102 
101-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>103+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
102-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::CopyInputGmToWs(104+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::CopyInputGmToWs(
103 GlobalTensor<VAR_T>& inGm, GlobalTensor<CAST_T>& outGm, int64_t offset, int64_t dataLen)105 GlobalTensor<VAR_T>& inGm, GlobalTensor<CAST_T>& outGm, int64_t offset, int64_t dataLen)
104{106{
105 DataCopyExtParams copyParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(VAR_T)),107 DataCopyExtParams copyParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(VAR_T)),
@@ -130,8 +132,8 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
130 varCastQue_.FreeTensor(dstLocal);132 varCastQue_.FreeTensor(dstLocal);
131}133}
132 134 
133-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>135+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
134-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::CopyWsToOutputGm(int64_t offset,136+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::CopyWsToOutputGm(int64_t offset,
135 int64_t dataLen)137 int64_t dataLen)
136{138{
137 DataCopyExtParams copyParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(CAST_T)),139 DataCopyExtParams copyParams = {static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(CAST_T)),
@@ -162,8 +164,8 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
162 varQue_.FreeTensor(dstLocal);164 varQue_.FreeTensor(dstLocal);
163}165}
164 166 
165-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>167+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
166-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::CastToInt32(168+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::CastToInt32(
167 LocalTensor<CAST_T>& dstLocal, LocalTensor<VAR_T>& srcLocal, uint32_t dataLen)169 LocalTensor<CAST_T>& dstLocal, LocalTensor<VAR_T>& srcLocal, uint32_t dataLen)
168{170{
169 __local_mem__ VAR_T* srcAddr = (__local_mem__ VAR_T*)srcLocal.GetPhyAddr();171 __local_mem__ VAR_T* srcAddr = (__local_mem__ VAR_T*)srcLocal.GetPhyAddr();
@@ -186,8 +188,8 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
186 }188 }
187}189}
188 190 
189-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>191+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
190-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::CastToOrigin(192+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::CastToOrigin(
191 LocalTensor<VAR_T>& dstLocal, LocalTensor<CAST_T>& srcLocal, uint32_t dataLen)193 LocalTensor<VAR_T>& dstLocal, LocalTensor<CAST_T>& srcLocal, uint32_t dataLen)
192{194{
193 __local_mem__ CAST_T* srcAddr = (__local_mem__ CAST_T*)srcLocal.GetPhyAddr();195 __local_mem__ CAST_T* srcAddr = (__local_mem__ CAST_T*)srcLocal.GetPhyAddr();
@@ -209,8 +211,8 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
209 }211 }
210}212}
211 213 
212-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>214+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
213-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::CopyVarToWs()215+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::CopyVarToWs()
214{216{
215 uint64_t offset = 0;217 uint64_t offset = 0;
216 for (uint64_t idx = 0; idx < static_cast<uint64_t>(curLoopNum_ - 1); idx++) {218 for (uint64_t idx = 0; idx < static_cast<uint64_t>(curLoopNum_ - 1); idx++) {
@@ -224,8 +226,8 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
224 }226 }
225}227}
226 228 
227-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>229+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
228-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::CopyWsToVar()230+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::CopyWsToVar()
229{231{
230 uint64_t offset = 0;232 uint64_t offset = 0;
231 for (uint64_t idx = 0; idx < static_cast<uint64_t>(curLoopNum_ - 1); idx++) {233 for (uint64_t idx = 0; idx < static_cast<uint64_t>(curLoopNum_ - 1); idx++) {
@@ -239,7 +241,7 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
239 }241 }
240}242}
241 243 
242-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>244+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
243__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM) inline void ScatterAddSimtCompute(245__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM) inline void ScatterAddSimtCompute(
244 ADDR_T totalCol, ADDR_T indicesSize, ADDR_T varFirstDimSize, ADDR_T magic, ADDR_T shift, __gm__ VAR_T* var,246 ADDR_T totalCol, ADDR_T indicesSize, ADDR_T varFirstDimSize, ADDR_T magic, ADDR_T shift, __gm__ VAR_T* var,
245 __gm__ IDX_T* indices, __gm__ VAR_T* updates, __gm__ CAST_T* varWorkspaceGm, ADDR_T blockIdx, ADDR_T blockNum)247 __gm__ IDX_T* indices, __gm__ VAR_T* updates, __gm__ CAST_T* varWorkspaceGm, ADDR_T blockIdx, ADDR_T blockNum)
@@ -255,24 +257,40 @@ __simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM) inline void ScatterAddSimtComput
255 257 
256 ADDR_T tailRowIdx = i - indiceRow * totalCol;258 ADDR_T tailRowIdx = i - indiceRow * totalCol;
257 ADDR_T varDataIdx = varRow * totalCol + tailRowIdx;259 ADDR_T varDataIdx = varRow * totalCol + tailRowIdx;
258- if constexpr (IsSameType<VAR_T, int8_t>::value || IsSameType<VAR_T, uint8_t>::value) {260+ if constexpr (scatterOp == ADD) {
259- if constexpr (isUpdateScalar) {261+ if constexpr (IsSameType<VAR_T, int8_t>::value || IsSameType<VAR_T, uint8_t>::value) {
260- Simt::AtomicAdd(varWorkspaceGm + varDataIdx, static_cast<CAST_T>(updates[0]));262+ if constexpr (isUpdateScalar) {
263+ Simt::AtomicAdd(varWorkspaceGm + varDataIdx, static_cast<CAST_T>(updates[0]));
264+ } else {
265+ Simt::AtomicAdd(varWorkspaceGm + varDataIdx, static_cast<CAST_T>(updates[i]));
266+ }
261 } else {267 } else {
262- Simt::AtomicAdd(varWorkspaceGm + varDataIdx, static_cast<CAST_T>(updates[i]));268+ if constexpr (isUpdateScalar) {
269+ Simt::AtomicAdd(var + varDataIdx, static_cast<VAR_T>(updates[0]));
270+ } else {
271+ Simt::AtomicAdd(var + varDataIdx, static_cast<VAR_T>(updates[i]));
272+ }
263 }273 }
264- } else {274+ } else if constexpr (scatterOp == SUB) {
265- if constexpr (isUpdateScalar) {275+ if constexpr (IsSameType<VAR_T, int8_t>::value || IsSameType<VAR_T, uint8_t>::value) {
266- Simt::AtomicAdd(var + varDataIdx, static_cast<VAR_T>(updates[0]));276+ if constexpr (isUpdateScalar) {
277+ Simt::AtomicAdd(varWorkspaceGm + varDataIdx, static_cast<CAST_T>(-updates[0]));
278+ } else {
279+ Simt::AtomicAdd(varWorkspaceGm + varDataIdx, static_cast<CAST_T>(-updates[i]));
280+ }
267 } else {281 } else {
268- Simt::AtomicAdd(var + varDataIdx, static_cast<VAR_T>(updates[i]));282+ if constexpr (isUpdateScalar) {
283+ Simt::AtomicAdd(var + varDataIdx, static_cast<VAR_T>(-updates[0]));
284+ } else {
285+ Simt::AtomicAdd(var + varDataIdx, static_cast<VAR_T>(-updates[i]));
286+ }
269 }287 }
270 }288 }
271 }289 }
272}290}
273 291 
274-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>292+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
275-__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>::Process()293+__aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>::Process()
276{294{
277 if constexpr (IsSameType<VAR_T, int8_t>::value || IsSameType<VAR_T, uint8_t>::value) {295 if constexpr (IsSameType<VAR_T, int8_t>::value || IsSameType<VAR_T, uint8_t>::value) {
278 if (blockIdx_ < usedCoreNum_) {296 if (blockIdx_ < usedCoreNum_) {
@@ -288,7 +306,7 @@ __aicore__ inline void ScatterAddSimt<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScal
288 ADDR_T shift = 0;306 ADDR_T shift = 0;
289 GetUintDivMagicAndShift(magic, shift, totalCol);307 GetUintDivMagicAndShift(magic, shift, totalCol);
290 308 
291- Simt::VF_CALL<ScatterAddSimtCompute<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>>(Simt::Dim3(THREAD_NUM), 309+ Simt::VF_CALL<ScatterAddSimtCompute<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>>(Simt::Dim3(THREAD_NUM),
292 totalCol, indicesSize, varFirstDimSize, magic, shift, (__gm__ VAR_T*)(var_.GetPhyAddr()),310 totalCol, indicesSize, varFirstDimSize, magic, shift, (__gm__ VAR_T*)(var_.GetPhyAddr()),
293 (__gm__ IDX_T*)(indices_.GetPhyAddr()), (__gm__ VAR_T*)(updates_.GetPhyAddr()),311 (__gm__ IDX_T*)(indices_.GetPhyAddr()), (__gm__ VAR_T*)(updates_.GetPhyAddr()),
294 (__gm__ CAST_T*)(varWorkspaceGm_.GetPhyAddr()), blockIdx_, blockNum_);312 (__gm__ CAST_T*)(varWorkspaceGm_.GetPhyAddr()), blockIdx_, blockNum_);
Mindex/scatter_add/op_kernel/arch35/scatter_add_simt_sort.h+33-32
@@ -18,6 +18,7 @@
18#include "kernel_operator.h"18#include "kernel_operator.h"
19#include "../inc/platform.h"19#include "../inc/platform.h"
20#include "../inc/kernel_utils.h"20#include "../inc/kernel_utils.h"
21+#include "scatter_add_common.h"
21 22 
22namespace ScatterAdd23namespace ScatterAdd
23{24{
@@ -31,7 +32,7 @@ constexpr uint32_t THREAD_NUM_SORT = 256;
31constexpr uint32_t THREAD_NUM_SORT = 1024;32constexpr uint32_t THREAD_NUM_SORT = 1024;
32#endif33#endif
33 34 
34-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType>35+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType, uint32_t scatterOp>
35class ScatterAddSimtSort36class ScatterAddSimtSort
36{37{
37public:38public:
@@ -51,7 +52,6 @@ private:
51 TQue<QuePosition::VECIN, 1> indicesInQueue_;52 TQue<QuePosition::VECIN, 1> indicesInQueue_;
52 TQue<QuePosition::VECIN, 1> updatesInQueue_;53 TQue<QuePosition::VECIN, 1> updatesInQueue_;
53 TBuf<QuePosition::VECCALC> castIndicesQue_;54 TBuf<QuePosition::VECCALC> castIndicesQue_;
54- TBuf<QuePosition::VECCALC> castTmpIndicesQue_;
55 TBuf<QuePosition::VECCALC> sortIndicesQue_;55 TBuf<QuePosition::VECCALC> sortIndicesQue_;
56 TBuf<QuePosition::VECCALC> updatesOriginIdexQue_;56 TBuf<QuePosition::VECCALC> updatesOriginIdexQue_;
57 TBuf<QuePosition::VECCALC> uniqueIdCountQue_;57 TBuf<QuePosition::VECCALC> uniqueIdCountQue_;
@@ -65,8 +65,8 @@ private:
65 static constexpr uint32_t shiftOffset_ = platform::GetUbBlockSize() / sizeof(CAST_T);65 static constexpr uint32_t shiftOffset_ = platform::GetUbBlockSize() / sizeof(CAST_T);
66};66};
67 67 
68-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType>68+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType, uint32_t scatterOp>
69-__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType>::ParseTilingData()69+__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType, scatterOp>::ParseTilingData()
70{70{
71 if (blockIdx_ == td_.sortCoreNum - 1) {71 if (blockIdx_ == td_.sortCoreNum - 1) {
72 currLoopCount_ = td_.tailBlockLoop;72 currLoopCount_ = td_.tailBlockLoop;
@@ -77,8 +77,8 @@ __aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdate
77 }77 }
78}78}
79 79 
80-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType>80+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType, uint32_t scatterOp>
81-__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType>::Init(81+__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType, scatterOp>::Init(
82 GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR workspace)82 GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR workspace)
83{83{
84 blockIdx_ = GetBlockIdx();84 blockIdx_ = GetBlockIdx();
@@ -94,17 +94,13 @@ __aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdate
94 pipe_.InitBuffer(uniqueIdCountQue_, ops::CeilAlign(td_.indicesFactor * sizeof(int32_t), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);94 pipe_.InitBuffer(uniqueIdCountQue_, ops::CeilAlign(td_.indicesFactor * sizeof(int32_t), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);
95 if constexpr (castType == CAST_0) {95 if constexpr (castType == CAST_0) {
96 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(IDX_T), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);96 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(IDX_T), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);
97- } else if constexpr (castType == CAST_3) {
98- pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(CAST_T), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);
99- pipe_.InitBuffer(castIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(CAST_T), UB_AGLIN_VALUE));
100- pipe_.InitBuffer(castTmpIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(int32_t), UB_AGLIN_VALUE));
101 } else {97 } else {
102 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(CAST_T), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);98 pipe_.InitBuffer(sortIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(CAST_T), UB_AGLIN_VALUE) + SORT_PAD_NUM * UB_AGLIN_VALUE);
103 pipe_.InitBuffer(castIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(CAST_T), UB_AGLIN_VALUE));99 pipe_.InitBuffer(castIndicesQue_, ops::CeilAlign(td_.indicesFactor * sizeof(CAST_T), UB_AGLIN_VALUE));
104 }100 }
105}101}
106 102 
107-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar>103+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t scatterOp>
108__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM_SORT) inline void ScatterAddSimtSortCompute(104__simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM_SORT) inline void ScatterAddSimtSortCompute(
109 ADDR_T outputOuterDimSize, __gm__ VAR_T* outputAddr, __local_mem__ VAR_T* inputAddr,105 ADDR_T outputOuterDimSize, __gm__ VAR_T* outputAddr, __local_mem__ VAR_T* inputAddr,
110 __local_mem__ CAST_T* sortedAddr, __local_mem__ uint32_t* sortedOriginIndexAddr, __local_mem__ int32_t* cumSumAddr,106 __local_mem__ CAST_T* sortedAddr, __local_mem__ uint32_t* sortedOriginIndexAddr, __local_mem__ int32_t* cumSumAddr,
@@ -120,15 +116,24 @@ __simt_vf__ __aicore__ LAUNCH_BOUND(THREAD_NUM_SORT) inline void ScatterAddSimtS
120 VAR_T result = 0;116 VAR_T result = 0;
121 for (int32_t tid = 0; tid < cumSumAddr[i + 1] - cumSumAddr[i]; tid++) {117 for (int32_t tid = 0; tid < cumSumAddr[i + 1] - cumSumAddr[i]; tid++) {
122 int32_t srcOffset = sortedOriginIndexAddr[cumSumAddr[i] + tid] * lastDim + innerOffset;118 int32_t srcOffset = sortedOriginIndexAddr[cumSumAddr[i] + tid] * lastDim + innerOffset;
123- result += inputAddr[srcOffset];119+ if constexpr (isUpdateScalar) {
120+ result += updateScalarValue;
121+ } else {
122+ result += inputAddr[srcOffset];
123+ }
124 }124 }
125 int64_t gmDstOffset = sortedAddr[cumSumAddr[i]] * lastDim + innerOffset;125 int64_t gmDstOffset = sortedAddr[cumSumAddr[i]] * lastDim + innerOffset;
126- Simt::AtomicAdd(outputAddr + gmDstOffset, result);126+ if constexpr (scatterOp == ADD) {
127+ Simt::AtomicAdd(outputAddr + gmDstOffset, result);
128+ } else if constexpr (scatterOp == SUB) {
129+ result = -result;
130+ Simt::AtomicAdd(outputAddr + gmDstOffset, result);
131+ }
127 }132 }
128}133}
129 134 
130-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType>135+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType, uint32_t scatterOp>
131-__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType>::CopyInIndicesUpdates(136+__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType, scatterOp>::CopyInIndicesUpdates(
132 uint32_t loopIdx, uint32_t indicesCount)137 uint32_t loopIdx, uint32_t indicesCount)
133{138{
134 LocalTensor<IDX_T> indicesLocal = indicesInQueue_.AllocTensor<IDX_T>();139 LocalTensor<IDX_T> indicesLocal = indicesInQueue_.AllocTensor<IDX_T>();
@@ -138,16 +143,19 @@ __aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdate
138 DataCopyPadExtParams<IDX_T> indicesPadParams { false, 0, 0, 0 };143 DataCopyPadExtParams<IDX_T> indicesPadParams { false, 0, 0, 0 };
139 DataCopyPad(indicesLocal, indices_[blockIdx_ * td_.normBlockIndices + loopIdx * td_.indicesFactor],144 DataCopyPad(indicesLocal, indices_[blockIdx_ * td_.normBlockIndices + loopIdx * td_.indicesFactor],
140 indicesCopyParams, indicesPadParams);145 indicesCopyParams, indicesPadParams);
141- DataCopyExtParams updatesCopyParams { 1, (uint32_t)(indicesCount * td_.varShape[1] * sizeof(VAR_T)), 0, 0, 0 };146+
142- DataCopyPadExtParams<VAR_T> updatesPadParams { false, 0, 0, 0 };147+ if constexpr (!isUpdateScalar) {
143- DataCopyPad(updatesLocal, updates_[(blockIdx_ * td_.normBlockIndices + loopIdx * td_.indicesFactor) * td_.varShape[1]],148+ DataCopyExtParams updatesCopyParams { 1, (uint32_t)(indicesCount * td_.varShape[1] * sizeof(VAR_T)), 0, 0, 0 };
144- updatesCopyParams, updatesPadParams);149+ DataCopyPadExtParams<VAR_T> updatesPadParams { false, 0, 0, 0 };
150+ DataCopyPad(updatesLocal, updates_[(blockIdx_ * td_.normBlockIndices + loopIdx * td_.indicesFactor) * td_.varShape[1]],
151+ updatesCopyParams, updatesPadParams);
152+ }
145 indicesInQueue_.EnQue<IDX_T>(indicesLocal);153 indicesInQueue_.EnQue<IDX_T>(indicesLocal);
146 updatesInQueue_.EnQue<VAR_T>(updatesLocal);154 updatesInQueue_.EnQue<VAR_T>(updatesLocal);
147}155}
148 156 
149-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType>157+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType, uint32_t scatterOp>
150-__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType>::Compute(158+__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType, scatterOp>::Compute(
151 uint32_t loopIdx, uint32_t indicesCount, VAR_T updateScalarValue)159 uint32_t loopIdx, uint32_t indicesCount, VAR_T updateScalarValue)
152{160{
153 ADDR_T totalCol = static_cast<ADDR_T>(td_.varShape[1]);161 ADDR_T totalCol = static_cast<ADDR_T>(td_.varShape[1]);
@@ -166,16 +174,9 @@ __aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdate
166 if constexpr (castType == CAST_0) {174 if constexpr (castType == CAST_0) {
167 uniqueIdNum = SortAndComputeUniqueIdx<IDX_T>(175 uniqueIdNum = SortAndComputeUniqueIdx<IDX_T>(
168 indicesCount, indicesLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);176 indicesCount, indicesLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);
169- } else if constexpr (castType == CAST_3) {
170- LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();
171- LocalTensor<int32_t> indicesCastTmpLocal = castTmpIndicesQue_.Get<int32_t>();
172- Cast<int32_t, IDX_T>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
173- Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
174- uniqueIdNum = SortAndComputeUniqueIdx<CAST_T>(
175- indicesCount, indicesCastLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);
176 } else {177 } else {
177 LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();178 LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();
178- Cast<CAST_T, IDX_T>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);179+ IndicesSortCast<IDX_T, CAST_T, castType>(indicesLocal, indicesCastLocal, uniqueIdCountLocal, indicesCount);
179 uniqueIdNum = SortAndComputeUniqueIdx<CAST_T>(180 uniqueIdNum = SortAndComputeUniqueIdx<CAST_T>(
180 indicesCount, indicesCastLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);181 indicesCount, indicesCastLocal, indicesSortedLocal, uniqueIdCountLocal, updatesOriginIdxLocal);
181 }182 }
@@ -186,7 +187,7 @@ __aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdate
186 int32_t threadBlock = currentMaxThread / totalCol;187 int32_t threadBlock = currentMaxThread / totalCol;
187 threadBlock = threadBlock < uniqueIdNum ? threadBlock : uniqueIdNum;188 threadBlock = threadBlock < uniqueIdNum ? threadBlock : uniqueIdNum;
188 189 
189- Simt::VF_CALL<ScatterAddSimtSortCompute<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar>>(190+ Simt::VF_CALL<ScatterAddSimtSortCompute<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, scatterOp>>(
190 Simt::Dim3({static_cast<uint32_t>(totalCol), static_cast<uint32_t>(threadBlock)}),191 Simt::Dim3({static_cast<uint32_t>(totalCol), static_cast<uint32_t>(threadBlock)}),
191 varFirstDimSize, (__gm__ VAR_T*)(var_.GetPhyAddr()), updatesLocalPtr, indicesSortedPtr,192 varFirstDimSize, (__gm__ VAR_T*)(var_.GetPhyAddr()), updatesLocalPtr, indicesSortedPtr,
192 (__local_mem__ uint32_t*)(updatesOriginIdxLocal.GetPhyAddr()),193 (__local_mem__ uint32_t*)(updatesOriginIdxLocal.GetPhyAddr()),
@@ -196,8 +197,8 @@ __aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdate
196 updatesInQueue_.FreeTensor(updatesLocal);197 updatesInQueue_.FreeTensor(updatesLocal);
197}198}
198 199 
199-template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType>200+template <typename IDX_T, typename VAR_T, typename CAST_T, typename ADDR_T, bool isUpdateScalar, uint32_t castType, uint32_t scatterOp>
200-__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType>::Process()201+__aicore__ inline void ScatterAddSimtSort<IDX_T, VAR_T, CAST_T, ADDR_T, isUpdateScalar, castType, scatterOp>::Process()
201{202{
202 uint32_t indicesCount = 0;203 uint32_t indicesCount = 0;
203 VAR_T updateScalarValue = ((__gm__ VAR_T*)(updates_.GetPhyAddr()))[0];204 VAR_T updateScalarValue = ((__gm__ VAR_T*)(updates_.GetPhyAddr()))[0];
Mindex/scatter_add/op_kernel/scatter_add_apt.cpp+404-343
@@ -24,348 +24,409 @@
24 using namespace AscendC;24 using namespace AscendC;
25 25 
26 #define TILING_KEY_0 026 #define TILING_KEY_0 0
27- #define TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR 10000000003333301000UL27+ #define TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR 10000000000333331000UL
28- #define TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR 10000000003333300000UL28+ #define TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR 10000000000333330000UL
29- #define TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR 10000000003333301100UL29+ #define TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR 10000000000333331100UL
30- #define TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR 10000000003333300100UL30+ #define TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR 10000000000333330100UL
31- #define TILING_KEY_SORT_NOCAST_SIMT_ADDR32_SCALAR 10000000003333301001UL31+ #define TILING_KEY_SORT_SIMT_ADDR32_SCALAR 10000000000333331001UL
32- #define TILING_KEY_SORT_NOCAST_SIMT_ADDR32_TENSOR 10000000003333300001UL32+ #define TILING_KEY_SORT_SIMT_ADDR32_TENSOR 10000000000333330001UL
33- #define TILING_KEY_SORT_NOCAST_SIMT_ADDR64_SCALAR 10000000003333301101UL33+ #define TILING_KEY_SORT_SIMT_ADDR64_SCALAR 10000000000333331101UL
34- #define TILING_KEY_SORT_NOCAST_SIMT_ADDR64_TENSOR 10000000003333300101UL34+ #define TILING_KEY_SORT_SIMT_ADDR64_TENSOR 10000000000333330101UL
35- #define TILING_KEY_SORT_CAST1_SIMT_ADDR32_SCALAR 10000000003333311001UL35+ #define TILING_KEY_UNSORT_SIMD_SCALAR 10000000000333331010UL
36- #define TILING_KEY_SORT_CAST1_SIMT_ADDR32_TENSOR 10000000003333310001UL36+ #define TILING_KEY_UNSORT_SIMD_TENSOR 10000000000333330010UL
37- #define TILING_KEY_SORT_CAST1_SIMT_ADDR64_SCALAR 10000000003333311101UL37+ #define TILING_KEY_SORT_SIMD_SCALAR 10000000000333331011UL
38- #define TILING_KEY_SORT_CAST1_SIMT_ADDR64_TENSOR 10000000003333310101UL38+ #define TILING_KEY_SORT_SIMD_TENSOR 10000000000333330011UL
39- #define TILING_KEY_SORT_CAST2_SIMT_ADDR32_SCALAR 10000000003333321001UL
40- #define TILING_KEY_SORT_CAST2_SIMT_ADDR32_TENSOR 10000000003333320001UL
41- #define TILING_KEY_SORT_CAST2_SIMT_ADDR64_SCALAR 10000000003333321101UL
42- #define TILING_KEY_SORT_CAST2_SIMT_ADDR64_TENSOR 10000000003333320101UL
43- #define TILING_KEY_SORT_CAST3_SIMT_ADDR32_SCALAR 10000000003333331001UL
44- #define TILING_KEY_SORT_CAST3_SIMT_ADDR32_TENSOR 10000000003333330001UL
45- #define TILING_KEY_SORT_CAST3_SIMT_ADDR64_SCALAR 10000000003333331101UL
46- #define TILING_KEY_SORT_CAST3_SIMT_ADDR64_TENSOR 10000000003333330101UL
47- #define TILING_KEY_UNSORT_SIMD_SCALAR 10000000003333301010UL
48- #define TILING_KEY_UNSORT_SIMD_TENSOR 10000000003333300010UL
49- #define TILING_KEY_SORT_NOCAST_SIMD_SCALAR 10000000003333301011UL
50- #define TILING_KEY_SORT_NOCAST_SIMD_TENSOR 10000000003333300011UL
51- #define TILING_KEY_SORT_CAST1_SIMD_SCALAR 10000000003333311011UL
52- #define TILING_KEY_SORT_CAST1_SIMD_TENSOR 10000000003333310011UL
53- #define TILING_KEY_SORT_CAST2_SIMD_SCALAR 10000000003333321011UL
54- #define TILING_KEY_SORT_CAST2_SIMD_TENSOR 10000000003333320011UL
55- #define TILING_KEY_SORT_CAST3_SIMD_SCALAR 10000000003333331011UL
56- #define TILING_KEY_SORT_CAST3_SIMD_TENSOR 10000000003333330011UL
57 39 
58 using namespace ScatterAdd;40 using namespace ScatterAdd;
59- extern "C" __global__ __aicore__ void scatter_add(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,41+ 
60- GM_ADDR workspace, GM_ADDR tiling)42+ __aicore__ inline void ScatterAddUnsortSimtAddr32Scalar(
61- {43+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
62- if (workspace == nullptr) {44+{
63- return;45+ GET_TILING_DATA(tilingData, tiling);
64- }46+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
65- SetSysWorkspace(workspace);47+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint32_t, true, ADD> op(tilingData, pipe);
66- GM_ADDR userWs = GetUserWorkspace(workspace);48+ op.Init(var, indices, updates, userWs);
67- if (userWs == nullptr) {49+ op.Process();
68- return;50+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
69- }51+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, true, ADD> op(tilingData, pipe);
70- GET_TILING_DATA(tilingData, tiling);52+ op.Init(var, indices, updates, userWs);
71- TPipe pipe;53+ op.Process();
72- KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);54+ } else {
73- 55+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint32_t, true, ADD> op(tilingData, pipe);
74- if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR)) {56+ op.Init(var, indices, updates, userWs);
75- if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {57+ op.Process();
76- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint32_t, true> op(tilingData, pipe);58+ }
77- op.Init(var, indices, updates, userWs);59+}
78- op.Process();60+ 
79- } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {61+ __aicore__ inline void ScatterAddUnsortSimtAddr32Tensor(
80- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, true> op(tilingData, pipe);62+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
81- op.Init(var, indices, updates, userWs);63+{
82- op.Process();64+ GET_TILING_DATA(tilingData, tiling);
83- } else {65+ if (tilingData.isDeterminTemplate) {
84- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint32_t, true> op(tilingData, pipe);66+ if constexpr (is_same<half, DTYPE_VAR>::value || is_same<float, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {
85- op.Init(var, indices, updates, userWs);67+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, ADD> op(tilingData, pipe);
86- op.Process();68+ op.Init(var, indices, updates, y, userWs);
87- }69+ op.Process();
88- } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR)) {70+ }
89- if (tilingData.isDeterminTemplate) {71+ } else {
90- if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {72+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
91- ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES> op(tilingData, pipe);73+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint32_t, false, ADD> op(tilingData, pipe);
92- op.Init(var, indices, updates, y, userWs);74+ op.Init(var, indices, updates, userWs);
93- op.Process();75+ op.Process();
94- }76+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
95- } else {77+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, false, ADD> op(tilingData, pipe);
96- if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {78+ op.Init(var, indices, updates, userWs);
97- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint32_t, false> op(tilingData, pipe);79+ op.Process();
98- op.Init(var, indices, updates, userWs);80+ } else {
99- op.Process();81+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint32_t, false, ADD> op(tilingData, pipe);
100- } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {82+ op.Init(var, indices, updates, userWs);
101- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, false> op(tilingData, pipe);83+ op.Process();
102- op.Init(var, indices, updates, userWs);84+ }
103- op.Process();85+ }
104- } else {86+}
105- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint32_t, false> op(tilingData, pipe);87+ 
106- op.Init(var, indices, updates, userWs);88+ __aicore__ inline void ScatterAddUnsortSimtAddr64Scalar(
107- op.Process();89+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
108- }90+{
109- }91+ GET_TILING_DATA(tilingData, tiling);
110- } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR)) {92+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
111- if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {93+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint64_t, true, ADD> op(tilingData, pipe);
112- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint64_t, true> op(tilingData, pipe);94+ op.Init(var, indices, updates, userWs);
113- op.Init(var, indices, updates, userWs);95+ op.Process();
114- op.Process();96+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
115- } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {97+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, true, ADD> op(tilingData, pipe);
116- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, true> op(tilingData, pipe);98+ op.Init(var, indices, updates, userWs);
117- op.Init(var, indices, updates, userWs);99+ op.Process();
118- op.Process();100+ } else {
119- } else {101+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint64_t, true, ADD> op(tilingData, pipe);
120- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint64_t, true> op(tilingData, pipe);102+ op.Init(var, indices, updates, userWs);
121- op.Init(var, indices, updates, userWs);103+ op.Process();
122- op.Process();104+ }
123- }105+}
124- } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR)) {106+ 
125- if (tilingData.isDeterminTemplate) {107+ __aicore__ inline void ScatterAddUnsortSimtAddr64Tensor(
126- if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {108+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
127- ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES> op(tilingData, pipe);109+{
128- op.Init(var, indices, updates, y, userWs);110+ GET_TILING_DATA(tilingData, tiling);
129- op.Process();111+ if (tilingData.isDeterminTemplate) {
130- }112+ if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {
131- } else {113+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, ADD> op(tilingData, pipe);
132- if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {114+ op.Init(var, indices, updates, y, userWs);
133- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint64_t, false> op(tilingData, pipe);115+ op.Process();
134- op.Init(var, indices, updates, userWs);116+ }
135- op.Process();117+ } else {
136- } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {118+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
137- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, false> op(tilingData, pipe);119+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint64_t, false, ADD> op(tilingData, pipe);
138- op.Init(var, indices, updates, userWs);120+ op.Init(var, indices, updates, userWs);
139- op.Process();121+ op.Process();
140- } else {122+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
141- ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint64_t, false> op(tilingData, pipe);123+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, false, ADD> op(tilingData, pipe);
142- op.Init(var, indices, updates, userWs);124+ op.Init(var, indices, updates, userWs);
143- op.Process();125+ op.Process();
144- }126+ } else {
145- }127+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint64_t, false, ADD> op(tilingData, pipe);
146- } else if (TILING_KEY_IS(TILING_KEY_SORT_NOCAST_SIMT_ADDR32_SCALAR)) {128+ op.Init(var, indices, updates, userWs);
147- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {129+ op.Process();
148- return;130+ }
149- } else {131+ }
150- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint32_t, true, CAST_0> op(tilingData, pipe);132+}
151- op.Init(var, indices, updates, userWs);133+ 
152- op.Process();134+ __aicore__ inline void ScatterAddSortSimtAddr32Scalar(
153- }135+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
154- } else if (TILING_KEY_IS(TILING_KEY_SORT_NOCAST_SIMT_ADDR32_TENSOR)) {136+{
155- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {137+ GET_TILING_DATA(tilingData, tiling);
156- return;138+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
157- } else {139+ return;
158- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint32_t, false, CAST_0> op(tilingData, pipe);140+ } else {
159- op.Init(var, indices, updates, userWs);141+ if (tilingData.indicesCastMode == CAST_0) {
160- op.Process();142+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint32_t, true, CAST_0, ADD> op(tilingData, pipe);
161- }143+ op.Init(var, indices, updates, userWs);
162- } else if (TILING_KEY_IS(TILING_KEY_SORT_NOCAST_SIMT_ADDR64_SCALAR)) {144+ op.Process();
163- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {145+ } else if (tilingData.indicesCastMode == CAST_1) {
164- return;146+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, true, CAST_1, ADD> op(tilingData, pipe);
165- } else {147+ op.Init(var, indices, updates, userWs);
166- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint64_t, true, CAST_0> op(tilingData, pipe);148+ op.Process();
167- op.Init(var, indices, updates, userWs);149+ } else if (tilingData.indicesCastMode == CAST_2) {
168- op.Process();150+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, true, CAST_2, ADD> op(tilingData, pipe);
169- }151+ op.Init(var, indices, updates, userWs);
170- } else if (TILING_KEY_IS(TILING_KEY_SORT_NOCAST_SIMT_ADDR64_TENSOR)) {152+ op.Process();
171- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {153+ } else if (tilingData.indicesCastMode == CAST_3) {
172- return;154+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, true, CAST_3, ADD> op(tilingData, pipe);
173- } else {155+ op.Init(var, indices, updates, userWs);
174- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint64_t, false, CAST_0> op(tilingData, pipe);156+ op.Process();
175- op.Init(var, indices, updates, userWs);157+ } else if (tilingData.indicesCastMode == CAST_4) {
176- op.Process();158+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, true, CAST_4, ADD> op(tilingData, pipe);
177- }159+ op.Init(var, indices, updates, userWs);
178- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST1_SIMT_ADDR32_SCALAR)) {160+ op.Process();
179- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {161+ } else if (tilingData.indicesCastMode == CAST_5) {
180- return;162+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, true, CAST_5, ADD> op(tilingData, pipe);
181- } else {163+ op.Init(var, indices, updates, userWs);
182- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, true, CAST_1> op(tilingData, pipe);164+ op.Process();
183- op.Init(var, indices, updates, userWs);165+ }
184- op.Process();166+ }
185- }167+}
186- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST1_SIMT_ADDR32_TENSOR)) {168+ 
187- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {169+ __aicore__ inline void ScatterAddSortSimtAddr32Tensor(
188- return;170+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
189- } else {171+{
190- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, false, CAST_1> op(tilingData, pipe);172+ GET_TILING_DATA(tilingData, tiling);
191- op.Init(var, indices, updates, userWs);173+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
192- op.Process();174+ return;
193- }175+ } else {
194- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST1_SIMT_ADDR64_SCALAR)) {176+ if (tilingData.indicesCastMode == CAST_0) {
195- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {177+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint32_t, false, CAST_0, ADD> op(tilingData, pipe);
196- return;178+ op.Init(var, indices, updates, userWs);
197- } else {179+ op.Process();
198- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, true, CAST_1> op(tilingData, pipe);180+ } else if (tilingData.indicesCastMode == CAST_1) {
199- op.Init(var, indices, updates, userWs);181+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, false, CAST_1, ADD> op(tilingData, pipe);
200- op.Process();182+ op.Init(var, indices, updates, userWs);
201- }183+ op.Process();
202- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST1_SIMT_ADDR64_TENSOR)) {184+ } else if (tilingData.indicesCastMode == CAST_2) {
203- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {185+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, false, CAST_2, ADD> op(tilingData, pipe);
204- return;186+ op.Init(var, indices, updates, userWs);
205- } else {187+ op.Process();
206- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, false, CAST_1> op(tilingData, pipe);188+ } else if (tilingData.indicesCastMode == CAST_3) {
207- op.Init(var, indices, updates, userWs);189+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, false, CAST_3, ADD> op(tilingData, pipe);
208- op.Process();190+ op.Init(var, indices, updates, userWs);
209- }191+ op.Process();
210- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST2_SIMT_ADDR32_SCALAR)) {192+ } else if (tilingData.indicesCastMode == CAST_4) {
211- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {193+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, false, CAST_4, ADD> op(tilingData, pipe);
212- return;194+ op.Init(var, indices, updates, userWs);
213- } else {195+ op.Process();
214- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, true, CAST_2> op(tilingData, pipe);196+ } else if (tilingData.indicesCastMode == CAST_5) {
215- op.Init(var, indices, updates, userWs);197+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, false, CAST_5, ADD> op(tilingData, pipe);
216- op.Process();198+ op.Init(var, indices, updates, userWs);
217- }199+ op.Process();
218- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST2_SIMT_ADDR32_TENSOR)) {200+ }
219- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {201+ }
220- return;202+}
221- } else {203+ 
222- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, false, CAST_2> op(tilingData, pipe);204+ __aicore__ inline void ScatterAddSortSimtAddr64Scalar(
223- op.Init(var, indices, updates, userWs);205+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
224- op.Process();206+{
225- }207+ GET_TILING_DATA(tilingData, tiling);
226- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST2_SIMT_ADDR64_SCALAR)) {208+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
227- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {209+ return;
228- return;210+ } else {
229- } else {211+ if (tilingData.indicesCastMode == CAST_0) {
230- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, true, CAST_2> op(tilingData, pipe);212+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint64_t, true, CAST_0, ADD> op(tilingData, pipe);
231- op.Init(var, indices, updates, userWs);213+ op.Init(var, indices, updates, userWs);
232- op.Process();214+ op.Process();
233- }215+ } else if (tilingData.indicesCastMode == CAST_1) {
234- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST2_SIMT_ADDR64_TENSOR)) {216+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, true, CAST_1, ADD> op(tilingData, pipe);
235- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {217+ op.Init(var, indices, updates, userWs);
236- return;218+ op.Process();
237- } else {219+ } else if (tilingData.indicesCastMode == CAST_2) {
238- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, false, CAST_2> op(tilingData, pipe);220+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, true, CAST_2, ADD> op(tilingData, pipe);
239- op.Init(var, indices, updates, userWs);221+ op.Init(var, indices, updates, userWs);
240- op.Process();222+ op.Process();
241- }223+ } else if (tilingData.indicesCastMode == CAST_3) {
242- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST3_SIMT_ADDR32_SCALAR)) {224+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, true, CAST_3, ADD> op(tilingData, pipe);
243- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {225+ op.Init(var, indices, updates, userWs);
244- return;226+ op.Process();
245- } else {227+ } else if (tilingData.indicesCastMode == CAST_4) {
246- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, true, CAST_3> op(tilingData, pipe);228+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, true, CAST_4, ADD> op(tilingData, pipe);
247- op.Init(var, indices, updates, userWs);229+ op.Init(var, indices, updates, userWs);
248- op.Process();230+ op.Process();
249- }231+ } else if (tilingData.indicesCastMode == CAST_5) {
250- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST3_SIMT_ADDR32_TENSOR)) {232+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, true, CAST_5, ADD> op(tilingData, pipe);
251- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {233+ op.Init(var, indices, updates, userWs);
252- return;234+ op.Process();
253- } else {235+ }
254- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, false, CAST_3> op(tilingData, pipe);236+ }
255- op.Init(var, indices, updates, userWs);237+}
256- op.Process();238+ 
257- }239+ __aicore__ inline void ScatterAddSortSimtAddr64Tensor(
258- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST3_SIMT_ADDR64_SCALAR)) {240+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
259- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {241+{
260- return;242+ GET_TILING_DATA(tilingData, tiling);
261- } else {243+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
262- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, true, CAST_3> op(tilingData, pipe);244+ return;
263- op.Init(var, indices, updates, userWs);245+ } else {
264- op.Process();246+ if (tilingData.indicesCastMode == CAST_0) {
265- }247+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint64_t, false, CAST_0, ADD> op(tilingData, pipe);
266- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST3_SIMT_ADDR64_TENSOR)) {248+ op.Init(var, indices, updates, userWs);
267- if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {249+ op.Process();
268- return;250+ } else if (tilingData.indicesCastMode == CAST_1) {
269- } else {251+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, false, CAST_1, ADD> op(tilingData, pipe);
270- ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, false, CAST_3> op(tilingData, pipe);252+ op.Init(var, indices, updates, userWs);
271- op.Init(var, indices, updates, userWs);253+ op.Process();
272- op.Process();254+ } else if (tilingData.indicesCastMode == CAST_2) {
273- }255+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, false, CAST_2, ADD> op(tilingData, pipe);
274- } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMD_SCALAR)) {256+ op.Init(var, indices, updates, userWs);
275- if (tilingData.isDeterminTemplate) {257+ op.Process();
276- if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {258+ } else if (tilingData.indicesCastMode == CAST_3) {
277- ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES> op(tilingData, pipe);259+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, false, CAST_3, ADD> op(tilingData, pipe);
278- op.Init(var, indices, updates, y, userWs);260+ op.Init(var, indices, updates, userWs);
279- op.Process();261+ op.Process();
280- }262+ } else if (tilingData.indicesCastMode == CAST_4) {
281- } else {263+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, false, CAST_4, ADD> op(tilingData, pipe);
282- ScatterAddSIMDImpl<DTYPE_VAR, DTYPE_INDICES, true> op(tilingData, pipe);264+ op.Init(var, indices, updates, userWs);
283- op.Init(var, indices, updates, y, userWs);265+ op.Process();
284- op.Process();266+ } else if (tilingData.indicesCastMode == CAST_5) {
285- }267+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, false, CAST_5, ADD> op(tilingData, pipe);
286- } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMD_TENSOR)) {268+ op.Init(var, indices, updates, userWs);
287- if (tilingData.isDeterminTemplate) {269+ op.Process();
288- if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {270+ }
289- ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES> op(tilingData, pipe);271+ }
290- op.Init(var, indices, updates, y, userWs);272+}
291- op.Process();273+ 
292- }274+ __aicore__ inline void ScatterAddUnsortSimdScalar(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
293- } else {275+ GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
294- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {276+{
295- ScatterAddSIMDSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, false> op(tilingData, pipe);277+ GET_TILING_DATA(tilingData, tiling);
296- op.Init(var, indices, updates, y, userWs);278+ if (tilingData.isDeterminTemplate) {
297- op.Process();279+ if constexpr (is_same<bfloat16_t, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<float, DTYPE_VAR>::value) {
298- } else {280+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, ADD> op(tilingData, pipe);
299- ScatterAddSIMDImpl<DTYPE_VAR, DTYPE_INDICES, false> op(tilingData, pipe);281+ op.Init(var, indices, updates, y, userWs);
300- op.Init(var, indices, updates, y, userWs);282+ op.Process();
301- op.Process();283+ }
302- }284+ } else {
303- }285+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
304- } else if (TILING_KEY_IS(TILING_KEY_SORT_NOCAST_SIMD_SCALAR)) {286+ ScatterAddSIMDSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, true, ADD> op(tilingData, pipe);
305- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {287+ op.Init(var, indices, updates, y, userWs);
306- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, DTYPE_INDICES, true, CAST_1> op(tilingData, pipe);288+ op.Process();
307- op.Init(var, indices, updates, y, userWs);289+ } else {
308- op.Process();290+ ScatterAddSIMDImpl<DTYPE_VAR, DTYPE_INDICES, true, ADD> op(tilingData, pipe);
309- } else {291+ op.Init(var, indices, updates, y, userWs);
310- return;292+ op.Process();
311- }293+ }
312- } else if (TILING_KEY_IS(TILING_KEY_SORT_NOCAST_SIMD_TENSOR)) {294+ }
313- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {295+}
314- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, DTYPE_INDICES, false, CAST_1> op(tilingData, pipe);296+ 
315- op.Init(var, indices, updates, y, userWs);297+ __aicore__ inline void ScatterAddUnsortSimdTensor(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
316- op.Process();298+ GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
317- } else {299+{
318- return;300+ GET_TILING_DATA(tilingData, tiling);
319- }301+ if (tilingData.isDeterminTemplate) {
320- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST1_SIMD_SCALAR)) {302+ if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {
321- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {303+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, ADD> op(tilingData, pipe);
322- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, true, CAST_1> op(tilingData, pipe);304+ op.Init(var, indices, updates, y, userWs);
323- op.Init(var, indices, updates, y, userWs);305+ op.Process();
324- op.Process();306+ }
325- } else {307+ } else {
326- return;308+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
327- }309+ ScatterAddSIMDSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, false, ADD> op(tilingData, pipe);
328- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST1_SIMD_TENSOR)) {310+ op.Init(var, indices, updates, y, userWs);
329- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {311+ op.Process();
330- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, false, CAST_1> op(tilingData, pipe);312+ } else {
331- op.Init(var, indices, updates, y, userWs);313+ ScatterAddSIMDImpl<DTYPE_VAR, DTYPE_INDICES, false, ADD> op(tilingData, pipe);
332- op.Process();314+ op.Init(var, indices, updates, y, userWs);
333- } else {315+ op.Process();
334- return;316+ }
335- }317+ }
336- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST2_SIMD_SCALAR)) {318+}
337- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {319+ 
338- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int32_t, true, CAST_2> op(tilingData, pipe);320+ __aicore__ inline void ScatterAddSortSimdScalar(
339- op.Init(var, indices, updates, y, userWs);321+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
340- op.Process();322+{
341- } else {323+ GET_TILING_DATA(tilingData, tiling);
342- return;324+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
343- }325+ if (tilingData.indicesCastMode == CAST_0) {
344- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST2_SIMD_TENSOR)) {326+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, DTYPE_INDICES, true, CAST_0, ADD> op(tilingData, pipe);
345- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {327+ op.Init(var, indices, updates, y, userWs);
346- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int32_t, false, CAST_2> op(tilingData, pipe);328+ op.Process();
347- op.Init(var, indices, updates, y, userWs);329+ } else if (tilingData.indicesCastMode == CAST_1) {
348- op.Process();330+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, true, CAST_1, ADD> op(tilingData, pipe);
349- } else {331+ op.Init(var, indices, updates, y, userWs);
350- return;332+ op.Process();
351- }333+ } else if (tilingData.indicesCastMode == CAST_2) {
352- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST3_SIMD_SCALAR)) {334+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int32_t, true, CAST_2, ADD> op(tilingData, pipe);
353- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {335+ op.Init(var, indices, updates, y, userWs);
354- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, true, CAST_3> op(tilingData, pipe);336+ op.Process();
355- op.Init(var, indices, updates, y, userWs);337+ } else if (tilingData.indicesCastMode == CAST_3) {
356- op.Process();338+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, true, CAST_3, ADD> op(tilingData, pipe);
357- } else {339+ op.Init(var, indices, updates, y, userWs);
358- return;340+ op.Process();
359- }341+ } else if (tilingData.indicesCastMode == CAST_4) {
360- } else if (TILING_KEY_IS(TILING_KEY_SORT_CAST3_SIMD_TENSOR)) {342+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, true, CAST_4, ADD> op(tilingData, pipe);
361- if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {343+ op.Init(var, indices, updates, y, userWs);
362- ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, false, CAST_3> op(tilingData, pipe);344+ op.Process();
363- op.Init(var, indices, updates, y, userWs);345+ } else if (tilingData.indicesCastMode == CAST_5) {
364- op.Process();346+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, true, CAST_5, ADD> op(tilingData, pipe);
365- } else {347+ op.Init(var, indices, updates, y, userWs);
366- return;348+ op.Process();
367- }349+ }
368- } else if (TILING_KEY_IS(TILING_KEY_0)) {350+ } else {
369- return;351+ return;
370- }352+ }
371- }353+}
354+ 
355+ __aicore__ inline void ScatterAddSortSimdTensor(
356+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
357+{
358+ GET_TILING_DATA(tilingData, tiling);
359+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
360+ if (tilingData.indicesCastMode == CAST_0) {
361+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, DTYPE_INDICES, false, CAST_0, ADD> op(tilingData, pipe);
362+ op.Init(var, indices, updates, y, userWs);
363+ op.Process();
364+ } else if (tilingData.indicesCastMode == CAST_1) {
365+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, false, CAST_1, ADD> op(tilingData, pipe);
366+ op.Init(var, indices, updates, y, userWs);
367+ op.Process();
368+ } else if (tilingData.indicesCastMode == CAST_2) {
369+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int32_t, false, CAST_2, ADD> op(tilingData, pipe);
370+ op.Init(var, indices, updates, y, userWs);
371+ op.Process();
372+ } else if (tilingData.indicesCastMode == CAST_3) {
373+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, false, CAST_3, ADD> op(tilingData, pipe);
374+ op.Init(var, indices, updates, y, userWs);
375+ op.Process();
376+ } else if (tilingData.indicesCastMode == CAST_4) {
377+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, false, CAST_4, ADD> op(tilingData, pipe);
378+ op.Init(var, indices, updates, y, userWs);
379+ op.Process();
380+ } else if (tilingData.indicesCastMode == CAST_5) {
381+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, false, CAST_5, ADD> op(tilingData, pipe);
382+ op.Init(var, indices, updates, y, userWs);
383+ op.Process();
384+ }
385+ } else {
386+ return;
387+ }
388+}
389+ 
390+extern "C" __global__ __aicore__ void scatter_add(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
391+ GM_ADDR workspace, GM_ADDR tiling)
392+{
393+ if (workspace == nullptr) {
394+ return;
395+ }
396+ SetSysWorkspace(workspace);
397+ GM_ADDR userWs = GetUserWorkspace(workspace);
398+ if (userWs == nullptr) {
399+ return;
400+ }
401+ GET_TILING_DATA(tilingData, tiling);
402+ TPipe pipe;
403+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
404+ 
405+ if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR)) {
406+ ScatterAddUnsortSimtAddr32Scalar(var, indices, updates, y, userWs, tiling, pipe);
407+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR)) {
408+ ScatterAddUnsortSimtAddr32Tensor(var, indices, updates, y, userWs, tiling, pipe);
409+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR)) {
410+ ScatterAddUnsortSimtAddr64Scalar(var, indices, updates, y, userWs, tiling, pipe);
411+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR)) {
412+ ScatterAddUnsortSimtAddr64Tensor(var, indices, updates, y, userWs, tiling, pipe);
413+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR32_SCALAR)) {
414+ ScatterAddSortSimtAddr32Scalar(var, indices, updates, y, userWs, tiling, pipe);
415+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR32_TENSOR)) {
416+ ScatterAddSortSimtAddr32Tensor(var, indices, updates, y, userWs, tiling, pipe);
417+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR64_SCALAR)) {
418+ ScatterAddSortSimtAddr64Scalar(var, indices, updates, y, userWs, tiling, pipe);
419+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR64_TENSOR)) {
420+ ScatterAddSortSimtAddr64Tensor(var, indices, updates, y, userWs, tiling, pipe);
421+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMD_SCALAR)) {
422+ ScatterAddUnsortSimdScalar(var, indices, updates, y, userWs, tiling, pipe);
423+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMD_TENSOR)) {
424+ ScatterAddUnsortSimdTensor(var, indices, updates, y, userWs, tiling, pipe);
425+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMD_SCALAR)) {
426+ ScatterAddSortSimdScalar(var, indices, updates, y, userWs, tiling, pipe);
427+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMD_TENSOR)) {
428+ ScatterAddSortSimdTensor(var, indices, updates, y, userWs, tiling, pipe);
429+ } else if (TILING_KEY_IS(TILING_KEY_0)) {
430+ return;
431+ }
432+}
Mindex/scatter_add/tests/ut/op_api/test_aclnn_scatter_add.cpp+45-0
@@ -32,6 +32,25 @@ public:
32 }32 }
33};33};
34 34 
35+TEST_F(l2_scatter_add_test, case_1)
36+{
37+ auto self_desc = TensorDesc({3, 3}, ACL_DOUBLE, ACL_FORMAT_ND)
38+ .Value(vector<double>{1.111131123123, 1, 1, 1, 1, 1, 1, 1, 1})
39+ .Precision(0.0001, 0.0001);
40+ int64_t dim = 0;
41+ auto index_desc = TensorDesc({1,3}, ACL_INT64, ACL_FORMAT_ND)
42+ .Value(vector<int>{0, 1, 2});
43+ auto src_desc = TensorDesc({3, 4}, ACL_DOUBLE, ACL_FORMAT_ND)
44+ .Value(vector<double>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
45+ auto ut = OP_API_UT(aclnnScatterAdd, INPUT(self_desc, dim, index_desc, src_desc), OUTPUT(self_desc));
46+ // SAMPLE: only test GetWorkspaceSize
47+ uint64_t workspaceSize = 0;
48+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
49+ EXPECT_EQ(aclRet, ACL_SUCCESS);
50+ // SAMPLE: precision simulate
51+ ut.TestPrecision();
52+}
53+ 
35// 空tensor54// 空tensor
36TEST_F(l2_scatter_add_test, case_2)55TEST_F(l2_scatter_add_test, case_2)
37{56{
@@ -172,3 +191,29 @@ TEST_F(l2_scatter_add_test, ascend910B2_case_fp16)
172 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);191 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
173 EXPECT_EQ(aclRet, ACL_SUCCESS);192 EXPECT_EQ(aclRet, ACL_SUCCESS);
174}193}
194+ 
195+TEST_F(l2_scatter_add_test, ascend910B2_case_expand) {
196+ auto selfTensor = CreateAclTensor({4096, 6144}, {6144, 1}, 0, {4096, 6144});
197+ auto indexTensor = CreateAclTensor({2076, 6144}, {1, 0}, 0, {2076,}, ACL_INT64);
198+ auto srcTensor = CreateAclTensor({2076, 6144}, {6144, 1}, 0, {2076, 6144});
199+ auto outTensor = CreateAclTensor({4096, 6144}, {6144, 1}, 0, {4096, 6144});
200+ 
201+ uint64_t workspaceSize = 0U;
202+ aclOpExecutor* exe = nullptr;
203+ auto aclRet = aclnnScatterAddGetWorkspaceSize(selfTensor, 0, indexTensor, srcTensor, outTensor, &workspaceSize, &exe);
204+ EXPECT_EQ(aclRet, ACL_SUCCESS);
205+ EXPECT_NE(exe, nullptr);
206+}
207+ 
208+TEST_F(l2_scatter_add_test, ascend910B2_case_expand_same) {
209+ auto selfTensor = CreateAclTensor({4096, 6144}, {6144, 1}, 0, {4096, 6144});
210+ auto indexTensor = CreateAclTensor({4096, 6144}, {1, 0}, 0, {4096,}, ACL_INT64);
211+ auto srcTensor = CreateAclTensor({4096, 6144}, {6144, 1}, 0, {4096, 6144});
212+ auto outTensor = CreateAclTensor({4096, 6144}, {6144, 1}, 0, {4096, 6144});
213+ 
214+ uint64_t workspaceSize = 0U;
215+ aclOpExecutor* exe = nullptr;
216+ auto aclRet = aclnnScatterAddGetWorkspaceSize(selfTensor, 0, indexTensor, srcTensor, outTensor, &workspaceSize, &exe);
217+ EXPECT_EQ(aclRet, ACL_SUCCESS);
218+ EXPECT_NE(exe, nullptr);
219+}
Aindex/scatter_sub/CMakeLists.txt+15-0
@@ -0,0 +1,15 @@
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+# 设置算子定义时支持的芯片类型
12+set(SUPPORT_COMPUTE_UNIT "ascend950")
13+# 设置每种芯片类型对应的tiling文件目录,即采用op_host目录下哪个文件夹下的tiling文件编译
14+set(SUPPORT_TILING_DIR "arch35")
15+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE scatter_sub ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE DEPENDENCIES linear_index scatter_elements_v2 scatter_update tf_scatter_add scatter_nd_add)
Aindex/scatter_sub/op_graph/fusion_pass/.gitkeep+0-0
The file is empty
Aindex/scatter_sub/op_graph/scatter_sub_proto.h+54-0
@@ -0,0 +1,54 @@
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 OPS_BUILT_IN_OP_PROTO_INC_SCATTER_SUB_H_
12+#define OPS_BUILT_IN_OP_PROTO_INC_SCATTER_SUB_H_
13+ 
14+#include "graph/operator_reg.h"
15+#include "graph/operator.h"
16+ 
17+namespace ge {
18+ 
19+ 
20+/**
21+* @brief Subs sparse "updates" to a variable reference.
22+ 
23+* @par Inputs:
24+* @li var: The rewritten tensor. An ND tensor. Support 1D ~ 8D. Must be one of the following types:
25+* float16, float32, int32, int8, uint8, bfloat16.
26+* @li indices: The index tensor. An ND tensor. Support 1D ~ 8D. Must be one of the following types: int32, int64.
27+* @li updates: The source tensor. An ND tensor. Support 1D ~ 8D. Shape should be equal to the shape of "indices" concats
28+* the shape of "var" except for the first dimension. Must have the same type of "var".
29+ 
30+* @par Attributes:
31+* use_locking: Ignore this attribute. This attribute does not take effect even if it is set. \n
32+ 
33+* @par Outputs:
34+* var: An ND tensor. Support 1D ~ 8D. Must have the same type, shape and format as input "var".
35+ 
36+* @attention Constraints:
37+* updates.shape = indices.shape + var.shape[1:] or updates.shape = []. \n
38+ 
39+* @par Third-party framework compatibility
40+* Compatible with the TensorFlow operator ScatterSub.
41+*/
42+ 
43+#define SCATTER_SUB_SUPPORT_TYPES {DT_FLOAT16,DT_FLOAT,DT_INT32,DT_INT8,DT_UINT8,DT_BF16}
44+REG_OP(ScatterSub)
CANN-robot
CANN-robotCANN-robot1月29日

代码结构与可维护性: REG_OP宏调用中,INPUT和OUTPUT的TensorType数据类型列表存在重复定义。var输入和var输出使用了完全相同的数据类型列表{DT_FLOAT16,DT_FLOAT,DT_INT32,DT_INT8,DT_UINT8,DT_BF16},updates输入也使用了相同列表。这种重复定义违反了DRY原则,当需要修改支持的数据类型时,需要在多个地方同步修改,容易遗漏导致不一致。

问题类型: 代码结构与可维护性 文件路径: index/scatter_sub/op_graph/scatter_sub_proto.h 行号: 42 问题代码:

REG_OP(ScatterSub)
    .INPUT(var, TensorType({DT_FLOAT16,DT_FLOAT,DT_INT32,DT_INT8,DT_UINT8,DT_BF16}))
    .INPUT(indices, TensorType::IndexNumberType())
    .INPUT(updates, TensorType({DT_FLOAT16,DT_FLOAT,DT_INT32,DT_INT8,DT_UINT8,DT_BF16}))
    .OUTPUT(var, TensorType({DT_FLOAT16,DT_FLOAT,DT_INT32,DT_INT8,DT_UINT8,DT_BF16}))

修改建议:

建议将数据类型列表定义为常量或宏,在多个地方复用。例如:
#define SCATTER_SUB_SUPPORTED_TYPES {DT_FLOAT16,DT_FLOAT,DT_INT32,DT_INT8,DT_UINT8,DT_BF16}
然后在REG_OP调用中使用:
.INPUT(var, TensorType(SCATTER_SUB_SUPPORTED_TYPES))
.INPUT(updates, TensorType(SCATTER_SUB_SUPPORTED_TYPES))
.OUTPUT(var, TensorType(SCATTER_SUB_SUPPORTED_TYPES))

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

likedislike
45+ .INPUT(var, TensorType(SCATTER_SUB_SUPPORT_TYPES))
46+ .INPUT(indices, TensorType::IndexNumberType())
47+ .INPUT(updates, TensorType(SCATTER_SUB_SUPPORT_TYPES))
48+ .OUTPUT(var, TensorType(SCATTER_SUB_SUPPORT_TYPES))
49+ .ATTR(use_locking, Bool, false)
CANN-robot
CANN-robotCANN-robot1月29日

代码结构与可维护性: use_locking属性被标记为'忽略此属性,即使设置也不会生效',但在代码中仍然定义了这个属性。这可能导致用户困惑,认为这个属性有实际作用。如果属性确实没有作用,应该考虑是否完全移除该属性定义,或者在文档中更明确地说明其无效性。

问题类型: 代码结构与可维护性 文件路径: index/scatter_sub/op_graph/scatter_sub_proto.h 行号: 47 问题代码:

.ATTR(use_locking, Bool, false)

修改建议:

建议:1) 如果确实不需要此属性,考虑移除.ATTR(use_locking, Bool, false)行;2) 如果为了保持与TensorFlow的API兼容性而保留,在注释中更明确地说明:'此属性仅为保持API兼容性而保留,实际不会影响任何锁定行为。'

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

likedislike
50+ .OP_END_FACTORY_REG(ScatterSub)
51+ 
52+}
53+ 
54+#endif // OPS_BUILT_IN_OP_PROTO_INC_SCATTER_SUB_H_
Aindex/scatter_sub/op_host/arch35/scatter_sub_tiling.cpp+36-0
@@ -0,0 +1,36 @@
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+/*!
12+ * \file scatter_sub_tiling.cpp
13+ * \brief scatter_sub_tiling
14+ */
15+#include "scatter_sub_tiling.h"
16+#include "../../../scatter_add/op_host/arch35/scatter_add_tiling_base.h"
17+ 
18+namespace optiling {
19+// -----------------ScatterSub Util START------------------
20+static ge::graphStatus TilingPrepare4ScatterSub(gert::TilingParseContext* context) {
21+ return ge::GRAPH_SUCCESS;
22+}
23+ 
24+static ge::graphStatus Tiling4ScatterSub(gert::TilingContext* context)
25+{
26+ OP_LOGD(context->GetNodeName(), "ScatterSubTiling running begin");
27+ auto compileInfo = reinterpret_cast<const ScatterSubCompileInfo*>(context->GetCompileInfo());
28+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
29+ 
30+ OP_LOGD(context->GetNodeName(), "ScatterSubTiling is ascendc. runing ascendc tiling.");
31+ return ScatterAddTilingForAscendC(context);
32+}
33+ 
34+// register tiling interface of the ScatterSubTiling op.
35+IMPL_OP_OPTILING(ScatterSub).Tiling(Tiling4ScatterSub).TilingParse<ScatterSubCompileInfo>(TilingPrepare4ScatterSub);
36+} // namespace optiling
Aindex/scatter_sub/op_host/arch35/scatter_sub_tiling.h+34-0
@@ -0,0 +1,34 @@
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+/*!
CANN-robot
CANN-robotCANN-robot1月29日

代码结构与可维护性: 文件头注释中的\brief部分为空,没有提供该文件的简要说明。这不符合良好的文档规范,不利于其他开发者快速理解文件的用途。

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

/*!
 * \file scatter_sub_tiling.h
 * \brief
 */

修改建议:

在\brief后面添加简要说明,描述该文件的主要内容和作用,例如:"\brief ScatterSub算子的tiling相关数据结构和声明"。

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

likedislike
12+ * \file scatter_sub_tiling.h
13+ * \brief scatter_sub_tiling
14+ */
15+#ifndef OPS_BUILT_IN_OP_TILING_RUNTIME_SCATTER_SUB_H_
16+#define OPS_BUILT_IN_OP_TILING_RUNTIME_SCATTER_SUB_H_
17+ 
18+#include <sstream>
19+#include <cctype>
20+#include "register/op_impl_registry.h"
21+#include "register/tilingdata_base.h"
22+#include "log/log.h"
23+ 
24+namespace optiling {
25+struct ScatterSubCompileInfo {
26+ int64_t core_num{1};
27+ int64_t ub_size{1};
28+ int64_t var_size{1};
29+ int64_t indices_size{1};
30+ int64_t support_atomic{1};
31+};
32+ 
33+} // namespace optiling
34+#endif // OPS_BUILT_IN_OP_TILING_RUNTIME_SCATTER_SUB_H_
Aindex/scatter_sub/op_host/config/ascend950/scatter_sub_binary.json+701-0
@@ -0,0 +1,701 @@
1+{
2+ "op_type": "ScatterSub",
3+ "op_list": [
4+ {
5+ "bin_filename": "ScatterSub_af73979639cedf6a56b77932ce5d26b6",
6+ "inputs": [
7+ {
8+ "name": "var",
9+ "index": 0,
10+ "dtype": "float16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ],
16+ "format_match_mode": "FormatAgnostic"
17+ },
18+ {
19+ "name": "indices",
20+ "index": 1,
21+ "dtype": "int32",
22+ "format": "ND",
23+ "paramType": "required",
24+ "shape": [
25+ -2
26+ ],
27+ "format_match_mode": "FormatAgnostic"
28+ },
29+ {
30+ "name": "updates",
31+ "index": 2,
32+ "dtype": "float16",
33+ "format": "ND",
34+ "paramType": "required",
35+ "shape": [
36+ -2
37+ ],
38+ "format_match_mode": "FormatAgnostic"
39+ }
40+ ],
41+ "outputs": [
42+ {
43+ "name": "var",
44+ "index": 0,
45+ "dtype": "float16",
46+ "format": "ND",
47+ "paramType": "required",
48+ "shape": [
49+ -2
50+ ],
51+ "format_match_mode": "FormatAgnostic"
52+ }
53+ ],
54+ "attrs": [
55+ {
56+ "name": "use_locking",
57+ "dtype": "bool",
58+ "value": false
59+ }
60+ ]
61+ },
62+ {
63+ "bin_filename": "ScatterSub_40a17d59649f2152504e77f5aa9c3bdb",
64+ "inputs": [
65+ {
66+ "name": "var",
67+ "index": 0,
68+ "dtype": "float16",
69+ "format": "ND",
70+ "paramType": "required",
71+ "shape": [
72+ -2
73+ ],
74+ "format_match_mode": "FormatAgnostic"
75+ },
76+ {
77+ "name": "indices",
78+ "index": 1,
79+ "dtype": "int64",
80+ "format": "ND",
81+ "paramType": "required",
82+ "shape": [
83+ -2
84+ ],
85+ "format_match_mode": "FormatAgnostic"
86+ },
87+ {
88+ "name": "updates",
89+ "index": 2,
90+ "dtype": "float16",
91+ "format": "ND",
92+ "paramType": "required",
93+ "shape": [
94+ -2
95+ ],
96+ "format_match_mode": "FormatAgnostic"
97+ }
98+ ],
99+ "outputs": [
100+ {
101+ "name": "var",
102+ "index": 0,
103+ "dtype": "float16",
104+ "format": "ND",
105+ "paramType": "required",
106+ "shape": [
107+ -2
108+ ],
109+ "format_match_mode": "FormatAgnostic"
110+ }
111+ ],
112+ "attrs": [
113+ {
114+ "name": "use_locking",
115+ "dtype": "bool",
116+ "value": false
117+ }
118+ ]
119+ },
120+ {
121+ "bin_filename": "ScatterSub_43ab6a2ee49b30cd7edd72cc50df16bd",
122+ "inputs": [
123+ {
124+ "name": "var",
125+ "index": 0,
126+ "dtype": "bfloat16",
127+ "format": "ND",
128+ "paramType": "required",
129+ "shape": [
130+ -2
131+ ],
132+ "format_match_mode": "FormatAgnostic"
133+ },
134+ {
135+ "name": "indices",
136+ "index": 1,
137+ "dtype": "int32",
138+ "format": "ND",
139+ "paramType": "required",
140+ "shape": [
141+ -2
142+ ],
143+ "format_match_mode": "FormatAgnostic"
144+ },
145+ {
146+ "name": "updates",
147+ "index": 2,
148+ "dtype": "bfloat16",
149+ "format": "ND",
150+ "paramType": "required",
151+ "shape": [
152+ -2
153+ ],
154+ "format_match_mode": "FormatAgnostic"
155+ }
156+ ],
157+ "outputs": [
158+ {
159+ "name": "var",
160+ "index": 0,
161+ "dtype": "bfloat16",
162+ "format": "ND",
163+ "paramType": "required",
164+ "shape": [
165+ -2
166+ ],
167+ "format_match_mode": "FormatAgnostic"
168+ }
169+ ],
170+ "attrs": [
171+ {
172+ "name": "use_locking",
173+ "dtype": "bool",
174+ "value": false
175+ }
176+ ]
177+ },
178+ {
179+ "bin_filename": "ScatterSub_984b3f93e00beabd733af7cd4fg27230",
180+ "inputs": [
181+ {
182+ "name": "var",
183+ "index": 0,
184+ "dtype": "bfloat16",
185+ "format": "ND",
186+ "paramType": "required",
187+ "shape": [
188+ -2
189+ ],
190+ "format_match_mode": "FormatAgnostic"
191+ },
192+ {
193+ "name": "indices",
194+ "index": 1,
195+ "dtype": "int64",
196+ "format": "ND",
197+ "paramType": "required",
198+ "shape": [
199+ -2
200+ ],
201+ "format_match_mode": "FormatAgnostic"
202+ },
203+ {
204+ "name": "updates",
205+ "index": 2,
206+ "dtype": "bfloat16",
207+ "format": "ND",
208+ "paramType": "required",
209+ "shape": [
210+ -2
211+ ],
212+ "format_match_mode": "FormatAgnostic"
213+ }
214+ ],
215+ "outputs": [
216+ {
217+ "name": "var",
218+ "index": 0,
219+ "dtype": "bfloat16",
220+ "format": "ND",
221+ "paramType": "required",
222+ "shape": [
223+ -2
224+ ],
225+ "format_match_mode": "FormatAgnostic"
226+ }
227+ ],
228+ "attrs": [
229+ {
230+ "name": "use_locking",
231+ "dtype": "bool",
232+ "value": false
233+ }
234+ ]
235+ },
236+ {
237+ "bin_filename": "ScatterSub_40ab8e237639e3d59531596577dd7cce",
238+ "inputs": [
239+ {
240+ "name": "var",
241+ "index": 0,
242+ "dtype": "float32",
243+ "format": "ND",
244+ "paramType": "required",
245+ "shape": [
246+ -2
247+ ],
248+ "format_match_mode": "FormatAgnostic"
249+ },
250+ {
251+ "name": "indices",
252+ "index": 1,
253+ "dtype": "int32",
254+ "format": "ND",
255+ "paramType": "required",
256+ "shape": [
257+ -2
258+ ],
259+ "format_match_mode": "FormatAgnostic"
260+ },
261+ {
262+ "name": "updates",
263+ "index": 2,
264+ "dtype": "float32",
265+ "format": "ND",
266+ "paramType": "required",
267+ "shape": [
268+ -2
269+ ],
270+ "format_match_mode": "FormatAgnostic"
271+ }
272+ ],
273+ "outputs": [
274+ {
275+ "name": "var",
276+ "index": 0,
277+ "dtype": "float32",
278+ "format": "ND",
279+ "paramType": "required",
280+ "shape": [
281+ -2
282+ ],
283+ "format_match_mode": "FormatAgnostic"
284+ }
285+ ],
286+ "attrs": [
287+ {
288+ "name": "use_locking",
289+ "dtype": "bool",
290+ "value": false
291+ }
292+ ]
293+ },
294+ {
295+ "bin_filename": "ScatterSub_a438a20ae7b53fee52ad749b6eb2d4bc",
296+ "inputs": [
297+ {
298+ "name": "var",
299+ "index": 0,
300+ "dtype": "float32",
301+ "format": "ND",
302+ "paramType": "required",
303+ "shape": [
304+ -2
305+ ],
306+ "format_match_mode": "FormatAgnostic"
307+ },
308+ {
309+ "name": "indices",
310+ "index": 1,
311+ "dtype": "int64",
312+ "format": "ND",
313+ "paramType": "required",
314+ "shape": [
315+ -2
316+ ],
317+ "format_match_mode": "FormatAgnostic"
318+ },
319+ {
320+ "name": "updates",
321+ "index": 2,
322+ "dtype": "float32",
323+ "format": "ND",
324+ "paramType": "required",
325+ "shape": [
326+ -2
327+ ],
328+ "format_match_mode": "FormatAgnostic"
329+ }
330+ ],
331+ "outputs": [
332+ {
333+ "name": "var",
334+ "index": 0,
335+ "dtype": "float32",
336+ "format": "ND",
337+ "paramType": "required",
338+ "shape": [
339+ -2
340+ ],
341+ "format_match_mode": "FormatAgnostic"
342+ }
343+ ],
344+ "attrs": [
345+ {
346+ "name": "use_locking",
347+ "dtype": "bool",
348+ "value": false
349+ }
350+ ]
351+ },
352+ {
353+ "bin_filename": "ScatterSub_92f6de01e6e3d10aed28c4bfac64cac5",
354+ "inputs": [
355+ {
356+ "name": "var",
357+ "index": 0,
358+ "dtype": "int32",
359+ "format": "ND",
360+ "paramType": "required",
361+ "shape": [
362+ -2
363+ ],
364+ "format_match_mode": "FormatAgnostic"
365+ },
366+ {
367+ "name": "indices",
368+ "index": 1,
369+ "dtype": "int32",
370+ "format": "ND",
371+ "paramType": "required",
372+ "shape": [
373+ -2
374+ ],
375+ "format_match_mode": "FormatAgnostic"
376+ },
377+ {
378+ "name": "updates",
379+ "index": 2,
380+ "dtype": "int32",
381+ "format": "ND",
382+ "paramType": "required",
383+ "shape": [
384+ -2
385+ ],
386+ "format_match_mode": "FormatAgnostic"
387+ }
388+ ],
389+ "outputs": [
390+ {
391+ "name": "var",
392+ "index": 0,
393+ "dtype": "int32",
394+ "format": "ND",
395+ "paramType": "required",
396+ "shape": [
397+ -2
398+ ],
399+ "format_match_mode": "FormatAgnostic"
400+ }
401+ ],
402+ "attrs": [
403+ {
404+ "name": "use_locking",
405+ "dtype": "bool",
406+ "value": false
407+ }
408+ ]
409+ },
410+ {
411+ "bin_filename": "ScatterSub_69bff26fadaf3ef1de5fefa5359eadd3",
412+ "inputs": [
413+ {
414+ "name": "var",
415+ "index": 0,
416+ "dtype": "int32",
417+ "format": "ND",
418+ "paramType": "required",
419+ "shape": [
420+ -2
421+ ],
422+ "format_match_mode": "FormatAgnostic"
423+ },
424+ {
425+ "name": "indices",
426+ "index": 1,
427+ "dtype": "int64",
428+ "format": "ND",
429+ "paramType": "required",
430+ "shape": [
431+ -2
432+ ],
433+ "format_match_mode": "FormatAgnostic"
434+ },
435+ {
436+ "name": "updates",
437+ "index": 2,
438+ "dtype": "int32",
439+ "format": "ND",
440+ "paramType": "required",
441+ "shape": [
442+ -2
443+ ],
444+ "format_match_mode": "FormatAgnostic"
445+ }
446+ ],
447+ "outputs": [
448+ {
449+ "name": "var",
450+ "index": 0,
451+ "dtype": "int32",
452+ "format": "ND",
453+ "paramType": "required",
454+ "shape": [
455+ -2
456+ ],
457+ "format_match_mode": "FormatAgnostic"
458+ }
459+ ],
460+ "attrs": [
461+ {
462+ "name": "use_locking",
463+ "dtype": "bool",
464+ "value": false
465+ }
466+ ]
467+ },
468+ {
469+ "bin_filename": "ScatterSub_505fe7a4ead48684de27a9edd2dcdf35",
470+ "inputs": [
471+ {
472+ "name": "var",
473+ "index": 0,
474+ "dtype": "int8",
475+ "format": "ND",
476+ "paramType": "required",
477+ "shape": [
478+ -2
479+ ],
480+ "format_match_mode": "FormatAgnostic"
481+ },
482+ {
483+ "name": "indices",
484+ "index": 1,
485+ "dtype": "int32",
486+ "format": "ND",
487+ "paramType": "required",
488+ "shape": [
489+ -2
490+ ],
491+ "format_match_mode": "FormatAgnostic"
492+ },
493+ {
494+ "name": "updates",
495+ "index": 2,
496+ "dtype": "int8",
497+ "format": "ND",
498+ "paramType": "required",
499+ "shape": [
500+ -2
501+ ],
502+ "format_match_mode": "FormatAgnostic"
503+ }
504+ ],
505+ "outputs": [
506+ {
507+ "name": "var",
508+ "index": 0,
509+ "dtype": "int8",
510+ "format": "ND",
511+ "paramType": "required",
512+ "shape": [
513+ -2
514+ ],
515+ "format_match_mode": "FormatAgnostic"
516+ }
517+ ],
518+ "attrs": [
519+ {
520+ "name": "use_locking",
521+ "dtype": "bool",
522+ "value": false
523+ }
524+ ]
525+ },
526+ {
527+ "bin_filename": "ScatterSub_845bdcd8aa42597f2bd3428c633241c7",
528+ "inputs": [
529+ {
530+ "name": "var",
531+ "index": 0,
532+ "dtype": "int8",
533+ "format": "ND",
534+ "paramType": "required",
535+ "shape": [
536+ -2
537+ ],
538+ "format_match_mode": "FormatAgnostic"
539+ },
540+ {
541+ "name": "indices",
542+ "index": 1,
543+ "dtype": "int64",
544+ "format": "ND",
545+ "paramType": "required",
546+ "shape": [
547+ -2
548+ ],
549+ "format_match_mode": "FormatAgnostic"
550+ },
551+ {
552+ "name": "updates",
553+ "index": 2,
554+ "dtype": "int8",
555+ "format": "ND",
556+ "paramType": "required",
557+ "shape": [
558+ -2
559+ ],
560+ "format_match_mode": "FormatAgnostic"
561+ }
562+ ],
563+ "outputs": [
564+ {
565+ "name": "var",
566+ "index": 0,
567+ "dtype": "int8",
568+ "format": "ND",
569+ "paramType": "required",
570+ "shape": [
571+ -2
572+ ],
573+ "format_match_mode": "FormatAgnostic"
574+ }
575+ ],
576+ "attrs": [
577+ {
578+ "name": "use_locking",
579+ "dtype": "bool",
580+ "value": false
581+ }
582+ ]
583+ },
584+ {
585+ "bin_filename": "ScatterSub_f1ee7ccc7a4c25bc333f873dd4bc4197",
586+ "inputs": [
587+ {
588+ "name": "var",
589+ "index": 0,
590+ "dtype": "uint8",
591+ "format": "ND",
592+ "paramType": "required",
593+ "shape": [
594+ -2
595+ ],
596+ "format_match_mode": "FormatAgnostic"
597+ },
598+ {
599+ "name": "indices",
600+ "index": 1,
601+ "dtype": "int32",
602+ "format": "ND",
603+ "paramType": "required",
604+ "shape": [
605+ -2
606+ ],
607+ "format_match_mode": "FormatAgnostic"
608+ },
609+ {
610+ "name": "updates",
611+ "index": 2,
612+ "dtype": "uint8",
613+ "format": "ND",
614+ "paramType": "required",
615+ "shape": [
616+ -2
617+ ],
618+ "format_match_mode": "FormatAgnostic"
619+ }
620+ ],
621+ "outputs": [
622+ {
623+ "name": "var",
624+ "index": 0,
625+ "dtype": "uint8",
626+ "format": "ND",
627+ "paramType": "required",
628+ "shape": [
629+ -2
630+ ],
631+ "format_match_mode": "FormatAgnostic"
632+ }
633+ ],
634+ "attrs": [
635+ {
636+ "name": "use_locking",
637+ "dtype": "bool",
638+ "value": false
639+ }
640+ ]
641+ },
642+ {
643+ "bin_filename": "ScatterSub_8ab00e6fe6e5d8a874f662ee720738b3",
644+ "inputs": [
645+ {
646+ "name": "var",
647+ "index": 0,
648+ "dtype": "uint8",
649+ "format": "ND",
650+ "paramType": "required",
651+ "shape": [
652+ -2
653+ ],
654+ "format_match_mode": "FormatAgnostic"
655+ },
656+ {
657+ "name": "indices",
658+ "index": 1,
659+ "dtype": "int64",
660+ "format": "ND",
661+ "paramType": "required",
662+ "shape": [
663+ -2
664+ ],
665+ "format_match_mode": "FormatAgnostic"
666+ },
667+ {
668+ "name": "updates",
669+ "index": 2,
670+ "dtype": "uint8",
671+ "format": "ND",
672+ "paramType": "required",
673+ "shape": [
674+ -2
675+ ],
676+ "format_match_mode": "FormatAgnostic"
677+ }
678+ ],
679+ "outputs": [
680+ {
681+ "name": "var",
682+ "index": 0,
683+ "dtype": "uint8",
684+ "format": "ND",
685+ "paramType": "required",
686+ "shape": [
687+ -2
688+ ],
689+ "format_match_mode": "FormatAgnostic"
690+ }
691+ ],
692+ "attrs": [
693+ {
694+ "name": "use_locking",
695+ "dtype": "bool",
696+ "value": false
697+ }
698+ ]
699+ }
700+ ]
701+ }
Aindex/scatter_sub/op_host/config/ascend950/scatter_sub_simplified_key.ini+13-0
@@ -0,0 +1,13 @@
1+; 该文件主要影响 opc 工具 编译二进制kernel时, --simplified_key_mode 选项中填写的值,格式如下所示:
2+; [某算子]
3+; default=xx
4+; ascendxx=xx
5+; 其中,default为默认mode,ascnedxx为可选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+[ScatterSub]
13+default=0
Aindex/scatter_sub/op_host/scatter_sub_def.cpp+71-0
@@ -0,0 +1,71 @@
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+/*!
12+ * \file scatter_sub_def.cpp
13+ * \brief scatter_sub_def
14+ */
15+#include "register/op_def_registry.h"
16+ 
17+namespace ops
18+{
19+class ScatterSub : public OpDef
20+{
21+public:
22+ explicit ScatterSub(const char* name) : OpDef(name)
23+ {
24+ this->Input("var")
25+ .ParamType(REQUIRED)
26+ .DataType({ge::DT_INT32, ge::DT_UINT8, ge::DT_INT8, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16,
27+ ge::DT_INT32, ge::DT_UINT8, ge::DT_INT8, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
28+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
29+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
30+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
31+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
32+ this->Input("indices")
33+ .ParamType(REQUIRED)
34+ .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
35+ ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
36+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
37+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
38+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
39+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
40+ this->Input("updates")
41+ .ParamType(REQUIRED)
42+ .DataType({ge::DT_INT32, ge::DT_UINT8, ge::DT_INT8, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16,
43+ ge::DT_INT32, ge::DT_UINT8, ge::DT_INT8, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
44+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
45+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
46+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
47+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
48+ this->Output("var")
49+ .ParamType(REQUIRED)
50+ .DataType({ge::DT_INT32, ge::DT_UINT8, ge::DT_INT8, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16,
51+ ge::DT_INT32, ge::DT_UINT8, ge::DT_INT8, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
52+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
53+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
54+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
55+ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
56+ 
57+ this->Attr("use_locking").AttrType(OPTIONAL).Bool(false);
58+
59+ OpAICoreConfig aicoreConfig;
60+ aicoreConfig.DynamicCompileStaticFlag(true)
61+ .DynamicFormatFlag(false)
62+ .DynamicRankSupportFlag(true)
63+ .DynamicShapeSupportFlag(true)
64+ .NeedCheckSupportFlag(false)
65+ .ExtendCfgInfo("opFile.value", "scatter_sub_apt");
66+ this->AICore().AddConfig("ascend950", aicoreConfig);
67+ }
68+};
69+ 
70+OP_ADD(ScatterSub);
71+} // namespace ops
Aindex/scatter_sub/op_kernel/scatter_sub_apt.cpp+432-0
@@ -0,0 +1,432 @@
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+/*!
12+ * \file scatter_sub.cpp
13+ * \brief scatter_sub
14+ */
15+ 
16+ #include "kernel_operator.h"
17+ #include "../scatter_add/arch35/scatter_add_simt.h"
18+ #include "../scatter_add/arch35/scatter_add_simd.h"
19+ #include "../scatter_add/arch35/scatter_add_simd_support_atomicadd.h"
20+ #include "../scatter_add/arch35/scatter_add_deterministic.h"
21+ #include "../scatter_add/arch35/scatter_add_simt_sort.h"
22+ #include "../scatter_add/arch35/scatter_add_simd_sort_support_atomicadd.h"
23+
24+ using namespace AscendC;
25+ 
26+ #define TILING_KEY_0 0
CANN-robot
CANN-robotCANN-robot1月29日

代码结构与可维护性: 魔数(Magic Numbers)的使用。第26-38行定义了一系列TILING_KEY常量,这些值是硬编码的非常大的无符号长整型。这些数字没有明显的含义,难以理解和维护。

问题类型: 代码结构与可维护性 文件路径: index/scatter_sub/op_kernel/scatter_sub_apt.cpp 行号: 26 问题代码:

#define TILING_KEY_0                                0
#define TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR        10000000000333331000UL
#define TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR        10000000000333330000UL
#define TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR        10000000000333331100UL
#define TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR        10000000000333330100UL
#define TILING_KEY_SORT_SIMT_ADDR32_SCALAR          10000000000333331001UL
#define TILING_KEY_SORT_SIMT_ADDR32_TENSOR          10000000000333330001UL
#define TILING_KEY_SORT_SIMT_ADDR64_SCALAR          10000000000333331101UL
#define TILING_KEY_SORT_SIMT_ADDR64_TENSOR          10000000000333330101UL
#define TILING_KEY_UNSORT_SIMD_SCALAR               10000000000333331010UL
#define TILING_KEY_UNSORT_SIMD_TENSOR               10000000000333330010UL
#define TILING_KEY_SORT_SIMD_SCALAR                 10000000000333331011UL
#define TILING_KEY_SORT_SIMD_TENSOR                 10000000000333330011UL

修改建议:

为这些常量添加注释,解释每个字段的含义(例如,哪些位表示排序、SIMT/SIMD、地址大小、标量/张量等)。如果可能,考虑使用枚举或constexpr常量,并赋予更有意义的名字。或者,使用一个结构体或位域来编码这些信息,而不是使用魔数。

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

likedislike
27+ #define TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR 10000000000333331000UL
28+ #define TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR 10000000000333330000UL
29+ #define TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR 10000000000333331100UL
30+ #define TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR 10000000000333330100UL
31+ #define TILING_KEY_SORT_SIMT_ADDR32_SCALAR 10000000000333331001UL
32+ #define TILING_KEY_SORT_SIMT_ADDR32_TENSOR 10000000000333330001UL
33+ #define TILING_KEY_SORT_SIMT_ADDR64_SCALAR 10000000000333331101UL
34+ #define TILING_KEY_SORT_SIMT_ADDR64_TENSOR 10000000000333330101UL
35+ #define TILING_KEY_UNSORT_SIMD_SCALAR 10000000000333331010UL
36+ #define TILING_KEY_UNSORT_SIMD_TENSOR 10000000000333330010UL
37+ #define TILING_KEY_SORT_SIMD_SCALAR 10000000000333331011UL
38+ #define TILING_KEY_SORT_SIMD_TENSOR 10000000000333330011UL
39+ 
40+ using namespace ScatterAdd;
41+ 
42+ __aicore__ inline void ScatterAddUnsortSimtSubr32Scalar(
43+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
44+{
45+ GET_TILING_DATA(tilingData, tiling);
46+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
47+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint32_t, true, SUB> op(tilingData, pipe);
48+ op.Init(var, indices, updates, userWs);
49+ op.Process();
50+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
51+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, true, SUB> op(tilingData, pipe);
52+ op.Init(var, indices, updates, userWs);
53+ op.Process();
54+ } else {
55+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint32_t, true, SUB> op(tilingData, pipe);
56+ op.Init(var, indices, updates, userWs);
57+ op.Process();
58+ }
59+}
60+ 
61+ __aicore__ inline void ScatterAddUnsortSimtSubr32Tensor(
62+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
63+{
64+ GET_TILING_DATA(tilingData, tiling);
65+ if (tilingData.isDeterminTemplate) {
66+ if constexpr (is_same<half, DTYPE_VAR>::value || is_same<float, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {
67+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, SUB> op(tilingData, pipe);
68+ op.Init(var, indices, updates, y, userWs);
69+ op.Process();
70+ }
71+ } else {
72+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
73+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint32_t, false, SUB> op(tilingData, pipe);
74+ op.Init(var, indices, updates, userWs);
75+ op.Process();
76+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
77+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, false, SUB> op(tilingData, pipe);
78+ op.Init(var, indices, updates, userWs);
79+ op.Process();
80+ } else {
81+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint32_t, false, SUB> op(tilingData, pipe);
82+ op.Init(var, indices, updates, userWs);
83+ op.Process();
84+ }
85+ }
86+}
87+ 
88+ __aicore__ inline void ScatterAddUnsortSimtSubr64Scalar(
89+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
90+{
91+ GET_TILING_DATA(tilingData, tiling);
92+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
93+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint64_t, true, SUB> op(tilingData, pipe);
94+ op.Init(var, indices, updates, userWs);
95+ op.Process();
96+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
97+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, true, SUB> op(tilingData, pipe);
98+ op.Init(var, indices, updates, userWs);
99+ op.Process();
100+ } else {
101+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint64_t, true, SUB> op(tilingData, pipe);
102+ op.Init(var, indices, updates, userWs);
103+ op.Process();
104+ }
105+}
106+ 
107+ __aicore__ inline void ScatterAddUnsortSimtSubr64Tensor(
108+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
109+{
110+ GET_TILING_DATA(tilingData, tiling);
111+ if (tilingData.isDeterminTemplate) {
112+ if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {
113+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, SUB> op(tilingData, pipe);
114+ op.Init(var, indices, updates, y, userWs);
115+ op.Process();
116+ }
117+ } else {
118+ if constexpr (is_same<uint8_t, DTYPE_VAR>::value) {
119+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, uint32_t, uint64_t, false, SUB> op(tilingData, pipe);
120+ op.Init(var, indices, updates, userWs);
121+ op.Process();
122+ } else if constexpr(is_same<int8_t, DTYPE_VAR>::value) {
123+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, false, SUB> op(tilingData, pipe);
124+ op.Init(var, indices, updates, userWs);
125+ op.Process();
126+ } else {
127+ ScatterAddSimt<DTYPE_INDICES, DTYPE_VAR, DTYPE_VAR, uint64_t, false, SUB> op(tilingData, pipe);
128+ op.Init(var, indices, updates, userWs);
129+ op.Process();
130+ }
131+ }
132+}
133+ 
134+ __aicore__ inline void ScatterAddSortSimtSubr32Scalar(
135+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
136+{
137+ GET_TILING_DATA(tilingData, tiling);
138+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
139+ return;
140+ } else {
141+ if (tilingData.indicesCastMode == CAST_0) {
142+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint32_t, true, CAST_0, SUB> op(tilingData, pipe);
143+ op.Init(var, indices, updates, userWs);
144+ op.Process();
145+ } else if (tilingData.indicesCastMode == CAST_1) {
146+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, true, CAST_1, SUB> op(tilingData, pipe);
147+ op.Init(var, indices, updates, userWs);
148+ op.Process();
149+ } else if (tilingData.indicesCastMode == CAST_2) {
150+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, true, CAST_2, SUB> op(tilingData, pipe);
151+ op.Init(var, indices, updates, userWs);
152+ op.Process();
153+ } else if (tilingData.indicesCastMode == CAST_3) {
154+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, true, CAST_3, SUB> op(tilingData, pipe);
155+ op.Init(var, indices, updates, userWs);
156+ op.Process();
157+ } else if (tilingData.indicesCastMode == CAST_4) {
158+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, true, CAST_4, SUB> op(tilingData, pipe);
159+ op.Init(var, indices, updates, userWs);
160+ op.Process();
161+ } else if (tilingData.indicesCastMode == CAST_5) {
162+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, true, CAST_5, SUB> op(tilingData, pipe);
163+ op.Init(var, indices, updates, userWs);
164+ op.Process();
165+ }
166+ }
167+}
168+ 
169+ __aicore__ inline void ScatterAddSortSimtSubr32Tensor(
170+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
171+{
172+ GET_TILING_DATA(tilingData, tiling);
173+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
174+ return;
175+ } else {
176+ if (tilingData.indicesCastMode == CAST_0) {
177+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint32_t, false, CAST_0, SUB> op(tilingData, pipe);
178+ op.Init(var, indices, updates, userWs);
179+ op.Process();
180+ } else if (tilingData.indicesCastMode == CAST_1) {
181+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, false, CAST_1, SUB> op(tilingData, pipe);
182+ op.Init(var, indices, updates, userWs);
183+ op.Process();
184+ } else if (tilingData.indicesCastMode == CAST_2) {
185+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint32_t, false, CAST_2, SUB> op(tilingData, pipe);
186+ op.Init(var, indices, updates, userWs);
187+ op.Process();
188+ } else if (tilingData.indicesCastMode == CAST_3) {
189+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint32_t, false, CAST_3, SUB> op(tilingData, pipe);
190+ op.Init(var, indices, updates, userWs);
191+ op.Process();
192+ } else if (tilingData.indicesCastMode == CAST_4) {
193+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, false, CAST_4, SUB> op(tilingData, pipe);
194+ op.Init(var, indices, updates, userWs);
195+ op.Process();
196+ } else if (tilingData.indicesCastMode == CAST_5) {
197+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint32_t, false, CAST_5, SUB> op(tilingData, pipe);
198+ op.Init(var, indices, updates, userWs);
199+ op.Process();
200+ }
201+ }
202+}
203+ 
204+ __aicore__ inline void ScatterAddSortSimtSubr64Scalar(
205+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
206+{
207+ GET_TILING_DATA(tilingData, tiling);
208+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
209+ return;
210+ } else {
211+ if (tilingData.indicesCastMode == CAST_0) {
212+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint64_t, true, CAST_0, SUB> op(tilingData, pipe);
213+ op.Init(var, indices, updates, userWs);
214+ op.Process();
215+ } else if (tilingData.indicesCastMode == CAST_1) {
216+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, true, CAST_1, SUB> op(tilingData, pipe);
217+ op.Init(var, indices, updates, userWs);
218+ op.Process();
219+ } else if (tilingData.indicesCastMode == CAST_2) {
220+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, true, CAST_2, SUB> op(tilingData, pipe);
221+ op.Init(var, indices, updates, userWs);
222+ op.Process();
223+ } else if (tilingData.indicesCastMode == CAST_3) {
224+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, true, CAST_3, SUB> op(tilingData, pipe);
225+ op.Init(var, indices, updates, userWs);
226+ op.Process();
227+ } else if (tilingData.indicesCastMode == CAST_4) {
228+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, true, CAST_4, SUB> op(tilingData, pipe);
229+ op.Init(var, indices, updates, userWs);
230+ op.Process();
231+ } else if (tilingData.indicesCastMode == CAST_5) {
232+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, true, CAST_5, SUB> op(tilingData, pipe);
233+ op.Init(var, indices, updates, userWs);
234+ op.Process();
235+ }
236+ }
237+}
238+ 
239+ __aicore__ inline void ScatterAddSortSimtSubr64Tensor(
240+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
241+{
242+ GET_TILING_DATA(tilingData, tiling);
243+ if constexpr (is_same<int8_t, DTYPE_VAR>::value || is_same<uint8_t, DTYPE_VAR>::value) {
244+ return;
245+ } else {
246+ if (tilingData.indicesCastMode == CAST_0) {
247+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, DTYPE_INDICES, uint64_t, false, CAST_0, SUB> op(tilingData, pipe);
248+ op.Init(var, indices, updates, userWs);
249+ op.Process();
250+ } else if (tilingData.indicesCastMode == CAST_1) {
251+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, false, CAST_1, SUB> op(tilingData, pipe);
252+ op.Init(var, indices, updates, userWs);
253+ op.Process();
254+ } else if (tilingData.indicesCastMode == CAST_2) {
255+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int32_t, uint64_t, false, CAST_2, SUB> op(tilingData, pipe);
256+ op.Init(var, indices, updates, userWs);
257+ op.Process();
258+ } else if (tilingData.indicesCastMode == CAST_3) {
259+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, int16_t, uint64_t, false, CAST_3, SUB> op(tilingData, pipe);
260+ op.Init(var, indices, updates, userWs);
261+ op.Process();
262+ } else if (tilingData.indicesCastMode == CAST_4) {
263+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, false, CAST_4, SUB> op(tilingData, pipe);
264+ op.Init(var, indices, updates, userWs);
265+ op.Process();
266+ } else if (tilingData.indicesCastMode == CAST_5) {
267+ ScatterAddSimtSort<DTYPE_INDICES, DTYPE_VAR, uint8_t, uint64_t, false, CAST_5, SUB> op(tilingData, pipe);
268+ op.Init(var, indices, updates, userWs);
269+ op.Process();
270+ }
271+ }
272+}
273+ 
274+ __aicore__ inline void ScatterAddUnsortSimdScalar(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
275+ GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
276+{
277+ GET_TILING_DATA(tilingData, tiling);
278+ if (tilingData.isDeterminTemplate) {
279+ if constexpr (is_same<bfloat16_t, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<float, DTYPE_VAR>::value) {
280+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, SUB> op(tilingData, pipe);
281+ op.Init(var, indices, updates, y, userWs);
282+ op.Process();
283+ }
284+ } else {
285+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
286+ ScatterAddSIMDSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, true, SUB> op(tilingData, pipe);
287+ op.Init(var, indices, updates, y, userWs);
288+ op.Process();
289+ } else {
290+ ScatterAddSIMDImpl<DTYPE_VAR, DTYPE_INDICES, true, SUB> op(tilingData, pipe);
291+ op.Init(var, indices, updates, y, userWs);
292+ op.Process();
293+ }
294+ }
295+}
296+ 
297+ __aicore__ inline void ScatterAddUnsortSimdTensor(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
298+ GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
299+{
300+ GET_TILING_DATA(tilingData, tiling);
301+ if (tilingData.isDeterminTemplate) {
302+ if constexpr (is_same<float, DTYPE_VAR>::value || is_same<half, DTYPE_VAR>::value || is_same<bfloat16_t, DTYPE_VAR>::value) {
303+ ScatterAddDeterministicImpl<DTYPE_VAR, DTYPE_INDICES, SUB> op(tilingData, pipe);
304+ op.Init(var, indices, updates, y, userWs);
305+ op.Process();
306+ }
307+ } else {
308+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
309+ ScatterAddSIMDSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, false, SUB> op(tilingData, pipe);
310+ op.Init(var, indices, updates, y, userWs);
311+ op.Process();
312+ } else {
313+ ScatterAddSIMDImpl<DTYPE_VAR, DTYPE_INDICES, false, SUB> op(tilingData, pipe);
314+ op.Init(var, indices, updates, y, userWs);
315+ op.Process();
316+ }
317+ }
318+}
319+ 
320+ __aicore__ inline void ScatterAddSortSimdScalar(
321+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
322+{
323+ GET_TILING_DATA(tilingData, tiling);
324+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
325+ if (tilingData.indicesCastMode == CAST_0) {
326+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, DTYPE_INDICES, true, CAST_0, SUB> op(tilingData, pipe);
327+ op.Init(var, indices, updates, y, userWs);
328+ op.Process();
329+ } else if (tilingData.indicesCastMode == CAST_1) {
330+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, true, CAST_1, SUB> op(tilingData, pipe);
331+ op.Init(var, indices, updates, y, userWs);
332+ op.Process();
333+ } else if (tilingData.indicesCastMode == CAST_2) {
334+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int32_t, true, CAST_2, SUB> op(tilingData, pipe);
335+ op.Init(var, indices, updates, y, userWs);
336+ op.Process();
337+ } else if (tilingData.indicesCastMode == CAST_3) {
338+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, true, CAST_3, SUB> op(tilingData, pipe);
339+ op.Init(var, indices, updates, y, userWs);
340+ op.Process();
341+ } else if (tilingData.indicesCastMode == CAST_4) {
342+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, true, CAST_4, SUB> op(tilingData, pipe);
343+ op.Init(var, indices, updates, y, userWs);
344+ op.Process();
345+ } else if (tilingData.indicesCastMode == CAST_5) {
346+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, true, CAST_5, SUB> op(tilingData, pipe);
347+ op.Init(var, indices, updates, y, userWs);
348+ op.Process();
349+ }
350+ } else {
351+ return;
352+ }
353+}
354+ 
355+ __aicore__ inline void ScatterAddSortSimdTensor(
356+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, GM_ADDR userWs, GM_ADDR tiling, TPipe &pipe)
357+{
358+ GET_TILING_DATA(tilingData, tiling);
359+ if constexpr (platform::IsSupportAtomicAddTypeSIMD<DTYPE_VAR>()) {
360+ if (tilingData.indicesCastMode == CAST_0) {
361+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, DTYPE_INDICES, false, CAST_0, SUB> op(tilingData, pipe);
362+ op.Init(var, indices, updates, y, userWs);
363+ op.Process();
364+ } else if (tilingData.indicesCastMode == CAST_1) {
365+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, false, CAST_1, SUB> op(tilingData, pipe);
366+ op.Init(var, indices, updates, y, userWs);
367+ op.Process();
368+ } else if (tilingData.indicesCastMode == CAST_2) {
369+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int32_t, false, CAST_2, SUB> op(tilingData, pipe);
370+ op.Init(var, indices, updates, y, userWs);
371+ op.Process();
372+ } else if (tilingData.indicesCastMode == CAST_3) {
373+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, int16_t, false, CAST_3, SUB> op(tilingData, pipe);
374+ op.Init(var, indices, updates, y, userWs);
375+ op.Process();
376+ } else if (tilingData.indicesCastMode == CAST_4) {
377+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, false, CAST_4, SUB> op(tilingData, pipe);
378+ op.Init(var, indices, updates, y, userWs);
379+ op.Process();
380+ } else if (tilingData.indicesCastMode == CAST_5) {
381+ ScatterAddSIMDSortSupportAtomicAdd<DTYPE_VAR, DTYPE_INDICES, uint8_t, false, CAST_5, SUB> op(tilingData, pipe);
382+ op.Init(var, indices, updates, y, userWs);
383+ op.Process();
384+ }
385+ } else {
386+ return;
387+ }
388+}
389+ 
390+extern "C" __global__ __aicore__ void scatter_sub(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
391+ GM_ADDR workspace, GM_ADDR tiling)
392+{
393+ if (workspace == nullptr) {
394+ return;
395+ }
396+ SetSysWorkspace(workspace);
397+ GM_ADDR userWs = GetUserWorkspace(workspace);
398+ if (userWs == nullptr) {
399+ return;
400+ }
401+ GET_TILING_DATA(tilingData, tiling);
402+ TPipe pipe;
403+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
404+ 
405+ if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR)) {
406+ ScatterAddUnsortSimtSubr32Scalar(var, indices, updates, y, userWs, tiling, pipe);
407+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR)) {
408+ ScatterAddUnsortSimtSubr32Tensor(var, indices, updates, y, userWs, tiling, pipe);
409+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR)) {
410+ ScatterAddUnsortSimtSubr64Scalar(var, indices, updates, y, userWs, tiling, pipe);
411+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR)) {
412+ ScatterAddUnsortSimtSubr64Tensor(var, indices, updates, y, userWs, tiling, pipe);
413+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR32_SCALAR)) {
414+ ScatterAddSortSimtSubr32Scalar(var, indices, updates, y, userWs, tiling, pipe);
415+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR32_TENSOR)) {
416+ ScatterAddSortSimtSubr32Tensor(var, indices, updates, y, userWs, tiling, pipe);
417+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR64_SCALAR)) {
418+ ScatterAddSortSimtSubr64Scalar(var, indices, updates, y, userWs, tiling, pipe);
419+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMT_ADDR64_TENSOR)) {
420+ ScatterAddSortSimtSubr64Tensor(var, indices, updates, y, userWs, tiling, pipe);
421+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMD_SCALAR)) {
422+ ScatterAddUnsortSimdScalar(var, indices, updates, y, userWs, tiling, pipe);
423+ } else if (TILING_KEY_IS(TILING_KEY_UNSORT_SIMD_TENSOR)) {
424+ ScatterAddUnsortSimdTensor(var, indices, updates, y, userWs, tiling, pipe);
425+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMD_SCALAR)) {
426+ ScatterAddSortSimdScalar(var, indices, updates, y, userWs, tiling, pipe);
427+ } else if (TILING_KEY_IS(TILING_KEY_SORT_SIMD_TENSOR)) {
428+ ScatterAddSortSimdTensor(var, indices, updates, y, userWs, tiling, pipe);
429+ } else if (TILING_KEY_IS(TILING_KEY_0)) {
430+ return;
431+ }
432+}
Aindex/scatter_sub/tests/CMakeLists.txt+18-0
@@ -0,0 +1,18 @@
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+message(STATUS "=== Debug: start ops.index.scatter_sub.tests.CMakeLists.txt ")
12+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
13+message(STATUS "=== Debug: CURRENT_DIRS =${CURRENT_DIRS} ")
14+foreach(SUB_DIR ${CURRENT_DIRS})
15+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
16+ add_subdirectory(${SUB_DIR})
17+ endif()
18+endforeach()
Aindex/scatter_sub/tests/ut/CMakeLists.txt+17-0
@@ -0,0 +1,17 @@
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+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+message(STATUS "=== Debug: CURRENT_DIRS =${CURRENT_DIRS} ")
13+foreach(SUB_DIR ${CURRENT_DIRS})
14+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
15+ add_subdirectory(${SUB_DIR})
16+ endif()
17+endforeach()
Aindex/scatter_sub/tests/ut/op_host/CMakeLists.txt+15-0
@@ -0,0 +1,15 @@
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+file(GLOB CURRENT_DIR RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+if(UT_TEST_ALL OR OP_HOST_UT)
13+ add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14+ add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
15+endif()
Aindex/scatter_sub/tests/ut/op_host/arch35/test_scatter_sub_tiling.cpp+197-0
@@ -0,0 +1,197 @@
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+ /* !
12+ * \file test_scatter_sub_tiling.cpp
13+ * \brief the ut of scatter_sub_tiling
14+ */
15+ 
16+#include <iostream>
17+#include <fstream>
18+#include <vector>
19+#include "log/log.h"
20+#include <gtest/gtest.h>
21+#include "register/op_impl_registry.h"
22+#include "platform/platform_infos_def.h"
23+#include "ut_op_common.h"
24+#include "ut_op_util.h"
25+#include "../../../../op_host/arch35/scatter_sub_tiling.h"
26+#include "kernel_run_context_facker.h"
27+#include "test_cube_util.h"
28+#include "exe_graph/runtime/storage_format.h"
29+#include "exe_graph/runtime/storage_shape.h"
30+ 
31+using namespace std;
32+using namespace ge;
33+using namespace ut_util;
34+ 
35+class ScatterSubTiling : public testing::Test {
36+ protected:
37+ static void SetUpTestCase() {
38+ std::cout << "ScatterSubTiling SetUp" << std::endl;
39+ }
40+ 
41+ static void TearDownTestCase() {
42+ std::cout << "ScatterSubTiling TearDown" << std::endl;
43+ }
44+};
45+ 
46+static string to_string(const std::stringstream& tiling_data) {
47+ auto data = tiling_data.str();
48+ string result;
49+ int64_t tmp = 0;
50+ for (size_t i = 0; i < data.length(); i += sizeof(int64_t)) {
51+ memcpy(&tmp, data.c_str() + i, sizeof(tmp));
52+ result += std::to_string(tmp);
53+ result += " ";
54+ }
55+ 
56+ return result;
57+}
58+ 
59+struct ScatterSubOpsParamInfos {
60+ ge::DataType indiceDtype;
61+ ge::DataType varDtype;
62+ gert::StorageShape indiceShape;
63+ gert::StorageShape updateShape;
64+ gert::StorageShape varShape;
65+};
66+ 
67+static void ExecuteTestCase(const ScatterSubOpsParamInfos& scatterSubOpsParamInfos, int32_t deterministic,
68+ ge::graphStatus status = ge::GRAPH_SUCCESS)
69+{
70+ string compileInfoString = R"({
71+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
72+ "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
73+ "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
74+ "UB_SIZE": 253952, "L2_SIZE": 33554432, "L1_SIZE": 524288,
75+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
76+ "CORE_NUM": 64}
77+ })";
78+ map<string, string> socInfos;
79+ map<string, string> aicoreSpec;
80+ map<string, string> intrinsics;
81+ 
82+ GetPlatFormInfos(compileInfoString.c_str(), socInfos, aicoreSpec, intrinsics);
83+ 
84+ // platform info
85+ fe::PlatFormInfos platformInfo;
86+ platformInfo.Init();
87+ 
88+ // compile info
89+ optiling::ScatterSubCompileInfo compileInfo;
90+ compileInfo.core_num = 64;
91+ compileInfo.ub_size = 253952;
92+ 
93+ std::string opType("ScatterSub");
94+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(opType.c_str()), nullptr);
95+ auto tilingFunc = gert::OpImplRegistry::GetInstance().GetOpImpl(opType.c_str())->tiling;
96+ 
97+ // tilingFunc simulate
98+ auto param = gert::TilingData::CreateCap(4096);
99+ auto workspaceSizeHoler = gert::ContinuousVector::Create<size_t>(4096);
100+ auto wsSize = reinterpret_cast<gert::ContinuousVector *>(workspaceSizeHoler.get());
101+ ASSERT_NE(param, nullptr);
102+ gert::StorageShape indiceShape = scatterSubOpsParamInfos.indiceShape;
103+ gert::StorageShape updateShape = scatterSubOpsParamInfos.updateShape;
104+ gert::StorageShape varShape = scatterSubOpsParamInfos.varShape;
105+ auto holder = gert::TilingContextFaker()
106+ .NodeIoNum(3, 1)
107+ .IrInstanceNum({1, 1, 1})
108+ .InputShapes({&varShape, &indiceShape, &updateShape})
109+ .OutputShapes({&varShape})
110+ .CompileInfo(&compileInfo)
111+ .PlatformInfo(reinterpret_cast<char *>(&platformInfo))
112+ .NodeInputTd(0, scatterSubOpsParamInfos.varDtype, ge::FORMAT_ND, ge::FORMAT_ND)
113+ .NodeInputTd(1, scatterSubOpsParamInfos.indiceDtype, ge::FORMAT_ND, ge::FORMAT_ND)
114+ .NodeInputTd(2, scatterSubOpsParamInfos.varDtype, ge::FORMAT_ND, ge::FORMAT_ND)
115+ .NodeOutputTd(0, scatterSubOpsParamInfos.varDtype, ge::FORMAT_ND, ge::FORMAT_ND)
116+ .DeterministicInfo(reinterpret_cast<int32_t*>(deterministic))
117+ .TilingData(param.get())
118+ .Workspace(wsSize)
119+ .Build();
120+ gert::TilingContext* tilingContext = holder.GetContext<gert::TilingContext>();
121+ ASSERT_NE(tilingContext, nullptr);
122+ auto infos = tilingContext->GetPlatformInfo();
123+ ASSERT_NE(infos, nullptr);
124+ infos->SetPlatformRes("SoCInfo", socInfos);
125+ infos->SetPlatformRes("AICoreSpec", aicoreSpec);
126+ infos->SetCoreNumByCoreType("AICore");
127+ infos->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
128+ 
129+ // workspaces nullptr return failed
130+ EXPECT_EQ(tilingFunc(tilingContext), status);
131+}
132+ 
133+TEST_F(ScatterSubTiling, ScatterSub_tiling_ascendc_varUint8_indiceInt32_simd_notSupportAtomicAdd) {
134+ ScatterSubOpsParamInfos scatterSubOpsParamInfos;
135+ scatterSubOpsParamInfos.indiceDtype = ge::DT_INT32;
136+ scatterSubOpsParamInfos.varDtype = ge::DT_UINT8;
137+ scatterSubOpsParamInfos.indiceShape = {{65, 133}, {65, 133}};
138+ scatterSubOpsParamInfos.updateShape = {{65, 133, 1, 130}, {65, 133, 1, 130}};
139+ scatterSubOpsParamInfos.varShape = {{191, 1, 130}, {191, 1, 130}};
140+ 
141+ ExecuteTestCase(scatterSubOpsParamInfos, 0);
142+}
143+ 
144+TEST_F(ScatterSubTiling, ScatterSub_tiling_ascendc_varUint8_indiceInt64_simd_notSupportAtomicAdd) {
145+ ScatterSubOpsParamInfos scatterSubOpsParamInfos;
146+ scatterSubOpsParamInfos.indiceDtype = ge::DT_INT64;
147+ scatterSubOpsParamInfos.varDtype = ge::DT_UINT8;
148+ scatterSubOpsParamInfos.indiceShape = {{65, 133}, {65, 133}};
149+ scatterSubOpsParamInfos.updateShape = {{65, 133, 1, 130}, {65, 133, 1, 130}};
150+ scatterSubOpsParamInfos.varShape = {{191, 1, 130}, {191, 1, 130}};
151+ 
152+ ExecuteTestCase(scatterSubOpsParamInfos, 0);
153+}
154+ 
155+TEST_F(ScatterSubTiling, ScatterSub_tiling_ascendc_varUint8_indiceInt64_simd_exceed_ubBoundary) {
156+ ScatterSubOpsParamInfos scatterSubOpsParamInfos;
157+ scatterSubOpsParamInfos.indiceDtype = ge::DT_INT64;
158+ scatterSubOpsParamInfos.varDtype = ge::DT_UINT8;
159+ scatterSubOpsParamInfos.indiceShape = {{6}, {6}};
160+ scatterSubOpsParamInfos.updateShape = {{6, 3, 3, 3, 3, 3, 3, 137}, {6, 3, 3, 3, 3, 3, 3, 137}};
161+ scatterSubOpsParamInfos.varShape = {{3, 3, 3, 3, 3, 3, 3, 137}, {3, 3, 3, 3, 3, 3, 3, 137}};
162+ 
163+ ExecuteTestCase(scatterSubOpsParamInfos, 0);
164+}
165+ 
166+TEST_F(ScatterSubTiling, ScatterSub_tiling_ascendc_varUint8_indiceInt32_simt_notSupportAtomicAdd) {
167+ ScatterSubOpsParamInfos scatterSubOpsParamInfos;
168+ scatterSubOpsParamInfos.indiceDtype = ge::DT_INT32;
169+ scatterSubOpsParamInfos.varDtype = ge::DT_UINT8;
170+ scatterSubOpsParamInfos.indiceShape = {{65, 133}, {65, 133}};
171+ scatterSubOpsParamInfos.updateShape = {{}, {}};
172+ scatterSubOpsParamInfos.varShape = {{191, 1, 2}, {191, 1, 2}};
173+ 
174+ ExecuteTestCase(scatterSubOpsParamInfos, 0);
175+}
176+ 
177+TEST_F(ScatterSubTiling, ScatterSub_tiling_ascendc_varUint8_indiceInt64_simt_notSupportAtomicAdd) {
178+ ScatterSubOpsParamInfos scatterSubOpsParamInfos;
179+ scatterSubOpsParamInfos.indiceDtype = ge::DT_INT64;
180+ scatterSubOpsParamInfos.varDtype = ge::DT_UINT8;
181+ scatterSubOpsParamInfos.indiceShape = {{65, 133}, {65, 133}};
182+ scatterSubOpsParamInfos.updateShape = {{}, {}};
183+ scatterSubOpsParamInfos.varShape = {{191, 1, 2}, {191, 1, 2}};
184+ 
185+ ExecuteTestCase(scatterSubOpsParamInfos, 0);
186+}
187+ 
188+TEST_F(ScatterSubTiling, ScatterSub_tiling_ascendc_varFp32_indiceInt32_deterministic) {
189+ ScatterSubOpsParamInfos scatterSubOpsParamInfos;
190+ scatterSubOpsParamInfos.indiceDtype = ge::DT_INT32;
191+ scatterSubOpsParamInfos.varDtype = ge::DT_FLOAT;
192+ scatterSubOpsParamInfos.indiceShape = {{19, 31}, {19, 31}};
193+ scatterSubOpsParamInfos.updateShape = {{19, 31, 2, 9}, {19, 31, 2, 9}};
194+ scatterSubOpsParamInfos.varShape = {{129, 2, 9}, {129, 2, 9}};
195+ 
196+ ExecuteTestCase(scatterSubOpsParamInfos, 1);
197+}
Mscripts/kernel/binary_config/ascendc_config.json+1-0
@@ -514,6 +514,7 @@
514 {"name":"InplaceScatterAdd", "compute_units": ["ascend910b", "ascend910_93"], "auto_sync" : false},514 {"name":"InplaceScatterAdd", "compute_units": ["ascend910b", "ascend910_93"], "auto_sync" : false},
515 {"name":"InplaceIndexAdd", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},515 {"name":"InplaceIndexAdd", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},
516 {"name":"GemmV2", "compute_units": ["ascend910b", "ascend910_93", "kirinx90"], "auto_sync" : false},516 {"name":"GemmV2", "compute_units": ["ascend910b", "ascend910_93", "kirinx90"], "auto_sync" : false},
517+ {"name":"ScatterSub", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "high_precision", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},
517 {"name":"AscendQuant", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},518 {"name":"AscendQuant", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},
518 {"name":"Quantize", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},519 {"name":"Quantize", "compute_units": ["ascend950"], "auto_sync": false, "impl_mode": "", "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},
519 {"name":"AscendQuantV2", "compute_units": ["ascend310p", "ascend910b", "ascend910_93", "ascend950", "kirinx90"], "auto_sync": false, "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},520 {"name":"AscendQuantV2", "compute_units": ["ascend310p", "ascend910b", "ascend910_93", "ascend950", "kirinx90"], "auto_sync": false, "compile_options": {"ascend950": ["-mllvm -cce-aicore-dcci-before-kernel-end=false"]}},