Star914
1830
代码
代码
Issues226
Pull Requests316
Actions
项目讨论
Wiki
分析
Star914
1830
已合并
新增ScatterNdMax/ScatterNdMin算子 #3899
z30075199创建于 4月16日
新增ScatterNdMax/ScatterNdMin算子 #3899
已合并
z30075199创建于 4月16日
从z30075199/ops-nn_zxscatterNdM合入到cann/ops-nnmaster
共 40 个文件变更+6736-1
Mdocs/zh/op_list.md+20-0
@@ -1846,6 +1846,26 @@
1846 <td>AI Core</td>1846 <td>AI Core</td>
1847 <td>该算子暂无Ascend C代码实现，欢迎开发者补充贡献，贡献方式参考<a href="../../CONTRIBUTING.md">贡献指南</a>。</td>1847 <td>该算子暂无Ascend C代码实现，欢迎开发者补充贡献，贡献方式参考<a href="../../CONTRIBUTING.md">贡献指南</a>。</td>
1848 </tr>1848 </tr>
1849+ <tr>
1850+ <td>index</td>
1851+ <td><a href="../../index/scatter_nd_max/README.md">scatter_nd_max</a></td>
1852+ <td>✓</td>
1853+ <td>✓</td>
1854+ <td>✗</td>
1855+ <td>✓</td>
1856+ <td>AI Core</td>
1857+ <td>算子功能：根据indices在给定变量内，在updates和单个值或切片之间求最大值。</td>
1858+ </tr>
1859+ <tr>
1860+ <td>index</td>
1861+ <td><a href="../../index/scatter_nd_min/README.md">scatter_nd_min</a></td>
1862+ <td>✓</td>
1863+ <td>✓</td>
1864+ <td>✗</td>
1865+ <td>✓</td>
1866+ <td>AI Core</td>
1867+ <td>算子功能：根据indices在给定变量内，在updates和单个值或切片之间求最小值。</td>
1868+ </tr>
1849 <tr>1869 <tr>
1850 <td>index</td>1870 <td>index</td>
1851 <td><a href="../../index/scatter_nd_update/README.md">scatter_nd_update</a></td>1871 <td><a href="../../index/scatter_nd_update/README.md">scatter_nd_update</a></td>
Mindex/gather_v2/op_graph/fusion_pass/gather_fusion_pass.cpp+1-1
@@ -48,7 +48,7 @@ bool IsTargetPlatform()
48 PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo) != SUCCESS,48 PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo) != SUCCESS,
49 false, PASS_NAME.c_str(), "Get platformInfo failed.");49 false, PASS_NAME.c_str(), "Get platformInfo failed.");
50 const std::string soc = platformInfo.str_info.short_soc_version;50 const std::string soc = platformInfo.str_info.short_soc_version;
51- bool isSupported = (soc == "Ascend950");51+ bool isSupported = (soc == "Ascend950" || soc == "MC62CM12A");
52 if (!isSupported) {52 if (!isSupported) {
53 OPS_LOG_D(PASS_NAME.c_str(), "Platform %s is not supported.", soc.c_str());53 OPS_LOG_D(PASS_NAME.c_str(), "Platform %s is not supported.", soc.c_str());
54 return false;54 return false;
Aindex/scatter_nd_common/CMakeLists.txt+15-0
@@ -0,0 +1,15 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+# 设置算子定义时支持的芯片类型
12+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} COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
Aindex/scatter_nd_common/op_host/arch35/scatter_nd_common_base_tiling.cpp+228-0
@@ -0,0 +1,228 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_base_tiling.cpp
13+ * \brief scatter_nd_common_base_tiling
14+ */
15+ 
16+#include "scatter_nd_common_base_tiling.h"
17+ 
18+using namespace AscendC;
19+using namespace ge;
20+ 
21+namespace optiling {
22+constexpr uint64_t ASCENDC_WORKSPACE = static_cast<uint64_t>(16) * 1024 * 1024;
23+static constexpr uint64_t CASTMODE1 = 1; // int32 Cast int16
24+static constexpr uint64_t CASTMODE2 = 2; // int64 Cast int32
25+static constexpr uint64_t CASTMODE3 = 3; // int64 Cast int16
26+static constexpr uint64_t CASTMODE4 = 4; // int32 Cast uint8
27+static constexpr uint64_t CASTMODE5 = 5; // int64 Cast uint8
28+ 
29+static constexpr uint16_t RANK_MIN_VALUE = 1;
30+static constexpr uint16_t RANK_MAX_VALUE = 7;
31+ 
32+ 
33+ge::graphStatus ScatterNdCommonBaseTiling::GetCastType()
34+{
35+ indiceCastDtype_ = indiceDtype_;
36+ 
37+ if (indiceDtype_ == ge::DT_INT32) {
38+ if (varInAxis_ < UINT8_MAX) {
39+ indiceCastMode_ = CASTMODE4; // int32 Cast uint8
40+ indiceCastDtype_ = ge::DT_UINT8;
41+ } else if (varInAxis_ < INT16_MAX) {
42+ indiceCastMode_ = CASTMODE1; // int32 Cast int16
43+ indiceCastDtype_ = ge::DT_INT16;
44+ }
45+ } else {
46+ if (varInAxis_ < UINT8_MAX) {
47+ indiceCastMode_ = CASTMODE5; // int64 Cast uint8
48+ indiceCastDtype_ = ge::DT_UINT8;
49+ } else if (varInAxis_ < INT16_MAX) {
50+ indiceCastMode_ = CASTMODE3; // int64 Cast int16
51+ indiceCastDtype_ = ge::DT_INT16;
52+ } else if (varInAxis_ < INT32_MAX) {
53+ indiceCastMode_ = CASTMODE2; // int64 Cast int32
54+ indiceCastDtype_ = ge::DT_INT32;
55+ }
56+ }
57+ 
58+ if (indiceCastMode_ != 0) {
59+ indiceCastDtypeSize_ = ge::GetSizeByDataType(indiceCastDtype_);
60+ }
61+ return ge::GRAPH_SUCCESS;
62+}
63+ 
64+void ScatterNdCommonBaseTiling::SetStride()
65+{
66+ auto outPutShape = context_->GetOutputShape(OUTPUT_IDX_SHAPE)->GetStorageShape();
67+ strideList_[rankSize_ - ONE] = static_cast<uint64_t>(1);
68+ for (int16_t dim = static_cast<int16_t>(rankSize_ - TWO); dim >= 0; --dim) {
69+ strideList_[dim] = strideList_[dim + 1] * outPutShape.GetDim(dim + 1);
70+ }
71+}
72+ 
73+ 
74+ge::graphStatus ScatterNdCommonBaseTiling::GetPlatformInfo()
75+{
76+ auto compileInfo = context_->GetCompileInfo<ScatterNdCommonCompileInfo>();
77+ OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);
78+ totalCoreNum_ = compileInfo->core_num;
79+ ubSize_ = compileInfo->ub_size;
80+ 
81+ OP_CHECK_IF(totalCoreNum_ <= 0, OP_LOGE(context_, "GetPlatformInfo get corenum <= 0"), return ge::GRAPH_FAILED);
82+ OP_CHECK_IF(ubSize_ <= 0, OP_LOGE(context_, "GetPlatformInfo get ub size <= 0"), return ge::GRAPH_FAILED);
83+ return ge::GRAPH_SUCCESS;
84+}
85+ 
86+ 
87+uint32_t ScatterNdCommonBaseTiling::GetSortTmpSize(ge::DataType dataType, uint32_t lastAxisNum, bool isDescend)
88+{
89+ std::vector<int64_t> shapeVec = {lastAxisNum};
90+ ge::Shape srcShape(shapeVec);
91+ AscendC::SortConfig config;
92+ config.type = AscendC::SortType::RADIX_SORT;
93+ config.isDescend = isDescend;
94+ config.hasSrcIndex = false;
95+ config.hasDstIndex = true;
96+ uint32_t maxValue = 0;
97+ uint32_t minValue = 0;
98+ AscendC::GetSortMaxMinTmpSize(srcShape, dataType, ge::DT_UINT32, false, config, maxValue, minValue);
99+ OP_LOGI("RadixSortTilingForAscendC", "Need tmp buffer %u byte for ac sort api", maxValue);
100+ return maxValue;
101+}
102+ 
103+ 
104+ge::graphStatus ScatterNdCommonBaseTiling::GetShapeAttrsInfo()
105+{
106+ auto opName = context_->GetNodeName();
107+ OP_LOGD(opName, "GetShapeAttrsInfo begin.");
108+ 
109+ auto var = context_->GetInputTensor(0);
110+ OP_CHECK_NULL_WITH_CONTEXT(context_, var);
111+ auto varShapeSize = var->GetShapeSize();
112+ OP_CHECK_IF((varShapeSize <= 0),
113+ OP_LOGE(opName, "var shape size is invalid(%ld)", varShapeSize),
114+ return ge::GRAPH_FAILED);
115+ auto varDesc = context_->GetInputDesc(INPUT_IDX_UPDATES);
116+ OP_CHECK_NULL_WITH_CONTEXT(context_, varDesc);
117+ auto varDtype = varDesc->GetDataType();
118+ varTypeSize_ = ge::GetSizeByDataType(varDtype);
119+ OP_CHECK_IF(
120+ varTypeSize_ <= 0,
121+ OP_LOGE(context_, "varTypeSize must be greater than 0, varTypeSize: %ld", varTypeSize_),
122+ return ge::GRAPH_FAILED);
123+ 
124+ auto indices = context_->GetInputTensor(INPUT_IDX_INDICES);
125+ OP_CHECK_NULL_WITH_CONTEXT(context_, indices);
126+ indiceShapeSize_ = indices->GetShapeSize();
127+ OP_CHECK_IF((indiceShapeSize_ < 0),
128+ OP_LOGE(opName,
129+ "update shape size is invalid(%ld)", indiceShapeSize_), return ge::GRAPH_FAILED);
130+ auto indicesDesc = context_->GetInputDesc(INPUT_IDX_INDICES);
131+ OP_CHECK_NULL_WITH_CONTEXT(context_, indicesDesc);
132+ indiceDtype_ = indicesDesc->GetDataType();
133+ indicesTypeSize_ = ge::GetSizeByDataType(indiceDtype_);
134+
135+ auto indiceShape = indices->GetStorageShape();
136+ auto indiceDims = indiceShape.GetDimNum();
137+ rankSize_ = indiceShape.GetDim(indiceDims - 1);
138+ OP_CHECK_IF(
139+ (RANK_MIN_VALUE > static_cast<uint16_t>(rankSize_) || static_cast<uint16_t>(rankSize_) > RANK_MAX_VALUE),
140+ OP_LOGE(opName,
141+ "rankSize_ %u out of range[1, 7], please check.", rankSize_),
142+ return ge::GRAPH_FAILED);
143+
144+ auto updates = context_->GetInputTensor(INPUT_IDX_UPDATES);
145+ OP_CHECK_NULL_WITH_CONTEXT(context_, updates);
146+ updateShapeSize_ = updates->GetShapeSize();
147+ OP_CHECK_IF((updateShapeSize_ < 0),
148+ OP_LOGE(opName,
149+ "update shape size is invalid(%ld)", updateShapeSize_), return ge::GRAPH_FAILED);
150+ 
151+ auto updateDesc = context_->GetInputDesc(INPUT_IDX_UPDATES);
152+ OP_CHECK_NULL_WITH_CONTEXT(context_, updateDesc);
153+ updateDtype_ = updateDesc->GetDataType();
154+ OP_CHECK_IF(
155+ (updateDtype_ != varDtype),
156+ OP_LOGE(opName, "updates [%s] and var [%s] must have the same dtype.",
157+ Ops::Base::ToString(updateDtype_).c_str(), Ops::Base::ToString(varDtype).c_str()),
158+ return ge::GRAPH_FAILED);
159+ 
160+ auto outputShape = context_->GetOutputShape(OUTPUT_IDX_SHAPE);
161+ OP_CHECK_NULL_WITH_CONTEXT(context_, outputShape);
162+ auto shapeValue = outputShape->GetStorageShape();
163+ uint64_t shapeRank = shapeValue.GetDimNum();
164+ OP_CHECK_IF((shapeRank < rankSize_),
165+ OP_LOGE(opName,
166+ "shapeRank %lu less than rank %u, please check.", shapeRank, rankSize_),
167+ return ge::GRAPH_FAILED);
168+ 
169+ for (uint64_t idx = 0; idx < shapeRank; idx++) {
170+ outPutShape_[idx] = shapeValue.GetDim(idx);
171+ outputShapeSize_ *= outPutShape_[idx];
172+ }
173+ 
174+ if (indiceShapeSize_ == 0 || updateShapeSize_ == 0) {
175+ return ge::GRAPH_SUCCESS;
176+ }
177+ // indicesAxis_ equal updatesInAxis
178+ indicesAxis_ = indiceShapeSize_ / rankSize_; // g
179+ afterAxis_ = updateShapeSize_ / indicesAxis_; // n
180+ varInAxis_ = varShapeSize / afterAxis_; // m
181+ if (varInAxis_ < INT32_MAX) { // rank维索引合一可能超过int32最大值
182+ outOfSetTypeSize_ = indicesTypeSize_;
183+ outOfSetDtype_ = indiceDtype_;
184+ } else {
185+ outOfSetTypeSize_ = sizeof(int64_t);
186+ outOfSetDtype_ = ge::DataType::DT_INT64;
187+ }
188+ 
189+ return ge::GRAPH_SUCCESS;
190+}
191+ 
192+ 
193+ 
194+bool ScatterNdCommonBaseTiling::IsCapable()
195+{
196+ return true;
197+}
198+ 
199+ge::graphStatus ScatterNdCommonBaseTiling::DoOpTiling()
200+{
201+ return ge::GRAPH_SUCCESS;
202+}
203+ 
204+ge::graphStatus ScatterNdCommonBaseTiling::DoLibApiTiling()
205+{
206+ return ge::GRAPH_SUCCESS;
207+}
208+ 
209+uint64_t ScatterNdCommonBaseTiling::GetTilingKey() const
210+{
211+ return 0;
212+}
213+ 
214+ge::graphStatus ScatterNdCommonBaseTiling::GetWorkspaceSize()
215+{
216+ size_t* workspaces = context_->GetWorkspaceSizes(1);
217+ OP_CHECK_NULL_WITH_CONTEXT(context_, workspaces);
218+ workspaces[0] = ASCENDC_WORKSPACE;
219+ 
220+ return ge::GRAPH_SUCCESS;
221+}
222+ 
223+ge::graphStatus ScatterNdCommonBaseTiling::PostTiling()
224+{
225+ return ge::GRAPH_SUCCESS;
226+}
227+ 
228+} // namespace optiling
Aindex/scatter_nd_common/op_host/arch35/scatter_nd_common_base_tiling.h+93-0
@@ -0,0 +1,93 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_base_tiling.h
13+ * \brief scatter_nd_common_base_tiling
14+ */
15+ 
16+#ifndef SCATTER_ND_COMMON_BASE_TILING_H
17+#define SCATTER_ND_COMMON_BASE_TILING_H
18+ 
19+ 
20+#include "register/op_def_registry.h"
21+#include "tiling/tiling_api.h"
22+#include "op_common/op_host/util/math_util.h"
23+#include "op_common/op_host/util/platform_util.h"
24+#include "../op_kernel/arch35/scatter_nd_common_struct.h"
25+#include "../op_kernel/arch35/scatter_nd_max_tiling_key.h"
26+#include "op_host/tiling_templates_registry.h"
27+ 
28+namespace optiling {
29+ 
30+constexpr uint16_t OPTILING_MAX_RANK_COUNT = 7;
31+constexpr uint16_t OPTILING_MAX_SHAPE_RANK = 8;
32+static constexpr uint16_t INPUT_IDX_INDICES = 1;
33+static constexpr uint16_t INPUT_IDX_UPDATES = 2;
34+static constexpr uint16_t OUTPUT_IDX_SHAPE = 0;
35+static constexpr uint32_t ONE = 1;
36+static constexpr int64_t TWO = 2;
37+ 
38+struct ScatterNdCommonCompileInfo {
39+ int64_t core_num;
40+ int64_t ub_size;
41+};
42+ 
43+class ScatterNdCommonBaseTiling : public Ops::NN::Optiling::TilingBaseClass
44+{
45+public:
46+ explicit ScatterNdCommonBaseTiling(gert::TilingContext* context) : TilingBaseClass(context)
47+ {
48+ }
49+ ~ScatterNdCommonBaseTiling() override
50+ {
51+ }
52+ 
53+protected:
54+ bool IsCapable() override;
55+ ge::graphStatus GetPlatformInfo() override;
56+ ge::graphStatus GetShapeAttrsInfo() override;
57+ ge::graphStatus DoOpTiling() override;
58+ ge::graphStatus DoLibApiTiling() override;
59+ uint64_t GetTilingKey() const override;
60+ ge::graphStatus GetWorkspaceSize() override;
61+ ge::graphStatus PostTiling() override;
62+ void DumpTilingInfo() override
63+ {}
64+ 
65+ uint32_t GetSortTmpSize(ge::DataType dataType, uint32_t lastAxisNum, bool isDescend);
66+ ge::graphStatus GetCastType();
67+ void SetStride();
68+ 
69+public:
70+ uint64_t ubSize_ = 0;
71+ int64_t totalCoreNum_ = 0;
72+ int64_t varTypeSize_ = 0;
73+ int64_t indicesTypeSize_ = 0;
74+ int64_t outOfSetTypeSize_ = 0;
75+ uint32_t rankSize_ = 0;
76+ int64_t updateShapeSize_ = 0;
77+ uint64_t outputShapeSize_ = 1;
78+ int64_t indiceShapeSize_ = 0;
79+ int64_t indicesAxis_ = 0;
80+ int64_t varInAxis_ = 1;
81+ int64_t afterAxis_ = 1;
82+ uint64_t indiceCastMode_ = 0; // 0: 不Cast; 1：int32 Cast int16; 2：int64 Cast int32; 3：int64 Cast int16; 4:int32 Cast uint8; 5:int64 Cast uint8.
83+ int64_t indiceCastDtypeSize_ = 0;
84+ uint64_t strideList_[OPTILING_MAX_RANK_COUNT] = {0};
85+ uint64_t outPutShape_[OPTILING_MAX_SHAPE_RANK] = {0};
86+ 
87+ ge::DataType updateDtype_ = ge::DT_UNDEFINED;
88+ ge::DataType indiceDtype_ = ge::DT_UNDEFINED;
89+ ge::DataType indiceCastDtype_ = ge::DT_UNDEFINED;
90+ ge::DataType outOfSetDtype_ = ge::DT_UNDEFINED;
91+};
92+} // namespace optiling
93+#endif // SCATTER_ND_COMMON_BASE_TILING_H
Aindex/scatter_nd_common/op_host/arch35/scatter_nd_common_simd_sort_tiling.cpp+415-0
@@ -0,0 +1,415 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_simd_sort_tiling.cpp
13+ * \brief scatter_nd_common_simd_sort_tiling
14+ */
15+ 
16+#include "scatter_nd_common_simd_sort_tiling.h"
17+ 
18+using namespace AscendC;
19+using namespace ge;
20+ 
21+namespace optiling {
22+ 
23+using namespace ScatterNdCommon;
24+ 
25+static constexpr int64_t MIN_SIZE_SIMD_NONDETERMINSTIC = 128;
26+static constexpr int64_t TEMPLATE_MODE = 2;
27+static constexpr int64_t SIMD_NONDETERMINSTIC_MIN_HANDLE_SIZE = 4096;
28+static constexpr int64_t SPLIT_THRESHOLD = 64;
29+static constexpr int64_t ALIGN_SIZE = 32;
30+static constexpr int64_t MIN_HANDLE_SIZE = 128;
31+static constexpr int64_t FP32_BYTES = 4;
32+static constexpr int64_t INT32_BYTES = 4;
33+static constexpr uint64_t RESERVE_SIZE = 256;
34+static constexpr int64_t CUT_THRESHOLD = 128;
35+static constexpr int64_t INDICES_MIN_BLOCK_SIZE = 1024;
36+static constexpr uint64_t DB_BUFFER = 2;
37+ 
38+static const std::set<ge::DataType> setAtomicNotSupport = {ge::DT_UINT32, ge::DT_INT64, ge::DT_UINT64};
39+ 
40+bool ScatterNdCommonSimdSortTiling::IsCapable()
41+{
42+ if (afterAxis_ * varTypeSize_ >= MIN_SIZE_SIMD_NONDETERMINSTIC &&
43+ setAtomicNotSupport.find(updateDtype_) == setAtomicNotSupport.end()) {
44+ return true;
45+ }
46+ if (updateDtype_ == ge::DT_INT8 || updateDtype_ == ge::DT_INT16) {
47+ return true;
48+ }
49+ return false;
50+}
51+ 
52+uint64_t ScatterNdCommonSimdSortTiling::GetTilingKey() const
53+{
54+ OP_LOGD("ScatterNdCommonSimdSortTiling::GetTilingKey begin");
55+ uint64_t addrMode = 1;
56+ if (varInAxis_ < INT32_MAX) { // rank维索引合一可能超过int32最大值
57+ addrMode = 0;
58+ }
59+ const uint64_t tilingKey = GET_TPL_TILING_KEY(TEMPLATE_MODE, indiceCastMode_, addrMode);
60+ OP_LOGD(context_->GetNodeName(), "tilingKey is: [%lu]", tilingKey);
61+ return tilingKey;
62+}
63+ 
64+void ScatterNdCommonSimdSortTiling::SetTilingData()
65+{
66+ ScatterNdCommon::ScatterNdCommonSimdSortTilingData *tilingData =
67+ context_->GetTilingData<ScatterNdCommon::ScatterNdCommonSimdSortTilingData>();
68+ 
69+ for (int32_t i = 0; i < OPTILING_MAX_RANK_COUNT; i++) {
70+ tilingData->strideList[i] = strideList_[i];
71+ }
72+ for (int32_t i = 0; i < OPTILING_MAX_SHAPE_RANK; i++) {
73+ tilingData->outPutShape[i] = outPutShape_[i];
74+ }
75+ tilingData->eachCoreAfterAxisCount = eachCoreAfterAxisCount_;
76+ tilingData->indexRankSize = rankSize_;
77+ tilingData->eachCoreIndexCount = eachCoreIndexCount_;
78+ tilingData->tailCoreIndexCount = tailCoreIndexCount_;
79+ tilingData->indicesFactor = indicesFactor_;
80+ tilingData->indiceTailNum = indiceTailNum_;
81+ tilingData->indicesLoopSize = indicesLoopSize_;
82+ tilingData->afterAxis = afterAxis_;
83+ tilingData->afterAxisFactor = afterAxisFactor_;
84+ tilingData->usedCoreNumBefore = usedCoreNumBefore_;
85+ tilingData->updateLoopSize = updateLoopSize_;
86+ tilingData->tailUpdateLoopSize = tailUpdateLoopSize_;
87+ tilingData->tailUpdateTailNum = tailUpdateTailNum_;
88+ tilingData->updateTailNum = updateTailNum_;
89+ tilingData->isSplitAfterAxis = isSplitAfterAxis_;
90+ tilingData->singleCol = singleCol_;
91+}
92+ 
93+ 
94+ 
95+ 
96+void ScatterNdCommonSimdSortTiling::DoOpTilingSplitIndicesSingleCol()
97+{
98+ int64_t indicesDtypeCastSize = indiceCastMode_ ? indiceCastDtypeSize_ : 0;
99+ int64_t indicesRealTypeSize = indiceCastMode_ ? indiceCastDtypeSize_ : outOfSetTypeSize_;
100+ ge::DataType sortTmpType = indiceCastMode_ ? indiceCastDtype_ : outOfSetDtype_;
101+ 
102+ int64_t alignNum = ALIGN_SIZE / varTypeSize_;
103+ int64_t halfUbSize = static_cast<int64_t>((ubSize_ - RESERVE_SIZE - MIN_HANDLE_SIZE * FP32_BYTES) / DB_BUFFER);
104+ eachCoreIndexCount_ = Ops::Base::CeilDiv(indicesAxis_, totalCoreNum_);
105+ usedCoreNumBefore_ = Ops::Base::CeilDiv(indicesAxis_, eachCoreIndexCount_);
106+ tailCoreIndexCount_ = indicesAxis_ - eachCoreIndexCount_ * (usedCoreNumBefore_ - 1);
107+ int64_t oneIndexSize = static_cast<int64_t>(rankSize_) * indicesTypeSize_;
108+ 
109+ auto ubBlock = static_cast<int64_t>(Ops::Base::GetUbBlockSize(context_));
110+ int64_t indicesAlignSize = Ops::Base::CeilAlign(oneIndexSize, ubBlock) + // indices
111+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock) + // outOfset
112+ Ops::Base::CeilAlign(indicesDtypeCastSize, ubBlock) + // cast_indices
113+ Ops::Base::CeilAlign(indicesRealTypeSize + TWO * ALIGN_SIZE, ubBlock) +
114+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
115+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
116+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock);
117+ 
118+ int64_t updateAlignSize = Ops::Base::CeilAlign(varTypeSize_ * afterAxis_, ubBlock) +
119+ Ops::Base::CeilAlign(varTypeSize_ * afterAxis_, ubBlock) +
120+ GetSortTmpSize(sortTmpType, SPLIT_THRESHOLD, false);
121+ //每次最少搬入indices数
122+ int64_t minIndicesFactorSize = indicesAlignSize * SPLIT_THRESHOLD;
123+ // 搬入多行indices，搬入一行updates
124+ if (minIndicesFactorSize + updateAlignSize > halfUbSize) { // 优先保证indices至少处理64行，实际 afterAxisFactor_ 按剩余ub给一行updates
125+ indicesFactor_ = SPLIT_THRESHOLD;
126+ int64_t sortTmpSize = GetSortTmpSize(sortTmpType, SPLIT_THRESHOLD, false);
127+ afterAxisFactor_ = (halfUbSize - indicesFactor_ * indicesAlignSize - sortTmpSize) / (varTypeSize_ + varTypeSize_); // (update + updateSum)
128+ afterAxisFactor_ = Ops::Base::FloorAlign(afterAxisFactor_, alignNum);
129+ } else { // 完整处理一行 afterAxis_,逐步调整 indicesFactor_ 以满足UB约束
130+ afterAxisFactor_ = Ops::Base::CeilAlign(afterAxis_, alignNum);
131+ indicesFactor_ = (halfUbSize - updateAlignSize) / indicesAlignSize;
132+ int64_t restSize = static_cast<int64_t>(-1);
133+ while (restSize <= 0) {
134+ int64_t occupy = Ops::Base::CeilAlign(indicesFactor_ * rankSize_ * indicesTypeSize_, ubBlock) +
135+ Ops::Base::CeilAlign(indicesFactor_ * outOfSetTypeSize_, ubBlock) +
136+ Ops::Base::CeilAlign(indicesFactor_ * indicesDtypeCastSize, ubBlock) +
137+ Ops::Base::CeilAlign(indicesFactor_ * (indicesRealTypeSize + TWO * ALIGN_SIZE), ubBlock) +
138+ Ops::Base::CeilAlign(indicesFactor_ * INT32_BYTES, ubBlock) +
139+ Ops::Base::CeilAlign(indicesFactor_ * (INT32_BYTES), ubBlock) +
140+ Ops::Base::CeilAlign(indicesFactor_ * outOfSetTypeSize_, ubBlock) +
141+ Ops::Base::CeilAlign(varTypeSize_ * afterAxisFactor_, ubBlock) +
142+ Ops::Base::CeilAlign(varTypeSize_ * afterAxisFactor_, ubBlock) +
143+ GetSortTmpSize(sortTmpType, indicesFactor_, false);
144+ restSize = halfUbSize - occupy;
145+ if (indicesFactor_ > indicesAxis_) {
146+ indicesFactor_ = indicesAxis_;
147+ break;
148+ }
149+ --indicesFactor_;
150+ }
151+ }
152+ updateLoopSize_ = Ops::Base::CeilDiv(afterAxis_, afterAxisFactor_);
153+ updateTailNum_ = afterAxis_ - (updateLoopSize_ - 1) * afterAxisFactor_;
154+ singleCol_ = 1;
155+}
156+ 
157+void ScatterNdCommonSimdSortTiling::DoOpTilingSplitAfter()
158+{
159+ int64_t ubSize = static_cast<int64_t>((ubSize_ - RESERVE_SIZE));
160+ ubSize = ubSize - MIN_HANDLE_SIZE * FP32_BYTES;//maxScore
161+ int64_t alignNum = ALIGN_SIZE / varTypeSize_;
162+ int64_t afterAxisSize = afterAxis_ * varTypeSize_;
163+ /* split afterAxis */
164+ if (afterAxisSize < CUT_THRESHOLD) { // 尾轴不大，索引也不大，开1个核
165+ eachCoreAfterAxisCount_ = afterAxis_;
166+ usedCoreNumBefore_ = ONE;
167+ tailCoreAfterAxisCount_ = afterAxis_;
168+ } else {
169+ eachCoreAfterAxisCount_ = Ops::Base::CeilDiv(afterAxis_, totalCoreNum_); // 正常核处理的尾轴个数
170+ usedCoreNumBefore_ = Ops::Base::CeilDiv(afterAxis_, eachCoreAfterAxisCount_); // 使用核数
171+ tailCoreAfterAxisCount_ = afterAxis_ - eachCoreAfterAxisCount_ * (usedCoreNumBefore_ - 1); // 尾核处理的尾轴个数
172+ }
173+ 
174+ auto ubBlock = static_cast<int64_t>(Ops::Base::GetUbBlockSize(context_));
175+ /* 同地址优化:搬入多少行indices,就搬入相同行数的updates, strideBuf放在RESERVE_SIZE中:
176+ * indicesFactor_: outOfsetBuf + indiecesQue + (sortIndicesQue + 2 * shiftOfset) + originIdxQue +
177+ * (uniqueIdCntQue_ + 1) + updateSumIdxQue_,
178+ * indicesFactor_ * eachCoreAfterAxisCount_: updatesQue_ + updateSumQue_
179+ */
180+ int64_t indicesSize = 0;
181+ if (!indiceCastMode_) {
182+ indicesSize = Ops::Base::CeilAlign(rankSize_ * indicesTypeSize_, ubBlock) + // indices
183+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock) + // outofset
184+ Ops::Base::CeilAlign(outOfSetTypeSize_ + TWO * ALIGN_SIZE, ubBlock) + // sortIndice
185+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) + // updateOriginIndices
186+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) + // uniqueIdCount
187+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock) + // updateSumIdx
188+ GetSortTmpSize(outOfSetDtype_, 1, false);
189+ } else {
190+ indicesSize = Ops::Base::CeilAlign(rankSize_ * indicesTypeSize_, ubBlock) + // indices
191+ Ops::Base::CeilAlign(indicesTypeSize_, ubBlock) + // outofset
192+ Ops::Base::CeilAlign(indiceCastDtypeSize_, ubBlock) + // cast_outofset
193+ Ops::Base::CeilAlign(indiceCastDtypeSize_ + TWO * ALIGN_SIZE, ubBlock) + // sortIndice
194+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) + // updateOriginIndices
195+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) + // uniqueIdCount
196+ Ops::Base::CeilAlign(indicesTypeSize_, ubBlock) + // updateSumIdx
197+ GetSortTmpSize(indiceCastDtype_, 1, false);
198+ }
199+ int64_t oneBlockSize = indicesSize + (varTypeSize_ + varTypeSize_) * eachCoreAfterAxisCount_; // (update + updateSum)
200+ indicesFactor_ = ubSize / oneBlockSize;
201+ 
202+ int64_t occupy = 0;
203+ if (!indiceCastMode_) {
204+ occupy = Ops::Base::CeilAlign(rankSize_ * indicesTypeSize_, ubBlock) +
205+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock) +
206+ Ops::Base::CeilAlign(outOfSetTypeSize_ + TWO * ALIGN_SIZE, ubBlock) +
207+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
208+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
209+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock) +
210+ Ops::Base::CeilAlign(varTypeSize_ * eachCoreAfterAxisCount_, ubBlock) +
211+ Ops::Base::CeilAlign(varTypeSize_ * eachCoreAfterAxisCount_, ubBlock) +
212+ GetSortTmpSize(outOfSetDtype_, 1, false);
213+ } else {
214+ occupy = Ops::Base::CeilAlign(rankSize_ * indicesTypeSize_, ubBlock) +
215+ Ops::Base::CeilAlign(indicesTypeSize_, ubBlock) +
216+ Ops::Base::CeilAlign(indiceCastDtypeSize_, ubBlock) +
217+ Ops::Base::CeilAlign(indiceCastDtypeSize_ + TWO * ALIGN_SIZE, ubBlock) +
218+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
219+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
220+ Ops::Base::CeilAlign(indicesTypeSize_, ubBlock) +
221+ Ops::Base::CeilAlign(varTypeSize_ * eachCoreAfterAxisCount_, ubBlock) +
222+ Ops::Base::CeilAlign(varTypeSize_ * eachCoreAfterAxisCount_, ubBlock) +
223+ GetSortTmpSize(indiceCastDtype_, 1, false);
224+ }
225+ 
226+ if (occupy > ubSize) {
227+ int64_t indicesUbSize = std::min(INDICES_MIN_BLOCK_SIZE, indicesAxis_ * indicesSize);
228+ /* indicesBuf_ + outOfstBuf_ */
229+ indicesFactor_ = Ops::Base::CeilAlign(indicesUbSize, ALIGN_SIZE) / indicesSize;
230+ int64_t updatesSize = Ops::Base::CeilAlign((varTypeSize_) * indicesFactor_, ubBlock) + Ops::Base::CeilAlign((varTypeSize_) * indicesFactor_, ubBlock);
231+ afterAxisFactor_ = (ubSize - indicesFactor_ * indicesSize) / updatesSize;
232+ afterAxisFactor_ = Ops::Base::FloorAlign(afterAxisFactor_, alignNum) / updateDtype_;
233+ } else {
234+ afterAxisFactor_ = Ops::Base::CeilAlign(eachCoreAfterAxisCount_, alignNum);
235+ indicesFactor_ = ubSize / (afterAxisFactor_ * (varTypeSize_ + varTypeSize_) + indicesSize);
236+ 
237+ int64_t restSize = static_cast<int64_t>(-1);
238+ int64_t indicesDtypeCastSize = indiceCastMode_ ? indiceCastDtypeSize_ : 0;
239+ int64_t indicesRealTypeSize = indiceCastMode_ ? indiceCastDtypeSize_ : outOfSetTypeSize_;
240+ ge::DataType sortTmpType = indiceCastMode_ ? indiceCastDtype_ : outOfSetDtype_;
241+ 
242+ while (restSize <= 0) {
243+ restSize = ubSize - (Ops::Base::CeilAlign(indicesFactor_ * rankSize_ * indicesTypeSize_, ubBlock) +
244+ Ops::Base::CeilAlign(indicesFactor_ * outOfSetTypeSize_, ubBlock) +
245+ Ops::Base::CeilAlign(indicesFactor_ * indicesDtypeCastSize, ubBlock) +
246+ Ops::Base::CeilAlign(indicesFactor_ * (indicesRealTypeSize + TWO * ALIGN_SIZE), ubBlock) +
247+ Ops::Base::CeilAlign(indicesFactor_ * INT32_BYTES, ubBlock) +
248+ Ops::Base::CeilAlign(indicesFactor_ * (INT32_BYTES), ubBlock) +
249+ Ops::Base::CeilAlign(indicesFactor_ * outOfSetTypeSize_, ubBlock) +
250+ indicesFactor_ * Ops::Base::CeilAlign((varTypeSize_) * eachCoreAfterAxisCount_, ubBlock) +
251+ indicesFactor_ * Ops::Base::CeilAlign((varTypeSize_) * eachCoreAfterAxisCount_, ubBlock) +
252+ GetSortTmpSize(sortTmpType, indicesFactor_, false));
253+ if (indicesFactor_ > indicesAxis_) {
254+ indicesFactor_ = indicesAxis_;
255+ break;
256+ }
257+ --indicesFactor_;
258+ }
259+ // indicesFactor_ = Ops::Base::CeilAlign(indicesFactor_ * indicesTypeSize_, ALIGN_SIZE) / indicesTypeSize_;
260+ }
261+ 
262+ /* 每个核分的indices相同 */
263+ indicesLoopSize_ = Ops::Base::CeilDiv(indicesAxis_, indicesFactor_);
264+ indiceTailNum_ = indicesAxis_ - (indicesLoopSize_ - 1) * indicesFactor_;
265+ 
266+ /* 主核循环次数 */
267+ updateLoopSize_ = Ops::Base::CeilDiv(eachCoreAfterAxisCount_, afterAxisFactor_);
268+ /* 主核尾loop处理afterAxis大小 */
269+ updateTailNum_ = eachCoreAfterAxisCount_ - (updateLoopSize_ - 1) * afterAxisFactor_;
270+ 
271+ /* 尾核循环次数 */
272+ tailUpdateLoopSize_ = Ops::Base::CeilDiv(tailCoreAfterAxisCount_, afterAxisFactor_);
273+ /* 尾核尾loop处理afterAxis大小 */
274+ tailUpdateTailNum_ = tailCoreAfterAxisCount_ - (tailUpdateLoopSize_ - 1) * afterAxisFactor_;
275+ isSplitAfterAxis_ = 1;
276+}
277+ 
278+void ScatterNdCommonSimdSortTiling::DoOpTilingSimdSplitIndices() // TODO: 没开double buffer 为啥 half ub ?
279+{
280+ int64_t indicesDtypeCastSize = indiceCastMode_ ? indiceCastDtypeSize_ : 0;
281+ int64_t indicesRealTypeSize = indiceCastMode_ ? indiceCastDtypeSize_ : outOfSetTypeSize_;
282+ ge::DataType sortTmpType = indiceCastMode_ ? indiceCastDtype_ : outOfSetDtype_;
283+ 
284+ int64_t alignNum = ALIGN_SIZE / varTypeSize_;
285+ // int64_t halfUbSize = static_cast<int64_t>((ubSize_ - RESERVE_SIZE) / DB_BUFFER);
286+ int64_t halfUbSize = static_cast<int64_t>(ubSize_ - RESERVE_SIZE);
287+ halfUbSize = halfUbSize - MIN_HANDLE_SIZE * FP32_BYTES;//maxScore
288+ 
289+ /* split indices分核 */
290+ eachCoreIndexCount_ = Ops::Base::CeilDiv(indicesAxis_, totalCoreNum_);
291+ usedCoreNumBefore_ = Ops::Base::CeilDiv(indicesAxis_, eachCoreIndexCount_);
292+ tailCoreIndexCount_ = indicesAxis_ - eachCoreIndexCount_ * (usedCoreNumBefore_ - 1);
293+ int64_t oneIndexSize = static_cast<int64_t>(rankSize_) * indicesTypeSize_;
294+
295+ /* 同地址优化:搬入多少行indices,就搬入相同行数的updates, strideBuf放在RESERVE_SIZE中:
296+ * indicesFactor_: indiecesQue + outOfsetBuf + (sortIndicesQue + 2 * shiftOfset) + originIdxQue +
297+ * (uniqueIdCntQue_ + 1) + updateSumIdxQue_,
298+ * indicesFactor_ * eachCoreAfterAxisCount_: updatesQue_ + updateSumQue_
299+ */
300+ auto ubBlock = static_cast<int64_t>(Ops::Base::GetUbBlockSize(context_));
301+ int64_t indicesAlignSize = Ops::Base::CeilAlign(oneIndexSize, ubBlock) +
302+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock) +
303+ Ops::Base::CeilAlign(indicesDtypeCastSize, ubBlock) +
304+ Ops::Base::CeilAlign(indicesRealTypeSize + TWO * ALIGN_SIZE, ubBlock) +
305+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
306+ // Ops::Base::CeilAlign(INT32_BYTES * TWO, ubBlock) +
307+ Ops::Base::CeilAlign(INT32_BYTES, ubBlock) +
308+ Ops::Base::CeilAlign(outOfSetTypeSize_, ubBlock);
309+ 
310+ int64_t updateAlignSize = Ops::Base::CeilAlign(varTypeSize_ * afterAxis_, ubBlock) +
311+ Ops::Base::CeilAlign(varTypeSize_ * afterAxis_, ubBlock) +
312+ GetSortTmpSize(sortTmpType, 1, false);
313+ if (indicesAlignSize + updateAlignSize > halfUbSize) { // 优先保证indices处理的行数，实际 afterAxisFactor_ 可能远小于afterAxis_
314+ int64_t indicesSize = std::min(INDICES_MIN_BLOCK_SIZE, indicesAxis_ * indicesAlignSize);
315+ /* indicesBuf_ + outOfstBuf_ */
316+ indicesFactor_ = Ops::Base::CeilAlign(indicesSize, ALIGN_SIZE) / indicesAlignSize;
317+ afterAxisFactor_ = (halfUbSize - indicesFactor_ * indicesAlignSize) / indicesFactor_;
318+ afterAxisFactor_ = Ops::Base::FloorAlign(afterAxisFactor_, alignNum);
319+ } else { // 全载 afterAxis_,逐步调整 indicesFactor_ 以满足UB约束
320+ afterAxisFactor_ = Ops::Base::CeilAlign(afterAxis_, alignNum);
321+ indicesFactor_ = halfUbSize / (updateAlignSize + indicesAlignSize);
322+ int64_t restSize = static_cast<int64_t>(-1);
323+ while (restSize <= 0) {
324+ int64_t occupy = Ops::Base::CeilAlign(indicesFactor_ * rankSize_ * indicesTypeSize_, ubBlock) +
325+ Ops::Base::CeilAlign(indicesFactor_ * outOfSetTypeSize_, ubBlock) +
326+ Ops::Base::CeilAlign(indicesFactor_ * indicesDtypeCastSize, ubBlock) +
327+ Ops::Base::CeilAlign(indicesFactor_ * (indicesRealTypeSize + TWO * ALIGN_SIZE), ubBlock) +
328+ Ops::Base::CeilAlign(indicesFactor_ * INT32_BYTES, ubBlock) +
329+ Ops::Base::CeilAlign(indicesFactor_ * (INT32_BYTES), ubBlock) +
330+ Ops::Base::CeilAlign(indicesFactor_ * outOfSetTypeSize_, ubBlock) +
331+ indicesFactor_ * Ops::Base::CeilAlign((varTypeSize_) * afterAxisFactor_, ubBlock) + // update
332+ indicesFactor_ * Ops::Base::CeilAlign((varTypeSize_) * afterAxisFactor_, ubBlock) + // updateSum
333+ GetSortTmpSize(sortTmpType, indicesFactor_, false);
334+ restSize = halfUbSize - occupy;
335+ if (indicesFactor_ > indicesAxis_) {
336+ indicesFactor_ = indicesAxis_;
337+ break;
338+ }
339+ --indicesFactor_;
340+ }
341+ }
342+ /* 每个核分的update相同 */
343+ updateLoopSize_ = Ops::Base::CeilDiv(afterAxis_, afterAxisFactor_);
344+ updateTailNum_ = afterAxis_ - (updateLoopSize_ - 1) * afterAxisFactor_;
345+}
346+ 
347+void ScatterNdCommonSimdSortTiling::DoOpTilingForSimdSort()
348+{
349+ GetCastType();
350+ // 列大于4k单行搬入搬出
351+ if (afterAxis_ * varTypeSize_ > SIMD_NONDETERMINSTIC_MIN_HANDLE_SIZE &&
352+ (indicesAxis_ > (totalCoreNum_ * SPLIT_THRESHOLD))) {
353+ DoOpTilingSplitIndicesSingleCol();
354+ return;
355+ }
356+ /* 优先分after */
357+ int64_t splitThresh = totalCoreNum_ * MIN_HANDLE_SIZE / varTypeSize_;
358+ if ((afterAxis_ > splitThresh) || (indicesAxis_ < (totalCoreNum_ / TWO))) {
359+ DoOpTilingSplitAfter(); // 如果尾轴大于 核数*128B 或者 索引(g)小于一半核数 则切尾轴
360+ return;
361+ }
362+ DoOpTilingSimdSplitIndices(); // 否则切索引
363+ return;
364+}
365+ 
366+ 
367+ge::graphStatus ScatterNdCommonSimdSortTiling::DoOpTiling()
368+{
369+ DoOpTilingForSimdSort();
370+ SetStride();
371+ SetTilingData();
372+ return ge::GRAPH_SUCCESS;
373+}
374+ 
375+ge::graphStatus ScatterNdCommonSimdSortTiling::PostTiling()
376+{
377+ OP_LOGD("ScatterNdCommonSimdSortTiling::PostTiling begin");
378+ context_->SetBlockDim(usedCoreNumBefore_);
379+ 
380+ return ge::GRAPH_SUCCESS;
381+}
382+ 
383+ 
384+std::string ScatterNdCommonSimdSortTiling::TilingDataToString()
385+{
386+ ScatterNdCommon::ScatterNdCommonSimdSortTilingData* tilingData =
387+ context_->GetTilingData<ScatterNdCommon::ScatterNdCommonSimdSortTilingData>();
388+ std::string str = " eachCoreAfterAxisCount:" + std::to_string(tilingData->eachCoreAfterAxisCount);
389+ str += " indexRankSize:" + std::to_string(tilingData->indexRankSize);
390+ str += " eachCoreIndexCount:" + std::to_string(tilingData->eachCoreIndexCount);
391+ str += " tailCoreIndexCount:" + std::to_string(tilingData->tailCoreIndexCount);
392+ str += " indicesFactor:" + std::to_string(tilingData->indicesFactor);
393+ str += " indiceTailNum:" + std::to_string(tilingData->indiceTailNum);
394+ str += " indicesLoopSize:" + std::to_string(tilingData->indicesLoopSize);
395+ str += " afterAxis:" + std::to_string(tilingData->afterAxis);
396+ str += " afterAxisFactor:" + std::to_string(tilingData->afterAxisFactor);
397+ str += " usedCoreNumBefore:" + std::to_string(tilingData->usedCoreNumBefore);
398+ str += " updateLoopSize:" + std::to_string(tilingData->updateLoopSize);
399+ str += " tailUpdateLoopSize:" + std::to_string(tilingData->tailUpdateLoopSize);
400+ str += " tailUpdateTailNum:" + std::to_string(tilingData->tailUpdateTailNum);
401+ str += " updateTailNum:" + std::to_string(tilingData->updateTailNum);
402+ str += " isSplitAfterAxis:" + std::to_string(tilingData->isSplitAfterAxis);
403+ str += " singleCol:" + std::to_string(tilingData->singleCol);
404+ for (int32_t i = 0; i < OPTILING_MAX_RANK_COUNT; i++) {
405+ str += " strideList[i]:" + std::to_string(tilingData->strideList[i]);
406+ }
407+ return str;
408+}
409+ 
410+void ScatterNdCommonSimdSortTiling::DumpTilingInfo()
411+{
412+ OP_LOGI(context_->GetNodeName(), "Tiling info is: %s", TilingDataToString().c_str());
413+}
414+ 
415+} // namespace optiling
Aindex/scatter_nd_common/op_host/arch35/scatter_nd_common_simd_sort_tiling.h+61-0
@@ -0,0 +1,61 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_simd_sort_tiling.h
13+ * \brief scatter_nd_common_simd_sort_tiling
14+ */
15+ 
16+#ifndef SCATTER_ND_COMMON_SIMD_SORT_TILING_H
17+#define SCATTER_ND_COMMON_SIMD_SORT_TILING_H
18+ 
19+#include "scatter_nd_common_base_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+class ScatterNdCommonSimdSortTiling : public ScatterNdCommonBaseTiling
24+{
25+public:
26+ explicit ScatterNdCommonSimdSortTiling(gert::TilingContext* context) : ScatterNdCommonBaseTiling(context)
27+ {}
28+ ~ScatterNdCommonSimdSortTiling() override
29+ {}
30+ 
31+protected:
32+ bool IsCapable() override;
33+ ge::graphStatus DoOpTiling() override;
34+ ge::graphStatus PostTiling() override;
35+ uint64_t GetTilingKey() const override;
36+ void DumpTilingInfo() override;
37+ std::string TilingDataToString();
38+ void SetTilingData();
39+ void DoOpTilingForSimdSort();
40+ void DoOpTilingSplitIndicesSingleCol();
41+ void DoOpTilingSplitAfter();
42+ void DoOpTilingSimdSplitIndices();
43+ 
44+ int64_t eachCoreIndexCount_ = 0;
45+ int64_t usedCoreNumBefore_ = 0;
46+ int64_t tailCoreIndexCount_ = 0;
47+ int64_t indicesFactor_ = 0;
48+ int64_t afterAxisFactor_ = 0;
49+ int64_t updateLoopSize_ = 0;
50+ int64_t updateTailNum_ = 0;
51+ int64_t singleCol_ = 0;
52+ int64_t eachCoreAfterAxisCount_ = 0;
53+ int64_t tailCoreAfterAxisCount_ = 0;
54+ int64_t indicesLoopSize_ = 0;
55+ int64_t indiceTailNum_ = 0;
56+ int64_t tailUpdateLoopSize_ = 0;
57+ int64_t tailUpdateTailNum_ = 0;
58+ int64_t isSplitAfterAxis_ = 0;
59+};
60+} // namespace optiling
61+#endif // SCATTER_ND_COMMON_SIMD_SORT_TILING_H
Aindex/scatter_nd_common/op_host/arch35/scatter_nd_common_simt_tiling.cpp+185-0
@@ -0,0 +1,185 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_simt_tiling.cpp
13+ * \brief scatter_nd_common_simt_tiling
14+ */
15+ 
16+#include "scatter_nd_common_simt_tiling.h"
17+ 
18+using namespace AscendC;
19+using namespace ge;
20+ 
21+namespace optiling {
22+ 
23+using namespace ScatterNdCommon;
24+ 
25+static constexpr uint64_t RESERVE_SIZE = 256;
26+static constexpr int64_t DCACHE_SIZE = 32768; // 32k
27+static constexpr int64_t TEMPLATE_MODE_SIMT = 8;
28+static constexpr uint64_t DB_BUFFER = 2;
29+static constexpr uint64_t MIN_TILING_SIZE = 128;
30+ 
31+ 
32+bool ScatterNdCommonSimtTiling::IsCapable()
33+{
34+ return true;
35+}
36+ 
37+uint64_t ScatterNdCommonSimtTiling::GetTilingKey() const
38+{
39+ OP_LOGD("ScatterNdCommonSimtTiling::GetTilingKey begin");
40+ uint64_t addrMode = 1;
41+ if (indiceShapeSize_ < UINT32_MAX && updateShapeSize_ < UINT32_MAX && outputShapeSize_ < UINT32_MAX) {
42+ addrMode = 0;
43+ }
44+ const uint64_t tilingKey = GET_TPL_TILING_KEY(TEMPLATE_MODE_SIMT, indiceCastMode_, addrMode);
45+ OP_LOGD(context_->GetNodeName(), "tilingKey is: [%lu]", tilingKey);
46+ return tilingKey;
47+}
48+ 
49+void ScatterNdCommonSimtTiling::SetTilingData()
50+{
51+ ScatterNdCommon::ScatterNdCommonSimtTilingData *tilingData =
52+ context_->GetTilingData<ScatterNdCommon::ScatterNdCommonSimtTilingData>();
53+ 
54+ for (int32_t i = 0; i < OPTILING_MAX_RANK_COUNT; i++) {
55+ tilingData->strideList[i] = strideList_[i];
56+ }
57+ for (int32_t i = 0; i < OPTILING_MAX_SHAPE_RANK; i++) {
58+ tilingData->outPutShape[i] = outPutShape_[i];
59+ }
60+ tilingData->blockNum = blockNum_;
61+ tilingData->rankSize = rankSize_;
62+ tilingData->blockTilingSize = blockTilingSize_;
63+ tilingData->tailBlockTilingSize = tailBlockTilingSize_;
64+ tilingData->ubTilingSize = ubTilingSize_;
65+ tilingData->sliceSize = sliceSize_;
66+ tilingData->varInAxis = varInAxis_;
67+}
68+ 
69+ 
70+ge::graphStatus ScatterNdCommonSimtTiling::BlockTiling()
71+{
72+ auto typeSize = ge::GetSizeByDataType(updateDtype_);
73+ OP_CHECK_IF(
74+ typeSize <= 0,
75+ OP_LOGE(context_, "update dtypeSize must be greater than 0, dtypeSize: %ld", typeSize),
76+ return ge::GRAPH_FAILED);
77+ alignFactor_ = Ops::Base::GetUbBlockSize(context_) / typeSize;
78+ auto blockFactor = Ops::Base::CeilDiv(updateShapeSize_, totalCoreNum_);
79+ auto blockAlignFactor = Ops::Base::CeilDiv(blockFactor, alignFactor_) * alignFactor_;
80+ blockTilingSize_ = std::max(static_cast<uint64_t>(blockAlignFactor), MIN_TILING_SIZE);
81+ blockNum_ = Ops::Base::CeilDiv(updateShapeSize_, blockTilingSize_);
82+ tailBlockTilingSize_ = updateShapeSize_ - blockTilingSize_ * (blockNum_ - 1UL);
83+ OP_LOGD(context_->GetNodeName(),
84+ "updateShapeSize = %lld, blockFactor = %lld, blockAlignFactor = %lld,"
85+ "blockTilingSize = %d, tailBlockTilingSize = %d", updateShapeSize_,
86+ blockFactor, blockAlignFactor, blockTilingSize_, tailBlockTilingSize_);
87+ return ge::GRAPH_SUCCESS;
88+}
89+ 
90+ 
91+ge::graphStatus ScatterNdCommonSimtTiling::UbTiling()
92+{
93+ if (indiceShapeSize_ == static_cast<uint64_t>(0) || updateShapeSize_ == static_cast<uint64_t>(0)) {
94+ return ge::GRAPH_SUCCESS;
95+ }
96+ // halfUbSize for double buffer
97+ auto halfUbSize = (ubSize_ - DCACHE_SIZE - RESERVE_SIZE) / DB_BUFFER;
98+ auto indiceNum = indiceShapeSize_ / rankSize_;
99+ sliceSize_ = updateShapeSize_ / indiceNum;
100+ OP_CHECK_IF(sliceSize_ == static_cast<uint64_t>(0),
101+ OP_LOGE(context_->GetNodeName(), "sliceSize %lu is zero. please check.", sliceSize_),
102+ return ge::GRAPH_FAILED);
103+ auto updateTypeSize = ge::GetSizeByDataType(updateDtype_);
104+ OP_CHECK_IF(
105+ updateTypeSize <= 0,
106+ OP_LOGE(context_, "updateTypeSize must be greater than 0, updateTypeSize: %ld", updateTypeSize),
107+ return ge::GRAPH_FAILED);
108+ indiceDtype_ = context_->GetInputDesc(INPUT_IDX_INDICES)->GetDataType();
109+ auto indiceTypeSize = ge::GetSizeByDataType(indiceDtype_);
110+ // sliceUb : the required size of UB for one scatter operation;
111+ auto sliceUb = sliceSize_ * updateTypeSize + rankSize_ * indiceTypeSize;
112+ sliceUb = Ops::Base::CeilDiv(sliceUb, static_cast<uint64_t>(alignFactor_)) * alignFactor_;
113+ if (sliceUb > halfUbSize) {
114+ // for scatter operator. At least rank size index need to be move in UB.
115+ ubTilingSize_ = (halfUbSize - rankSize_ * indiceTypeSize) / updateTypeSize;
116+ } else {
117+ // calculate the size of updates that need to be move in UB
118+ auto maxIndiceCnt = halfUbSize / sliceUb;
119+ ubTilingSize_ = maxIndiceCnt * sliceSize_;
120+ }
121+ OP_LOGD(context_->GetNodeName(), "sliceUb = %lu, halfUbSize = %u, ubTilingSize = %u", sliceUb, halfUbSize,
122+ ubTilingSize_);
123+ return ge::GRAPH_SUCCESS;
124+}
125+ 
126+ 
127+ge::graphStatus ScatterNdCommonSimtTiling::DoOpTiling()
128+{
129+ ge::graphStatus res = BlockTiling();
130+ if (res == ge::GRAPH_FAILED) {
131+ return ge::GRAPH_FAILED;
132+ }
133+ res = UbTiling();
134+ if (res == ge::GRAPH_FAILED) {
135+ return ge::GRAPH_FAILED;
136+ }
137+ 
138+ SetStride();
139+ SetTilingData();
140+ return ge::GRAPH_SUCCESS;
141+}
142+ 
143+ge::graphStatus ScatterNdCommonSimtTiling::PostTiling()
144+{
145+ OP_LOGD("ScatterNdCommonSimtTiling::PostTiling begin");
146+ context_->SetBlockDim(blockNum_);
147+ if (indiceShapeSize_ == 0 || updateShapeSize_ == 0) {
148+ // 输入为空tensor时，设置blockNum为1，在kernel中直接返回
149+ context_->SetBlockDim(1);
150+ }
151+ 
152+ auto res = context_->SetLocalMemorySize(ubSize_ - DCACHE_SIZE);
153+ OP_CHECK_IF((res != ge::GRAPH_SUCCESS),
154+ OP_LOGE(context_->GetNodeName(), "SetLocalMemorySize ubSize = %lu failed.", ubSize_), return ge::GRAPH_FAILED);
155+ 
156+ return ge::GRAPH_SUCCESS;
157+}
158+ 
159+ 
160+std::string ScatterNdCommonSimtTiling::TilingDataToString()
161+{
162+ ScatterNdCommon::ScatterNdCommonSimtTilingData* tilingData =
163+ context_->GetTilingData<ScatterNdCommon::ScatterNdCommonSimtTilingData>();
164+ std::string str = " blockNum:" + std::to_string(tilingData->blockNum);
165+ str += " rankSize:" + std::to_string(tilingData->rankSize);
166+ str += " blockTilingSize:" + std::to_string(tilingData->blockTilingSize);
167+ str += " tailBlockTilingSize:" + std::to_string(tilingData->tailBlockTilingSize);
168+ str += " ubTilingSize:" + std::to_string(tilingData->ubTilingSize);
169+ str += " sliceSize:" + std::to_string(tilingData->sliceSize);
170+ str += " varInAxis:" + std::to_string(tilingData->varInAxis);
171+ for (int32_t i = 0; i < OPTILING_MAX_SHAPE_RANK; i++) {
172+ str += " outPutShape[i]:" + std::to_string(tilingData->outPutShape[i]);
173+ }
174+ for (int32_t i = 0; i < OPTILING_MAX_RANK_COUNT; i++) {
175+ str += " strideList[i]:" + std::to_string(tilingData->strideList[i]);
176+ }
177+ return str;
178+}
179+ 
180+void ScatterNdCommonSimtTiling::DumpTilingInfo()
181+{
182+ OP_LOGI(context_->GetNodeName(), "Tiling info is: %s", TilingDataToString().c_str());
183+}
184+ 
185+} // namespace optiling
Aindex/scatter_nd_common/op_host/arch35/scatter_nd_common_simt_tiling.h+50-0
@@ -0,0 +1,50 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_simt_tiling.h
13+ * \brief scatter_nd_common_simt_tiling
14+ */
15+ 
16+#ifndef SCATTER_ND_COMMON_SIMT_TILING_H
17+#define SCATTER_ND_COMMON_SIMT_TILING_H
18+ 
19+#include "scatter_nd_common_base_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+class ScatterNdCommonSimtTiling : public ScatterNdCommonBaseTiling
24+{
25+public:
26+ explicit ScatterNdCommonSimtTiling(gert::TilingContext* context) : ScatterNdCommonBaseTiling(context)
27+ {}
28+ ~ScatterNdCommonSimtTiling() override
29+ {}
30+ 
31+protected:
32+ bool IsCapable() override;
33+ ge::graphStatus DoOpTiling() override;
34+ ge::graphStatus PostTiling() override;
35+ uint64_t GetTilingKey() const override;
36+ void DumpTilingInfo() override;
37+ std::string TilingDataToString();
38+ void SetTilingData();
39+ ge::graphStatus UbTiling();
40+ ge::graphStatus BlockTiling();
41+ 
42+ int64_t blockNum_ = 0;
43+ int64_t alignFactor_ = 0;
44+ int64_t blockTilingSize_ = 0;
45+ uint64_t tailBlockTilingSize_ = 0;
46+ uint32_t ubTilingSize_ = 0;
47+ uint64_t sliceSize_ = 0;
48+};
49+} // namespace optiling
50+#endif // SCATTER_ND_COMMON_SIMT_TILING_H
Aindex/scatter_nd_common/op_kernel/arch35/indices_sort_utils.h+184-0
@@ -0,0 +1,184 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License")
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file indices_sort_utils.h
13+ * \brief indices_sort_utils
14+ */
15+#ifndef ASCENDC_SCATTER_COMMON_INDICES_SORT_UTILS_H_
16+#define ASCENDC_SCATTER_COMMON_INDICES_SORT_UTILS_H_
17+ 
18+#define LAST_DIM_SIZE_LIMIT 512
19+#define INDICES_BUCKETS_SIZE 256
20+#define FNV_PRIME_B32 0x01000193UL
21+#define FNV_PRIME_B64 0x0100000001B3UL
22+#define FNV_OFFSET_BIASIS_B32 0x811C9DC5UL
23+#define FNV_OFFSET_BIASIS_B64 0xCBF29CE484222325UL
24+
25+using namespace AscendC;
26+static constexpr MicroAPI::CastTrait castTraitInt16ToFp32 = {
27+ MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
28+/*
29+* 用于计算
30+*/
31+template<typename INDICE_CAST_TYPE>
32+__aicore__ void IndexStatisticInt32(
33+ LocalTensor<INDICE_CAST_TYPE>& srcLocal, LocalTensor<float>& dstLocal, float& maxScore, int64_t rowLen, int64_t lastDim)
34+{
35+ __local_mem__ uint32_t* srcLocalAddr = (__local_mem__ uint32_t*)srcLocal.GetPhyAddr();
36+ __local_mem__ uint32_t* srcM = srcLocalAddr;
37+ __local_mem__ float* dstLocalAddr = (__local_mem__ float*)dstLocal.GetPhyAddr();
38+
39+ uint64_t lastDimSize = lastDim * sizeof(INDICE_CAST_TYPE);
40+ int32_t lastDimShift = 0;
41+ if (lastDimSize < LAST_DIM_SIZE_LIMIT) {
42+ uint64_t divisor = LAST_DIM_SIZE_LIMIT / lastDimSize;
43+ auto lz = 64 - ScalarCountLeadingZero(divisor);
44+ auto bc1 = ScalarGetCountOfValue<1>(divisor);
45+ lastDimShift = (bc1 != 1) ? lz : lz - 1;
46+ }
47+ uint16_t mainLoop = rowLen / INDICES_BUCKETS_SIZE;
48+ uint32_t tailNum = rowLen % INDICES_BUCKETS_SIZE;
49+ uint16_t tailLoop = ops::CeilDiv(tailNum, static_cast<uint32_t>(64));
50+ __VEC_SCOPE__
51+ {
52+ using namespace AscendC::MicroAPI;
53+ MaskReg patAllB32 = CreateMask<uint32_t, MaskPattern::ALL>();
54+ MaskReg patAllB16 = CreateMask<uint16_t, MaskPattern::ALL>();
55+ MaskReg patAllB8 = CreateMask<uint8_t, MaskPattern::ALL>();
56+ 
57+ RegTensor<float> maxCntFp32;
58+ RegTensor<uint16_t> histVector0;
59+ RegTensor<uint16_t> histVector1;
60+ Duplicate(histVector0, (uint16_t)0);
61+ Duplicate(histVector1, (uint16_t)0);
62+ 
63+ RegTensor<uint32_t> xorOffset;
64+ Duplicate(xorOffset, FNV_OFFSET_BIASIS_B32);
65+ RegTensor<int32_t> lastDimShiftAmount;
66+ Duplicate(lastDimShiftAmount, lastDimShift);
67+ 
68+ for (uint16_t i = 0; i < mainLoop; i++) {
69+ RegTensor<uint32_t> vectorIndex0, vectorIndex1, vectorIndex2, vectorIndex3;
70+ RegTensor<uint16_t> vectorB16Tmp0, vectorB16Tmp1, vectorB16Tmp2, vectorB16Tmp3;
71+ RegTensor<uint8_t> vectorB8Hash0, vectorB8Hash1;
72+ 
73+ DataCopy<uint32_t, PostLiteral::POST_MODE_UPDATE, LoadDist::DIST_NORM>(vectorIndex0, srcM, 64);
74+ DataCopy<uint32_t, PostLiteral::POST_MODE_UPDATE, LoadDist::DIST_NORM>(vectorIndex1, srcM, 64);
75+ DataCopy<uint32_t, PostLiteral::POST_MODE_UPDATE, LoadDist::DIST_NORM>(vectorIndex2, srcM, 64);
76+ DataCopy<uint32_t, PostLiteral::POST_MODE_UPDATE, LoadDist::DIST_NORM>(vectorIndex3, srcM, 64);
77+ 
78+ ShiftRight(vectorIndex0, vectorIndex0, lastDimShiftAmount, patAllB32);
79+ ShiftRight(vectorIndex1, vectorIndex1, lastDimShiftAmount, patAllB32);
80+ ShiftRight(vectorIndex2, vectorIndex2, lastDimShiftAmount, patAllB32);
81+ ShiftRight(vectorIndex3, vectorIndex3, lastDimShiftAmount, patAllB32);
82+ 
83+ Xor(vectorIndex0, vectorIndex0, xorOffset, patAllB32);
84+ Xor(vectorIndex1, vectorIndex1, xorOffset, patAllB32);
85+ Xor(vectorIndex2, vectorIndex2, xorOffset, patAllB32);
86+ Xor(vectorIndex3, vectorIndex3, xorOffset, patAllB32);
87+ 
88+ Muls(vectorIndex0, vectorIndex0, FNV_PRIME_B32, patAllB32);
89+ Muls(vectorIndex1, vectorIndex1, FNV_PRIME_B32, patAllB32);
90+ Muls(vectorIndex2, vectorIndex2, FNV_PRIME_B32, patAllB32);
91+ Muls(vectorIndex3, vectorIndex3, FNV_PRIME_B32, patAllB32);
92+ 
93+ DeInterleave(vectorB16Tmp0, vectorB16Tmp1, (RegTensor<uint16_t> &)vectorIndex0, (RegTensor<uint16_t> &)vectorIndex1);
94+ DeInterleave(vectorB16Tmp2, vectorB16Tmp3, (RegTensor<uint16_t> &)vectorIndex2, (RegTensor<uint16_t> &)vectorIndex3);
95+ DeInterleave(vectorB8Hash0, vectorB8Hash1, (RegTensor<uint8_t> &)vectorB16Tmp0, (RegTensor<uint8_t> &)vectorB16Tmp2);
96+ 
97+ Histograms<uint8_t, uint16_t, HistogramsBinType::BIN0, HistogramsType::FREQUENCY>(histVector0, vectorB8Hash0, patAllB8);
98+ Histograms<uint8_t, uint16_t, HistogramsBinType::BIN1, HistogramsType::FREQUENCY>(histVector1, vectorB8Hash0, patAllB8);
99+ }
100+ 
101+ for (uint16_t i = 0; i < tailLoop; i++) {
102+ MaskReg maskReg = UpdateMask<uint32_t>(tailNum);
103+ RegTensor<uint32_t> vectorIndex0;
104+ DataCopy<uint32_t, PostLiteral::POST_MODE_UPDATE, LoadDist::DIST_NORM>(vectorIndex0, srcM, 64);
105+ ShiftRight(vectorIndex0, vectorIndex0, lastDimShiftAmount, patAllB32);
106+ Xor(vectorIndex0, vectorIndex0, xorOffset, patAllB32);
107+ Muls(vectorIndex0, vectorIndex0, FNV_PRIME_B32, patAllB32);
108+ 
109+ Histograms<uint8_t, uint16_t, HistogramsBinType::BIN0, HistogramsType::FREQUENCY>(histVector0, (RegTensor<uint8_t> &)vectorIndex0, maskReg);
110+ Histograms<uint8_t, uint16_t, HistogramsBinType::BIN1, HistogramsType::FREQUENCY>(histVector1, (RegTensor<uint8_t> &)vectorIndex0, maskReg);
111+ }
112+ 
113+ Max(histVector0, histVector0, histVector1, patAllB16);
114+ ReduceMax(histVector0, histVector0, patAllB16);
115+ 
116+ Cast<float, int16_t, castTraitInt16ToFp32>(maxCntFp32, (RegTensor<int16_t> &)histVector0, patAllB16);
117+ DataCopy<float, PostLiteral::POST_MODE_UPDATE, StoreDist::DIST_FIRST_ELEMENT_B32>(dstLocalAddr, maxCntFp32, 1, patAllB32);
118+ }
119+ 
120+ event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
121+ SetFlag<HardEvent::V_S>(eventIdVToS);
122+ WaitFlag<HardEvent::V_S>(eventIdVToS);
123+ 
124+ maxScore = dstLocalAddr[0] / rowLen;
125+}
126+ 
127+template<typename INDICE_CAST_TYPE>
128+__aicore__ void IndexStatisticInt64(
129+ LocalTensor<INDICE_CAST_TYPE>& srcLocal, LocalTensor<float>& dstLocal, float& maxScore, int64_t rowLen, int64_t lastDim)
130+{
131+ __local_mem__ uint64_t* srcLocalAddr = (__local_mem__ uint64_t*)srcLocal.GetPhyAddr();
132+ __local_mem__ uint64_t* srcM = srcLocalAddr;
133+ __local_mem__ float* dstLocalAddr = (__local_mem__ float*)dstLocal.GetPhyAddr();
134+
135+ int32_t lastDimSize = lastDim * sizeof(INDICE_CAST_TYPE);
136+ int64_t lastDimShift = 0;
137+ if (lastDimSize < LAST_DIM_SIZE_LIMIT) {
138+ uint64_t divisor = LAST_DIM_SIZE_LIMIT / lastDimSize;
139+ auto lz = 64 - ScalarCountLeadingZero(divisor);
140+ auto bc1 = ScalarGetCountOfValue<1>(divisor);
141+ lastDimShift = (bc1 != 1) ? lz : lz - 1;
142+ }
143+ uint32_t dataLen = static_cast<uint32_t>(rowLen);
144+ uint16_t loopSize = ops::CeilDiv(rowLen, 32L);
145+ __VEC_SCOPE__
146+ {
147+ using namespace AscendC::MicroAPI;
148+ MaskReg patAllB32 = CreateMask<uint32_t, MaskPattern::ALL>();
149+ MaskReg patAllB16 = CreateMask<uint16_t, MaskPattern::ALL>();
150+ 
151+ RegTensor<float> maxCntFp32;
152+ RegTensor<uint16_t> histVector0, histVector1, maxValue;
153+ Duplicate(histVector0, (uint16_t)0);
154+ Duplicate(histVector1, (uint16_t)0);
155+ 
156+ RegTensor<uint64_t> xorOffset;
157+ RegTensor<int64_t> lastDimShiftAmount;
158+ Duplicate(xorOffset, FNV_OFFSET_BIASIS_B64);
159+ Duplicate(lastDimShiftAmount, lastDimShift);
160+ 
161+ RegTensor<uint64_t> vectorIndex0;
162+ for (uint16_t i = 0; i < loopSize; i++) {
163+ MaskReg maskReg = UpdateMask<uint64_t>(dataLen);
164+ DataCopy(vectorIndex0, srcM);
165+ ShiftRight(vectorIndex0, vectorIndex0, lastDimShiftAmount, maskReg);
166+ Xor(vectorIndex0, vectorIndex0, xorOffset, maskReg);
167+ Muls(vectorIndex0, vectorIndex0, FNV_PRIME_B64, maskReg);
168+ 
169+ Histograms<uint8_t, uint16_t, HistogramsBinType::BIN0, HistogramsType::FREQUENCY>(histVector0, (RegTensor<uint8_t> &)vectorIndex0, maskReg);
170+ Histograms<uint8_t, uint16_t, HistogramsBinType::BIN1, HistogramsType::FREQUENCY>(histVector1, (RegTensor<uint8_t> &)vectorIndex0, maskReg);
171+ Max(maxValue, histVector0, histVector1, patAllB16);
172+ }
173+ ReduceMax(maxValue, maxValue, patAllB16);
174+ Cast<float, int16_t, castTraitInt16ToFp32>(maxCntFp32, (RegTensor<int16_t> &)maxValue, patAllB16);
175+ DataCopy<float, PostLiteral::POST_MODE_UPDATE, StoreDist::DIST_FIRST_ELEMENT_B32>(dstLocalAddr, maxCntFp32, 1, patAllB32);
176+ }
177+ 
178+ event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
179+ SetFlag<HardEvent::V_S>(eventIdVToS);
180+ WaitFlag<HardEvent::V_S>(eventIdVToS);
181+ maxScore = dstLocalAddr[0] / rowLen;
182+}
183+ 
184+#endif
Aindex/scatter_nd_common/op_kernel/arch35/scatter_nd_common_base.h+888-0
@@ -0,0 +1,888 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_base.h
13+ * \brief scatter_nd_common_base
14+ */
15+#ifndef ASCENDC_SCATTER_ND_COMMON_BASE_H_
16+#define ASCENDC_SCATTER_ND_COMMON_BASE_H_
17+ 
18+#include "kernel_operator.h"
19+#include "../inc/platform.h"
20+#include "scatter_nd_common_struct.h"
21+#include "../inc/load_store_utils.h"
22+#include "./indices_sort_utils.h"
23+ 
24+namespace ScatterNdCommon {
25+using namespace AscendC;
26+ 
27+#ifdef __DAV_FPGA__
28+constexpr uint32_t THREAD_NUM = 128;
29+constexpr uint32_t THREAD_NUM_LAUNCH_BOUND = 512;
30+#else
31+constexpr uint32_t THREAD_NUM = 1024;
32+constexpr uint32_t THREAD_NUM_LAUNCH_BOUND = 1024;
33+#endif
34+constexpr int64_t DOUBLE_BUFFER = 2;
35+constexpr uint64_t UB_AGLIN_VALUE = 32;
36+constexpr uint16_t MAX_RANK_COUNT = 7;
37+constexpr uint64_t SORT_PAD_NUM = 2;
38+constexpr uint16_t MIN_SAME_IDX_ACCM_COUNT = 256;
39+constexpr uint16_t MAX_SHAPE_RANK = 8;
40+constexpr uint16_t INDICE_RANK_TWO = 2;
41+constexpr float SORT_HIST_THRESHOLD = 0.03f;
42+constexpr uint32_t HASH_SCORE_BUF_SIZE = 128;
43+constexpr uint32_t CAST_0 = 0;
44+constexpr uint32_t CAST_1 = 1;
45+constexpr uint32_t CAST_2 = 2;
46+constexpr uint32_t CAST_3 = 3;
47+constexpr uint32_t CAST_4 = 4;
48+constexpr uint32_t CAST_5 = 5;
49+constexpr uint8_t MODE_MAX = 1; // 1: scatter_nd_max;
50+constexpr uint8_t MODE_MIN = 0; // 0: scatter_nd_min;
51+constexpr float FLOAT32_MAX = 3.4028235e+38f;
52+constexpr half FLOAT16_MAX = 65504.0f;
53+constexpr bfloat16_t BFLOAT16_MAX = 3.3895314e+38f;
54+constexpr float FLOAT32_MIN = -3.4028235e+38f;
55+constexpr half FLOAT16_MIN = -65504.0f;
56+constexpr bfloat16_t BFLOAT16_MIN = -3.3895314e+38f;
57+constexpr int64_t VFLEN_INT64 = platform::GetVRegSize() / sizeof(int64_t);
58+constexpr int64_t VFLEN_INT32 = platform::GetVRegSize() / sizeof(int32_t);
59+constexpr int64_t VFLEN_INT16 = platform::GetVRegSize() / sizeof(int16_t);
60+constexpr int64_t VFLEN_INT16HALF = platform::GetVRegSize() / sizeof(int16_t) / TWO;
61+constexpr int64_t VFLEN_UINT8 = platform::GetVRegSize() / sizeof(uint8_t);
62+constexpr int64_t VFLEN_UINT8HALFHALF = platform::GetVRegSize() / sizeof(uint8_t) / FOUR;
63+constexpr uint32_t U8_MAX = 255;
64+ 
65+constexpr MicroAPI::CastTrait castTraitB322B64 = {
66+ MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
67+ 
68+static constexpr SortConfig sortConfig{SortType::RADIX_SORT, false};
69+ 
70+__aicore__ inline uint32_t ROUND_UP32(uint32_t x)
71+{
72+ if (x % UB_AGLIN_VALUE != 0) {
73+ return (x / UB_AGLIN_VALUE + 1) * UB_AGLIN_VALUE;
74+ }
75+ return x;
76+}
77+ 
78+template <typename TYPE>
79+__aicore__ inline constexpr TYPE GetDtypeMax()
80+{
81+ TYPE dtypeMax = 0;
82+ if constexpr (IsSameType<TYPE, int32_t>::value) {
83+ dtypeMax = INT32_MAX;
84+ } else if constexpr (IsSameType<TYPE, int64_t>::value) {
85+ dtypeMax = INT64_MAX;
86+ } else if constexpr (IsSameType<TYPE, uint32_t>::value) {
87+ dtypeMax = UINT32_MAX;
88+ } else if constexpr (IsSameType<TYPE, uint64_t>::value) {
89+ dtypeMax = UINT64_MAX;
90+ } else if constexpr (IsSameType<TYPE, bfloat16_t>::value) {
91+ dtypeMax = BFLOAT16_MAX;
92+ } else if constexpr (IsSameType<TYPE, half>::value) {
93+ dtypeMax = FLOAT16_MAX;
94+ } else if constexpr (IsSameType<TYPE, float>::value) {
95+ dtypeMax = FLOAT32_MAX;
96+ } else if constexpr (IsSameType<TYPE, int8_t>::value) {
97+ dtypeMax = INT8_MAX;
98+ } else if constexpr (IsSameType<TYPE, int16_t>::value) {
99+ dtypeMax = INT16_MAX;
100+ }
101+ return dtypeMax;
102+}
103+ 
104+template <typename TYPE>
105+__aicore__ inline constexpr TYPE GetDtypeMin()
106+{
107+ TYPE dtypeMin = 0;
108+ if constexpr (IsSameType<TYPE, int32_t>::value) {
109+ dtypeMin = INT32_MIN;
110+ } else if constexpr (IsSameType<TYPE, int64_t>::value) {
111+ dtypeMin = INT64_MIN;
112+ } else if constexpr (IsSameType<TYPE, uint32_t>::value) {
113+ dtypeMin = 0;
114+ } else if constexpr (IsSameType<TYPE, uint64_t>::value) {
115+ dtypeMin = 0;
116+ } else if constexpr (IsSameType<TYPE, bfloat16_t>::value) {
117+ dtypeMin = BFLOAT16_MIN;
118+ } else if constexpr (IsSameType<TYPE, half>::value) {
119+ dtypeMin = FLOAT16_MIN;
120+ } else if constexpr (IsSameType<TYPE, float>::value) {
121+ dtypeMin = FLOAT32_MIN;
122+ } else if constexpr (IsSameType<TYPE, int8_t>::value) {
123+ dtypeMin = INT8_MIN;
124+ } else if constexpr (IsSameType<TYPE, int16_t>::value) {
125+ dtypeMin = INT16_MIN;
126+ }
127+ return dtypeMin;
128+}
129+ 
130+template <typename T, typename U, typename CAST_T = U, typename OFFSET_T = U, uint32_t castType = CAST_0, uint8_t Mode = MODE_MAX>
131+class ScatterNdCommonBase
132+{
133+public:
134+ int64_t indicesFactor_ = 0;
135+ int64_t afterAxisFactor_ = 0;
136+ int64_t afterAxis_ = 0;
137+ int64_t indexRankSize_ = 0;
138+ int64_t eachCoreAfterAxisCount_ = 0;
139+ int64_t eachCoreIndexCount_ = 0;
140+ int64_t shiftOffset_ = UB_AGLIN_VALUE / sizeof(CAST_T);
141+ uint32_t uniqueIdNum_ = 0;
142+ float maxScore_ = static_cast<float>(0);
143+ 
144+ AscendC::GlobalTensor<U> indicesGm_;
145+ AscendC::GlobalTensor<T> updatesGm_;
146+ AscendC::GlobalTensor<T> yGm_;
147+ 
148+ TBuf<QuePosition::VECCALC> indicesBuf_;
149+ TBuf<QuePosition::VECCALC> outOfstBuf_;
150+ TBuf<QuePosition::VECCALC> strideBuf_;
151+ TBuf<QuePosition::VECCALC> outputShapeBuf_;
152+ TBuf<QuePosition::VECCALC> maxScoreBuf_;
153+ TQueBind<QuePosition::VECIN, QuePosition::VECOUT, 1> dataQueue_;
154+ 
155+ TBuf<QuePosition::VECCALC> sortIndicesQue_;
156+ TBuf<QuePosition::VECCALC> castIndicesQue_;
157+ TBuf<QuePosition::VECCALC> castTmpIndicesQue_;
158+ TQue<QuePosition::VECIN, 1> updatesOriginIdexQue_;
159+ TQue<QuePosition::VECIN, 1> uniqueIdCountQue_;
160+ TQue<QuePosition::VECOUT, 1> updateSumIdxQue_;
161+ TQue<QuePosition::VECOUT, 1> updateSumQue_;
162+ 
163+ using IndexRegType = typename std::conditional<
164+ IsSameType<U, int64_t>::value,
165+ typename AscendC::MicroAPI::RegTensor<uint64_t, AscendC::MicroAPI::RegTraitNumTwo>,
166+ typename AscendC::MicroAPI::RegTensor<uint32_t>>::type;
167+ using InnerRegType = typename std::conditional<
168+ IsSameType<OFFSET_T, int64_t>::value,
169+ typename AscendC::MicroAPI::RegTensor<int64_t, AscendC::MicroAPI::RegTraitNumTwo>,
170+ typename AscendC::MicroAPI::RegTensor<int32_t>>::type;
171+ 
172+ using selRegType = typename std::conditional<IsSameType<T, bool>::value, int8_t, T>::type;
173+ 
174+ using COMPUTE_TYPE = typename std::conditional<IsSameType<T, bool>::value, int8_t, T>::type;
175+ 
176+ __aicore__ inline void InitBaseBuffer(
177+ TPipe& pipe, uint32_t indicesNumber, GM_ADDR indices, GM_ADDR updates, GM_ADDR y, int64_t singleCol = 0)
178+ {
179+ indicesGm_.SetGlobalBuffer((__gm__ U*)(indices));
180+ updatesGm_.SetGlobalBuffer((__gm__ T*)(updates));
181+ yGm_.SetGlobalBuffer((__gm__ T*)(y));
182+ 
183+ pipe.InitBuffer(strideBuf_, MAX_RANK_COUNT * sizeof(OFFSET_T));
184+ pipe.InitBuffer(outputShapeBuf_, MAX_SHAPE_RANK * sizeof(U));
185+ pipe.InitBuffer(outOfstBuf_, indicesFactor_ * sizeof(OFFSET_T));
186+ pipe.InitBuffer(indicesBuf_, indicesFactor_ * indexRankSize_ * sizeof(U));
187+ pipe.InitBuffer(maxScoreBuf_, HASH_SCORE_BUF_SIZE * sizeof(float));
188+ 
189+ pipe.InitBuffer(updatesOriginIdexQue_, 1, indicesFactor_ * sizeof(uint32_t));
190+ pipe.InitBuffer(
191+ uniqueIdCountQue_, 1, ops::CeilAlign((indicesFactor_ + 1) * sizeof(int32_t), UB_AGLIN_VALUE));
192+ pipe.InitBuffer(
193+ updateSumIdxQue_, 1, ops::CeilAlign((indicesFactor_ + 1) * sizeof(OFFSET_T), UB_AGLIN_VALUE));
194+ if (singleCol) {
195+ pipe.InitBuffer(dataQueue_, DOUBLE_BUFFER, afterAxisFactor_ * sizeof(T));
196+ pipe.InitBuffer(updateSumQue_, DOUBLE_BUFFER, afterAxisFactor_ * sizeof(T));
197+ } else {
198+ pipe.InitBuffer(dataQueue_, 1, indicesFactor_ * afterAxisFactor_ * sizeof(T));
199+ pipe.InitBuffer(updateSumQue_, 1, indicesFactor_ * afterAxisFactor_ * sizeof(T));
200+ }
201+ 
202+ if constexpr (castType == CAST_0) {
203+ pipe.InitBuffer(sortIndicesQue_, ops::CeilAlign(indicesFactor_ * sizeof(OFFSET_T) + SORT_PAD_NUM * UB_AGLIN_VALUE, UB_AGLIN_VALUE));
204+ } else {
205+ pipe.InitBuffer(sortIndicesQue_, ops::CeilAlign(indicesFactor_ * sizeof(CAST_T) + SORT_PAD_NUM * UB_AGLIN_VALUE, UB_AGLIN_VALUE));
206+ pipe.InitBuffer(castIndicesQue_, ops::CeilAlign(indicesFactor_ * sizeof(CAST_T), UB_AGLIN_VALUE));
207+ }
208+ }
209+ 
210+ 
211+ template <typename PARAM_T>
212+ __aicore__ inline void CopyIn(
213+ const LocalTensor<PARAM_T>& dstTensor, const GlobalTensor<PARAM_T>& srcTensor, int64_t dataLen)
214+ {
215+ DataCopyExtParams copyParams = {
216+ static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(PARAM_T)), static_cast<uint32_t>(0),
217+ static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
218+ DataCopyPadExtParams<PARAM_T> padParams = {
219+ false, static_cast<uint8_t>(0), static_cast<uint8_t>(0), static_cast<PARAM_T>(0)};
220+ DataCopyPad(dstTensor, srcTensor, copyParams, padParams);
221+ }
222+ 
223+ template <typename PARAM_T>
224+ __aicore__ inline void CopyOut(
225+ const GlobalTensor<PARAM_T>& dstTensor, const LocalTensor<PARAM_T>& srcTensor, int64_t dataLen)
226+ {
227+ DataCopyExtParams copyParams = {
228+ static_cast<uint16_t>(1), static_cast<uint32_t>(dataLen * sizeof(PARAM_T)), static_cast<uint32_t>(0),
229+ static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
230+ DataCopyPad(dstTensor, srcTensor, copyParams);
231+ }
232+ 
233+ __aicore__ inline void ComputeOutOfset(
234+ const LocalTensor<U> indicesLocal, const LocalTensor<OFFSET_T> outOfstLocal, int32_t indicesLen, int32_t rankSize)
235+ {
236+ LocalTensor<OFFSET_T> strideLocal = strideBuf_.Get<OFFSET_T>();
237+ LocalTensor<U> outputShapeLocal = outputShapeBuf_.Get<U>();
238+ 
239+ __local_mem__ U* indicesLocalPtr = ((__local_mem__ U*)indicesLocal.GetPhyAddr());
240+ __local_mem__ OFFSET_T* outOfstLocalPtr = ((__local_mem__ OFFSET_T*)outOfstLocal.GetPhyAddr());
241+ 
242+ uint32_t dataLen = indicesLen;
243+ uint32_t vfLen = platform::GetVRegSize() / sizeof(int32_t);
244+ uint32_t indicesLenTimes = ops::CeilDiv(dataLen, vfLen);
245+ uint16_t loopCnt = static_cast<uint16_t>(indicesLenTimes);
246+ uint16_t rankSizeLoops = static_cast<uint16_t>(rankSize);
247+ 
248+ __VEC_SCOPE__
249+ {
250+ InnerRegType inReg;
251+ InnerRegType outReg;
252+ InnerRegType orderReg;
253+ InnerRegType selectReg;
254+ IndexRegType indexReg;
255+ AscendC::MicroAPI::MaskReg pregLoop;
256+ AscendC::MicroAPI::MaskReg cmpMask;
257+ AscendC::MicroAPI::MaskReg invalidMask;
258+ 
259+ for (uint16_t i = 0; i < loopCnt; i++) {
260+ if constexpr (IsSameType<OFFSET_T, int64_t>::value) {
261+ pregLoop = AscendC::MicroAPI::UpdateMask<OFFSET_T, AscendC::MicroAPI::RegTraitNumTwo>(dataLen);
262+ invalidMask = AscendC::MicroAPI::CreateMask<OFFSET_T, MicroAPI::MaskPattern::ALLF, AscendC::MicroAPI::RegTraitNumTwo>();
263+ } else {
264+ pregLoop = AscendC::MicroAPI::UpdateMask<OFFSET_T>(dataLen);
265+ invalidMask = AscendC::MicroAPI::CreateMask<OFFSET_T, MicroAPI::MaskPattern::ALLF>();
266+ }
267+ AscendC::MicroAPI::Duplicate(outReg, 0, pregLoop);
268+ AscendC::MicroAPI::Arange(orderReg, i * vfLen);
269+ AscendC::MicroAPI::Muls(orderReg, orderReg, rankSize, pregLoop);
270+ for (uint16_t dim = 0; dim < rankSizeLoops; dim++) {
271+ OFFSET_T strideValue = strideLocal(dim);
272+ U outputShapeValue = outputShapeLocal(dim);
273+ indexReg = (IndexRegType&)orderReg;
274+ 
275+ if constexpr (IsSameType<U, int32_t>::value && IsSameType<OFFSET_T, int64_t>::value) {
276+ AscendC::MicroAPI::RegTensor<int32_t> castReg;
277+ AscendC::MicroAPI::DataCopyGather(castReg, indicesLocalPtr, indexReg, pregLoop);
278+ MicroAPI::Cast<int64_t, int32_t, castTraitB322B64>(inReg, castReg, pregLoop);
279+ } else {
280+ AscendC::MicroAPI::DataCopyGather(inReg, indicesLocalPtr, indexReg, pregLoop);
281+ }
282+
283+ 
284+ AscendC::MicroAPI::CompareScalar<OFFSET_T, CMPMODE::LT>(cmpMask, inReg, static_cast<OFFSET_T>(0), pregLoop);
285+ AscendC::MicroAPI::Or(invalidMask, invalidMask, cmpMask, pregLoop);
286+ AscendC::MicroAPI::CompareScalar<OFFSET_T, CMPMODE::GE>(cmpMask, inReg, static_cast<OFFSET_T>(outputShapeValue), pregLoop);
287+ AscendC::MicroAPI::Or(invalidMask, invalidMask, cmpMask, pregLoop);
288+ 
289+ AscendC::MicroAPI::Muls(inReg, inReg, strideValue, pregLoop);
290+ AscendC::MicroAPI::Add(outReg, inReg, outReg, pregLoop);
291+ AscendC::MicroAPI::Adds(orderReg, orderReg, (OFFSET_T)(1), pregLoop);
292+ }
293+ AscendC::MicroAPI::Duplicate(selectReg, static_cast<OFFSET_T>(-2), pregLoop);
294+ AscendC::MicroAPI::Select(outReg, selectReg, outReg, invalidMask);
295+ auto outOfstAddr = outOfstLocalPtr + i * vfLen;
296+ AscendC::MicroAPI::DataCopy(outOfstAddr, outReg, pregLoop);
297+ }
298+ }
299+ return;
300+ }
301+ 
302+ __aicore__ inline void ComputeUniqueIdNumInt64(__local_mem__ CAST_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
303+ {
304+ uint32_t counter = dataLen + 1;
305+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, selReg;
306+ AscendC::MicroAPI::RegTensor<CAST_T> sortedIdxReg, sortedIdxShiftOneReg;
307+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskHalf;
308+ AscendC::MicroAPI::UnalignReg u0, uOut;
309+ for (uint16_t i = 0; i < loopCnt; ++i) {
310+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT64);
311+ maskReg = AscendC::MicroAPI::UpdateMask<CAST_T>(counter);
312+ auto startAddr = indicesAddr + i * VFLEN_INT64;
313+ DataCopy(sortedIdxReg, startAddr);
314+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
315+ AscendC::MicroAPI::DataCopyUnAlign<CAST_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
316+ AscendC::MicroAPI::Compare<CAST_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
317+ AscendC::MicroAPI::MaskPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskHalf, cmpMask);
318+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskHalf);
319+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
320+ uniqueIdCountsAddr, selReg, uOut);
321+ }
322+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
323+ }
324+ 
325+ __aicore__ inline void ComputeUniqueIdNumInt32(__local_mem__ CAST_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
326+ {
327+ uint32_t counter = dataLen + 1;
328+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, selReg;
329+ AscendC::MicroAPI::RegTensor<CAST_T> sortedIdxReg, sortedIdxShiftOneReg;
330+ AscendC::MicroAPI::MaskReg cmpMask, maskReg;
331+ AscendC::MicroAPI::UnalignReg u0, uOut;
332+ for (uint16_t i = 0; i < loopCnt; ++i) {
333+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT32);
334+ maskReg = AscendC::MicroAPI::UpdateMask<CAST_T>(counter);
335+ auto startAddr = indicesAddr + i * VFLEN_INT32;
336+ DataCopy(sortedIdxReg, startAddr);
337+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
338+ AscendC::MicroAPI::DataCopyUnAlign<CAST_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
339+ AscendC::MicroAPI::Compare<CAST_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
340+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, cmpMask);
341+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
342+ uniqueIdCountsAddr, selReg, uOut);
343+ }
344+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
345+ }
346+ 
347+ __aicore__ inline void ComputeUniqueIdNumInt16(__local_mem__ CAST_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
348+ {
349+ uint32_t counter = dataLen + 1;
350+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, orderReg2, selReg, selReg2;
351+ AscendC::MicroAPI::RegTensor<CAST_T> sortedIdxReg, sortedIdxShiftOneReg;
352+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskDouble1, maskDouble2;
353+ AscendC::MicroAPI::UnalignReg u0, uOut;
354+ for (uint16_t i = 0; i < loopCnt; ++i) {
355+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_INT16);
356+ AscendC::MicroAPI::Arange(orderReg2, i * VFLEN_INT16 + VFLEN_INT16HALF);
357+ maskReg = AscendC::MicroAPI::UpdateMask<CAST_T>(counter);
358+ auto startAddr = indicesAddr + i * VFLEN_INT16;
359+ DataCopy(sortedIdxReg, startAddr);
360+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
361+ AscendC::MicroAPI::DataCopyUnAlign<CAST_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
362+ AscendC::MicroAPI::Compare<CAST_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
363+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskDouble1, cmpMask);
364+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskDouble2, cmpMask);
365+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskDouble1);
366+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg, uOut);
367+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskDouble2);
368+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg2, uOut);
369+ }
370+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
371+ }
372+ 
373+ __aicore__ inline void ComputeUniqueIdNumUint8(__local_mem__ CAST_T* indicesAddr, __local_mem__ int32_t* uniqueIdCountsAddr, uint16_t loopCnt, int64_t dataLen)
374+ {
375+ uint32_t counter = dataLen + 1;
376+ AscendC::MicroAPI::RegTensor<int32_t> orderReg, orderReg2, orderReg3, orderReg4;
377+ AscendC::MicroAPI::RegTensor<int32_t> selReg, selReg2, selReg3, selReg4;
378+ AscendC::MicroAPI::RegTensor<CAST_T> sortedIdxReg, sortedIdxShiftOneReg;
379+ AscendC::MicroAPI::MaskReg cmpMask, maskReg, maskFour1, maskFour2, maskFour3, maskFour4;
380+ AscendC::MicroAPI::UnalignReg u0, uOut;
381+ for (uint16_t i = 0; i < loopCnt; ++i) {
382+ AscendC::MicroAPI::Arange(orderReg, i * VFLEN_UINT8);
383+ AscendC::MicroAPI::Arange(orderReg2, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF);
384+ AscendC::MicroAPI::Arange(orderReg3, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF * TWO);
385+ AscendC::MicroAPI::Arange(orderReg4, i * VFLEN_UINT8 + VFLEN_UINT8HALFHALF * THREE);
386+ maskReg = AscendC::MicroAPI::UpdateMask<CAST_T>(counter);
387+ auto startAddr = indicesAddr + i * VFLEN_UINT8;
388+ DataCopy(sortedIdxReg, startAddr);
389+ AscendC::MicroAPI::DataCopyUnAlignPre(u0, startAddr - 1);
390+ AscendC::MicroAPI::DataCopyUnAlign<CAST_T>(sortedIdxShiftOneReg, u0, startAddr - 1);
391+ AscendC::MicroAPI::Compare<CAST_T, CMPMODE::NE>(cmpMask, sortedIdxReg, sortedIdxShiftOneReg, maskReg);
392+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour3, cmpMask);
393+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour4, cmpMask);
394+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour1, maskFour3);
395+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour2, maskFour3);
396+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::LOWEST>(maskFour3, maskFour4);
397+ AscendC::MicroAPI::MaskUnPack<AscendC::MicroAPI::HighLowPart::HIGHEST>(maskFour4, maskFour4);
398+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg, orderReg, maskFour1);
399+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg, uOut);
400+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg2, orderReg2, maskFour2);
401+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg2, uOut);
402+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg3, orderReg3, maskFour3);
403+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg3, uOut);
404+ AscendC::MicroAPI::GatherMask<int32_t, AscendC::MicroAPI::GatherMaskMode::STORE_REG>(selReg4, orderReg4, maskFour4);
405+ AscendC::MicroAPI::DataCopyUnAlign<int32_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(uniqueIdCountsAddr, selReg4, uOut);
406+ }
407+ AscendC::MicroAPI::DataCopyUnAlignPost(uniqueIdCountsAddr, uOut);
408+ }
409+ 
410+ __aicore__ inline uint32_t
411+ ComputeUniqueIdNum(LocalTensor<CAST_T> indicesLocal, LocalTensor<int32_t> uniqueIdCountLocal, int64_t dataLen)
412+ {
413+ __local_mem__ CAST_T* indicesAddr = (__local_mem__ CAST_T*)indicesLocal[(UB_AGLIN_VALUE / sizeof(CAST_T))].GetPhyAddr();
414+ __local_mem__ int32_t* uniqueIdCountsAddr = (__local_mem__ int32_t*)uniqueIdCountLocal.GetPhyAddr();
415+ 
416+ int64_t vfLen = platform::GetVRegSize() / sizeof(CAST_T);
417+ uint16_t loopCnt = ops::CeilDiv(dataLen + 1, vfLen);
418+ __VEC_SCOPE__
419+ {
420+ AscendC::MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
421+ if constexpr (std::is_same<int64_t, CAST_T>::value) {
422+ ComputeUniqueIdNumInt64(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
423+ } else if constexpr (std::is_same<int32_t, CAST_T>::value) {
424+ ComputeUniqueIdNumInt32(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
425+ } else if constexpr (std::is_same<int16_t, CAST_T>::value) {
426+ ComputeUniqueIdNumInt16(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
427+ } else {
428+ ComputeUniqueIdNumUint8(indicesAddr, uniqueIdCountsAddr, loopCnt, dataLen);
429+ }
430+ }
431+ uint32_t uniqueIdNum = ((AscendC::MicroAPI::GetSpr<AscendC::SpecialPurposeReg::AR>()) / sizeof(int32_t)) - 1;
432+ return uniqueIdNum;
433+ }
434+ 
435+ __aicore__ inline void ComputeUinqueIdTimes(LocalTensor<int32_t> uniqueIdCountLocal, uint32_t uniqueIdNum)
436+ {
437+ __local_mem__ int32_t* uniqueIdCountsAddr = (__local_mem__ int32_t*)uniqueIdCountLocal.GetPhyAddr();
438+ uint32_t vfLen = platform::GetVRegSize() / sizeof(int32_t);
439+ uint16_t loopSize = ops::CeilDiv(uniqueIdNum, vfLen);
440+ __VEC_SCOPE__
441+ {
442+ AscendC::MicroAPI::RegTensor<int32_t> preReg;
443+ AscendC::MicroAPI::RegTensor<int32_t> postReg;
444+ AscendC::MicroAPI::RegTensor<int32_t> subReg;
445+ AscendC::MicroAPI::UnalignReg uIn;
446+ AscendC::MicroAPI::MaskReg maskReg;
447+ for (uint16_t i = 0; i < loopSize; ++i) {
448+ maskReg = AscendC::MicroAPI::UpdateMask<int32_t>(uniqueIdNum);
449+ auto startAddr = uniqueIdCountsAddr + i * vfLen;
450+ auto startAddrOfstOne = startAddr + 1;
451+ DataCopy(preReg, startAddr);
452+ AscendC::MicroAPI::DataCopyUnAlignPre(uIn, startAddrOfstOne);
453+ AscendC::MicroAPI::DataCopyUnAlign<int32_t>(postReg, uIn, startAddrOfstOne, vfLen);
454+ AscendC::MicroAPI::Sub(subReg, postReg, preReg, maskReg);
455+ DataCopy(startAddr, subReg, maskReg);
456+ }
457+ }
458+ }
459+ 
460+ __aicore__ inline void IndicesSortCast(LocalTensor<U> indicesLocal, LocalTensor<CAST_T> indicesCastLocal,
461+ LocalTensor<int32_t> indicesCastTmpLocal, uint32_t indicesCount)
462+ {
463+ if constexpr (castType == CAST_4) { // int32 Cast uint8
464+ CompareScalar(indicesCastLocal, indicesLocal, static_cast<U>(0), CMPMODE::GE, indicesCount);
465+ Select(indicesLocal, indicesCastLocal, indicesLocal, static_cast<U>(U8_MAX), SELMODE::VSEL_TENSOR_SCALAR_MODE, indicesCount);
466+ Cast<CAST_T, U>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
467+ } else if constexpr (castType == CAST_3) { // int64 Cast int16
468+ Cast<int32_t, U>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
469+ Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
470+ } else if constexpr (castType == CAST_5) { // int64 Cast uint8
471+ CompareScalar(indicesCastLocal, indicesLocal, static_cast<U>(0), CMPMODE::GE, indicesCount);
472+ Select(indicesLocal, indicesCastLocal, indicesLocal, static_cast<U>(U8_MAX), SELMODE::VSEL_TENSOR_SCALAR_MODE, indicesCount);
473+ Cast<int32_t, U>(indicesCastTmpLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
474+ Cast<CAST_T, int32_t>(indicesCastLocal, indicesCastTmpLocal, RoundMode::CAST_NONE, indicesCount);
475+ } else { // CAST_1 + CAST_2, int32 Cast int16 + int64 Cast int32
476+ Cast<CAST_T, U>(indicesCastLocal, indicesLocal, RoundMode::CAST_NONE, indicesCount);
477+ }
478+ }
479+ 
480+ __aicore__ inline void ComputeSumWithOutCast(
481+ LocalTensor<int32_t> uniqueIdCountLocal, LocalTensor<uint32_t> updatesOriginIdexLocal,
482+ LocalTensor<T> updatesLocal, LocalTensor<T> updateSumLocal, uint32_t uniqueIdNum, int64_t colLen)
483+ {
484+ __local_mem__ selRegType* updatesAddr = (__local_mem__ selRegType*)updatesLocal.GetPhyAddr();
485+ __local_mem__ selRegType* updateSumAddr = (__local_mem__ selRegType*)updateSumLocal.GetPhyAddr();
486+ 
487+ uint32_t vfLen = platform::GetVRegSize() / sizeof(selRegType);
488+ int32_t loopSize = (colLen + vfLen - 1) / vfLen;
489+ int32_t idLocation = 0;
490+ int64_t colLenAlignSize = ops::CeilAlign(colLen * sizeof(selRegType), UB_AGLIN_VALUE) / sizeof(selRegType);
491+
492+ selRegType dupNum = 0;
493+ if constexpr (Mode == MODE_MAX) {
494+ dupNum = GetDtypeMin<selRegType>();
495+ } else {
496+ dupNum = GetDtypeMax<selRegType>();
497+ }
498+ 
499+ __VEC_SCOPE__
500+ {
501+ for (uint16_t i = 0; i < static_cast<uint16_t>(uniqueIdNum); i++) {
502+ AscendC::MicroAPI::RegTensor<selRegType> sumReg;
503+ AscendC::MicroAPI::RegTensor<selRegType> updateReg;
504+ AscendC::MicroAPI::MaskReg maskReg;
505+ AscendC::MicroAPI::MaskReg zeroMask = AscendC::MicroAPI::CreateMask<selRegType>();
506+ uint32_t maskLen = static_cast<uint32_t>(colLen);
507+ uint16_t idRepeatTimes = static_cast<uint16_t>(uniqueIdCountLocal(i));
508+ for (uint16_t j = 0; j < static_cast<uint16_t>(loopSize); j++) {
509+ maskReg = AscendC::MicroAPI::UpdateMask<selRegType>(maskLen);
510+ AscendC::MicroAPI::Duplicate(sumReg, dupNum, zeroMask);
511+ for (uint16_t k = 0; k < idRepeatTimes; k++) {
512+ auto updatesOffet = updatesOriginIdexLocal(idLocation + k) * colLenAlignSize + j * vfLen;
513+ auto startAddr = updatesAddr + updatesOffet;
514+ AscendC::MicroAPI::DataCopy(updateReg, startAddr);
515+ if constexpr (Mode == MODE_MAX) {
516+ AscendC::MicroAPI::Max(sumReg, sumReg, updateReg, maskReg);
517+ } else {
518+ AscendC::MicroAPI::Min(sumReg, sumReg, updateReg, maskReg);
519+ }
520+ }
521+ auto updateSumAddrOfst = updateSumAddr + i * colLenAlignSize + j * vfLen;
522+ AscendC::MicroAPI::DataCopy(updateSumAddrOfst, sumReg, maskReg);
523+ }
524+ idLocation += idRepeatTimes;
525+ }
526+ }
527+ }
528+ 
529+ __aicore__ inline void ComputeUpdateSum(int64_t rowLen, int64_t colLen) // rowLen 不用
530+ {
531+ LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountQue_.DeQue<int32_t>(); // 存放每个排好序的 非重复的 indices的个数
532+ LocalTensor<uint32_t> updatesOriginIdexLocal = updatesOriginIdexQue_.DeQue<uint32_t>(); // 排序之后的indices的索引
533+ LocalTensor<T> updatesLocal = dataQueue_.DeQue<T>();
534+ 
535+ LocalTensor<T> updateSumLocal = updateSumQue_.AllocTensor<T>();
536+ 
537+ ComputeSumWithOutCast(uniqueIdCountLocal, updatesOriginIdexLocal, updatesLocal, updateSumLocal, uniqueIdNum_, colLen);
538+ 
539+ updatesOriginIdexQue_.EnQue(updatesOriginIdexLocal);
540+ uniqueIdCountQue_.EnQue(uniqueIdCountLocal);
541+ 
542+ updateSumQue_.EnQue(updateSumLocal);
543+ dataQueue_.EnQue(updatesLocal);
544+ }
545+ 
546+ __aicore__ inline void SortIndices(
547+ LocalTensor<OFFSET_T> outOfstLocal, int64_t rowLen)
548+ {
549+ LocalTensor<CAST_T> sortIndicesLocal = sortIndicesQue_.Get<CAST_T>(); // CAST_0 情况下 CAST_T = OFFSET_T
550+ LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountQue_.AllocTensor<int32_t>();
551+ LocalTensor<uint32_t> updatesOriginIdexLocal = updatesOriginIdexQue_.AllocTensor<uint32_t>();
552+ LocalTensor<CAST_T> shiftSortLocal = sortIndicesLocal[(UB_AGLIN_VALUE / sizeof(CAST_T))];
553+ if constexpr (castType == CAST_0) {
554+ AscendC::Sort<CAST_T, false, sortConfig>( // 排序之后的indices, 排序之后的indices的索引, 原始indices
555+ shiftSortLocal, updatesOriginIdexLocal, outOfstLocal, static_cast<uint32_t>(rowLen));
556+ Duplicate(sortIndicesLocal, (CAST_T)-1, shiftOffset_);
557+ } else {
558+ LocalTensor<CAST_T> indicesCastLocal = castIndicesQue_.Get<CAST_T>();
559+ IndicesSortCast(outOfstLocal, indicesCastLocal, uniqueIdCountLocal, rowLen);
560+ AscendC::Sort<CAST_T, false, sortConfig>(
561+ shiftSortLocal, updatesOriginIdexLocal, indicesCastLocal, static_cast<uint32_t>(rowLen));
562+ Duplicate(sortIndicesLocal, (CAST_T)-1, shiftOffset_);
563+ }
564+ 
565+ event_t eventIdV2S = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
566+ SetFlag<HardEvent::V_S>(eventIdV2S);
567+ WaitFlag<HardEvent::V_S>(eventIdV2S);
568+ 
569+ shiftSortLocal(rowLen) = -1;
570+ PipeBarrier<PIPE_V>();
571+ 
572+ LocalTensor<OFFSET_T> updateSumIdxLocal = updateSumIdxQue_.AllocTensor<OFFSET_T>();
573+
574+ uniqueIdNum_ = ComputeUniqueIdNum(sortIndicesLocal, uniqueIdCountLocal, rowLen);
575+ event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
576+ SetFlag<HardEvent::V_S>(eventIdVToS);
577+ WaitFlag<HardEvent::V_S>(eventIdVToS);
578+ 
579+ for (uint32_t idx = 0; idx < uniqueIdNum_; idx++) {
580+ auto offset = uniqueIdCountLocal(idx); // uniqueIdCountLocal 存放排好序的 非重复的 indices的索引
581+ updateSumIdxLocal(idx) = static_cast<OFFSET_T>(shiftSortLocal(offset)); // 存放排好序的 非重复的 indices值
582+ }
583+ event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
584+ SetFlag<HardEvent::S_V>(eventIdSToV);
585+ WaitFlag<HardEvent::S_V>(eventIdSToV);
586+ ComputeUinqueIdTimes(uniqueIdCountLocal, uniqueIdNum_); // uniqueIdCountLocal 存放每个排好序的 非重复的 indices的个数
587+ PipeBarrier<PIPE_V>();
588+ 
589+ updatesOriginIdexQue_.EnQue(updatesOriginIdexLocal);
590+ uniqueIdCountQue_.EnQue(uniqueIdCountLocal);
591+ 
592+ updateSumIdxQue_.EnQue(updateSumIdxLocal);
593+ }
594+ 
595+ __aicore__ inline void CopyOutSplitIndices(
596+ LocalTensor<OFFSET_T> ofstLocal, LocalTensor<T> dataLocal, int64_t rowLen, int64_t colLen, int64_t colIdx)
597+ {
598+ int64_t colLenAlignSize = ops::CeilAlign(colLen * sizeof(T), UB_AGLIN_VALUE) / sizeof(T);
599+ for (int64_t i = 0; i < rowLen; i++) {
600+ if (ofstLocal(i) < 0) {
601+ continue;
602+ }
603+ int64_t rowOfset = ofstLocal(i) * afterAxis_;
604+ int64_t outOfset = rowOfset + colIdx * afterAxisFactor_;
605+ if constexpr (IsSameType<T, bool>::value) {
606+ if constexpr (Mode == MODE_MAX) {
607+ SetAtomicMax<int8_t>();
608+ } else {
609+ SetAtomicMin<int8_t>();
610+ }
611+ CopyOut<bool>(yGm_[outOfset], dataLocal[i * colLenAlignSize], colLen);
612+ SetAtomicNone();
613+ } else {
614+ if constexpr (Mode == MODE_MAX) {
615+ SetAtomicMax<T>();
616+ } else {
617+ SetAtomicMin<T>();
618+ }
619+ CopyOut<T>(yGm_[outOfset], dataLocal[i * colLenAlignSize], colLen);
620+ SetAtomicNone();
621+ }
622+ }
623+ }
624+ 
625+ __aicore__ inline void CopyIndiceInSplitIndices(int64_t rowIdx, int64_t rowLen)
626+ {
627+ LocalTensor<U> indicesLocal = indicesBuf_.Get<U>();
628+ LocalTensor<OFFSET_T> outOfstLocal = outOfstBuf_.Get<OFFSET_T>();
629+ LocalTensor<float> dstLocal = maxScoreBuf_.Get<float>();
630+ 
631+ int64_t rankSize = indexRankSize_;
632+ int64_t indicesOfset = GetBlockIdx() * eachCoreIndexCount_ + rowIdx * indicesFactor_;
633+ this->template CopyIn<U>(indicesLocal, indicesGm_[indicesOfset * rankSize], rowLen * rankSize);
634+ event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
635+ SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
636+ WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
637+ this->ComputeOutOfset(indicesLocal, outOfstLocal, rowLen, rankSize);
638+ if constexpr (IsSameType<OFFSET_T, int32_t>::value) {
639+ IndexStatisticInt32(outOfstLocal, dstLocal, this->maxScore_, rowLen, afterAxis_);
640+ } else {
641+ IndexStatisticInt64(outOfstLocal, dstLocal, this->maxScore_, rowLen, afterAxis_);
642+ }
643+ }
644+ 
645+ __aicore__ inline void SingleColAdd(LocalTensor<COMPUTE_TYPE> updateSumLocal, int64_t colLen) {
646+ LocalTensor<T> updatesLocal = dataQueue_.DeQue<T>();
647+ __local_mem__ selRegType* updatesAddr = (__local_mem__ selRegType*)updatesLocal.GetPhyAddr();
648+ __local_mem__ COMPUTE_TYPE* updateSumAddr = (__local_mem__ COMPUTE_TYPE*)updateSumLocal.GetPhyAddr();
649+ uint32_t vfLen = platform::GetVRegSize() / sizeof(COMPUTE_TYPE);
650+ int32_t loopSize = (colLen + vfLen - 1) / vfLen;
651+ __VEC_SCOPE__
652+ {
653+ AscendC::MicroAPI::RegTensor<COMPUTE_TYPE> sumReg;
654+ AscendC::MicroAPI::RegTensor<COMPUTE_TYPE> updateReg;
655+ AscendC::MicroAPI::MaskReg maskReg;
656+ uint32_t maskLen = static_cast<uint32_t>(colLen);
657+ for (uint16_t j = 0; j < static_cast<uint16_t>(loopSize); j++) {
658+ maskReg = AscendC::MicroAPI::UpdateMask<COMPUTE_TYPE>(maskLen);
659+ AscendC::MicroAPI::DataCopy(sumReg, updateSumAddr + j * vfLen);
660+ AscendC::MicroAPI::DataCopy(updateReg, updatesAddr + j * vfLen);
661+ if constexpr (Mode == MODE_MAX) {
662+ AscendC::MicroAPI::Max(sumReg, sumReg, updateReg, maskReg);
663+ } else {
664+ AscendC::MicroAPI::Min(sumReg, sumReg, updateReg, maskReg);
665+ }
666+ auto updateSumAddrOfst = updateSumAddr + j * vfLen;
667+ AscendC::MicroAPI::DataCopy(updateSumAddrOfst, sumReg, maskReg);
668+ }
669+ }
670+ dataQueue_.FreeTensor(updatesLocal);
671+ }
672+ 
673+ __aicore__ inline void CopyOutSingleCol(LocalTensor<T> outLocal, int64_t outOffset, int64_t colLen)
674+ {
675+ if constexpr (IsSameType<T, bool>::value) {
676+ if constexpr (Mode == MODE_MAX) {
677+ SetAtomicMax<int8_t>();
678+ } else {
679+ SetAtomicMin<int8_t>();
680+ }
681+ CopyOut<bool>(yGm_[outOffset], outLocal, colLen);
682+ SetAtomicNone();
683+ } else {
684+ if constexpr (Mode == MODE_MAX) {
685+ SetAtomicMax<T>();
686+ } else {
687+ SetAtomicMin<T>();
688+ }
689+ CopyOut<T>(yGm_[outOffset], outLocal, colLen);
690+ SetAtomicNone();
691+ }
692+ }
693+ 
694+ __aicore__ inline void CopyInUpdates(int64_t updatesOffset, int64_t rowLen, int64_t colLen) {
695+ LocalTensor<T> updatesLocal = dataQueue_.AllocTensor<T>();
696+ DataCopyExtParams copyParams = {
697+ static_cast<uint16_t>(rowLen), static_cast<uint32_t>(colLen * sizeof(T)),
698+ static_cast<uint32_t>((afterAxis_ - colLen) * sizeof(T)), static_cast<uint32_t>(0),
699+ static_cast<uint32_t>(0)};
700+ DataCopyPadExtParams<T> updatePadParams = {false, 0, 0, 0};
701+ DataCopyPad(updatesLocal, updatesGm_[updatesOffset], copyParams, updatePadParams);
702+ dataQueue_.EnQue(updatesLocal);
703+ }
704+ 
705+ __aicore__ inline void ComputeSortInOutSingleCol(
706+ int64_t rowIdx, int64_t colIdx, int64_t rowLen, int64_t colLen) // rowLen 不用
707+ {
708+ // rowIdx 可以改为indicesOffset， splitAfter和SplitBefore通用 colIdx 同理
709+ int64_t indicesOffset = GetBlockIdx() * eachCoreIndexCount_ + rowIdx * indicesFactor_;
710+ event_t eventIdMte2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
711+ SetFlag<HardEvent::MTE2_S>(eventIdMte2ToS);
712+ WaitFlag<HardEvent::MTE2_S>(eventIdMte2ToS);
713+ 
714+ event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
715+ SetFlag<HardEvent::V_S>(eventIdVToS);
716+ WaitFlag<HardEvent::V_S>(eventIdVToS);
717+ 
718+ LocalTensor<int32_t> uniqueIdCountLocal = uniqueIdCountQue_.DeQue<int32_t>();
719+ LocalTensor<uint32_t> updatesOriginIdexLocal = updatesOriginIdexQue_.DeQue<uint32_t>(); // 排序之后的indices的索引
720+ LocalTensor<OFFSET_T> ofstLocal = updateSumIdxQue_.DeQue<OFFSET_T>();
721+ int64_t updatesOffset = indicesOffset * afterAxis_ + colIdx * afterAxisFactor_;
722+ int32_t idLocation = 0;
723+ COMPUTE_TYPE dupNum = 0;
724+ if constexpr (Mode == MODE_MAX) {
725+ dupNum = GetDtypeMin<COMPUTE_TYPE>();
726+ } else {
727+ dupNum = GetDtypeMax<COMPUTE_TYPE>();
728+ }
729+ for (int32_t i = 0; i < uniqueIdNum_; i++) {
730+ if (ofstLocal(i) < 0) {
731+ continue;
732+ }
733+ LocalTensor<COMPUTE_TYPE> updateSumLocal = updateSumQue_.AllocTensor<COMPUTE_TYPE>();
734+ AscendC::Duplicate<COMPUTE_TYPE>(updateSumLocal, dupNum, colLen);
735+ int32_t idRepeatTimes = uniqueIdCountLocal(i);
736+ for (int32_t k = 0; k < idRepeatTimes; k++) {
737+ // 兼容splitAfter
738+ int64_t curOffset = updatesOffset + updatesOriginIdexLocal(idLocation) * afterAxis_;
739+ event_t eventIdVToMte2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
740+ SetFlag<HardEvent::V_MTE2>(eventIdVToMte2);
741+ WaitFlag<HardEvent::V_MTE2>(eventIdVToMte2);
742+ CopyInUpdates(curOffset, 1, colLen); // 每个索引单行搬入update
743+ event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
744+ SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
745+ WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
746+ SingleColAdd(updateSumLocal, colLen);
747+ idLocation += 1;
748+ }
749+ 
750+ updateSumQue_.EnQue(updateSumLocal);
751+ int64_t outOffset = ofstLocal(i) * afterAxis_ + colIdx * afterAxisFactor_;
752+ LocalTensor<T> outLocal = updateSumQue_.DeQue<T>();
753+ CopyOutSingleCol(outLocal, outOffset, colLen);
754+ updateSumQue_.FreeTensor(outLocal);
755+ }
756+ updatesOriginIdexQue_.EnQue(updatesOriginIdexLocal);
757+ uniqueIdCountQue_.EnQue(uniqueIdCountLocal);
758+ updateSumIdxQue_.EnQue(ofstLocal);
759+ }
760+ 
761+ __aicore__ inline void CopyInAndOutSingleCol(LocalTensor<OFFSET_T> outOffsetLocal, int64_t rowIdx, int64_t colIdx, int64_t rowLen, int64_t colLen)
762+ {
763+ for (int32_t i = 0; i < rowLen; i++) {
764+ if (outOffsetLocal(i) < 0) {
765+ continue;
766+ }
767+ int64_t indicesOfset = GetBlockIdx() * eachCoreIndexCount_ + rowIdx * indicesFactor_ + i;
768+ int64_t updatesOffset = indicesOfset * afterAxis_ + colIdx * afterAxisFactor_;
769+ CopyInUpdates(updatesOffset, 1, colLen);
770+ LocalTensor<T> dataLocal = dataQueue_.DeQue<T>();
771+ int64_t outOffset = outOffsetLocal(i) * afterAxis_ + colIdx * afterAxisFactor_;
772+ CopyOutSingleCol(dataLocal, outOffset, colLen);
773+ dataQueue_.FreeTensor(dataLocal);
774+ }
775+ }
776+ 
777+ __aicore__ inline void CopyUpdatesInSplitIndices(
778+ int64_t rowIdx, int64_t colIdx, int64_t rowLen, int64_t colLen)
779+ {
780+ LocalTensor<T> updatesLocal = dataQueue_.AllocTensor<T>();
781+ int64_t indicesOfset = GetBlockIdx() * eachCoreIndexCount_ + rowIdx * indicesFactor_;
782+ DataCopyExtParams copyParams = {
783+ static_cast<uint16_t>(rowLen), static_cast<uint32_t>(colLen * sizeof(T)),
784+ static_cast<uint32_t>((afterAxis_ - colLen) * sizeof(T)), static_cast<uint32_t>(0),
785+ static_cast<uint32_t>(0)};
786+ DataCopyPadExtParams<T> updatePadParams = {false, 0, 0, 0};
787+ int64_t rowOfset = indicesOfset * afterAxis_;
788+ int64_t updatesOfset = rowOfset + colIdx * afterAxisFactor_;
789+ DataCopyPad(updatesLocal, updatesGm_[updatesOfset], copyParams, updatePadParams);
790+ dataQueue_.EnQue(updatesLocal);
791+ }
792+ 
793+ __aicore__ inline void ComputeOutSplitAfter(int64_t colIdx, int64_t rowLen, int64_t colLen) // rowLen 不用
794+ {
795+ LocalTensor<T> updatesLocal = dataQueue_.DeQue<T>();
796+ LocalTensor<OFFSET_T> updateSumIdxLocal = updateSumIdxQue_.DeQue<OFFSET_T>();
797+ 
798+ LocalTensor<T> updateSumLocal = updateSumQue_.DeQue<T>();
799+ CopyOutSplitAfter(updateSumIdxLocal, updateSumLocal, uniqueIdNum_, colLen, colIdx);
800+ updateSumQue_.FreeTensor(updateSumLocal);
801+ 
802+ dataQueue_.FreeTensor(updatesLocal);
803+ updateSumIdxQue_.EnQue(updateSumIdxLocal);
804+ }
805+ 
806+ __aicore__ inline void ComputeOutSplitIndices(int64_t colIdx, int64_t rowLen, int64_t colLen) // rowLen 不用
807+ {
808+ LocalTensor<OFFSET_T> updateSumIdxLocal = updateSumIdxQue_.DeQue<OFFSET_T>();
809+ LocalTensor<T> updatesLocal = dataQueue_.DeQue<T>();
810+ 
811+ LocalTensor<T> updateSumLocal = updateSumQue_.DeQue<T>();
812+ CopyOutSplitIndices(updateSumIdxLocal, updateSumLocal, uniqueIdNum_, colLen, colIdx);
813+ updateSumQue_.FreeTensor(updateSumLocal);
814+ 
815+ dataQueue_.FreeTensor(updatesLocal);
816+ updateSumIdxQue_.EnQue(updateSumIdxLocal);
817+ }
818+ 
819+ __aicore__ inline void CopyOutSplitAfter(
820+ LocalTensor<OFFSET_T> ofstLocal, LocalTensor<T> dataLocal, int64_t rowLen, int64_t colLen, int64_t colIdx)
821+ {
822+ int64_t colLenAlignSize = ops::CeilAlign(colLen * sizeof(T), UB_AGLIN_VALUE) / sizeof(T);
823+ for (int64_t i = 0; i < rowLen; i++) {
824+ if (ofstLocal(i) < 0) {
825+ continue;
826+ }
827+ int64_t rowOfset = ofstLocal(i) * afterAxis_;
828+ int64_t outOfset = rowOfset + GetBlockIdx() * eachCoreAfterAxisCount_ + colIdx * afterAxisFactor_;
829+ if constexpr (IsSameType<T, bool>::value) {
830+ if constexpr (Mode == MODE_MAX) {
831+ SetAtomicMax<int8_t>();
832+ } else {
833+ SetAtomicMin<int8_t>();
834+ }
835+ CopyOut<bool>(yGm_[outOfset], dataLocal[i * colLenAlignSize], colLen);
836+ SetAtomicNone();
837+ } else {
838+ if constexpr (Mode == MODE_MAX) {
839+ SetAtomicMax<T>();
840+ } else {
841+ SetAtomicMin<T>();
842+ }
843+ CopyOut<T>(yGm_[outOfset], dataLocal[i * colLenAlignSize], colLen);
844+ SetAtomicNone();
845+ }
846+ }
847+ }
848+ 
849+ __aicore__ inline void CopyIndiceInSplitAfter(int64_t rowIdx, int64_t rowLen)
850+ {
851+ LocalTensor<U> indicesLocal = indicesBuf_.Get<U>();
852+ LocalTensor<OFFSET_T> outOfstLocal = outOfstBuf_.Get<OFFSET_T>();
853+ LocalTensor<float> dstLocal = maxScoreBuf_.Get<float>();
854+ 
855+ int64_t rankSize = indexRankSize_;
856+ int64_t indicesOfset = rowIdx * indicesFactor_;
857+ CopyIn<U>(indicesLocal, indicesGm_[indicesOfset * rankSize], rowLen * rankSize);
858+ event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
859+ SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
860+ WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
861+ ComputeOutOfset(indicesLocal, outOfstLocal, rowLen, rankSize);
862+ if constexpr (IsSameType<OFFSET_T, int32_t>::value) {
863+ IndexStatisticInt32(outOfstLocal, dstLocal, maxScore_, rowLen, afterAxis_);
864+ } else {
865+ IndexStatisticInt64(outOfstLocal, dstLocal, maxScore_, rowLen, afterAxis_);
866+ }
867+ }
868+ 
869+ __aicore__ inline void CopyUpdatesInSplitAfter(int64_t rowIdx, int64_t colIdx, int64_t rowLen, int64_t colLen)
870+ {
871+ LocalTensor<T> updatesLocal = dataQueue_.AllocTensor<T>();
872+ int64_t indicesOfset = rowIdx * indicesFactor_;
873+ DataCopyExtParams copyParams = {
874+ static_cast<uint16_t>(rowLen), static_cast<uint32_t>(colLen * sizeof(T)),
875+ static_cast<uint32_t>((afterAxis_ - colLen) * sizeof(T)), static_cast<uint32_t>(0),
876+ static_cast<uint32_t>(0)};
877+ DataCopyPadExtParams<T> updatePadParams = {false, 0, 0, 0};
878+ int64_t rowOfset = indicesOfset * afterAxis_;
879+ int64_t updatesOfset = rowOfset + GetBlockIdx() * eachCoreAfterAxisCount_ + colIdx * afterAxisFactor_;
880+ DataCopyPad(updatesLocal, updatesGm_[updatesOfset], copyParams, updatePadParams);
881+ dataQueue_.EnQue(updatesLocal);
882+ }
883+ 
884+};
885+ 
886+} // namespace ScatterNdCommon
887+ 
888+#endif
Aindex/scatter_nd_common/op_kernel/arch35/scatter_nd_common_simd_sort.h+234-0
@@ -0,0 +1,234 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file scatter_nd_common_simd_sort.h
13+ * \brief scatter_nd_common_simd_sort
14+ */
15+ 
16+#ifndef SCATTER_ND_COMMON_SIMD_SORT_H
17+#define SCATTER_ND_COMMON_SIMD_SORT_H
18+ 
19+#include "scatter_nd_common_base.h"
20+#include "../inc/kernel_utils.h"
21+ 
22+namespace ScatterNdCommon {
23+using namespace AscendC;
24+ 
25+template <typename T, typename U, typename CAST_T, typename OFFSET_T, uint32_t castType, uint8_t Mode>
26+class ScatterNdCommonSimdSort : public ScatterNdCommonBase<T, U, CAST_T, OFFSET_T, castType, Mode>
27+{
28+public:
29+ __aicore__ inline ScatterNdCommonSimdSort(const ScatterNdCommonSimdSortTilingData& tilingData, TPipe& pipe)
30+ : tilingData_(tilingData), pipe_(pipe){};
31+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y);
32+ __aicore__ inline void ProcessSplitAfter();
33+ __aicore__ inline void ProcessSplitIndices();
34+ __aicore__ inline void ProcessSplitIndicesSingleCol();
35+ __aicore__ inline void Process();
36+ 
37+private:
38+ AscendC::GlobalTensor<T> varGm_;
39+ AscendC::GlobalTensor<U> indicesGm_;
40+ AscendC::GlobalTensor<T> updatesGm_;
41+ AscendC::GlobalTensor<T> yGm_;
42+ 
43+ TPipe& pipe_;
44+ const ScatterNdCommonSimdSortTilingData& tilingData_;
45+ 
46+ int64_t curCoreIndexCount_{0};
47+ uint64_t strideList[MAX_RANK_COUNT];
48+};
49+ 
50+template <typename T, typename U, typename CAST_T, typename OFFSET_T, uint32_t castType, uint8_t Mode>
51+__aicore__ inline void ScatterNdCommonSimdSort<T, U, CAST_T, OFFSET_T, castType, Mode>::Init(
52+ GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y)
53+{
54+ 
55+ this->eachCoreAfterAxisCount_ = tilingData_.eachCoreAfterAxisCount;
56+ this->indexRankSize_ = tilingData_.indexRankSize;
57+ this->eachCoreIndexCount_ = tilingData_.eachCoreIndexCount;
58+ this->indicesFactor_ = tilingData_.indicesFactor;
59+ this->afterAxis_ = tilingData_.afterAxis;
60+ this->afterAxisFactor_ = tilingData_.afterAxisFactor;
61+ this->InitBaseBuffer(pipe_, tilingData_.indicesFactor, indices, updates, y, tilingData_.singleCol);
62+ 
63+ curCoreIndexCount_ =
64+ (GetBlockIdx() != (tilingData_.usedCoreNumBefore - 1) ? tilingData_.eachCoreIndexCount :
65+ tilingData_.tailCoreIndexCount);
66+}
67+ 
68+template <typename T, typename U, typename CAST_T, typename OFFSET_T, uint32_t castType, uint8_t Mode>
69+__aicore__ inline void ScatterNdCommonSimdSort<T, U, CAST_T, OFFSET_T, castType, Mode>::ProcessSplitAfter()
70+{
71+ if (GetBlockIdx() >= tilingData_.usedCoreNumBefore) {
72+ return;
73+ }
74+ int64_t rowMainDataLen = tilingData_.indicesFactor;
75+ int64_t rowTailDataLen = tilingData_.indiceTailNum;
76+ int64_t rowLoopNum = tilingData_.indicesLoopSize;
77+ int64_t colLoopNum = (GetBlockIdx() == tilingData_.usedCoreNumBefore - 1) ? tilingData_.tailUpdateLoopSize
78+ : tilingData_.updateLoopSize;
79+ int64_t colMainDataLen = tilingData_.afterAxisFactor;
80+ int64_t colTailDataLen = (GetBlockIdx() == tilingData_.usedCoreNumBefore - 1) ? tilingData_.tailUpdateTailNum
81+ : tilingData_.updateTailNum;
82+ for (int64_t rowIdx = 0; rowIdx < rowLoopNum; rowIdx++) {
83+ int64_t rowDataLen = (rowIdx == rowLoopNum - 1) ? rowTailDataLen : rowMainDataLen;
84+ this->CopyIndiceInSplitAfter(rowIdx, rowDataLen);
85+ LocalTensor<OFFSET_T> outOfstLocal = this->outOfstBuf_.template Get<OFFSET_T>();
86+ if (this->maxScore_ > SORT_HIST_THRESHOLD) {
87+ this->SortIndices(outOfstLocal, rowDataLen);
88+ for (int64_t colIdx = 0; colIdx < colLoopNum; colIdx++) {
89+ int64_t colDataLen = (colIdx == colLoopNum - 1) ? colTailDataLen : colMainDataLen;
90+ this->CopyUpdatesInSplitAfter(rowIdx, colIdx, rowDataLen, colDataLen);
91+ event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
92+ SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
93+ WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
94+ this->ComputeUpdateSum(rowDataLen, colDataLen);
95+ this->ComputeOutSplitAfter(colIdx, rowDataLen, colDataLen);
96+ }
97+ LocalTensor<int32_t> uniqueIdCountLocal = this->uniqueIdCountQue_.template DeQue<int32_t>();
98+ LocalTensor<uint32_t> updatesOriginIdexLocal = this->updatesOriginIdexQue_.template DeQue<uint32_t>();
99+ LocalTensor<OFFSET_T> updateSumIdxLocal = this->updateSumIdxQue_.template DeQue<OFFSET_T>();
100+ this->updatesOriginIdexQue_.template FreeTensor(updatesOriginIdexLocal);
101+ this->updateSumIdxQue_.template FreeTensor(updateSumIdxLocal);
102+ this->uniqueIdCountQue_.template FreeTensor(uniqueIdCountLocal);
103+ } else {
104+ for (int64_t colIdx = 0; colIdx < colLoopNum; colIdx++) {
105+ int64_t colDataLen = (colIdx == colLoopNum - 1) ? colTailDataLen : colMainDataLen;
106+ this->CopyUpdatesInSplitAfter(rowIdx, colIdx, rowDataLen, colDataLen);
107+ LocalTensor<T> updatesLocal = this->dataQueue_.template DeQue<T>();
108+ event_t eventIdMte2ToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_MTE3));
109+ SetFlag<HardEvent::MTE2_MTE3>(eventIdMte2ToMte3);
110+ WaitFlag<HardEvent::MTE2_MTE3>(eventIdMte2ToMte3);
111+ this->CopyOutSplitAfter(outOfstLocal, updatesLocal, rowDataLen, colDataLen, colIdx);
112+ this->dataQueue_.template FreeTensor(updatesLocal);
113+ }
114+ }
115+ }
116+}
117+ 
118+template <typename T, typename U, typename CAST_T, typename OFFSET_T, uint32_t castType, uint8_t Mode>
119+__aicore__ inline void ScatterNdCommonSimdSort<T, U, CAST_T, OFFSET_T, castType, Mode>::ProcessSplitIndices()
120+{
121+ if (GetBlockIdx() >= tilingData_.usedCoreNumBefore) {
122+ return;
123+ }
124+ 
125+ int64_t colLoopNum = tilingData_.updateLoopSize;
126+ int64_t colMainDataLen = tilingData_.afterAxisFactor;
127+ int64_t colTailDataLen = tilingData_.updateTailNum;
128+ int64_t rowLoopNum = ops::CeilDiv(curCoreIndexCount_, tilingData_.indicesFactor);
129+ int64_t rowMainDataLen = tilingData_.indicesFactor;
130+ int64_t rowTailDataLen = curCoreIndexCount_ - tilingData_.indicesFactor * (rowLoopNum - 1);
131+ 
132+ for (int64_t rowIdx = 0; rowIdx < rowLoopNum; rowIdx++) {
133+ int64_t rowDataLen = (rowIdx == rowLoopNum - 1) ? rowTailDataLen : rowMainDataLen;
134+ this->CopyIndiceInSplitIndices(rowIdx, rowDataLen);
135+ LocalTensor<OFFSET_T> outOfstLocal = this->outOfstBuf_.template Get<OFFSET_T>();
136+ if (this->maxScore_ > SORT_HIST_THRESHOLD) {
137+ this->SortIndices(outOfstLocal, rowDataLen);
138+ for (int64_t colIdx = 0; colIdx < colLoopNum; colIdx++) {
139+ int64_t colDataLen = (colIdx == colLoopNum - 1) ? colTailDataLen : colMainDataLen;
140+ this->CopyUpdatesInSplitIndices(rowIdx, colIdx, rowDataLen, colDataLen);
141+ // 累加相同索引的updates
142+ event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
143+ SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
144+ WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
145+ this->ComputeUpdateSum(rowDataLen, colDataLen);
146+ this->ComputeOutSplitIndices(colIdx, rowDataLen, colDataLen);
147+ }
148+ 
149+ //本次行数处理完成,释放资源
150+ LocalTensor<int32_t> uniqueIdCountLocal = this->uniqueIdCountQue_.template DeQue<int32_t>();
151+ LocalTensor<uint32_t> updatesOriginIdexLocal = this->updatesOriginIdexQue_.template DeQue<uint32_t>();
152+ LocalTensor<OFFSET_T> updateSumIdxLocal = this->updateSumIdxQue_.template DeQue<OFFSET_T>();
153+ this->uniqueIdCountQue_.template FreeTensor(uniqueIdCountLocal);
154+ this->updatesOriginIdexQue_.template FreeTensor(updatesOriginIdexLocal);
155+ this->updateSumIdxQue_.template FreeTensor(updateSumIdxLocal);
156+ } else {
157+ for (int64_t colIdx = 0; colIdx < colLoopNum; colIdx++) {
158+ int64_t colDataLen = (colIdx == colLoopNum - 1) ? colTailDataLen : colMainDataLen;
159+ this->CopyUpdatesInSplitIndices(rowIdx, colIdx, rowDataLen, colDataLen);
160+ LocalTensor<T> updatesLocal = this->dataQueue_.template DeQue<T>();
161+ event_t eventIdMte2ToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_MTE3));
162+ SetFlag<HardEvent::MTE2_MTE3>(eventIdMte2ToMte3);
163+ WaitFlag<HardEvent::MTE2_MTE3>(eventIdMte2ToMte3);
164+ this->CopyOutSplitIndices(outOfstLocal, updatesLocal, rowDataLen, colDataLen, colIdx);
165+ this->dataQueue_.template FreeTensor(updatesLocal);
166+ }
167+ }
168+ }
169+}
170+ 
171+ 
172+template <typename T, typename U, typename CAST_T, typename OFFSET_T, uint32_t castType, uint8_t Mode>
173+__aicore__ inline void ScatterNdCommonSimdSort<T, U, CAST_T, OFFSET_T, castType, Mode>::ProcessSplitIndicesSingleCol()
174+{
175+ if (GetBlockIdx() >= tilingData_.usedCoreNumBefore) {
176+ return;
177+ }
178+ 
179+ int64_t colLoopNum = tilingData_.updateLoopSize;
180+ int64_t colMainDataLen = tilingData_.afterAxisFactor;
181+ int64_t colTailDataLen = tilingData_.updateTailNum;
182+ int64_t rowLoopNum = ops::CeilDiv(curCoreIndexCount_, tilingData_.indicesFactor);
183+ int64_t rowMainDataLen = tilingData_.indicesFactor;
184+ int64_t rowTailDataLen = curCoreIndexCount_ - tilingData_.indicesFactor * (rowLoopNum - 1);
185+ 
186+ for (int64_t rowIdx = 0; rowIdx < rowLoopNum; rowIdx++) {
187+ int64_t rowDataLen = (rowIdx == rowLoopNum - 1) ? rowTailDataLen : rowMainDataLen;
188+ this->CopyIndiceInSplitIndices(rowIdx, rowDataLen);
189+ LocalTensor<OFFSET_T> outOfstLocal = this->outOfstBuf_.template Get<OFFSET_T>();
190+ if (this->maxScore_ > SORT_HIST_THRESHOLD) {
191+ this->SortIndices(outOfstLocal, rowDataLen);
192+ for (int64_t colIdx = 0; colIdx < colLoopNum; colIdx++) {
193+ int64_t colLen = (colIdx == colLoopNum - 1) ? colTailDataLen : colMainDataLen;
194+ this->ComputeSortInOutSingleCol(rowIdx, colIdx, rowDataLen, colLen);
195+ }
196+ LocalTensor<int32_t> uniqueIdCountLocal = this->uniqueIdCountQue_.template DeQue<int32_t>();
197+ LocalTensor<uint32_t> updatesOriginIdexLocal = this->updatesOriginIdexQue_.template DeQue<uint32_t>();
198+ LocalTensor<OFFSET_T> updateSumIdxLocal = this->updateSumIdxQue_.template DeQue<OFFSET_T>();
199+ this->uniqueIdCountQue_.template FreeTensor(uniqueIdCountLocal);
200+ this->updatesOriginIdexQue_.template FreeTensor(updatesOriginIdexLocal);
201+ this->updateSumIdxQue_.template FreeTensor(updateSumIdxLocal);
202+ } else {
203+ for (int64_t colIdx = 0; colIdx < colLoopNum; colIdx++) {
204+ int64_t colDataLen = (colIdx == colLoopNum - 1) ? colTailDataLen : colMainDataLen;
205+ this->CopyInAndOutSingleCol(outOfstLocal, rowIdx, colIdx, rowDataLen, colDataLen);
206+ }
207+ }
208+ }
209+}
210+ 
211+template <typename T, typename U, typename CAST_T, typename OFFSET_T, uint32_t castType, uint8_t Mode>
212+__aicore__ inline void ScatterNdCommonSimdSort<T, U, CAST_T, OFFSET_T, castType, Mode>::Process()
213+{
214+ LocalTensor<OFFSET_T> strideLocal = this->strideBuf_.template Get<OFFSET_T>();
215+ LocalTensor<U> outputShapeLocal = this->outputShapeBuf_.template Get<U>();
216+ for (int32_t i = 0; i < MAX_RANK_COUNT; i++) {
217+ strideLocal(i) = tilingData_.strideList[i];
218+ }
219+ for (int32_t i = 0; i < MAX_SHAPE_RANK; i++) {
220+ outputShapeLocal(i) = tilingData_.outPutShape[i];
221+ }
222+ 
223+ if (tilingData_.isSplitAfterAxis == 1) {
224+ ProcessSplitAfter();
225+ } else {
226+ if (tilingData_.singleCol) {
227+ ProcessSplitIndicesSingleCol();
228+ } else {
229+ ProcessSplitIndices();
230+ }
231+ }
232+}
233+} // namespace ScatterNdCommon
234+#endif // SCATTER_ND_COMMON_SIMD_SORT_H
Aindex/scatter_nd_common/op_kernel/arch35/scatter_nd_common_simt.h+296-0
@@ -0,0 +1,296 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+
11+/*!
12+ * \file scatter_nd_common_simt.h
13+ * \brief scatter_nd_common_simt
14+ */
15+ 
16+#ifndef SCATTER_ND_COMMON_SIMT_H
17+#define SCATTER_ND_COMMON_SIMT_H
18+ 
19+#include "scatter_nd_common_base.h"
20+#include "kernel_operator.h"
21+#include "simt_api/asc_simt.h"
22+ 
23+namespace ScatterNdCommon {
24+using namespace AscendC;
25+ 
26+template <typename INDICES_T, typename PARAMS_T, typename TYPE_T, uint8_t Mode>
27+__simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM_LAUNCH_BOUND) inline void SimtCompute(
28+ __local_mem__ INDICES_T* idxLocalAddr, __local_mem__ PARAMS_T* xLocalAddr, __gm__ PARAMS_T* outputGmAddr,
29+ const __local_mem__ TYPE_T* strideListAddr, const __local_mem__ TYPE_T* outputShapeAddr, const uint32_t currUbTilingSize,
30+ const TYPE_T xOffSet, const TYPE_T sliceSize, const uint32_t rankSize, const TYPE_T indiceOffSet,
31+ const TYPE_T magic, const TYPE_T shift)
32+{
33+ for (uint32_t index = threadIdx.x; index < currUbTilingSize; index += blockDim.x) {
34+ TYPE_T globalIdx = xOffSet + index;
35+ TYPE_T quotient = Simt::UintDiv(globalIdx, magic, shift); // GM第几行update
36+ TYPE_T currIndiceIdx = quotient * rankSize; // GM第几个索引
37+ TYPE_T scatterAxisIdx = globalIdx - quotient * sliceSize; // update尾轴第几个元素(globalIdx % sliceSize)
38+ TYPE_T idx = 0;
39+ bool outOfBound = false;
40+ for (TYPE_T dim = 0; dim < rankSize; ++dim) {
41+ INDICES_T indiceVal = idxLocalAddr[currIndiceIdx + dim - indiceOffSet];
42+ outOfBound |= (indiceVal < 0 || indiceVal >= outputShapeAddr[dim]);
Z
Zz300751995月20日

确认有没有负索引处理

likedislike
43+ idx += indiceVal * strideListAddr[dim];
44+ }
45+ if (!outOfBound) {
46+ TYPE_T dstIndex = idx * sliceSize + scatterAxisIdx;
47+ if constexpr (IsSameType<PARAMS_T, bool>::value) {
48+ PARAMS_T value = xLocalAddr[index];
49+ if constexpr (Mode == MODE_MAX) {
50+ if (value > 0) {
51+ outputGmAddr[dstIndex] = value;
52+ }
53+ } else {
54+ if (value <= 0) {
55+ outputGmAddr[dstIndex] = value;
56+ }
57+ }
58+ } else {
59+ if constexpr (Mode == MODE_MAX) {
60+ Simt::AtomicMax(outputGmAddr + dstIndex, xLocalAddr[index]);
61+ } else {
62+ Simt::AtomicMin(outputGmAddr + dstIndex, xLocalAddr[index]);
63+ }
64+ }
65+ }
66+ }
67+}
68+ 
69+template <typename INDICES_T, typename PARAMS_T, typename TYPE_T, uint8_t Mode>
70+__simt_vf__ __aicore__ __launch_bounds__(THREAD_NUM_LAUNCH_BOUND) inline void SimtComputeDimensionOne(
71+ __gm__ INDICES_T* idxGmAddr, __local_mem__ PARAMS_T* xLocalAddr, __gm__ PARAMS_T* outputGmAddr,
72+ uint32_t currUbTilingSize, TYPE_T xOffSet, TYPE_T sliceSize, TYPE_T indiceOffSet,
73+ TYPE_T varInAxis, TYPE_T magic, TYPE_T shift)
74+{
75+ for (uint32_t index = threadIdx.x; index < currUbTilingSize; index += blockDim.x) {
76+ TYPE_T globalIdx = xOffSet + index;
77+ TYPE_T currIndiceIdx = Simt::UintDiv(globalIdx, magic, shift);
78+ TYPE_T scatterAxisIdx = globalIdx - currIndiceIdx * sliceSize;
79+ INDICES_T idx = idxGmAddr[currIndiceIdx];
80+ 
81+ if (idx >= 0 && idx < varInAxis) {
82+ TYPE_T dstOffet = idx * sliceSize + scatterAxisIdx;
83+ if constexpr (IsSameType<PARAMS_T, bool>::value) {
84+ PARAMS_T value = xLocalAddr[index];
85+ if constexpr (Mode == MODE_MAX) {
86+ if (value > 0) {
87+ outputGmAddr[dstOffet] = value;
88+ }
89+ } else {
90+ if (value <= 0) {
91+ outputGmAddr[dstOffet] = value;
92+ }
93+ }
94+ } else {
95+ if constexpr (Mode == MODE_MAX) {
96+ Simt::AtomicMax(outputGmAddr + dstOffet, xLocalAddr[index]);
97+ } else {
98+ Simt::AtomicMin(outputGmAddr + dstOffet, xLocalAddr[index]);
99+ }
100+ }
101+ }
102+ }
103+}
104+ 
105+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
106+class ScatterNdCommonSimt
107+{
108+public:
109+ __aicore__ inline ScatterNdCommonSimt(const ScatterNdCommonSimtTilingData& tilingData, TPipe& pipe)
110+ : pipe_(pipe), tiling_(tilingData){};
111+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR indices, GM_ADDR updates, GM_ADDR y);
112+ __aicore__ inline void Process();
113+ 
114+private:
115+ __aicore__ inline void ComputeData();
116+ __aicore__ inline void CopyIn(LocalTensor<INDICES_T>& idxLocal, LocalTensor<PARAMS_T>& xLocal);
117+ __aicore__ inline void ComputeDimensionOther();
118+ __aicore__ inline void ComputeDimensionOne();
119+ __aicore__ inline void CopyInUpdate(LocalTensor<PARAMS_T>& xLocal);
120+ 
121+private:
122+ TPipe& pipe_;
123+ const ScatterNdCommonSimtTilingData& tiling_;
124+ GlobalTensor<INDICES_T> idxGm;
125+ GlobalTensor<PARAMS_T> xGm;
126+ GlobalTensor<PARAMS_T> outputGm;
127+ 
128+ TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueIdx, inQueX;
129+ TBuf<TPosition::VECCALC> strideListBuf;
130+ TBuf<TPosition::VECCALC> outputShapeBuf;
131+ 
132+ uint32_t blockIdx_;
133+ TYPE_T currBlockTilingSize_ = 0; // 当前核中要更新的update元素个数
134+ 
135+ uint32_t ubTilingSize_ = 0;
136+ uint32_t currUbTilingSize_ = 0; // 当前ub循环处理的update元素个数
137+ 
138+ TYPE_T xBlockOffSet_ = 0;
139+ TYPE_T xOffSet_ = 0; // 当前ub循环update的GM偏移
140+ TYPE_T indiceBlockOffSet_ = 0;
141+ TYPE_T indiceOffSet_ = 0;
142+ uint32_t currIdxTilingSize_ = 0;
143+ TYPE_T ubLoopCnt_ = 0;
144+};
145+ 
146+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
147+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::Init(
148+ GM_ADDR x, GM_ADDR indices, GM_ADDR updates, GM_ADDR y)
149+{
150+ if (tiling_.sliceSize == 0) {
151+ return;
152+ }
153+ 
154+ blockIdx_ = GetBlockIdx();
155+ 
156+ this->xBlockOffSet_ = tiling_.blockTilingSize * blockIdx_; // update的GM偏移
157+ // calculate indice offset size by x offset size
158+ this->indiceBlockOffSet_ = this->xBlockOffSet_ / tiling_.sliceSize * tiling_.rankSize; // 索引的偏移
159+ 
160+ if (blockIdx_ == tiling_.blockNum - 1) {
161+ this->currBlockTilingSize_ = tiling_.tailBlockTilingSize;
162+ } else {
163+ this->currBlockTilingSize_ = tiling_.blockTilingSize;
164+ }
165+ 
166+ this->ubTilingSize_ = tiling_.ubTilingSize; // ub 中要处理的update元素个数
167+ if (this->currBlockTilingSize_ <= tiling_.ubTilingSize) {
168+ this->ubTilingSize_ = this->currBlockTilingSize_;
169+ }
170+ 
171+ auto indiceUbTilingSize = (this->ubTilingSize_ + tiling_.sliceSize - 1) / tiling_.sliceSize * tiling_.rankSize; // ub 中indices元素个数
172+ idxGm.SetGlobalBuffer((__gm__ INDICES_T*)indices);
173+ xGm.SetGlobalBuffer((__gm__ PARAMS_T*)updates);
174+ outputGm.SetGlobalBuffer((__gm__ PARAMS_T*)y);
175+ 
176+ pipe_.InitBuffer(inQueX, DOUBLE_BUFFER, ROUND_UP32(this->ubTilingSize_ * sizeof(PARAMS_T)));
177+ if (tiling_.rankSize >= INDICE_RANK_TWO) {
178+ pipe_.InitBuffer(inQueIdx, DOUBLE_BUFFER, ROUND_UP32(indiceUbTilingSize * sizeof(INDICES_T)));
179+ pipe_.InitBuffer(strideListBuf, MAX_RANK_COUNT * sizeof(TYPE_T));
180+ pipe_.InitBuffer(outputShapeBuf, MAX_SHAPE_RANK * sizeof(TYPE_T));
181+ }
182+}
183+ 
184+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
185+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::Process()
186+{
187+ if (blockIdx_ < tiling_.blockNum) {
188+ this->ubLoopCnt_ = (this->currBlockTilingSize_ + this->ubTilingSize_ - 1) / this->ubTilingSize_;
189+ for (TYPE_T idx = 0; idx < this->ubLoopCnt_ - 1; idx++) {
190+ this->currUbTilingSize_ = this->ubTilingSize_;
191+ this->xOffSet_ = this->xBlockOffSet_ + idx * this->ubTilingSize_;
192+ ComputeData();
193+ }
194+ this->xOffSet_ = this->xBlockOffSet_ + (this->ubLoopCnt_ - 1) * this->ubTilingSize_;
195+ this->currUbTilingSize_ = this->currBlockTilingSize_ - this->ubTilingSize_ * (this->ubLoopCnt_ - 1);
196+ ComputeData();
197+ }
198+}
199+ 
200+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
201+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::ComputeData()
202+{
203+ auto currEnd = this->xOffSet_ + this->currUbTilingSize_;
204+ auto indiceBegin = this->xOffSet_ / tiling_.sliceSize * tiling_.rankSize;
205+ auto indiceEnd = (currEnd + tiling_.sliceSize - 1) / tiling_.sliceSize * tiling_.rankSize;
206+ this->currIdxTilingSize_ = indiceEnd - indiceBegin; // 左闭右开
207+ this->indiceOffSet_ = indiceBegin;
208+ if (tiling_.rankSize >= INDICE_RANK_TWO) {
209+ ComputeDimensionOther();
210+ } else {
211+ ComputeDimensionOne();
212+ }
213+}
214+ 
215+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
216+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::ComputeDimensionOther()
217+{
218+ LocalTensor<INDICES_T> idxLocal = inQueIdx.AllocTensor<INDICES_T>();
219+ LocalTensor<PARAMS_T> xLocal = inQueX.AllocTensor<PARAMS_T>();
220+ CopyIn(idxLocal, xLocal);
221+ uint32_t currUbTilingSize = this->currUbTilingSize_;
222+ TYPE_T sliceSize = tiling_.sliceSize;
223+ uint32_t rankSize = tiling_.rankSize;
224+ 
225+ LocalTensor<TYPE_T> strideList = strideListBuf.Get<TYPE_T>();
226+ LocalTensor<TYPE_T> outputShape = outputShapeBuf.Get<TYPE_T>();
227+ for (uint32_t i = 0; i < MAX_RANK_COUNT; i++) {
228+ strideList(i) = tiling_.strideList[i];
229+ }
230+ for (uint32_t i = 0; i < MAX_SHAPE_RANK; i++) {
231+ outputShape(i) = tiling_.outPutShape[i];
232+ }
233+ DataSyncBarrier<MemDsbT::UB>(); //
234+ TYPE_T magic = 0;
235+ TYPE_T shift = 0;
236+ GetUintDivMagicAndShift(magic, shift, sliceSize);
237+ asc_vf_call<SimtCompute<INDICES_T, PARAMS_T, TYPE_T, Mode>>(
238+ dim3(THREAD_NUM), (__local_mem__ INDICES_T*)idxLocal.GetPhyAddr(), (__local_mem__ PARAMS_T*)xLocal.GetPhyAddr(),
239+ (__gm__ PARAMS_T*)(outputGm.GetPhyAddr()), (__local_mem__ TYPE_T*)strideList.GetPhyAddr(),
240+ (__local_mem__ TYPE_T*)outputShape.GetPhyAddr(), currUbTilingSize, this->xOffSet_, sliceSize, rankSize, this->indiceOffSet_,
241+ magic, shift);
242+ 
243+ inQueIdx.FreeTensor(idxLocal);
244+ inQueX.FreeTensor(xLocal);
245+}
246+ 
247+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
248+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::ComputeDimensionOne()
249+{
250+ LocalTensor<PARAMS_T> xLocal = inQueX.AllocTensor<PARAMS_T>();
251+ CopyInUpdate(xLocal);
252+ uint32_t currUbTilingSize = this->currUbTilingSize_;
253+ TYPE_T sliceSize = tiling_.sliceSize;
254+ TYPE_T varInAxis = tiling_.varInAxis;
255+
256+ TYPE_T magic = 0;
257+ TYPE_T shift = 0;
258+ GetUintDivMagicAndShift(magic, shift, sliceSize);
259+ asc_vf_call<SimtComputeDimensionOne<INDICES_T, PARAMS_T, TYPE_T, Mode>>(
260+ dim3(THREAD_NUM), (__gm__ INDICES_T*)(idxGm.GetPhyAddr()), (__local_mem__ PARAMS_T*)xLocal.GetPhyAddr(),
261+ (__gm__ PARAMS_T*)(outputGm.GetPhyAddr()), currUbTilingSize, this->xOffSet_, sliceSize, this->indiceOffSet_,
262+ varInAxis, magic, shift);
263+ inQueX.FreeTensor(xLocal);
264+}
265+ 
266+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
267+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::CopyIn(
268+ LocalTensor<INDICES_T>& idxLocal, LocalTensor<PARAMS_T>& xLocal)
269+{
270+ DataCopyExtParams idxCopyParams{1, static_cast<uint32_t>(this->currIdxTilingSize_ * sizeof(INDICES_T)), 0, 0, 0};
271+ DataCopyPadExtParams<INDICES_T> idxPadParams{false, 0, 0, 0};
272+ DataCopyPad(idxLocal, idxGm[this->indiceOffSet_], idxCopyParams, idxPadParams);
273+ 
274+ DataCopyExtParams xCopyParams{1, static_cast<uint32_t>(this->currUbTilingSize_ * sizeof(PARAMS_T)), 0, 0, 0};
275+ DataCopyPadExtParams<PARAMS_T> xPadParams{false, 0, 0, 0};
276+ DataCopyPad(xLocal, xGm[this->xOffSet_], xCopyParams, xPadParams);
277+ 
278+ inQueIdx.EnQue(idxLocal);
279+ inQueX.EnQue(xLocal);
280+ inQueIdx.DeQue<INDICES_T>();
281+ inQueX.DeQue<PARAMS_T>();
282+}
283+ 
284+template <typename PARAMS_T, typename INDICES_T, typename TYPE_T, uint8_t Mode>
285+__aicore__ inline void ScatterNdCommonSimt<PARAMS_T, INDICES_T, TYPE_T, Mode>::CopyInUpdate(LocalTensor<PARAMS_T>& xLocal)
286+{
287+ DataCopyExtParams xCopyParams{1, static_cast<uint32_t>(this->currUbTilingSize_ * sizeof(PARAMS_T)), 0, 0, 0};
288+ DataCopyPadExtParams<PARAMS_T> xPadParams{false, 0, 0, 0};
289+ DataCopyPad(xLocal, xGm[this->xOffSet_], xCopyParams, xPadParams);
290+ 
291+ inQueX.EnQue(xLocal);
292+ inQueX.DeQue<PARAMS_T>();
293+}
294+ 
295+} // namespace ScatterNdCommon
296+#endif // SCATTER_ND_COMMON_SIMT_H
Aindex/scatter_nd_common/op_kernel/arch35/scatter_nd_common_struct.h+77-0
@@ -0,0 +1,77 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file scatter_nd_common_struct.h
13+ * \brief tiling base data
14+ */
15+ 
16+#ifndef SCATTER_ND_COMMON_STRUCT_H
17+#define SCATTER_ND_COMMON_STRUCT_H
18+ 
19+#include <cstdint>
20+#include "kernel_tiling/kernel_tiling.h"
21+ 
22+namespace ScatterNdCommon {
23+ 
24+constexpr uint16_t MAX_RANK_COUNT_NUM = 7;
25+constexpr uint16_t MAX_SHAPE_RANK_NUM = 8;
26+ 
27+struct ScatterNdCommonSimtTilingData{
28+ uint64_t blockNum;
29+ uint32_t rankSize;
30+ uint64_t blockTilingSize;
31+ uint64_t tailBlockTilingSize;
32+ uint32_t ubTilingSize;
33+ uint64_t sliceSize;
34+ uint64_t outPutShape[MAX_SHAPE_RANK_NUM];
35+ uint64_t strideList[MAX_RANK_COUNT_NUM];
36+ int64_t varInAxis;
37+};
38+ 
39+struct ScatterNdCommonSimtSortTilingData{
40+ uint64_t strideList[MAX_RANK_COUNT_NUM];
41+ int64_t indicesFactor;
42+ int64_t afterAxis;
43+ int64_t varInAxis;
44+ int64_t afterAxisFactor;
45+ int64_t indexRankSize;
46+ int64_t eachCoreAfterAxisCount;
47+ int64_t eachCoreIndexCount;
48+ int64_t tailCoreIndexCount;
49+ int64_t usedCoreNumBefore;
50+ int64_t updateLoopSize;
51+ int64_t updateTailNum;
52+};
53+ 
54+struct ScatterNdCommonSimdSortTilingData{
55+ uint64_t strideList[MAX_RANK_COUNT_NUM];
56+ uint64_t outPutShape[MAX_SHAPE_RANK_NUM];
57+ int64_t eachCoreAfterAxisCount;
58+ int64_t indexRankSize;
59+ int64_t eachCoreIndexCount;
60+ int64_t tailCoreIndexCount;
61+ int64_t indicesFactor;
62+ int64_t indiceTailNum;
63+ int64_t indicesLoopSize;
64+ int64_t afterAxis;
65+ int64_t afterAxisFactor;
66+ int64_t usedCoreNumBefore;
67+ int64_t updateLoopSize;
68+ int64_t tailUpdateLoopSize;
69+ int64_t tailUpdateTailNum;
70+ int64_t updateTailNum;
71+ int64_t isSplitAfterAxis;
72+ int64_t singleCol;
73+};
74+ 
75+ 
76+}// namespace ScatterNdCommon
77+#endif
Aindex/scatter_nd_common/op_kernel/arch35/scatter_nd_max_tiling_key.h+63-0
@@ -0,0 +1,63 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file scatter_nd_max_tiling_key.h
13+ * \brief
14+ */
15+#ifndef SCATTER_ND_MAX_TILING_KEY_H_
16+#define SCATTER_ND_MAX_TILING_KEY_H_
17+ 
18+#include "ascendc/host_api/tiling/template_argument.h"
19+ 
20+#define TPL_MODE_ADDR_INT32 0
21+#define TPL_MODE_ADDR_INT64 1
22+ 
23+#define CAST_0 0
24+#define CAST_1 1
25+#define CAST_2 2
26+#define CAST_3 3
27+#define CAST_4 4
28+#define CAST_5 5
29+ 
30+#define TPL_MODE_TEMPLATE_SIMD_SORT 2
31+ 
32+#define TPL_MODE_TEMPLATE_SIMT 8
33+ 
34+namespace ScatterNdCommon {
35+ 
36+ASCENDC_TPL_ARGS_DECL(
Z
Zz300751995月20日

问下ascendc，能否写两个算子名

likedislike
37+ ScatterNdMax,
38+ ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, 2, ASCENDC_TPL_UI_LIST, TPL_MODE_TEMPLATE_SIMD_SORT, TPL_MODE_TEMPLATE_SIMT),
39+ ASCENDC_TPL_UINT_DECL(CAST_MODE, 3, ASCENDC_TPL_UI_LIST, CAST_0, CAST_1, CAST_2, CAST_3, CAST_4, CAST_5),
40+ ASCENDC_TPL_UINT_DECL(ADDR_MODE, 1, ASCENDC_TPL_UI_LIST, TPL_MODE_ADDR_INT32, TPL_MODE_ADDR_INT64)
41+);
42+ 
43+ 
44+ASCENDC_TPL_SEL(
45+ ASCENDC_TPL_ARGS_SEL(
46+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
47+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_TEMPLATE_SIMD_SORT),
48+ ASCENDC_TPL_UINT_SEL(CAST_MODE, ASCENDC_TPL_UI_LIST, CAST_0, CAST_1, CAST_2, CAST_3, CAST_4, CAST_5),
49+ ASCENDC_TPL_UINT_SEL(ADDR_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_ADDR_INT32, TPL_MODE_ADDR_INT64),
50+ ASCENDC_TPL_TILING_STRUCT_SEL(ScatterNdCommonSimdSortTilingData)
51+ ),
52+ ASCENDC_TPL_ARGS_SEL(
53+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
54+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_TEMPLATE_SIMT),
55+ ASCENDC_TPL_UINT_SEL(CAST_MODE, ASCENDC_TPL_UI_LIST, CAST_0),
56+ ASCENDC_TPL_UINT_SEL(ADDR_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_ADDR_INT32, TPL_MODE_ADDR_INT64),
57+ ASCENDC_TPL_TILING_STRUCT_SEL(ScatterNdCommonSimtTilingData)
58+ )
59+);
60+ 
61+} // namespace ScatterNdCommon
62+ 
63+#endif // SCATTER_ND_MAX_TILING_KEY_H_
Aindex/scatter_nd_common/op_kernel/arch35/scatter_nd_min_tiling_key.h+63-0
@@ -0,0 +1,63 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file scatter_nd_min_tiling_key.h
13+ * \brief
14+ */
15+#ifndef SCATTER_ND_MIN_TILING_KEY_H_
16+#define SCATTER_ND_MIN_TILING_KEY_H_
17+ 
18+#include "ascendc/host_api/tiling/template_argument.h"
19+ 
20+#define TPL_MODE_ADDR_INT32 0
21+#define TPL_MODE_ADDR_INT64 1
22+ 
23+#define CAST_0 0
24+#define CAST_1 1
25+#define CAST_2 2
26+#define CAST_3 3
27+#define CAST_4 4
28+#define CAST_5 5
29+ 
30+#define TPL_MODE_TEMPLATE_SIMD_SORT 2
31+ 
32+#define TPL_MODE_TEMPLATE_SIMT 8
33+ 
34+namespace ScatterNdCommon {
35+ 
36+ASCENDC_TPL_ARGS_DECL(
37+ ScatterNdMin,
38+ ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, 2, ASCENDC_TPL_UI_LIST, TPL_MODE_TEMPLATE_SIMD_SORT, TPL_MODE_TEMPLATE_SIMT),
39+ ASCENDC_TPL_UINT_DECL(CAST_MODE, 3, ASCENDC_TPL_UI_LIST, CAST_0, CAST_1, CAST_2, CAST_3, CAST_4, CAST_5),
40+ ASCENDC_TPL_UINT_DECL(ADDR_MODE, 1, ASCENDC_TPL_UI_LIST, TPL_MODE_ADDR_INT32, TPL_MODE_ADDR_INT64)
41+);
42+ 
43+ 
44+ASCENDC_TPL_SEL(
45+ ASCENDC_TPL_ARGS_SEL(
46+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
47+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_TEMPLATE_SIMD_SORT),
48+ ASCENDC_TPL_UINT_SEL(CAST_MODE, ASCENDC_TPL_UI_LIST, CAST_0, CAST_1, CAST_2, CAST_3, CAST_4, CAST_5),
49+ ASCENDC_TPL_UINT_SEL(ADDR_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_ADDR_INT32, TPL_MODE_ADDR_INT64),
50+ ASCENDC_TPL_TILING_STRUCT_SEL(ScatterNdCommonSimdSortTilingData)
51+ ),
52+ ASCENDC_TPL_ARGS_SEL(
53+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
54+ ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_TEMPLATE_SIMT),
55+ ASCENDC_TPL_UINT_SEL(CAST_MODE, ASCENDC_TPL_UI_LIST, CAST_0),
56+ ASCENDC_TPL_UINT_SEL(ADDR_MODE, ASCENDC_TPL_UI_LIST, TPL_MODE_ADDR_INT32, TPL_MODE_ADDR_INT64),
57+ ASCENDC_TPL_TILING_STRUCT_SEL(ScatterNdCommonSimtTilingData)
58+ )
59+);
60+ 
61+} // namespace ScatterNdCommon
62+ 
63+#endif // SCATTER_ND_MIN_TILING_KEY_H_
Aindex/scatter_nd_max/CMakeLists.txt+15-0
@@ -0,0 +1,15 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+# 设置算子定义时支持的芯片类型
12+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_nd_max ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE DEPENDENCIES scatter_nd_common)
Aindex/scatter_nd_max/README.md+88-0
@@ -0,0 +1,88 @@
1+# ScatterNdMax
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| :----------------------------------------------------------- | :------: |
7+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8+| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | × |
9+| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | × |
10+| <term>Atlas 200I/500 A2 推理产品</term> | × |
11+| <term>Atlas 推理系列产品</term> | × |
12+| <term>Atlas 训练系列产品</term> | × |
13+ 
14+## 功能说明
15+ 
16+- 算子功能：根据indices在给定变量内，在updates和单个值或切片之间求最大值。
17+- 计算公式：ref[indices[k]]=max(ref[indices[k]], updates[k])
18+ - ref是一个维度为 P的张量Tensor；
19+ - indices是一个维度为 Q的整型张量Tensor；
20+ - indices的shape一定是[d0,..., d_Q-2, K]，此处0<K<=P；
21+ - updates 是一个维度为Q−1+P−K 的张量：[d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]]
22+ 
23+## 参数说明
24+ 
25+<table style="undefined;table-layout: fixed; width: 1576px"><colgroup>
26+ <col style="width: 170px">
27+ <col style="width: 170px">
28+ <col style="width: 310px">
29+ <col style="width: 212px">
30+ <col style="width: 100px">
31+ </colgroup>
32+ <thead>
33+ <tr>
34+ <th>参数名</th>
35+ <th>输入/输出/属性</th>
36+ <th>描述</th>
37+ <th>数据类型</th>
38+ <th>数据格式</th>
39+ </tr></thead>
40+ <tbody>
41+ <tr>
42+ <td>var</td>
43+ <td>输入</td>
44+ <td>表示一个待被更新的张量, 等同于公式中的`ref`</td>
45+ <td>DT_INT8、DT_INT16、DT_INT64、DT_UINT64、DT_BOOL、DT_FLOAT16、DT_BF16、DT_FLOAT、DT_INT32、DT_UINT32</td>
46+ <td>ND</td>
47+ </tr>
48+ <tr>
49+ <td>indices</td>
50+ <td>输入</td>
51+ <td>一个索引张量，索引到公式中的`ref`</td>
52+ <td>INT32、INT64。</td>
53+ <td>ND</td>
54+ </tr>
55+ <tr>
56+ <td>updates</td>
57+ <td>输入</td>
58+ <td>使用此张量来更新var张量，必须跟输入var张量保持一样的数据类型</td>
59+ <td>DT_INT8、DT_INT16、DT_INT64、DT_UINT64、DT_BOOL、DT_FLOAT16、DT_BF16、DT_FLOAT、DT_INT32、DT_UINT32</td>
60+ <td>ND</td>
61+ </tr>
62+ <tr>
63+ <td>var</td>
64+ <td>输出</td>
65+ <td>表示更新后的张量</td>
66+ <td>DT_INT8、DT_INT16、DT_INT64、DT_UINT64、DT_BOOL、DT_FLOAT16、DT_BF16、DT_FLOAT、DT_INT32、DT_UINT32</td>
67+ <td>ND</td>
68+ </tr>
69+ <tr>
70+ <td>use_locking</td>
71+ <td>可选属性</td>
72+ <td>可选属性，默认值为`false`, 如果为`true`, 这次操作将会被一个lock保护</td>
73+ <td>BOOL</td>
74+ <td>-</td>
75+ </tr>
76+ </tbody></table>
77+ 
78+## 约束说明
C
Cchenjiao5月28日

算子调用说明缺失

likedislike
Z
z30075199
5月28日 评论：
Z
z30075199
5月28日 评论：
79+ 
80+- 输入shape限制：
81+ - indices至少是2维，其最后1维的大小不能超过varRef的维度大小。
82+ - 假设indices最后1维的大小是a，则updates的shape等于indices除最后1维外的shape加上varRef除前a维外的shape。举例：varRef的shape是(4, 5, 6)，indices的shape是(3, 2)，则updates的shape必须是(3, 6)。
83+ 
84+## 调用说明
85+ 
86+| 调用方式 | 样例代码 | 说明 |
87+| ---------------- | --------------------------- | --------------------------------------------------- |
88+| 图模式 | [test_geir_scatter_nd_max](examples/test_geir_scatter_nd_max.cpp) | 通过 GE IR 构图方式调用 ScatterNdMax 算子。 |