已合并
新增算子ScatterSub #1164
gengxueru创建于 1月29日
新增算子ScatterSub #1164
已合并
共 27 个文件变更+2606-597
| @@ -58,7 +58,6 @@ TILING_DATA_FIELD_DEF(uint64_t, isDeterministic); | |||
| 58 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleCol); | 58 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleCol); |
| 59 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsLoopSize); | 59 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsLoopSize); |
| 60 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsTailUbFactor); | 60 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreColsTailUbFactor); |
| 61 | -TILING_DATA_FIELD_DEF(uint64_t, ubSize); | ||
| 62 | TILING_DATA_FIELD_DEF(uint64_t, perCoreHandleRows); //反量化按var[0]分核或indices分核 | 61 | TILING_DATA_FIELD_DEF(uint64_t, perCoreHandleRows); //反量化按var[0]分核或indices分核 |
| 63 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleRows); | 62 | TILING_DATA_FIELD_DEF(uint64_t, tailCoreHandleRows); |
| 64 | TILING_DATA_FIELD_DEF(uint64_t, rowsInUb); | 63 | TILING_DATA_FIELD_DEF(uint64_t, rowsInUb); |
| @@ -83,9 +82,11 @@ TILING_DATA_FIELD_DEF(uint64_t, normBlockCol); // 整核分块列数 | |||
| 83 | TILING_DATA_FIELD_DEF(uint64_t, tailBlockCol); // 列尾核分块列数 | 82 | TILING_DATA_FIELD_DEF(uint64_t, tailBlockCol); // 列尾核分块列数 |
| 84 | TILING_DATA_FIELD_DEF(uint64_t, ubFactorRow); // UB每次循环搬运的行数 | 83 | TILING_DATA_FIELD_DEF(uint64_t, ubFactorRow); // UB每次循环搬运的行数 |
| 85 | TILING_DATA_FIELD_DEF(uint64_t, ubFactorCol); // UB每次循环搬运的列数 | 84 | TILING_DATA_FIELD_DEF(uint64_t, ubFactorCol); // UB每次循环搬运的列数 |
| 85 | +TILING_DATA_FIELD_DEF(uint64_t, indicesCastMode); | ||
| 86 | END_TILING_DATA_DEF; | 86 | END_TILING_DATA_DEF; |
| 87 | 87 | ||
| 88 | REGISTER_TILING_DATA_CLASS(ScatterAdd, ScatterAddTilingData) | 88 | REGISTER_TILING_DATA_CLASS(ScatterAdd, ScatterAddTilingData) |
| 89 | +REGISTER_TILING_DATA_CLASS(ScatterSub, ScatterAddTilingData) | ||
| 89 | 90 | ||
| 90 | ge::graphStatus ScatterAddTilingForAscendC(gert::TilingContext* context); | 91 | ge::graphStatus ScatterAddTilingForAscendC(gert::TilingContext* context); |
| 91 | 92 | ||
| @@ -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 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +extern "C" { | ||
| 18 | + | ||
| 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 | + | ||
| 58 | +} | ||
| 59 | + | ||
| 60 | + | ||
| 61 | + | ||
| @@ -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 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 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 | + | ||
| @@ -28,10 +28,23 @@ constexpr uint64_t UB_AGLIN_VALUE = 32; | |||
| 28 | constexpr uint64_t SORT_PAD_NUM = 2; | 28 | constexpr uint64_t SORT_PAD_NUM = 2; |
| 29 | constexpr uint64_t HASH_BUCKER_BUFFER_SIZE = 128 * sizeof(float); | 29 | constexpr uint64_t HASH_BUCKER_BUFFER_SIZE = 128 * sizeof(float); |
| 30 | constexpr int64_t DOUBLE_BUFFER = 2; | 30 | constexpr int64_t DOUBLE_BUFFER = 2; |
| 31 | +constexpr uint32_t TWO = 2; | ||
| 32 | +constexpr uint32_t THREE = 3; | ||
| 33 | +constexpr uint32_t FOUR = 4; | ||
| 31 | constexpr uint32_t CAST_0 = 0; | 34 | constexpr uint32_t CAST_0 = 0; |
| 32 | constexpr uint32_t CAST_1 = 1; | 35 | constexpr uint32_t CAST_1 = 1; |
| 33 | constexpr uint32_t CAST_2 = 2; | 36 | constexpr uint32_t CAST_2 = 2; |
| 34 | constexpr uint32_t CAST_3 = 3; | 37 | constexpr 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 | ||
| 36 | constexpr SortConfig sortConfig{SortType::RADIX_SORT, false}; | 49 | constexpr SortConfig sortConfig{SortType::RADIX_SORT, false}; |
| 37 | static constexpr MicroAPI::CastTrait castTraitU82Int32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, | 50 | static 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 | + | ||
| 74 | template <typename T> | 132 | template <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 | + | ||
| 113 | template<typename IDX_T> | 302 | template<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; |
| @@ -19,9 +19,11 @@ | |||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | 21 | ||
| 22 | + | ||
| 22 | 23 | ||
| 23 | namespace ScatterAdd { | 24 | namespace ScatterAdd { |
| 24 | using namespace AscendC; | 25 | using namespace AscendC; |
| 26 | +using namespace ScatterAddCommon; | ||
| 25 | 27 | ||
| 26 | constexpr uint64_t DOUBLE_BUF = 2; | 28 | constexpr uint64_t DOUBLE_BUF = 2; |
| 27 | constexpr uint64_t BLOCK_SIZE = 32; | 29 | constexpr uint64_t BLOCK_SIZE = 32; |
| @@ -48,7 +50,7 @@ AscendC::RoundMode::UNKNOWN, | |||
| 48 | 50 | ||
| 49 | static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false}; | 51 | static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false}; |
| 50 | 52 | ||
| 51 | -template<typename T, typename U> | 53 | +template<typename T, typename U, uint32_t scatterOp> |
| 52 | class ScatterAddDeterministicImpl { | 54 | class ScatterAddDeterministicImpl { |
| 53 | public: | 55 | public: |
| 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; |
| @@ -34,7 +34,7 @@ struct UpdatesDstPosSelector { constexpr static TPosition pos = TPosition::VECIN | |||
| 34 | template<> struct UpdatesDstPosSelector<true, false> { constexpr static TPosition pos = TPosition::VECOUT; }; // 支持atomic_add且updates不是标量,直接搬入搬出 | 34 | template<> struct UpdatesDstPosSelector<true, false> { constexpr static TPosition pos = TPosition::VECOUT; }; // 支持atomic_add且updates不是标量,直接搬入搬出 |
| 35 | template<> struct UpdatesDstPosSelector<true, true> { constexpr static TPosition pos = TPosition::GM; }; // 支持atomic_add且updates是标量,duplicate后搬出 | 35 | template<> 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> |
| 38 | class ScatterAddSIMDImpl { | 38 | class ScatterAddSIMDImpl { |
| 39 | public: | 39 | public: |
| 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; |
| @@ -24,7 +24,7 @@ using namespace ScatterAddCommon; | |||
| 24 | 24 | ||
| 25 | constexpr uint32_t SORT_PADDING = 64; | 25 | constexpr 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> |
| 28 | class ScatterAddSIMDSortSupportAtomicAdd { | 28 | class ScatterAddSIMDSortSupportAtomicAdd { |
| 29 | public: | 29 | public: |
| 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的重复度存放于uniqueIdCountLocal | 103 | 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; |
| @@ -22,7 +22,7 @@ namespace ScatterAdd { | |||
| 22 | using namespace AscendC; | 22 | using namespace AscendC; |
| 23 | using namespace ScatterAddCommon; | 23 | using namespace ScatterAddCommon; |
| 24 | 24 | ||
| 25 | -template<typename T, typename U, bool updatesIsScalar> | 25 | +template<typename T, typename U, bool updatesIsScalar, uint32_t scatterOp> |
| 26 | class ScatterAddSIMDSupportAtomicAdd { | 26 | class ScatterAddSIMDSupportAtomicAdd { |
| 27 | public: | 27 | public: |
| 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 | ||
| @@ -18,10 +18,12 @@ | |||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | + | ||
| 21 | 22 | ||
| 22 | namespace ScatterAdd | 23 | namespace ScatterAdd |
| 23 | { | 24 | { |
| 24 | using namespace AscendC; | 25 | using namespace AscendC; |
| 26 | +using namespace ScatterAddCommon; | ||
| 25 | 27 | ||
| 26 | constexpr uint32_t VECTOR_LENGTH = platform::GetVRegSize(); | 28 | constexpr uint32_t VECTOR_LENGTH = platform::GetVRegSize(); |
| 27 | constexpr uint32_t VL_B32 = VECTOR_LENGTH / sizeof(uint32_t); | 29 | constexpr uint32_t VL_B32 = VECTOR_LENGTH / sizeof(uint32_t); |
| @@ -38,7 +40,7 @@ constexpr uint32_t THREAD_NUM_LAUNCH_BOUND = 1024; | |||
| 38 | static constexpr MicroAPI::CastTrait castTraitB8B162B32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, | 40 | static 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> |
| 42 | class ScatterAddSimt | 44 | class ScatterAddSimt |
| 43 | { | 45 | { |
| 44 | public: | 46 | public: |
| @@ -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_); |
| @@ -18,6 +18,7 @@ | |||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | + | ||
| 21 | 22 | ||
| 22 | namespace ScatterAdd | 23 | namespace ScatterAdd |
| 23 | { | 24 | { |
| @@ -31,7 +32,7 @@ constexpr uint32_t THREAD_NUM_SORT = 256; | |||
| 31 | constexpr uint32_t THREAD_NUM_SORT = 1024; | 32 | constexpr uint32_t THREAD_NUM_SORT = 1024; |
| 32 | 33 | ||
| 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> |
| 35 | class ScatterAddSimtSort | 36 | class ScatterAddSimtSort |
| 36 | { | 37 | { |
| 37 | public: | 38 | public: |
| @@ -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]; |
| @@ -24,348 +24,409 @@ | |||
| 24 | using namespace AscendC; | 24 | using namespace AscendC; |
| 25 | 25 | ||
| 26 | 26 | ||
| 27 | - #define TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR 10000000003333301000UL | 27 | + #define TILING_KEY_UNSORT_SIMT_ADDR32_SCALAR 10000000000333331000UL |
| 28 | - #define TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR 10000000003333300000UL | 28 | + #define TILING_KEY_UNSORT_SIMT_ADDR32_TENSOR 10000000000333330000UL |
| 29 | - #define TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR 10000000003333301100UL | 29 | + #define TILING_KEY_UNSORT_SIMT_ADDR64_SCALAR 10000000000333331100UL |
| 30 | - #define TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR 10000000003333300100UL | 30 | + #define TILING_KEY_UNSORT_SIMT_ADDR64_TENSOR 10000000000333330100UL |
| 31 | - #define TILING_KEY_SORT_NOCAST_SIMT_ADDR32_SCALAR 10000000003333301001UL | 31 | + #define TILING_KEY_SORT_SIMT_ADDR32_SCALAR 10000000000333331001UL |
| 32 | - #define TILING_KEY_SORT_NOCAST_SIMT_ADDR32_TENSOR 10000000003333300001UL | 32 | + #define TILING_KEY_SORT_SIMT_ADDR32_TENSOR 10000000000333330001UL |
| 33 | - #define TILING_KEY_SORT_NOCAST_SIMT_ADDR64_SCALAR 10000000003333301101UL | 33 | + #define TILING_KEY_SORT_SIMT_ADDR64_SCALAR 10000000000333331101UL |
| 34 | - #define TILING_KEY_SORT_NOCAST_SIMT_ADDR64_TENSOR 10000000003333300101UL | 34 | + #define TILING_KEY_SORT_SIMT_ADDR64_TENSOR 10000000000333330101UL |
| 35 | - #define TILING_KEY_SORT_CAST1_SIMT_ADDR32_SCALAR 10000000003333311001UL | 35 | + #define TILING_KEY_UNSORT_SIMD_SCALAR 10000000000333331010UL |
| 36 | - #define TILING_KEY_SORT_CAST1_SIMT_ADDR32_TENSOR 10000000003333310001UL | 36 | + #define TILING_KEY_UNSORT_SIMD_TENSOR 10000000000333330010UL |
| 37 | - #define TILING_KEY_SORT_CAST1_SIMT_ADDR64_SCALAR 10000000003333311101UL | 37 | + #define TILING_KEY_SORT_SIMD_SCALAR 10000000000333331011UL |
| 38 | - #define TILING_KEY_SORT_CAST1_SIMT_ADDR64_TENSOR 10000000003333310101UL | 38 | + #define TILING_KEY_SORT_SIMD_TENSOR 10000000000333330011UL |
| 39 | - | ||
| 40 | - | ||
| 41 | - | ||
| 42 | - | ||
| 43 | - | ||
| 44 | - | ||
| 45 | - | ||
| 46 | - | ||
| 47 | - | ||
| 48 | - | ||
| 49 | - | ||
| 50 | - | ||
| 51 | - | ||
| 52 | - | ||
| 53 | - | ||
| 54 | - | ||
| 55 | - | ||
| 56 | - | ||
| 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 | +} | ||
| @@ -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 | // 空tensor | 54 | // 空tensor |
| 36 | TEST_F(l2_scatter_add_test, case_2) | 55 | TEST_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 | +} | ||
| @@ -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) | ||
| @@ -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 | + | ||
| 16 | + | ||
| 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 | ||
| @@ -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 | +/*! | ||
| 12 | + * \file scatter_sub_tiling.h | ||
| 13 | + * \brief scatter_sub_tiling | ||
| 14 | + */ | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 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 | + | ||
| @@ -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 | + } | ||
| @@ -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 | ||
| @@ -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 | + | ||
| 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 | ||
| @@ -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() | ||
| @@ -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() | ||
| @@ -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() | ||
| @@ -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 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 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 | +} | ||
| @@ -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"]}}, |


代码结构与可维护性: 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 问题代码:修改建议:
此评论由代码审查工具自动生成