Star906
1826
代码
代码
Issues226
Pull Requests306
Actions
项目讨论
Wiki
分析
Star906
1826
已合并
新增算子：MaxPool3DWithArgmaxV2 #1525
小王！创建于 2月4日
新增算子：MaxPool3DWithArgmaxV2 #1525
已合并
小王！创建于 2月4日
从已删除 :migrate合入到cann/ops-nnmaster
共 48 个文件变更+70469-100
Mcommon/stub/op_api/level0/max_pool3d_with_argmax_v2.h+3-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -17,7 +17,8 @@ namespace l0op {
17const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(const aclTensor* self, const aclIntArray* kernelSize,17const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(const aclTensor* self, const aclIntArray* kernelSize,
18 const aclIntArray* stride, const aclIntArray* padding,18 const aclIntArray* stride, const aclIntArray* padding,
19 const aclIntArray* dilation, bool ceilMode,19 const aclIntArray* dilation, bool ceilMode,
20- std::string dataFormat, aclOpExecutor* executor);20+ std::string dataFormat, aclOpExecutor* executor,
21+ op::DataType indicesDtype = op::DataType::DT_INT32);
21} // namespace l0op22} // namespace l0op
22 23 
23#endif // OP_API_INC_LEVEL0_MAX_POOL3D_WITH_ARGMAX_V2_H_24#endif // OP_API_INC_LEVEL0_MAX_POOL3D_WITH_ARGMAX_V2_H_
Mcommon/stub/op_api/op_api_stub.cpp+2-2文件内容审核中，请稍后刷新重试
Mpooling/max_pool3d_with_argmax_v2/op_api/aclnn_max_pool3d_with_argmax.cpp+37-18
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -24,6 +24,7 @@
24#include "opdev/make_op_executor.h"24#include "opdev/make_op_executor.h"
25#include "opdev/platform.h"25#include "opdev/platform.h"
26#include "opdev/framework_op.h"26#include "opdev/framework_op.h"
27+#include "op_api/aclnn_util.h"
27 28 
28using namespace op;29using namespace op;
29#ifdef __cplusplus30#ifdef __cplusplus
@@ -34,6 +35,7 @@ static const std::initializer_list<DataType> NULL_SUPPORT_LIST = {};
34static const std::initializer_list<DataType> OUT_DTYPE_SUPPORT_LIST = {35static const std::initializer_list<DataType> OUT_DTYPE_SUPPORT_LIST = {
35 DataType::DT_BF16, DataType::DT_FLOAT16, DataType::DT_FLOAT};36 DataType::DT_BF16, DataType::DT_FLOAT16, DataType::DT_FLOAT};
36static const std::initializer_list<op::DataType> INDICES_DTYPE_SUPPORT_LIST = {op::DataType::DT_INT32};37static const std::initializer_list<op::DataType> INDICES_DTYPE_SUPPORT_LIST = {op::DataType::DT_INT32};
38+static const std::initializer_list<op::DataType> INDICES_DTYPE_SUPPORT_LIST_REGBASE = {op::DataType::DT_INT32, op::DataType::DT_INT64};
37 39 
38static const size_t CDHW_DIMS = 4;40static const size_t CDHW_DIMS = 4;
39static const size_t NCDHW_DIMS = 5;41static const size_t NCDHW_DIMS = 5;
@@ -55,9 +57,8 @@ static bool CheckNotNullPtr(
55 57 
56static const std::initializer_list<DataType>& GetDtypeSupportList()58static const std::initializer_list<DataType>& GetDtypeSupportList()
57{59{
58- if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B ||60+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
59- GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910_93 ||61+ if (curArch == NpuArch::DAV_2201 || Ops::NN::AclnnUtil::IsRegbase(curArch)) {
60- Ops::NN::AclnnUtil::IsRegbase()) {
61 return OUT_DTYPE_SUPPORT_LIST;62 return OUT_DTYPE_SUPPORT_LIST;
62 } else {63 } else {
63 return NULL_SUPPORT_LIST;64 return NULL_SUPPORT_LIST;
@@ -68,7 +69,11 @@ static bool CheckDtypeValid(const aclTensor* self, const aclTensor* out, const a
68{69{
69 auto dtypeSupportList = GetDtypeSupportList();70 auto dtypeSupportList = GetDtypeSupportList();
70 OP_CHECK_DTYPE_NOT_SUPPORT(self, dtypeSupportList, return false);71 OP_CHECK_DTYPE_NOT_SUPPORT(self, dtypeSupportList, return false);
71- OP_CHECK_DTYPE_NOT_SUPPORT(indices, INDICES_DTYPE_SUPPORT_LIST, return false);72+ if (Ops::NN::AclnnUtil::IsRegbase()) {
73+ OP_CHECK_DTYPE_NOT_SUPPORT(indices, INDICES_DTYPE_SUPPORT_LIST_REGBASE, return false);
74+ } else {
75+ OP_CHECK_DTYPE_NOT_SUPPORT(indices, INDICES_DTYPE_SUPPORT_LIST, return false);
76+ }
72 OP_CHECK_DTYPE_NOT_SAME(self, out, return false);77 OP_CHECK_DTYPE_NOT_SAME(self, out, return false);
73 return true;78 return true;
74}79}
@@ -95,6 +100,16 @@ static bool CheckFormat(const aclTensor* self, const aclTensor* out, const aclTe
95 return false;100 return false;
96 }101 }
97 102 
103+ if (Ops::NN::AclnnUtil::IsRegbase()) {
104+ if (self->GetStorageFormat() != ge::FORMAT_NDHWC && self->GetStorageFormat() != ge::FORMAT_NCDHW &&
105+ self->GetStorageFormat() != ge::FORMAT_NCHW && self->GetStorageFormat() != ge::FORMAT_ND) {
106+ OP_LOGE(
107+ ACLNN_ERR_PARAM_INVALID, "Format of input is not supported, self [%s].",
108+ op::ToString(self->GetStorageFormat()).GetString());
109+ return false;
110+ }
111+ }
112+ 
98 return true;113 return true;
99}114}
100 115 
@@ -168,13 +183,16 @@ static bool CheckParamsValid(
168 const int64_t dilationD = dilationRef[0];183 const int64_t dilationD = dilationRef[0];
169 const int64_t dilationH = (dilationRef.Size() == 1) ? dilationD : dilationRef[1];184 const int64_t dilationH = (dilationRef.Size() == 1) ? dilationD : dilationRef[1];
170 const int64_t dilationW = (dilationRef.Size() == 1) ? dilationD : dilationRef[2];185 const int64_t dilationW = (dilationRef.Size() == 1) ? dilationD : dilationRef[2];
171- OP_CHECK(186+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
172- ((dilationD == 1) && (dilationH == 1) && (dilationW == 1)),187+ if (curArch == NpuArch::DAV_2201 || curArch == NpuArch::DAV_1001) {
173- OP_LOGE(188+ OP_CHECK(
174- ACLNN_ERR_PARAM_INVALID,189+ ((dilationD == 1) && (dilationH == 1) && (dilationW == 1)),
175- "The value of dilation should be equal 1, but got dilationD:%ld, dilationH:%ld, dilationW:%ld", dilationD,190+ OP_LOGE(
176- dilationH, dilationW),191+ ACLNN_ERR_PARAM_INVALID,
177- return false);192+ "The value of dilation should be equal 1, but got dilationD:%ld, dilationH:%ld, dilationW:%ld", dilationD,
193+ dilationH, dilationW),
194+ return false);
195+ }
178 196 
179 return true;197 return true;
180}198}
@@ -192,9 +210,8 @@ static aclnnStatus CheckParams(
192 210 
193static bool CheckPlatform()211static bool CheckPlatform()
194{212{
195- if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B ||213+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
196- GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910_93 ||214+ if (curArch == NpuArch::DAV_2201 || Ops::NN::AclnnUtil::IsRegbase(curArch)) {
197- Ops::NN::AclnnUtil::IsRegbase()) {
198 return true;215 return true;
199 } else {216 } else {
200 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnMaxPool3dWithArgmax is not supported on this platform");217 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnMaxPool3dWithArgmax is not supported on this platform");
@@ -251,7 +268,8 @@ aclnnStatus aclnnMaxPool3dWithArgmaxGetWorkspaceSize(
251 auto ret = CheckParams(self, kernelSize, stride, padding, dilation, out, indices);268 auto ret = CheckParams(self, kernelSize, stride, padding, dilation, out, indices);
252 CHECK_RET(ret == ACLNN_SUCCESS, ret);269 CHECK_RET(ret == ACLNN_SUCCESS, ret);
253 270 
254- std::string dataFormat = "NCDHW";271+ // Default to NCDHW
272+ std::string dataFormat = (self->GetStorageFormat() == ge::FORMAT_NDHWC) ? "NDHWC" : "NCDHW";
255 273 
256 // Check whether the tensor is empty (the operator does not support empty tensors)274 // Check whether the tensor is empty (the operator does not support empty tensors)
257 if (self->IsEmpty() || out->IsEmpty()) {275 if (self->IsEmpty() || out->IsEmpty()) {
@@ -272,8 +290,9 @@ aclnnStatus aclnnMaxPool3dWithArgmaxGetWorkspaceSize(
272 CHECK_RET(selfUnsqueezed != nullptr, ACLNN_ERR_INNER_NULLPTR);290 CHECK_RET(selfUnsqueezed != nullptr, ACLNN_ERR_INNER_NULLPTR);
273 291 
274 // Returns a tuple containing out and indices292 // Returns a tuple containing out and indices
293+ auto indicesDtype = indices->GetDataType();
275 auto [outResult, indicesResult] = l0op::MaxPool3DWithArgmaxV2Ncdhw(294 auto [outResult, indicesResult] = l0op::MaxPool3DWithArgmaxV2Ncdhw(
276- selfUnsqueezed, kernelSize, stride, padding, dilation, ceilMode, dataFormat, uniqueExecutor.get());295+ selfUnsqueezed, kernelSize, stride, padding, dilation, ceilMode, dataFormat, uniqueExecutor.get(), indicesDtype);
277 296 
278 CHECK_RET(outResult != nullptr, ACLNN_ERR_INNER_NULLPTR);297 CHECK_RET(outResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
279 CHECK_RET(indicesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);298 CHECK_RET(indicesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
@@ -313,4 +332,4 @@ aclnnStatus aclnnMaxPool3dWithArgmax(
313 332 
314#ifdef __cplusplus333#ifdef __cplusplus
315}334}
316-#endif335+#endif
Mpooling/max_pool3d_with_argmax_v2/op_api/aclnn_max_pool3d_with_argmax.h+1-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
Mpooling/max_pool3d_with_argmax_v2/op_api/max_pool3d_with_argmax_v2.cpp+35-15
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -34,12 +34,15 @@ static const uint32_t NC = 2;
34static const int64_t TWO_SIDED = 2;34static const int64_t TWO_SIDED = 2;
35static const size_t DHW_DIMS = 3;35static const size_t DHW_DIMS = 3;
36static const int64_t NCDHW_DIMS = 5;36static const int64_t NCDHW_DIMS = 5;
37+static const int64_t NCDHW_D_OFFSET = 3; // Get D-index by totalDim - offset
38+static const int64_t NDHWC_D_OFFSET = 4;
39+static const int32_t DTYPE_INT32 = 3;
40+static const int32_t DTYPE_INT64 = 9;
37 41 
38static const std::initializer_list<DataType> GetDtypeSupportListBySocVersion()42static const std::initializer_list<DataType> GetDtypeSupportListBySocVersion()
39{43{
40- if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B ||44+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
41- GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910_93 ||45+ if (curArch == NpuArch::DAV_2201 || Ops::NN::AclnnUtil::IsRegbase(curArch)) {
42- Ops::NN::AclnnUtil::IsRegbase()) {
43 return SELF_DTYPE_SUPPORT_LIST;46 return SELF_DTYPE_SUPPORT_LIST;
44 } else {47 } else {
45 return NULL_SUPPORT_LIST;48 return NULL_SUPPORT_LIST;
@@ -55,6 +58,9 @@ static bool IsMaxPool3DWithArgmaxV2NcdhwAiCoreSupported(const aclTensor* self)
55 op::ToString(dtypeSupportList).GetString());58 op::ToString(dtypeSupportList).GetString());
56 return false;59 return false;
57 }60 }
61+ if (Ops::NN::AclnnUtil::IsRegbase()) {
62+ return true; // Support for DHW exceeding max int32
63+ }
58 op::Shape input_shape = self->GetViewShape();64 op::Shape input_shape = self->GetViewShape();
59 int64_t dims = input_shape.GetDimNum();65 int64_t dims = input_shape.GetDimNum();
60 const uint32_t batchesDims = dims == NCDHW_DIMS ? NC : C;66 const uint32_t batchesDims = dims == NCDHW_DIMS ? NC : C;
@@ -76,9 +82,16 @@ static const std::tuple<aclTensor*, aclTensor*> MaxPool3DWithArgmaxV2NcdhwAiCore
76 MaxPool3DWithArgmaxV2NcdhwAiCore, self, kernelSize, stride, padding, dilation, ceilMode, dataFormat, out,82 MaxPool3DWithArgmaxV2NcdhwAiCore, self, kernelSize, stride, padding, dilation, ceilMode, dataFormat, out,
77 indices);83 indices);
78 84 
79- ADD_TO_LAUNCHER_LIST_AICORE(85+ if (Ops::NN::AclnnUtil::IsRegbase()) {
80- MaxPool3DWithArgmaxV2, OP_INPUT(self), OP_OUTPUT(out, indices),86+ int32_t dtype = (indices->GetDataType() == op::DataType::DT_INT64) ? DTYPE_INT64 : DTYPE_INT32;
81- OP_ATTR(kernelSize, stride, padding, dilation, ceilMode, dataFormat));87+ ADD_TO_LAUNCHER_LIST_AICORE(
88+ MaxPool3DWithArgmaxV2, OP_INPUT(self), OP_OUTPUT(out, indices),
89+ OP_ATTR(kernelSize, stride, padding, dilation, ceilMode, dataFormat, dtype));
90+ } else {
91+ ADD_TO_LAUNCHER_LIST_AICORE(
92+ MaxPool3DWithArgmaxV2, OP_INPUT(self), OP_OUTPUT(out, indices),
93+ OP_ATTR(kernelSize, stride, padding, dilation, ceilMode, dataFormat));
94+ }
82 return std::tuple<aclTensor*, aclTensor*>(out, indices);95 return std::tuple<aclTensor*, aclTensor*>(out, indices);
83}96}
84 97 
@@ -114,7 +127,7 @@ static int64_t PoolingOutShape(
114 127 
115static op::Shape GetOutputShape(128static op::Shape GetOutputShape(
116 const aclTensor* self, const aclIntArray* kernelSize, const aclIntArray* stride, const aclIntArray* padding,129 const aclTensor* self, const aclIntArray* kernelSize, const aclIntArray* stride, const aclIntArray* padding,
117- const aclIntArray* dilation, bool ceilMode)130+ const aclIntArray* dilation, bool ceilMode, std::string dataFormat)
118{131{
119 op::Shape input_shape = self->GetViewShape();132 op::Shape input_shape = self->GetViewShape();
120 int64_t dims = input_shape.GetDimNum();133 int64_t dims = input_shape.GetDimNum();
@@ -123,7 +136,7 @@ static op::Shape GetOutputShape(
123 const aclIntArray& paddingRef = *padding;136 const aclIntArray& paddingRef = *padding;
124 const aclIntArray& dilationRef = *dilation;137 const aclIntArray& dilationRef = *dilation;
125 op::Shape outputShape;138 op::Shape outputShape;
126- const uint32_t batchesDims = dims == NCDHW_DIMS ? NC : C;139+ const uint32_t batchesDims = dataFormat == "NDHWC" ? dims - NDHWC_D_OFFSET : dims - NCDHW_D_OFFSET;
127 int64_t curDim[DHW_DIMS];140 int64_t curDim[DHW_DIMS];
128 141 
129 curDim[0] = PoolingOutShape(142 curDim[0] = PoolingOutShape(
@@ -133,17 +146,24 @@ static op::Shape GetOutputShape(
133 curDim[NC] = PoolingOutShape(146 curDim[NC] = PoolingOutShape(
134 input_shape.GetDim(batchesDims + NC), kernelRef[NC], strideRef[NC], paddingRef[NC], dilationRef[NC], ceilMode);147 input_shape.GetDim(batchesDims + NC), kernelRef[NC], strideRef[NC], paddingRef[NC], dilationRef[NC], ceilMode);
135 148 
136- if (dims == NCDHW_DIMS) {149+ if (dataFormat == "NDHWC") {
137- outputShape = {input_shape.GetDim(0), input_shape.GetDim(1), curDim[0], curDim[1], curDim[NC]};150+ if (dims == NCDHW_DIMS) {
151+ outputShape = {input_shape.GetDim(0), curDim[0], curDim[1], curDim[NC], input_shape.GetDim(4)};
152+ }
138 } else {153 } else {
139- outputShape = {input_shape.GetDim(0), curDim[0], curDim[1], curDim[NC]};154+ if (dims == NCDHW_DIMS) {
155+ outputShape = {input_shape.GetDim(0), input_shape.GetDim(1), curDim[0], curDim[1], curDim[NC]};
156+ } else {
157+ outputShape = {input_shape.GetDim(0), curDim[0], curDim[1], curDim[NC]};
158+ }
140 }159 }
141 return outputShape;160 return outputShape;
142}161}
143 162 
144const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(163const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(
145 const aclTensor* self, const aclIntArray* kernelSize, const aclIntArray* stride, const aclIntArray* padding,164 const aclTensor* self, const aclIntArray* kernelSize, const aclIntArray* stride, const aclIntArray* padding,
146- const aclIntArray* dilation, bool ceilMode, std::string dataFormat, aclOpExecutor* executor)165+ const aclIntArray* dilation, bool ceilMode, std::string dataFormat, aclOpExecutor* executor,
166+ op::DataType indicesDtype)
147{167{
148 const aclIntArray& kernelRef = *kernelSize;168 const aclIntArray& kernelRef = *kernelSize;
149 const int64_t kernelD = kernelRef[0];169 const int64_t kernelD = kernelRef[0];
@@ -175,10 +195,10 @@ const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(
175 aclIntArray* padding3 = executor->AllocIntArray(paddingSizeData.data(), DHW_DIMS);195 aclIntArray* padding3 = executor->AllocIntArray(paddingSizeData.data(), DHW_DIMS);
176 aclIntArray* dilation3 = executor->AllocIntArray(dilationSizeData.data(), DHW_DIMS);196 aclIntArray* dilation3 = executor->AllocIntArray(dilationSizeData.data(), DHW_DIMS);
177 197 
178- op::Shape outShape = GetOutputShape(self, kernelSize3, stride3, padding3, dilation3, ceilMode);198+ op::Shape outShape = GetOutputShape(self, kernelSize3, stride3, padding3, dilation3, ceilMode, dataFormat);
179 op::DataType outType = self->GetDataType();199 op::DataType outType = self->GetDataType();
180 auto out = executor->AllocTensor(outShape, outType, self->GetViewFormat());200 auto out = executor->AllocTensor(outShape, outType, self->GetViewFormat());
181- auto indices = executor->AllocTensor(outShape, op::DataType::DT_INT32, self->GetViewFormat());201+ auto indices = executor->AllocTensor(outShape, indicesDtype, self->GetViewFormat());
182 202 
183 if (IsMaxPool3DWithArgmaxV2NcdhwAiCoreSupported(self)) {203 if (IsMaxPool3DWithArgmaxV2NcdhwAiCoreSupported(self)) {
184 return MaxPool3DWithArgmaxV2NcdhwAiCore(204 return MaxPool3DWithArgmaxV2NcdhwAiCore(
Mpooling/max_pool3d_with_argmax_v2/op_api/max_pool3d_with_argmax_v2.h+3-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -16,7 +16,8 @@
16namespace l0op {16namespace l0op {
17const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(17const std::tuple<const aclTensor*, const aclTensor*> MaxPool3DWithArgmaxV2Ncdhw(
18 const aclTensor* self, const aclIntArray* kernelSize, const aclIntArray* stride, const aclIntArray* padding,18 const aclTensor* self, const aclIntArray* kernelSize, const aclIntArray* stride, const aclIntArray* padding,
19- const aclIntArray* dilation, bool ceilMode, std::string dataFormat, aclOpExecutor* executor);19+ const aclIntArray* dilation, bool ceilMode, std::string dataFormat, aclOpExecutor* executor,
20+ op::DataType indicesDtype = op::DataType::DT_INT32);
20} // namespace l0op21} // namespace l0op
21 22 
22#endif // OP_API_INC_LEVEL0_MAX_POOL3D_WITH_ARGMAX_V2_H_23#endif // OP_API_INC_LEVEL0_MAX_POOL3D_WITH_ARGMAX_V2_H_
Apooling/max_pool3d_with_argmax_v2/op_graph/max_pool3d_with_argmax_v2_proto.h+70-0
@@ -0,0 +1,70 @@
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 max_pool3d_with_argmax_v2_proto.h
13+ * \brief
14+ */
15+ 
16+#ifndef OPS_POOLING_MAX_POOL3D_WITH_ARGMAX_V2_PROTO_H_
17+#define OPS_POOLING_MAX_POOL3D_WITH_ARGMAX_V2_PROTO_H_
18+ 
19+#include "graph/operator_reg.h"
20+ 
21+namespace ge {
22+/**
23+* @brief Performs max pooling on the input and outputs both max values and indices.
24+ 
25+* @par Inputs:
26+* One input:
27+* x: A tensor of type bfloat16, float16, float32, the shape is [batch, channels, depth_in, height_in, width_in] or
28+ [batch, depth_in, height_in, width_in, channels].
29+ 
30+* @par Attributes:
31+* @li ksize: A required list of int64 values,
32+* specifying the size of the window for each dimension of the input tensor.
33+* A list that has length 3.
34+* @li strides: A required list of int64 values,
35+* specifying the stride of the sliding window for each dimension of the input tensor.
36+* A list that has length 3.
37+* @li pads: A required list of int64 values,
38+* specifying the pad of the input feature map.
39+* A list that has length 3:
40+* 0 <= pads[0] <= (ksize[0]//2), 0 <= pads[1] <= (ksize[1]//2), 0 <= pads[2] <= (ksize[2]//2).
41+* @li dilation: A list that has length 3, default value is {1,1,1}.
42+* @li ceil_mode: When true, will use ceil instead of floor to compute the output shape, defaults to false.
43+* @li data_format: The value can be "NCDHW" or "NDHWC", defaults to "NCDHW".
44+* @li dtype: An optional int, default value is 3. (3 is int32, 9 is int64)
45+ 
46+* @par Outputs:
47+* @li y: A tensor has the same type and format as input "x", the shape is [batch, channels, depth_out, height_out, width_out] or
48+ [batch, depth_out, height_out, width_out, channels].
49+* @li argmax: A tensor of type is int64 or int32, the shape is [batch, channels, depth_out, height_out, width_out] or
50+ [batch, depth_out, height_out, width_out, channels].
51+ 
CANN-robotCANN-robot
CANN-robotCANN-robot2月4日

代码结构与可维护性: 属性dtype的默认值注释中使用了魔数3和9，分别表示int32和int64。这种硬编码的魔数降低了代码的可读性和可维护性。其他开发者需要查看注释才能理解这些数字的含义，容易导致误解或错误使用。

问题类型: 代码结构与可维护性 文件路径: pooling/max_pool3d_with_argmax_v2/op_graph/max_pool3d_with_argmax_v2_proto.h 行号: 51 问题代码:

* @li dtype: An optional int, default value is 3.  (3 is int32, 9 is int64)

修改建议:

建议定义枚举常量或宏定义来代替这些魔数，例如：
#define DTYPE_INT32 3
#define DTYPE_INT64 9
并在注释中引用这些常量，或者在头文件中包含相关的数据类型定义。

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

likedislike
CANN-robotCANN-robot2月4日

枚举值文档不完整: 第51行注释中说明dtype的默认值是3，并注明3是int32，9是int64，但没有提供完整的类型映射表。用户可能不知道其他可能的取值及其含义，也没有说明如果传入无效值会发生什么。

问题类型: 枚举值文档不完整 文件路径: pooling/max_pool3d_with_argmax_v2/op_graph/max_pool3d_with_argmax_v2_proto.h 行号: 51 问题代码:

* @li dtype: An optional int, default value is 3.  (3 is int32, 9 is int64)

修改建议:

建议：1) 提供完整的类型映射表，列出所有支持的dtype值及其对应的数据类型；2) 说明无效值的处理方式（是否抛出异常、是否使用默认值）；3) 考虑使用框架提供的类型枚举常量而不是数字值。

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

likedislike
52+* @par Third-party framework compatibility
53+* Compatible with the PyTorch operator max_pool3d_with_indices.
54+*/
55+ 
56+REG_OP(MaxPool3DWithArgmaxV2)
57+ .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
58+ .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
59+ .OUTPUT(argmax, TensorType({DT_INT32, DT_INT64}))
60+ .REQUIRED_ATTR(ksize, ListInt)
61+ .REQUIRED_ATTR(strides, ListInt)
62+ .REQUIRED_ATTR(pads, ListInt)
63+ .ATTR(dilation, ListInt, {1, 1, 1})
64+ .ATTR(ceil_mode, Bool, false)
65+ .ATTR(data_format, String, "NCDHW")
66+ .ATTR(dtype, Int, 3)
67+ .OP_END_FACTORY_REG(MaxPool3DWithArgmaxV2)
68+} //namespace ge
69+ 
70+#endif //OPS_POOLING_MAX_POOL3D_WITH_ARGMAX_V2_PROTO_H_
Apooling/max_pool3d_with_argmax_v2/op_host/common_dtype.h+45-0
@@ -0,0 +1,45 @@
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 common_dtype.h
13+ * \brief
14+ */
15+ 
16+#ifndef AIR_CXX_RUNTIME_V2_COMMON_DTYPE_H_
17+#define AIR_CXX_RUNTIME_V2_COMMON_DTYPE_H_
18+ 
19+#include <string>
20+#include <sstream>
21+ 
22+namespace optiling
23+{
24+/**
25+ ** function: ConcatString
26+ */
27+template <typename T>
28+std::string ConcatString(const T& arg)
29+{
30+ std::ostringstream oss;
31+ oss << arg;
32+ return oss.str();
33+}
34+ 
35+template <typename T, typename... Ts>
36+std::string ConcatString(const T& arg, const Ts&... argLeft)
37+{
38+ std::ostringstream oss;
39+ oss << arg;
40+ oss << ConcatString(argLeft...);
41+ return oss.str();
42+}
43+} //namespace optiling
44+ 
45+#endif
Mpooling/max_pool3d_with_argmax_v2/op_host/config/ascend910_93/max_pool3d_with_argmax_v2_binary.json+15-0
@@ -67,6 +67,11 @@
67 "name": "data_format",67 "name": "data_format",
68 "dtype": "string",68 "dtype": "string",
69 "value": "NCDHW"69 "value": "NCDHW"
70+ },
71+ {
72+ "name": "dtype",
73+ "dtype": "int",
74+ "value": 3
70 }75 }
71 ]76 ]
72 },77 },
@@ -136,6 +141,11 @@
136 "name": "data_format",141 "name": "data_format",
137 "dtype": "string",142 "dtype": "string",
138 "value": "NCDHW"143 "value": "NCDHW"
144+ },
145+ {
146+ "name": "dtype",
zengjuan
zengjuanzengjuan2月5日

有新增attr，所有版本的二进制json都得对应修改，否则kernel会编译失败

likedislike
147+ "dtype": "int",
148+ "value": 3
139 }149 }
140 ]150 ]
141 },151 },
@@ -205,6 +215,11 @@
205 "name": "data_format",215 "name": "data_format",
206 "dtype": "string",216 "dtype": "string",
207 "value": "NCDHW"217 "value": "NCDHW"
218+ },
219+ {
220+ "name": "dtype",
221+ "dtype": "int",
222+ "value": 3
208 }223 }
209 ]224 ]
210 }225 }
Mpooling/max_pool3d_with_argmax_v2/op_host/config/ascend910b/max_pool3d_with_argmax_v2_binary.json+15-0
@@ -67,6 +67,11 @@
67 "name": "data_format",67 "name": "data_format",
68 "dtype": "string",68 "dtype": "string",
69 "value": "NCDHW"69 "value": "NCDHW"
70+ },
71+ {
72+ "name": "dtype",
73+ "dtype": "int",
74+ "value": 3
70 }75 }
71 ]76 ]
72 },77 },
@@ -136,6 +141,11 @@
136 "name": "data_format",141 "name": "data_format",
137 "dtype": "string",142 "dtype": "string",
138 "value": "NCDHW"143 "value": "NCDHW"
144+ },
145+ {
146+ "name": "dtype",
147+ "dtype": "int",
148+ "value": 3
139 }149 }
140 ]150 ]
141 },151 },
@@ -205,6 +215,11 @@
205 "name": "data_format",215 "name": "data_format",
206 "dtype": "string",216 "dtype": "string",
207 "value": "NCDHW"217 "value": "NCDHW"
218+ },
219+ {
220+ "name": "dtype",
221+ "dtype": "int",
222+ "value": 3
208 }223 }
209 ]224 ]
210 }225 }
Mpooling/max_pool3d_with_argmax_v2/op_host/config/kirin9030/max_pool3d_with_argmax_v2_binary.json+10-0
@@ -67,6 +67,11 @@
67 "name": "data_format",67 "name": "data_format",
68 "dtype": "string",68 "dtype": "string",
69 "value": "NCDHW"69 "value": "NCDHW"
70+ },
71+ {
72+ "name": "dtype",
73+ "dtype": "int",
74+ "value": 3
70 }75 }
71 ]76 ]
72 },77 },
@@ -136,6 +141,11 @@
136 "name": "data_format",141 "name": "data_format",
137 "dtype": "string",142 "dtype": "string",
138 "value": "NCDHW"143 "value": "NCDHW"
144+ },
145+ {
146+ "name": "dtype",
147+ "dtype": "int",
148+ "value": 3
139 }149 }
140 ]150 ]
141 }151 }
Mpooling/max_pool3d_with_argmax_v2/op_host/config/kirinx90/max_pool3d_with_argmax_v2_binary.json+10-0
@@ -67,6 +67,11 @@
67 "name": "data_format",67 "name": "data_format",
68 "dtype": "string",68 "dtype": "string",
69 "value": "NCDHW"69 "value": "NCDHW"
70+ },
71+ {
72+ "name": "dtype",
73+ "dtype": "int",
74+ "value": 3
70 }75 }
71 ]76 ]
72 },77 },
@@ -136,6 +141,11 @@
136 "name": "data_format",141 "name": "data_format",
137 "dtype": "string",142 "dtype": "string",
138 "value": "NCDHW"143 "value": "NCDHW"
144+ },
145+ {
146+ "name": "dtype",
147+ "dtype": "int",
148+ "value": 3
139 }149 }
140 ]150 ]
141 }151 }
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_def.cpp+36-5文件内容审核中，请稍后刷新重试
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_gather_tiling.cpp+581-0
@@ -0,0 +1,581 @@
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 max_pool_with_argmax_v3_gather_tiling.cpp
CANN-robot
CANN-robotCANN-robot2月4日

代码可维护性问题: 文件头注释中的文件名与实际文件名不一致。第12行显示为"max_pool_with_argmax_v3_gather_tiling.cpp"，但实际文件名是"max_pool3d_with_argmax_v2_gather_tiling.cpp"。这种不一致可能导致开发者困惑。

问题类型: 代码可维护性问题 文件路径: pooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_gather_tiling.cpp 行号: 12 问题代码:

 * \file max_pool_with_argmax_v3_gather_tiling.cpp

修改建议:

将文件头注释中的文件名修正为实际文件名："max_pool3d_with_argmax_v2_gather_tiling.cpp"。

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

likedislike
13+ * \brief
14+ */
15+ 
16+#include <cctype>
17+#include <algorithm>
18+#include "log/log.h"
19+#include "util/math_util.h"
20+#include "error_util.h"
21+#include "tiling_base/tiling_base.h"
22+#include "tiling_base/tiling_templates_registry.h"
23+#include "max_pool3d_with_argmax_v2_gather_tiling.h"
24+#include "op_common/op_host/util/platform_util.h"
25+#include "platform/platform_ascendc.h"
26+#include "register/op_def_registry.h"
27+#include "platform/platform_info.h"
28+#include "tiling_base/tiling_util.h"
29+#include <iostream>
30+namespace optiling
31+{
32+static constexpr int64_t FLOAT16_OR_BF16_SIZE = 2;
33+static constexpr int64_t FLOAT32_SIZE = 4;
34+static constexpr int64_t INT32_SIZE = 4;
35+static constexpr int64_t INT64_SIZE = 8;
36+static constexpr int64_t UB_RESVERVED_SIZE = 0;
CANN-robot
CANN-robotCANN-robot2月4日

拼写错误: 第36行常量UB_RESVERVED_SIZE拼写错误，应为UB_RESERVED_SIZE。

问题类型: 拼写错误 文件路径: pooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_gather_tiling.cpp 行号: 36 问题代码:

static constexpr int64_t UB_RESVERVED_SIZE = 0;

修改建议:

修正拼写错误：UB_RESERVED_SIZE。

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

likedislike
37+static constexpr int64_t HELPER_BUFFER_SIZE = 1024;
38+static constexpr int64_t NO_PADDING_TILING_KEY = 400001;
39+static constexpr int64_t PADDING_TILING_KEY = 400002;
40+static constexpr int64_t MAX_BANDWIDTH_COEFFICIENTS = 2;
41+static constexpr int64_t DOUBLE = 2;
42+static constexpr int64_t CACHE_LINE_SIZE = 128;
43+static constexpr int64_t MIN_OUTPUT_THRESHOLD = 16;
44+static constexpr int64_t DILATION_THRESHOLD = 1;
45+ 
46+static constexpr int64_t NCDHW_DIMS = 5;
47+static constexpr int64_t NUM_64 = 64;
48+static constexpr int64_t INPUT_IDX_X = 0;
49+static constexpr int64_t KERNEL_POS = 0;
50+static constexpr int64_t STRIDE_POS = 1;
51+static constexpr int64_t PADDING_POS = 2;
52+static constexpr int64_t DTYPE_POS = 6;
53+static constexpr int64_t DILATION_POS = 3;
54+static constexpr int64_t CEIL_POS = 4;
55+static constexpr int64_t FORMAT_POS = 5;
56+ 
57+static const int32_t MP_MAX_3D_DIM_ZERO = 0;
58+static const int32_t MP_MAX_3D_DIM_ONE = 1;
59+static const int32_t MP_MAX_3D_DIM_TWO = 2;
60+static const int32_t MP_MAX_3D_DIM_THREE = 3;
61+static const int32_t MP_MAX_3D_DIM_FOUR = 4;
62+static const int64_t MP_MAX_3D_TYPE_INT32 = 3;
63+static const int64_t MP_MAX_3D_TYPE_INT64 = 9;
64+ 
65+static const gert::Shape g_vec_1_shape = {1};
66+ 
67+static const gert::Shape& EnsureNotScalar(const gert::Shape& in_shape)
68+{
69+ if (in_shape.IsScalar()) {
70+ return g_vec_1_shape;
71+ }
72+ return in_shape;
73+}
74+ 
75+ge::graphStatus MaxPool3DWithArgmaxV2GatherTiling::GetPlatformInfo()
76+{
77+ auto platformPtr = context_->GetPlatformInfo();
78+ if (platformPtr == nullptr) {
79+ auto compileInfoPtr = reinterpret_cast<const MaxPool3DWithArgmaxV2CompileInfo*>(context_->GetCompileInfo());
80+ OP_CHECK_IF(
81+ compileInfoPtr == nullptr, OP_LOGE(context_->GetNodeName(), "compile info is null"),
82+ return ge::GRAPH_FAILED);
83+ coreNum = compileInfoPtr->coreNum;
84+ 
85+ ubSize = compileInfoPtr->ubSize;
86+ } else {
87+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformPtr);
88+ coreNum = ascendcPlatform.GetCoreNumAiv();
89+ 
90+ uint64_t ubSizePlatform;
91+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
92+ ubSize = ubSizePlatform;
93+ }
94+ OP_CHECK_IF(
95+ coreNum == 0, OP_LOGE(context_->GetNodeName(), "coreNum is 0"),
96+ return ge::GRAPH_FAILED);
97+ 
98+ return ge::GRAPH_SUCCESS;
99+}
100+ 
101+ge::graphStatus MaxPool3DWithArgmaxV2GatherTiling::GetShapeAttrsInfo()
102+{
103+ auto platformInfo = context_->GetPlatformInfo();
104+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
105+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
106+ if (!Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
107+ return ge::GRAPH_PARAM_INVALID;
108+ }
109+
110+ auto inputX = context_->GetInputShape(0);
111+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputX);
112+ auto inputShape = EnsureNotScalar(inputX->GetStorageShape());
113+ OP_CHECK_IF(inputShape.GetDimNum() != NCDHW_DIMS,
114+ OP_LOGE(context_->GetNodeName(),
115+ "MaxPool3DWithArgmaxV2: input shape dim = %zu, should be equal 5",
116+ inputShape.GetDimNum()),
117+ return ge::GRAPH_FAILED);
118+ OP_CHECK_IF(inputShape.GetShapeSize() <= 0,
119+ OP_LOGE(context_->GetNodeName(),
120+ "MaxPool3DWithArgmaxV2: input shape size %ld less than zero failed",
121+ inputShape.GetShapeSize()),
122+ return ge::GRAPH_FAILED);
123+ auto inputDesc = context_->GetInputDesc(0);
124+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputDesc);
125+ dtype = inputDesc->GetDataType();
126+ if (dtype != ge::DataType::DT_BF16 && dtype != ge::DataType::DT_FLOAT16 && dtype != ge::DataType::DT_FLOAT) {
127+ OP_LOGE(context_->GetNodeName(), "MaxPool3DWithArgmaxV2: invalid dtype");
128+ return ge::GRAPH_FAILED;
129+ }
130+ 
131+ auto outX = context_->GetOutputShape(0);
132+ OP_CHECK_NULL_WITH_CONTEXT(context_, outX);
133+ auto outShape = EnsureNotScalar(outX->GetStorageShape());
134+ auto indicesX = context_->GetOutputShape(1);
135+ OP_CHECK_NULL_WITH_CONTEXT(context_, indicesX);
136+ auto indicesShape = EnsureNotScalar(indicesX->GetStorageShape());
137+ if (indicesShape != outShape) {
138+ OP_LOGE(context_->GetNodeName(),
139+ "MaxPool3DWithArgmaxV2: indices shape and values shape is different");
140+ return ge::GRAPH_FAILED;
141+ }
142+ auto runtimeAttrs = context_->GetAttrs();
143+ OP_CHECK_NULL_WITH_CONTEXT(context_, runtimeAttrs);
144+ std::string inputFormatStr("NCDHW");
145+ const char* inputFormat = runtimeAttrs->GetAttrPointer<char>(FORMAT_POS);
146+ if (inputFormat != nullptr) {
147+ inputFormatStr = inputFormat;
148+ }
149+ int d_dim = MP_MAX_3D_DIM_TWO;
150+ int h_dim = MP_MAX_3D_DIM_THREE;
151+ int w_dim = MP_MAX_3D_DIM_FOUR;
152+ 
153+ if (inputFormatStr == "NCDHW") {
154+ inputData.inputFormat = ge::Format::FORMAT_NCDHW;
155+ inputData.batches = inputShape.GetDim(MP_MAX_3D_DIM_ZERO) * inputShape.GetDim(MP_MAX_3D_DIM_ONE);
156+ inputData.nInput = inputShape.GetDim(MP_MAX_3D_DIM_ZERO);
157+ inputData.cInput = inputShape.GetDim(MP_MAX_3D_DIM_ONE);
158+ } else {
159+ return ge::GRAPH_PARAM_INVALID;
160+ }
161+ 
162+ OP_CHECK_IF(outShape.GetDim(d_dim) < 1 || outShape.GetDim(h_dim) < 1 || outShape.GetDim(w_dim) < 1 ,
163+ OP_LOGE(context_->GetNodeName(),
164+ "MaxPool3DWithArgmaxV2: output shape [%ld, %ld, %ld] not support",
165+ outShape.GetDim(d_dim), outShape.GetDim(h_dim), outShape.GetDim(w_dim)),
166+ return ge::GRAPH_FAILED);
167+ 
168+ inputData.inputShape =
169+ array<uint64_t, DHW_DIMS>{uint64_t(inputShape.GetDim(d_dim)), uint64_t(inputShape.GetDim(h_dim)), uint64_t(inputShape.GetDim(w_dim))};
170+ inputData.outShape =
171+ array<uint64_t, DHW_DIMS>{uint64_t(outShape.GetDim(d_dim)), uint64_t(outShape.GetDim(h_dim)), uint64_t(outShape.GetDim(w_dim))};
172+ int32_t dValue = 0;
173+ int32_t hValue = 0;
174+ int32_t wValue = 0;
175+ const gert::TypedContinuousVector<int64_t>* kernelSize = runtimeAttrs->GetListInt(KERNEL_POS);
176+ OP_CHECK_NULL_WITH_CONTEXT(context_, kernelSize);
177+ dValue = *(kernelSize->GetData());
178+ hValue = *(kernelSize->GetData() + 1);
179+ wValue = *(kernelSize->GetData() + 2);
180+ inputData.kernelSize = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
181+ OP_CHECK_IF(
182+ dValue <= 0 || hValue <= 0 || wValue <= 0,
183+ OP_LOGE(context_->GetNodeName(),
184+ "MaxPool3DWithArgmaxV2: not support kernel shape [%d,%d, %d]", dValue, hValue, wValue),
185+ return ge::GRAPH_FAILED);
186+ 
187+ int32_t kdValue = dValue;
188+ int32_t khValue = hValue;
189+ int32_t kwValue = wValue;
190+ const gert::TypedContinuousVector<int64_t>* stride = runtimeAttrs->GetListInt(STRIDE_POS);
191+ OP_CHECK_NULL_WITH_CONTEXT(context_, stride);
192+ dValue = *(stride->GetData());
193+ hValue = *(stride->GetData() + 1);
194+ wValue = *(stride->GetData() + 2);
195+ inputData.stride = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
196+ OP_CHECK_IF(
197+ hValue <= 0 || wValue <= 0 || dValue <=0,
198+ OP_LOGE(context_->GetNodeName(),
199+ "MaxPool3DWithArgmaxV2: not support stride shape [%d, %d, %d]", dValue, hValue, wValue),
200+ return ge::GRAPH_FAILED);
201+
202+ const gert::TypedContinuousVector<int64_t>* padding = runtimeAttrs->GetListInt(PADDING_POS);
203+ OP_CHECK_NULL_WITH_CONTEXT(context_, padding);
204+ dValue = *(padding->GetData());
205+ hValue = *(padding->GetData() + 1);
206+ wValue = *(padding->GetData() + 2);
207+ inputData.pad = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
208+ OP_CHECK_IF(
209+ hValue > khValue / 2 || wValue > kwValue / 2 || dValue > kdValue / 2,
210+ OP_LOGE(context_->GetNodeName(),
211+ "MaxPool3DWithArgmaxV2: not support pad shape [%d, %d, %d] kernel shape [%d, %d, %d]",
212+ dValue, hValue, wValue, kdValue, khValue, kwValue),
213+ return ge::GRAPH_FAILED);
214+ 
215+ inputData.dilation = array<uint64_t, DHW_DIMS>{1, 1, 1};
216+ dValue = 1;
217+ hValue = 1;
218+ wValue = 1;
219+ const gert::TypedContinuousVector<int64_t>* dilation = runtimeAttrs->GetListInt(DILATION_POS);
220+ if (dilation != nullptr) {
221+ dValue = *(dilation->GetData());
222+ hValue = *(dilation->GetData() + 1);
223+ wValue = *(dilation->GetData() + 2);
224+ inputData.dilation = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
225+ OP_CHECK_IF(
226+ dValue <= 0 || hValue <= 0 || wValue <= 0,
227+ OP_LOGE(context_->GetNodeName(),
228+ "MaxPool3DWithArgmaxV2: not support dilation shape [%d, %d, %d]", dValue, hValue, wValue),
229+ return ge::GRAPH_FAILED);
230+ }
231+ 
232+ inputData.ceilMode = false;
233+ const bool* ceilModePtr = runtimeAttrs->GetAttrPointer<bool>(CEIL_POS);
234+ if (ceilModePtr != nullptr) {
235+ inputData.ceilMode = *ceilModePtr;
236+ }
237+ 
238+ int indexDtype = 3;
239+ const int* indexDtypePtr = runtimeAttrs->GetAttrPointer<int>(DTYPE_POS);
240+ if (indexDtypePtr != nullptr) {
241+ indexDtype = *indexDtypePtr;
242+ }
243+ switch (indexDtype) {
244+ case MP_MAX_3D_TYPE_INT32:
245+ inputData.indexDtype = ge::DataType::DT_INT32;
246+ break;
247+ case MP_MAX_3D_TYPE_INT64:
248+ inputData.indexDtype = ge::DataType::DT_INT64;
249+ break;
250+ default:
251+ inputData.indexDtype = ge::DataType::DT_INT32;
252+ break;
253+ }
254+ return ge::GRAPH_SUCCESS;
255+}
256+ 
257+void MaxPool3DWithArgmaxV2GatherTiling::InitializationVars()
258+{
259+ baseData_.inputBytes = dtype == ge::DataType::DT_FLOAT ? FLOAT32_SIZE : FLOAT16_OR_BF16_SIZE;
260+ baseData_.indexBytes = inputData.indexDtype == ge::DataType::DT_INT32 ? INT32_SIZE : INT64_SIZE;
261+ baseData_.availableUb = ubSize - UB_RESVERVED_SIZE;
262+ baseData_.totalCoreNum = coreNum;
263+ baseData_.coreUsedForBestPerformance = baseData_.inputBytes == FLOAT32_SIZE
264+ ? baseData_.totalCoreNum / MAX_BANDWIDTH_COEFFICIENTS
265+ : baseData_.totalCoreNum;
266+ baseData_.coreUsedForBestPerformance = NUM_64;
267+ baseData_.padFront = inputData.pad[D_DIM];
268+ baseData_.padTop = inputData.pad[H_DIM];
269+ baseData_.padLeft = inputData.pad[W_DIM];
270+ baseData_.dInput = inputData.inputShape[D_DIM];
271+ baseData_.hInput = inputData.inputShape[H_DIM];
272+ baseData_.wInput = inputData.inputShape[W_DIM];
273+ baseData_.dOutput = inputData.outShape[D_DIM];
274+ baseData_.hOutput = inputData.outShape[H_DIM];
275+ baseData_.wOutput = inputData.outShape[W_DIM];
276+ baseData_.dStride = inputData.stride[D_DIM];
277+ baseData_.hStride = inputData.stride[H_DIM];
278+ baseData_.wStride = inputData.stride[W_DIM];
279+ baseData_.dKernel = inputData.kernelSize[D_DIM];
280+ baseData_.hKernel = inputData.kernelSize[H_DIM];
281+ baseData_.wKernel = inputData.kernelSize[W_DIM];
282+ baseData_.highAxisTotal = inputData.batches;
283+ baseData_.dDilation = inputData.dilation[D_DIM];
284+ baseData_.hDilation = inputData.dilation[H_DIM];
285+ baseData_.wDilation = inputData.dilation[W_DIM];
286+ baseData_.isPad = 0;
287+ if (baseData_.padTop != 0 || baseData_.padFront != 0 || baseData_.padLeft != 0) {
288+ baseData_.isPad = 1;
289+ }
290+ if (inputData.ceilMode && baseData_.isPad == 0) {
291+ if (((baseData_.wOutput - 1) * baseData_.wStride + baseData_.wKernel) != baseData_.wInput ||
292+ ((baseData_.hOutput - 1) * baseData_.hStride + baseData_.hKernel) != baseData_.hInput ||
293+ ((baseData_.dOutput - 1) * baseData_.dStride + baseData_.dKernel) != baseData_.dInput ) {
294+ baseData_.isPad = 1;
295+ }
296+ }
297+ 
298+ baseData_.oneBlockNumT1 = Ops::Base::GetUbBlockSize(context_) / baseData_.inputBytes;
299+ baseData_.oneBlockNumT2 = Ops::Base::GetUbBlockSize(context_) / baseData_.indexBytes;
300+}
301+ 
302+bool MaxPool3DWithArgmaxV2GatherTiling::IsCapable()
303+{
304+ if (inputData.dilation[D_DIM] > DILATION_THRESHOLD || inputData.dilation[H_DIM] > DILATION_THRESHOLD
305+ || inputData.dilation[W_DIM] > DILATION_THRESHOLD ||inputData.inputFormat != ge::Format::FORMAT_NCDHW) {
306+ return false;
307+ }
308+ 
309+ InitializationVars();
310+ if (baseData_.wKernel * baseData_.inputBytes >= CACHE_LINE_SIZE) {
311+ return false;
312+ }
313+ 
314+ splitData_.dOutputInner = 1;
315+ splitData_.hOutputInner = 1;
316+ splitData_.wOutputInner = 1;
317+ splitData_.highAxisInner = 1;
318+ DoBufferCalculate();
319+ int64_t dRate = (baseData_.dOutput == 1 || (baseData_.dStride / baseData_.dKernel) < 1)
320+ ? 1
321+ : (baseData_.dStride / baseData_.dKernel);
322+ int64_t hRate = (baseData_.hOutput == 1 || (baseData_.hStride / baseData_.hKernel) < 1)
323+ ? 1
324+ : (baseData_.hStride / baseData_.hKernel);
325+ int64_t wRate = (baseData_.wOutput == 1 || (baseData_.wStride / baseData_.wKernel) < 1)
326+ ? 1
327+ : (baseData_.wStride / baseData_.wKernel);
328+ return splitData_.totalBufferSize <= baseData_.availableUb / (MIN_OUTPUT_THRESHOLD * dRate * hRate * wRate);
329+}
330+ 
331+uint64_t MaxPool3DWithArgmaxV2GatherTiling::GetTilingKey() const
332+{
333+ uint64_t tilingKey = NO_PADDING_TILING_KEY;
334+ if (baseData_.isPad == 1) {
335+ tilingKey = PADDING_TILING_KEY;
336+ }
337+ return tilingKey;
338+}
339+ 
340+void MaxPool3DWithArgmaxV2GatherTiling::DoBufferCalculate()
341+{
342+ splitData_.dInputInner =
343+ (splitData_.dOutputInner - 1) * baseData_.dStride + (baseData_.dKernel - 1) * baseData_.dDilation + 1;
344+ splitData_.hInputInner =
345+ (splitData_.hOutputInner - 1) * baseData_.hStride + (baseData_.hKernel - 1) * baseData_.hDilation + 1;
346+ splitData_.wInputInner =
347+ (splitData_.wOutputInner - 1) * baseData_.wStride + (baseData_.wKernel - 1) * baseData_.wDilation + 1;
348+ int64_t maxDataNumInOneBlock = std::max(baseData_.oneBlockNumT1, baseData_.oneBlockNumT2);
349+ int64_t wInputInnerAligned = Ops::Base::CeilAlign(splitData_.wInputInner, baseData_.oneBlockNumT1);
350+ int64_t wOutputInnerAligned = Ops::Base::CeilAlign(splitData_.wOutputInner, maxDataNumInOneBlock);
351+ int64_t inputBufferSize =
352+ splitData_.highAxisInner * splitData_.dInputInner * splitData_.hInputInner * wInputInnerAligned * baseData_.inputBytes;
353+ splitData_.inputBufferSize = inputBufferSize;
354+ 
355+ if (baseData_.isPad == 1) {
356+ inputBufferSize *= DOUBLE;
357+ }
358+ int64_t outputDataSize = splitData_.highAxisInner * splitData_.dOutputInner * splitData_.hOutputInner * wOutputInnerAligned;
359+ splitData_.maxValueBufferSize = outputDataSize * baseData_.inputBytes;
360+ splitData_.argmaxBufferSize = outputDataSize * baseData_.indexBytes;
361+ 
362+ int64_t tmpTotalBufferSize =
363+ inputBufferSize + splitData_.maxValueBufferSize + splitData_.argmaxBufferSize + HELPER_BUFFER_SIZE;
364+
365+ splitData_.totalBufferSize = tmpTotalBufferSize * DOUBLE;
366+ if (baseData_.isPad == 1) {
367+ splitData_.totalBufferSize -= splitData_.inputBufferSize;
368+ }
369+}
370+ 
371+bool MaxPool3DWithArgmaxV2GatherTiling::IsMeetTargetCoreNum() const
372+{
373+ int64_t tmpWOutputOuter = Ops::Base::CeilDiv(baseData_.wOutput, splitData_.wOutputInner);
374+ int64_t tmpHOutputOuter = Ops::Base::CeilDiv(baseData_.hOutput, splitData_.hOutputInner);
375+ int64_t tmpDOutputOuter = Ops::Base::CeilDiv(baseData_.dOutput, splitData_.dOutputInner);
376+ int64_t tmpNCOutputOuter = Ops::Base::CeilDiv(baseData_.highAxisTotal, splitData_.highAxisInner);
377+ return tmpWOutputOuter * tmpHOutputOuter * tmpDOutputOuter * tmpNCOutputOuter >= baseData_.coreUsedForBestPerformance;
378+}
379+ 
380+bool MaxPool3DWithArgmaxV2GatherTiling::IsMeetUBSize()
381+{
382+ DoBufferCalculate();
383+ return splitData_.totalBufferSize <= baseData_.availableUb;
384+}
385+ 
386+void MaxPool3DWithArgmaxV2GatherTiling::BinarySearch(int64_t start, int64_t end, int64_t* value)
387+{
388+ int64_t left = start;
389+ int64_t right = end;
390+ int64_t bestSplit = 1;
391+ 
392+ while (left <= right) {
393+ int64_t mid = left + (right - left) / DOUBLE;
394+ *value = mid;
395+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
396+ bestSplit = mid;
397+ left = mid + 1;
398+ } else {
399+ right = mid - 1;
400+ }
401+ }
402+ 
403+ *value = bestSplit;
404+}
405+ 
406+bool MaxPool3DWithArgmaxV2GatherTiling::TrySplitNC()
407+{
408+ splitData_.dOutputInner = baseData_.dOutput;
409+ splitData_.hOutputInner = baseData_.hOutput;
410+ splitData_.wOutputInner = baseData_.wOutput;
411+ splitData_.highAxisInner = Ops::Base::CeilDiv(baseData_.highAxisTotal, baseData_.coreUsedForBestPerformance);
CANN-robot
CANN-robotCANN-robot2月4日

逻辑错误: 在TrySplitNC函数中，第411行计算highAxisInner时使用了CeilDiv(baseData_.highAxisTotal, baseData_.coreUsedForBestPerformance)，但如果highAxisTotal小于coreUsedForBestPerformance，计算结果可能为1或更小。然而第413-421行的逻辑存在问题：如果第411行的计算满足条件，函数返回true；否则将highAxisInner设为1，再次检查条件。但第二次检查时，highAxisInner=1，IsMeetUBSize()和IsMeetTargetCoreNum()的结果可能与第一次不同，这可能导致逻辑不一致。

问题类型: 逻辑错误 文件路径: pooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_gather_tiling.cpp 行号: 411 问题代码:

    splitData_.highAxisInner = Ops::Base::CeilDiv(baseData_.highAxisTotal, baseData_.coreUsedForBestPerformance);
    if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
        return true;
    }
    splitData_.highAxisInner = 1;
    if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
        BinarySearch(1, baseData_.highAxisTotal, &splitData_.highAxisInner);
        return true;
    } else {
        return false;
    }

修改建议:

重新审视这个逻辑：如果第一次计算不满足条件，将highAxisInner设为1后再次检查，如果满足条件则进行二分搜索。这个逻辑可能正确，但建议添加注释说明为什么需要这个回退逻辑。

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

likedislike
412+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
413+ return true;
414+ }
415+ splitData_.highAxisInner = 1;
416+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
417+ BinarySearch(1, baseData_.highAxisTotal, &splitData_.highAxisInner);
418+ return true;
419+ } else {
420+ return false;
421+ }
422+}
423+ 
424+bool MaxPool3DWithArgmaxV2GatherTiling::TrySplitD()
425+{
426+ splitData_.highAxisInner = 1;
427+ splitData_.dOutputInner = 1;
428+ splitData_.hOutputInner = baseData_.hOutput;
429+ splitData_.wOutputInner = baseData_.wOutput;
430+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
431+ BinarySearch(1, baseData_.dOutput, &splitData_.dOutputInner);
432+ return true;
433+ } else {
434+ return false;
435+ }
436+}
437+ 
438+bool MaxPool3DWithArgmaxV2GatherTiling::TrySplitH()
439+{
440+ splitData_.highAxisInner = 1;
441+ splitData_.dOutputInner = 1;
442+ splitData_.hOutputInner = 1;
443+ splitData_.wOutputInner = baseData_.wOutput;
444+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
445+ BinarySearch(1, baseData_.hOutput, &splitData_.hOutputInner);
446+ return true;
447+ } else {
448+ return false;
449+ }
450+}
451+ 
452+bool MaxPool3DWithArgmaxV2GatherTiling::TrySplitW()
453+{
454+ splitData_.highAxisInner = 1;
455+ splitData_.dOutputInner = 1;
456+ splitData_.hOutputInner = 1;
457+ splitData_.wOutputInner = 1;
458+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
459+ BinarySearch(1, baseData_.wOutput, &splitData_.wOutputInner);
460+ return true;
461+ } else {
462+ return false;
463+ }
464+}
465+ 
466+void MaxPool3DWithArgmaxV2GatherTiling::SearchBestTiling()
467+{
468+ if (TrySplitNC()) {
469+ return;
470+ }
471+ if (TrySplitD()) {
472+ return;
473+ }
474+ if (TrySplitH()) {
475+ return;
476+ }
477+ if (TrySplitW()) {
478+ return;
479+ }
480+}
481+ 
482+void MaxPool3DWithArgmaxV2GatherTiling::DoUBTiling()
483+{
484+ SearchBestTiling();
485+ DoBufferCalculate();
486+ splitData_.wOutputOuter = Ops::Base::CeilDiv(baseData_.wOutput, splitData_.wOutputInner);
487+ int64_t tempWOutputTail = baseData_.wOutput % splitData_.wOutputInner;
488+ splitData_.wOutputTail = tempWOutputTail == 0 ? splitData_.wOutputInner : tempWOutputTail;
489+ 
490+ splitData_.hOutputOuter = Ops::Base::CeilDiv(baseData_.hOutput, splitData_.hOutputInner);
491+ int64_t tempHOutputTail = baseData_.hOutput % splitData_.hOutputInner;
492+ splitData_.hOutputTail = tempHOutputTail == 0 ? splitData_.hOutputInner : tempHOutputTail;
493+ 
494+ splitData_.dOutputOuter = Ops::Base::CeilDiv(baseData_.dOutput, splitData_.dOutputInner);
495+ int64_t tempDOutputTail = baseData_.dOutput % splitData_.dOutputInner;
496+ splitData_.dOutputTail = tempDOutputTail == 0 ? splitData_.dOutputInner : tempDOutputTail;
497+ 
498+ splitData_.highAxisOuter = Ops::Base::CeilDiv(baseData_.highAxisTotal, splitData_.highAxisInner);
499+ int64_t tempNOutputTail = baseData_.highAxisTotal % splitData_.highAxisInner;
500+ splitData_.highAxisTail = tempNOutputTail == 0 ? splitData_.highAxisInner : tempNOutputTail;
501+}
502+ 
503+void MaxPool3DWithArgmaxV2GatherTiling::DoBlockTiling()
504+{
505+ splitData_.totalBaseBlockNum = splitData_.highAxisOuter * splitData_.dOutputOuter * splitData_.hOutputOuter * splitData_.wOutputOuter ;
506+ splitData_.normalCoreProcessNum = Ops::Base::CeilDiv(splitData_.totalBaseBlockNum, baseData_.totalCoreNum);
507+ splitData_.usedCoreNum = Ops::Base::CeilDiv(splitData_.totalBaseBlockNum, splitData_.normalCoreProcessNum);
508+ splitData_.tailCoreProcessNum =
509+ splitData_.totalBaseBlockNum - splitData_.normalCoreProcessNum * (splitData_.usedCoreNum - 1);
510+}
511+ 
512+void MaxPool3DWithArgmaxV2GatherTiling::PrintBaseData() const
513+{
514+ OP_LOGI("PrintBaseData", "%s", baseData_.ToString().c_str());
515+}
516+ 
517+void MaxPool3DWithArgmaxV2GatherTiling::PrintSplitData() const
518+{
519+ OP_LOGI("PrintSplitData", "%s", splitData_.ToString().c_str());
520+}
521+ 
522+void MaxPool3DWithArgmaxV2GatherTiling::SetTilingData()
523+{
524+ tilingData_->dInput = baseData_.dInput;
525+ tilingData_->hInput = baseData_.hInput;
526+ tilingData_->wInput = baseData_.wInput;
527+ tilingData_->dOutput = baseData_.dOutput;
528+ tilingData_->hOutput = baseData_.hOutput;
529+ tilingData_->wOutput = baseData_.wOutput;
530+ tilingData_->dKernel = baseData_.dKernel;
531+ tilingData_->hKernel = baseData_.hKernel;
532+ tilingData_->wKernel = baseData_.wKernel;
533+ tilingData_->dStride = baseData_.dStride;
534+ tilingData_->hStride = baseData_.hStride;
535+ tilingData_->wStride = baseData_.wStride;
536+ tilingData_->padFront = baseData_.padFront;
537+ tilingData_->padTop = baseData_.padTop;
538+ tilingData_->padLeft = baseData_.padLeft;
539+ tilingData_->highAxisInner = splitData_.highAxisInner;
540+ tilingData_->highAxisTail = splitData_.highAxisTail;
541+ tilingData_->highAxisOuter = splitData_.highAxisOuter;
542+ tilingData_->dOutputInner = splitData_.dOutputInner;
543+ tilingData_->dOutputTail = splitData_.dOutputTail;
544+ tilingData_->dOutputOuter = splitData_.dOutputOuter;
545+ tilingData_->hOutputInner = splitData_.hOutputInner;
546+ tilingData_->hOutputTail = splitData_.hOutputTail;
547+ tilingData_->hOutputOuter = splitData_.hOutputOuter;
548+ tilingData_->wOutputInner = splitData_.wOutputInner;
549+ tilingData_->wOutputTail = splitData_.wOutputTail;
550+ tilingData_->wOutputOuter = splitData_.wOutputOuter;
551+ tilingData_->normalCoreProcessNum = splitData_.normalCoreProcessNum;
552+ tilingData_->tailCoreProcessNum = splitData_.tailCoreProcessNum;
553+ tilingData_->usedCoreNum = splitData_.usedCoreNum;
554+ tilingData_->inputBufferSize = splitData_.inputBufferSize;
555+ tilingData_->maxValueBufferSize = splitData_.maxValueBufferSize;
556+ tilingData_->argmaxBufferSize = splitData_.argmaxBufferSize;
557+ tilingData_->isPad = baseData_.isPad;
558+ tilingData_->dDilation = baseData_.dDilation;
559+ tilingData_->hDilation = baseData_.hDilation;
560+ tilingData_->wDilation = baseData_.wDilation;
561+}
562+ 
563+ge::graphStatus MaxPool3DWithArgmaxV2GatherTiling::DoOpTiling()
564+{
565+ DoUBTiling();
566+ DoBlockTiling();
567+ SetTilingData();
568+ PrintBaseData();
569+ PrintSplitData();
570+ return ge::GRAPH_SUCCESS;
571+}
572+ 
573+ge::graphStatus MaxPool3DWithArgmaxV2GatherTiling::PostTiling()
574+{
575+ context_->SetBlockDim(tilingData_->usedCoreNum);
576+ return ge::GRAPH_SUCCESS;
577+}
578+ 
579+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2GatherTiling, 0);
580+ 
581+} // namespace optiling
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_gather_tiling.h+211-0
@@ -0,0 +1,211 @@
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 max_pool3d_with_argmax_v2_gather_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL3D_WITH_AGRMAX_V2_GATHER_TILING_H_
17+#define MAX_POOL3D_WITH_AGRMAX_V2_GATHER_TILING_H_
18+ 
19+#include "max_pool3d_with_argmax_v2_tiling_base.h"
20+#include "../op_kernel/arch35/max_pool3d_with_argmax_v2_tiling_struct.h"
21+ 
22+namespace optiling
23+{
24+struct MaxPool3DWithArgmaxV2GatherBaseInfo {
25+ int64_t inputBytes = 0;
26+ int64_t indexBytes = 0;
27+ int64_t availableUb = 0;
28+ int64_t totalCoreNum = 0;
29+ int64_t oneBlockNumT1 = 0;
30+ int64_t oneBlockNumT2 = 0;
31+ int64_t coreUsedForBestPerformance = 0;
32+ 
33+ int64_t padFront = 0;
34+ int64_t padTop = 0;
35+ int64_t padLeft = 0;
36+ int64_t dStride = 0;
37+ int64_t hStride = 0;
38+ int64_t wStride = 0;
39+ int64_t dKernel = 0;
40+ int64_t hKernel = 0;
41+ int64_t wKernel = 0;
42+ int64_t dInput = 0;
43+ int64_t hInput = 0;
44+ int64_t wInput = 0;
45+ int64_t dOutput = 0;
46+ int64_t hOutput = 0;
47+ int64_t wOutput = 0;
48+ int64_t highAxisTotal = 0;
49+ int64_t isPad = 0;
50+ int64_t dDilation = 0;
51+ int64_t hDilation = 0;
52+ int64_t wDilation = 0;
53+ std::string ToString() const
54+ {
55+ std::stringstream info;
56+ info << "MaxPool3DWithArgmaxV2GatherBaseInfo {";
57+ 
58+ info << "inputBytes:" << inputBytes
59+ << ", indexBytes:" << indexBytes
60+ << ", availableUb:" << availableUb
61+ << ", totalCoreNum:" << totalCoreNum
62+ << ", oneBlockNumT1:" << oneBlockNumT1
63+ << ", oneBlockNumT2:" << oneBlockNumT2
64+ << ", coreUsedForBestPerformance:" << coreUsedForBestPerformance;
65+ 
66+ info << ", padFront:" << padFront
67+ << ", padTop:" << padTop
68+ << ", padLeft:" << padLeft;
69+ 
70+ info << ", dStride:" << dStride
71+ << ", hStride:" << hStride
72+ << ", wStride:" << wStride;
73+ 
74+ info << ", dKernel:" << dKernel
75+ << ", hKernel:" << hKernel
76+ << ", wKernel:" << wKernel;
77+ 
78+ info << ", dInput:" << dInput
79+ << ", hInput:" << hInput
80+ << ", wInput:" << wInput;
81+ 
82+ info << ", dOutput:" << dOutput
83+ << ", hOutput:" << hOutput
84+ << ", wOutput:" << wOutput;
85+ 
86+ info << ", highAxisTotal:" << highAxisTotal
87+ << ", isPad:" << isPad;
88+ 
89+ info << ", dDilation:" << dDilation
90+ << ", hDilation:" << hDilation
91+ << ", wDilation:" << wDilation;
92+ info << " }";
93+ return info.str();
94+ }
95+};
96+ 
97+struct MaxPool3DWithArgmaxV2GatherSplitInfo {
98+ // InitializationVars
99+ int64_t highAxisInner = 0;
100+ int64_t highAxisTail = 0;
101+ int64_t highAxisOuter = 0;
102+ int64_t highAxisAligned = 0;
103+ 
104+ // DoUBTiling
105+ int64_t dOutputInner = 0;
106+ int64_t dOutputTail = 0;
107+ int64_t dOutputOuter = 0;
108+ int64_t hOutputInner = 0;
109+ int64_t hOutputTail = 0;
110+ int64_t hOutputOuter = 0;
111+ int64_t wOutputInner = 0;
112+ int64_t wOutputTail = 0;
113+ int64_t wOutputOuter = 0;
114+ 
115+ // DoBlockTiling
116+ int64_t normalCoreProcessNum = 0;
117+ int64_t tailCoreProcessNum = 0;
118+ int64_t usedCoreNum = 0;
119+ int64_t totalBaseBlockNum = 0;
120+ 
121+ // DoBufferCalculate
122+ int64_t dInputInner = 0;
123+ int64_t hInputInner = 0;
124+ int64_t wInputInner = 0;
125+ int64_t baseBlockPlaneSizeAligned = 0;
126+ int64_t inputBufferSize = 0;
127+ int64_t maxValueBufferSize = 0;
128+ int64_t argmaxBufferSize = 0;
129+ int64_t totalBufferSize = 0;
130+ std::string ToString() const
131+ {
132+ std::stringstream info;
133+ info << "MaxPool3DWithArgmaxV2GatherSplitInfo {";
134+
135+ info << " highAxisInner:" << highAxisInner
136+ << ", highAxisTail:" << highAxisTail
137+ << ", highAxisOuter:" << highAxisOuter
138+ << ", highAxisAligned:" << highAxisAligned;
139+
140+ info << ", dOutputInner:" << dOutputInner
141+ << ", dOutputTail:" << dOutputTail
142+ << ", dOutputOuter:" << dOutputOuter
143+ << ", hOutputInner:" << hOutputInner
144+ << ", hOutputTail:" << hOutputTail
145+ << ", hOutputOuter:" << hOutputOuter
146+ << ", wOutputInner:" << wOutputInner
147+ << ", wOutputTail:" << wOutputTail
148+ << ", wOutputOuter:" << wOutputOuter;
149+
150+ info << ", normalCoreProcessNum:" << normalCoreProcessNum
151+ << ", tailCoreProcessNum:" << tailCoreProcessNum
152+ << ", usedCoreNum:" << usedCoreNum
153+ << ", totalBaseBlockNum:" << totalBaseBlockNum;
154+
155+ info << ", dInputInner:" << dInputInner
156+ << ", hInputInner:" << hInputInner
157+ << ", wInputInner:" << wInputInner
158+ << ", baseBlockPlaneSizeAligned:" << baseBlockPlaneSizeAligned
159+ << ", inputBufferSize:" << inputBufferSize
160+ << ", maxValueBufferSize:" << maxValueBufferSize
161+ << ", argmaxBufferSize:" << argmaxBufferSize
162+ << ", totalBufferSize:" << totalBufferSize;
163+
164+ info << " }";
165+ return info.str();
166+ }
167+};
168+ 
169+class MaxPool3DWithArgmaxV2GatherTiling : public MaxPool3DWithArgmaxV2BaseTiling
170+{
171+public:
172+ explicit MaxPool3DWithArgmaxV2GatherTiling(gert::TilingContext* context) : MaxPool3DWithArgmaxV2BaseTiling(context)
173+ {
174+ }
175+ 
176+ ~MaxPool3DWithArgmaxV2GatherTiling() override
177+ {
178+ }
179+ 
180+private:
181+ void DoUBTiling();
182+ void InitializationVars();
183+ bool IsMeetTargetCoreNum() const;
184+ void SearchBestTiling();
185+ bool IsMeetUBSize();
186+ void SetTilingData();
187+ void BinarySearch(int64_t start, int64_t end, int64_t* value);
188+ bool TrySplitNC();
189+ bool TrySplitD();
190+ bool TrySplitH();
191+ bool TrySplitW();
192+ uint64_t GetTilingKey() const override;
193+ void PrintBaseData() const;
194+ void PrintSplitData() const;
195+ void DoBlockTiling();
196+ void DoBufferCalculate();
197+ bool IsCapable() override;
198+ ge::graphStatus GetPlatformInfo() override;
199+ ge::graphStatus DoOpTiling() override;
200+ ge::graphStatus PostTiling() override;
201+ ge::graphStatus GetShapeAttrsInfo() override;
202+ 
203+ MaxPool3DWithArgmaxV2Tiling::MaxPool3DWithArgmaxV2GatherTilingData* tilingData_ =
204+ context_->GetTilingData<MaxPool3DWithArgmaxV2Tiling::MaxPool3DWithArgmaxV2GatherTilingData>();
205+ MaxPool3DWithArgmaxV2GatherBaseInfo baseData_;
206+ MaxPool3DWithArgmaxV2GatherSplitInfo splitData_;
207+};
208+ 
209+} // namespace optiling
210+ 
211+#endif
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_infershape.cpp+217-0
@@ -0,0 +1,217 @@
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 max_pool3d_with_argmax_v2_infershape.cpp
13+ * \brief
14+ */
15+ 
16+#include "log/log.h"
17+#include "register/op_impl_registry.h"
18+#include "error_util.h"
19+#include "util/shape_util.h"
20+#include "graph/utils/type_utils.h"
21+#include "exe_graph/runtime/infer_shape_context.h"
22+#include <string>
23+#include "common_dtype.h"
24+ 
25+using namespace ge;
26+ 
27+namespace ops
28+{
29+static constexpr size_t INDEX_KSIZE = 0;
30+static constexpr size_t INDEX_STRIDES = 1;
31+static constexpr size_t INDEX_PADS = 2;
32+static constexpr size_t INDEX_DILATION = 3;
33+static constexpr size_t INDEX_CEIL_MODE = 4;
34+static constexpr size_t INDEX_DATA_FORMAT = 5;
35+static constexpr size_t INDEX_DTYPE = 6;
36+static constexpr size_t ATTR_LIST_SHAPE_SIZE = 3;
37+static constexpr size_t INDEX_OUT_MAX = 0;
38+static constexpr size_t INDEX_OUT_INDICES = 1;
39+static constexpr size_t PARAM_NUM = 4;
40+static constexpr size_t PARAM_D_DIM = 0;
41+static constexpr size_t PARAM_H_DIM = 1;
42+static constexpr size_t PARAM_W_DIM = 2;
43+static constexpr size_t SHAPE_D_DIM = 2;
44+static constexpr size_t SHAPE_H_DIM = 3;
45+static constexpr size_t SHAPE_W_DIM = 4;
46+static constexpr size_t CDHW_DIM = 4;
47+static constexpr size_t INT32_DTYPE = 3;
48+static constexpr size_t INT64_DTYPE = 9;
49+ 
50+static int64_t DivRtn(int64_t x, int64_t y)
51+{
52+ if (y == 0) {
53+ OP_LOGE("MaxPool3DWithArgmaxV2", "strides value cannot be zero.");
54+ return GRAPH_FAILED;
55+ }
56+ if (x < 0) {
57+ OP_LOGE("MaxPool3DWithArgmaxV2", "x value cannot small than zero.");
58+ return GRAPH_FAILED;
59+ }
60+ int64_t q = x / y;
61+ return q;
62+}
63+ 
64+static void UpdateMaxShape(const int64_t (&param)[PARAM_NUM], bool ceil_mode, const int64_t& dim_size,
65+ int64_t& out_max_shape)
66+{
67+ int64_t ksize = param[INDEX_KSIZE];
68+ int64_t strides = param[INDEX_STRIDES];
69+ int64_t pad = param[INDEX_PADS];
70+ int64_t dilation = param[PARAM_NUM - 1];
71+ int64_t exact_size = dim_size + 2 * pad - dilation * (ksize - 1) - 1 + (ceil_mode ? (strides - 1) : 0);
72+ out_max_shape = DivRtn(exact_size, strides) + 1;
73+ if (ceil_mode) {
74+ if ((out_max_shape - 1) * strides >= dim_size + pad) {
75+ out_max_shape = out_max_shape - 1;
76+ }
77+ }
78+}
79+ 
80+ge::graphStatus InferShape4MaxPool3DWithArgmaxV2(gert::InferShapeContext* context)
81+{
82+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPool3DWithArgmaxV2 infershape running");
83+ auto src_td = context->GetInputDesc(0);
84+ OPS_CHECK_NULL_WITH_CONTEXT(context, src_td);
85+ auto input_format = src_td->GetOriginFormat();
86+ auto indices_td = context->GetOutputDesc(INDEX_OUT_INDICES);
87+ OPS_CHECK_NULL_WITH_CONTEXT(context, indices_td);
88+ auto indices_dtype = indices_td->GetDataType();
89+ OP_LOGD(context->GetNodeName(), "indices_dtype = %d", indices_dtype);
90+ OP_CHECK_IF(input_format != FORMAT_ND && input_format != FORMAT_NCDHW && input_format != FORMAT_NDHWC,
91+ OP_LOGE(context->GetNodeName(), "format only supports ND, NCDHW, NDHWC"),
92+ return GRAPH_FAILED);
93+ 
94+ size_t param_d_dim = PARAM_D_DIM;
95+ size_t param_h_dim = PARAM_H_DIM;
96+ size_t param_w_dim = PARAM_W_DIM;
97+ size_t input_d_dim = SHAPE_D_DIM;
98+ size_t input_h_dim = SHAPE_H_DIM;
99+ size_t input_w_dim = SHAPE_W_DIM;
100+ 
101+ auto attrs = context->GetAttrs();
102+ OPS_CHECK_NULL_WITH_CONTEXT(context, attrs);
103+ 
104+ auto ksize = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_KSIZE);
105+ OPS_CHECK_NULL_WITH_CONTEXT(context, ksize);
106+ std::string errMsg4Ksize = optiling::ConcatString("Length of ksize ", ksize->GetSize(), " must be 3!");
107+ OP_CHECK_IF(ksize->GetSize() != ATTR_LIST_SHAPE_SIZE,
108+ OP_LOGE(context->GetNodeName(), "%s", errMsg4Ksize.c_str()), return GRAPH_FAILED);
109+ auto ksize_data = reinterpret_cast<const int64_t*>(ksize->GetData());
110+ 
111+ auto strides = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_STRIDES);
112+ OPS_CHECK_NULL_WITH_CONTEXT(context, strides);
113+ std::string errMsg4Strides = optiling::ConcatString("Length of strides ", strides->GetSize(), " must be 3!");
114+ OP_CHECK_IF(strides->GetSize() != ATTR_LIST_SHAPE_SIZE,
115+ OP_LOGE(context->GetNodeName(), "%s", errMsg4Strides.c_str()), return GRAPH_FAILED);
116+ auto strides_data = reinterpret_cast<const int64_t*>(strides->GetData());
117+ 
118+ auto pads = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_PADS);
119+ OPS_CHECK_NULL_WITH_CONTEXT(context, pads);
120+ std::string errMsg4Pads = optiling::ConcatString("Length of pads ", pads->GetSize(), " must be 3!");
121+ OP_CHECK_IF(pads->GetSize() != ATTR_LIST_SHAPE_SIZE,
122+ OP_LOGE(context->GetNodeName(), "%s", errMsg4Pads.c_str()), return GRAPH_FAILED);
123+ auto pads_data = reinterpret_cast<const int64_t*>(pads->GetData());
124+ 
125+ auto dilation = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_DILATION);
126+ OPS_CHECK_NULL_WITH_CONTEXT(context, dilation);
127+ std::string errMsg4Dilation = optiling::ConcatString("Length of dilation ", dilation->GetSize(), " must be 3!");
128+ OP_CHECK_IF(dilation->GetSize() != ATTR_LIST_SHAPE_SIZE,
129+ OP_LOGE(context->GetNodeName(), "%s", errMsg4Dilation.c_str()), return GRAPH_FAILED);
130+ auto dilation_data = reinterpret_cast<const int64_t*>(dilation->GetData());
131+ 
132+ auto ceil_mode = attrs->GetAttrPointer<bool>(INDEX_CEIL_MODE);
133+ OPS_CHECK_NULL_WITH_CONTEXT(context, ceil_mode);
134+ 
135+ const char* data_format = attrs->GetAttrPointer<char>(INDEX_DATA_FORMAT);
136+ OPS_CHECK_NULL_WITH_CONTEXT(context, data_format);
137+ 
138+ std::string data_format_str = data_format;
139+ if (data_format_str == "NDHWC") {
140+ input_d_dim = input_d_dim - 1;
141+ input_h_dim = input_h_dim - 1;
142+ input_w_dim = input_w_dim - 1;
143+ }
144+
145+ const gert::Shape* in_shape = context->GetInputShape(0);
146+ OPS_CHECK_NULL_WITH_CONTEXT(context, in_shape);
147+ gert::Shape* out_max_shape = context->GetOutputShape(INDEX_OUT_MAX);
148+ OPS_CHECK_NULL_WITH_CONTEXT(context, out_max_shape);
149+ *out_max_shape = *in_shape;
150+ gert::Shape* out_indices_shape = context->GetOutputShape(INDEX_OUT_INDICES);
151+ OPS_CHECK_NULL_WITH_CONTEXT(context, out_indices_shape);
152+ *out_indices_shape = *in_shape;
153+ 
154+ if (Ops::Base::IsUnknownRank(*in_shape) || Ops::Base::IsUnknownShape(*in_shape)) {
155+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPool3DWithArgmaxV2 infershape handle unknown rank or shape.");
156+ return ge::GRAPH_SUCCESS;
157+ }
158+ 
159+ size_t dim_num = in_shape->GetDimNum();
160+ int64_t max_dim = 0;
161+ if (dim_num == CDHW_DIM) {
162+ input_d_dim = input_d_dim - 1;
163+ input_h_dim = input_h_dim - 1;
164+ input_w_dim = input_w_dim - 1;
165+ }
166+ for (size_t i = 0; i < dim_num; i++) {
167+ int64_t input_dim = in_shape->GetDim(i);
168+ if (i == input_d_dim) {
169+ int64_t param[PARAM_NUM] = {ksize_data[param_d_dim], strides_data[param_d_dim], pads_data[param_d_dim],
170+ dilation_data[param_d_dim]};
171+ UpdateMaxShape(param, *ceil_mode, input_dim, max_dim);
172+ out_max_shape->SetDim(i, max_dim);
173+ out_indices_shape->SetDim(i, max_dim);
174+ } else if (i == input_h_dim) {
175+ int64_t param[PARAM_NUM] = {ksize_data[param_h_dim], strides_data[param_h_dim], pads_data[param_h_dim],
176+ dilation_data[param_h_dim]};
177+ UpdateMaxShape(param, *ceil_mode, input_dim, max_dim);
178+ out_max_shape->SetDim(i, max_dim);
179+ out_indices_shape->SetDim(i, max_dim);
180+ } else if (i == input_w_dim) {
181+ int64_t param[PARAM_NUM] = {ksize_data[param_w_dim], strides_data[param_w_dim], pads_data[param_w_dim],
182+ dilation_data[param_w_dim]};
183+ UpdateMaxShape(param, *ceil_mode, input_dim, max_dim);
184+ out_max_shape->SetDim(i, max_dim);
185+ out_indices_shape->SetDim(i, max_dim);
186+ } else {
187+ out_max_shape->SetDim(i, input_dim);
188+ out_indices_shape->SetDim(i, input_dim);
189+ }
190+ }
191+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPool3DWithArgmaxV2 infershape run success.");
192+ return GRAPH_SUCCESS;
193+}
194+ 
195+static ge::graphStatus InferDataType4MaxPool3DWithArgmaxV2(gert::InferDataTypeContext* context)
196+{
197+ if (context == nullptr) {
198+ return GRAPH_FAILED;
199+ }
200+ const ge::DataType x = context->GetInputDataType(0);
201+ context->SetOutputDataType(INDEX_OUT_MAX, x);
202+ 
203+ auto attrsPtr = context->GetAttrs();
204+ OPS_CHECK_NULL_WITH_CONTEXT(context, attrsPtr);
205+ const int64_t* dstDtype = attrsPtr->GetAttrPointer<int64_t>(INDEX_DTYPE);
206+ OPS_CHECK_NULL_WITH_CONTEXT(context, dstDtype);
207+ ge::DataType indicesDtype = *dstDtype == INT32_DTYPE ? ge::DT_INT32 : ge::DT_INT64;
208+ 
209+ context->SetOutputDataType(INDEX_OUT_INDICES, indicesDtype);
210+ 
211+ return GRAPH_SUCCESS;
212+}
213+ 
214+IMPL_OP_INFERSHAPE(MaxPool3DWithArgmaxV2)
215+ .InferShape(InferShape4MaxPool3DWithArgmaxV2)
216+ .InferDataType(InferDataType4MaxPool3DWithArgmaxV2);
217+} // namespace ops
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_no_expand_indices_tiling.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18namespace optiling {19namespace optiling {
@@ -122,6 +123,12 @@ void MaxPool3DWithArgmaxV2NoExpandIndicesTiling::DoOutputPadAdjustment()
122 123 
123bool MaxPool3DWithArgmaxV2NoExpandIndicesTiling::IsCapable()124bool MaxPool3DWithArgmaxV2NoExpandIndicesTiling::IsCapable()
124{125{
126+ auto platformInfo = context_->GetPlatformInfo();
127+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
128+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
129+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
130+ return false;
131+ }
125 if (inputData.dilation[D_DIM] != 1 || inputData.dilation[H_DIM] != 1 || inputData.dilation[W_DIM] != 1) {132 if (inputData.dilation[D_DIM] != 1 || inputData.dilation[H_DIM] != 1 || inputData.dilation[W_DIM] != 1) {
126 return false;133 return false;
127 }134 }
@@ -371,6 +378,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2NoExpandIndicesTiling::PostTiling()
371 return ge::GRAPH_SUCCESS;378 return ge::GRAPH_SUCCESS;
372}379}
373 380 
374-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2NoExpandIndicesTiling, 0);381+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2NoExpandIndicesTiling, 10);
375 382 
376} // namespace optiling383} // namespace optiling
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_simt_tiling.cpp+260-0
@@ -0,0 +1,260 @@
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 max_pool3d_with_argmax_v2_simt_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include <cctype>
17+#include <algorithm>
18+#include "log/log.h"
19+#include "util/math_util.h"
20+#include "error_util.h"
21+#include "tiling_base/tiling_base.h"
22+#include "tiling_base/tiling_templates_registry.h"
23+#include "max_pool3d_with_argmax_v2_simt_tiling.h"
24+#include "op_common/op_host/util/platform_util.h"
25+#include "platform/platform_ascendc.h"
26+#include "register/op_def_registry.h"
27+#include "platform/platform_info.h"
28+#include "tiling_base/tiling_util.h"
29+#include "register/op_impl_registry.h"
30+ 
31+using namespace ge;
32+ 
33+namespace optiling{
34+ 
35+static const gert::Shape g_vec_1_shape = {1};
36+ 
37+static const gert::Shape& EnsureNotScalar(const gert::Shape &inShape) {
38+ if (inShape.IsScalar()) {
39+ return g_vec_1_shape;
40+ }
41+ return inShape;
42+}
43+ 
44+bool MaxPool3DWithArgmaxV2TilingSIMT::IsCapable()
45+{
46+ return true;
47+}
48+ 
49+ge::graphStatus MaxPool3DWithArgmaxV2TilingSIMT::GetPlatformInfo()
50+{
51+ auto platformPtr = context_->GetPlatformInfo();
52+ if (platformPtr == nullptr) {
53+ auto compileInfoPtr = reinterpret_cast<const MaxPool3DWithArgmaxV2CompileInfo*>(context_->GetCompileInfo());
54+ OP_CHECK_IF(
55+ compileInfoPtr == nullptr, OP_LOGE(context_->GetNodeName(), "compile info is null"),
56+ return ge::GRAPH_FAILED);
57+ coreNum = compileInfoPtr->coreNum;
58+ 
59+ ubSize = compileInfoPtr->ubSize;
60+ } else {
61+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformPtr);
62+ coreNum = ascendcPlatform.GetCoreNumAiv();
63+ 
64+ uint64_t ubSizePlatform;
65+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
66+ ubSize = ubSizePlatform;
67+ }
68+ OP_CHECK_IF(
69+ coreNum == 0, OP_LOGE(context_->GetNodeName(), "coreNum is 0"),
70+ return ge::GRAPH_FAILED);
71+ 
72+ return ge::GRAPH_SUCCESS;
73+}
74+ 
75+ge::graphStatus MaxPool3DWithArgmaxV2TilingSIMT::GetShapeAttrsInfo()
76+{
77+ auto platformInfo = context_->GetPlatformInfo();
78+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
79+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
80+ if (!Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
81+ return ge::GRAPH_PARAM_INVALID;
82+ }
83+ auto runtimeAttrs = context_->GetAttrs();
84+ const char* data_format = runtimeAttrs->GetAttrPointer<char>(FORMAT_POS);
85+ OPS_CHECK_NULL_WITH_CONTEXT(context_, data_format);
86+ inputData.data_format = data_format;
87+ std::transform(inputData.data_format.begin(), inputData.data_format.end(), inputData.data_format.begin(),
88+ [](unsigned char c) { return std::tolower(c); });
89+ OP_CHECK_IF(!(inputData.data_format == "ndhwc" || inputData.data_format == "ncdhw"),
90+ OP_LOGE(context_, "ATTR data_format is %s ,expect [NDHWC] or [NCDHW].", data_format),
91+ return ge::GRAPH_FAILED);
92+ auto inputX = context_->GetInputShape(FIRPOS);
93+ OPS_CHECK_NULL_WITH_CONTEXT(context_, inputX);
94+ auto inputShape = EnsureNotScalar(inputX->GetStorageShape());
95+ 
96+ auto outX = context_->GetOutputShape(FIRPOS);
97+ OPS_CHECK_NULL_WITH_CONTEXT(context_, outX);
98+ auto outShape = EnsureNotScalar(outX->GetStorageShape());
99+ 
100+ auto indicesX = context_->GetOutputShape(SECPOS);
101+ OPS_CHECK_NULL_WITH_CONTEXT(context_, indicesX);
102+ auto indicesShape = EnsureNotScalar(indicesX->GetStorageShape());
103+ 
104+ if (inputShape.GetDimNum() != NCDHW_DIMS) {
105+ VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(),
106+ "MaxPool3DWithArgmaxV2: input shape dim = %zu, should be equal 5",
107+ inputShape.GetDimNum());
108+ return ge::GRAPH_FAILED;
109+ }
110+ 
111+ if (inputData.data_format == "ndhwc") {
112+ nDimPos = 0;
113+ cDimPos = 4;
114+ dDimPos = 1;
115+ hDimPos = 2;
116+ wDimPos = 3;
117+ }
118+ inputData.inputShape =
119+ array<uint64_t, NCDHW_DIMS>{uint64_t(inputShape.GetDim(nDimPos)), uint64_t(inputShape.GetDim(cDimPos)),
120+ uint64_t(inputShape.GetDim(dDimPos)), uint64_t(inputShape.GetDim(hDimPos)), uint64_t(inputShape.GetDim(wDimPos))};
121+ inputData.outShape =
122+ array<uint64_t, NCDHW_DIMS>{uint64_t(inputShape.GetDim(nDimPos)), uint64_t(inputShape.GetDim(cDimPos)),
123+ uint64_t(outShape.GetDim(dDimPos)), uint64_t(outShape.GetDim(hDimPos)), uint64_t(outShape.GetDim(wDimPos))};
124+ auto inputDesc = context_->GetInputDesc(0);
125+ OPS_CHECK_NULL_WITH_CONTEXT(context_, inputDesc);
126+ dtype = inputDesc->GetDataType();
127+ if (dtype != ge::DataType::DT_BF16 && dtype != ge::DataType::DT_FLOAT16 && dtype != ge::DataType::DT_FLOAT) {
128+ VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "MaxPool3DWithArgmaxV2: invalid dtype %s, should be BFloat16、Float16 or Float32", Ops::Base::ToString(dtype).c_str());
129+ return ge::GRAPH_FAILED;
130+ }
131+ if (indicesShape != outShape) {
132+ VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(),
133+ "MaxPool3DWithArgmaxV2: indices shape and values shape is different");
134+ return ge::GRAPH_FAILED;
135+ }
136+ OPS_CHECK_NULL_WITH_CONTEXT(context_, runtimeAttrs);
137+ const gert::TypedContinuousVector<int64_t>* kernelSize = runtimeAttrs->GetListInt(KERNEL_POS);
138+ OPS_CHECK_NULL_WITH_CONTEXT(context_, kernelSize);
139+ int32_t kSizeD = *(kernelSize->GetData());
CANN-robot
CANN-robotCANN-robot2月4日

边界检查缺失: kernelSize->GetData()返回的指针可能没有足够的元素。代码假设至少有3个元素，但没有检查kernelSize的大小。

问题类型: 边界检查缺失 文件路径: pooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_simt_tiling.cpp 行号: 139 问题代码:

int32_t kSizeD = *(kernelSize->GetData());
int32_t kSizeH = *(kernelSize->GetData() + 1);
int32_t kSizeW = *(kernelSize->GetData() + 2);

修改建议:

应该检查kernelSize的大小是否至少为3。建议：if (kernelSize->GetSize() < 3) { return ge::GRAPH_FAILED; }

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

likedislike
140+ int32_t kSizeH = *(kernelSize->GetData() + 1);
141+ int32_t kSizeW = *(kernelSize->GetData() + 2);
142+ inputData.kernelSize = array<uint64_t, DHW_DIMS>{uint64_t(kSizeD), uint64_t(kSizeH), uint64_t(kSizeW)};
143+ const gert::TypedContinuousVector<int64_t>* stride = runtimeAttrs->GetListInt(STRIDE_POS);
144+ OPS_CHECK_NULL_WITH_CONTEXT(context_, stride);
145+ uint64_t strideD = *(stride->GetData());
146+ uint64_t strideH = *(stride->GetData() + 1);
147+ uint64_t strideW = *(stride->GetData() + 2);
148+ inputData.stride = array<uint64_t, DHW_DIMS>{strideD, strideH, strideW};
149+ const gert::TypedContinuousVector<int64_t>* padding = runtimeAttrs->GetListInt(PADDING_POS);
150+ OPS_CHECK_NULL_WITH_CONTEXT(context_, padding);
151+ int32_t padsD = *(padding->GetData());
152+ int32_t padsH = *(padding->GetData() + 1);
153+ int32_t padsW = *(padding->GetData() + 2);
154+ OP_CHECK_IF((padsD * DOUB > kSizeD || padsH * DOUB > kSizeH || padsW * DOUB > kSizeW),
155+ OP_LOGE(context_, "pad should be smaller than or equal to half of kernel size, pad shape is [%d, %d, %d], kernel shape [%d, %d, %d]",
156+ padsD, padsH, padsW, kSizeD, kSizeH, kSizeW),
157+ return ge::GRAPH_FAILED);
158+ inputData.pad = array<uint64_t, DHW_DIMS>{uint64_t(padsD), uint64_t(padsH), uint64_t(padsW)};
159+ const gert::TypedContinuousVector<int64_t>* dilation = runtimeAttrs->GetListInt(DILATION_POS);
160+ OPS_CHECK_NULL_WITH_CONTEXT(context_, dilation);
161+ uint64_t dilationD = *(dilation->GetData());
162+ uint64_t dilationH = *(dilation->GetData() + 1);
163+ uint64_t dilationW = *(dilation->GetData() + 2);
164+ inputData.dilation =
165+ array<uint64_t, DHW_DIMS>{dilationD, dilationH, dilationW};
166+ inputData.ceilMode = *runtimeAttrs->GetAttrPointer<bool>(CEIL_POS);
167+ return ge::GRAPH_SUCCESS;
168+}
169+ 
170+ge::graphStatus MaxPool3DWithArgmaxV2TilingSIMT::DoOpTiling()
171+{
172+ tilingData_->nDim = inputData.inputShape[N_DIM_];
173+ tilingData_->cDim = inputData.inputShape[C_DIM_];
174+ tilingData_->dInDim = inputData.inputShape[D_DIM_];
175+ tilingData_->hInDim = inputData.inputShape[H_DIM_];
176+ tilingData_->wInDim = inputData.inputShape[W_DIM_];
177+ tilingData_->dOutDim = inputData.outShape[D_DIM_];
178+ tilingData_->hOutDim = inputData.outShape[H_DIM_];
179+ tilingData_->wOutDim = inputData.outShape[W_DIM_];
180+ tilingData_->kSizeD = inputData.kernelSize[D_IDX_];
181+ tilingData_->kSizeH = inputData.kernelSize[H_IDX_];
182+ tilingData_->kSizeW = inputData.kernelSize[W_IDX_];
183+ tilingData_->stridesD = inputData.stride[D_IDX_];
184+ tilingData_->stridesH = inputData.stride[H_IDX_];
185+ tilingData_->stridesW = inputData.stride[W_IDX_];
186+ tilingData_->padD = inputData.pad[D_IDX_];
187+ tilingData_->padH = inputData.pad[H_IDX_];
188+ tilingData_->padW = inputData.pad[W_IDX_];
189+ tilingData_->dilationD = inputData.dilation[D_IDX_];
190+ tilingData_->dilationH = inputData.dilation[H_IDX_];
191+ tilingData_->dilationW = inputData.dilation[W_IDX_];
192+ tilingData_->ceilMode = inputData.ceilMode;
193+ outputDataCount = tilingData_->nDim * tilingData_->cDim * tilingData_->dOutDim * tilingData_->hOutDim * tilingData_->wOutDim;
194+ int64_t threads = std::min(outputDataCount, MAX_THREAD_NUM);
195+ int64_t blockNum = Ops::Base::CeilDiv(outputDataCount, threads);
196+ blockNum = std::min(blockNum, static_cast<int64_t>(coreNum));
197+ context_->SetBlockDim(blockNum);
198+ tilingData_->threadNums = threads;
199+ tilingData_->blockNums = blockNum;
200+ return ge::GRAPH_SUCCESS;
201+}
202+ 
203+uint64_t MaxPool3DWithArgmaxV2TilingSIMT::GetTilingKey() const
204+{
205+ if (inputData.data_format == "ncdhw" && outputDataCount <= MAX_INT32) {
206+ return SIMT_NCDHW_TILING_KEY_INT32;
207+ } else if (inputData.data_format == "ndhwc" && outputDataCount <= MAX_INT32) {
208+ return SIMT_NDHWC_TILING_KEY_INT32;
209+ } else if (inputData.data_format == "ncdhw" && outputDataCount > MAX_INT32) {
210+ return SIMT_NCDHW_TILING_KEY_INT64;
211+ } else if (inputData.data_format == "ndhwc" && outputDataCount > MAX_INT32) {
212+ return SIMT_NDHWC_TILING_KEY_INT64;
213+ }
214+ return SIMT_NCDHW_TILING_KEY_INT32;
215+}
216+ 
217+ge::graphStatus MaxPool3DWithArgmaxV2TilingSIMT::GetWorkspaceSize()
218+{
219+ auto sys_workspace = SYS_WORKSPACE_SIZE;
220+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
221+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
222+ currentWorkspace[0] = static_cast<size_t>(sys_workspace);
223+ return ge::GRAPH_SUCCESS;
224+}
225+ 
226+ge::graphStatus MaxPool3DWithArgmaxV2TilingSIMT::PostTiling()
227+{
228+ return ge::GRAPH_SUCCESS;
229+}
230+ 
231+void MaxPool3DWithArgmaxV2TilingSIMT::DumpTilingInfo()
232+{
233+ std::string str;
234+ str += " threadNums:" + std::to_string(tilingData_->threadNums);
235+ str += " blockNums:" + std::to_string(tilingData_->blockNums);
236+ str += " nDim:" + std::to_string(tilingData_->nDim);
237+ str += " cDim:" + std::to_string(tilingData_->cDim);
238+ str += " dInDim:" + std::to_string(tilingData_->dInDim);
239+ str += " hInDim:" + std::to_string(tilingData_->hInDim);
240+ str += " wInDim:" + std::to_string(tilingData_->wInDim);
241+ str += " dOutDim:" + std::to_string(tilingData_->dOutDim);
242+ str += " hOutDim:" + std::to_string(tilingData_->hOutDim);
243+ str += " wOutDim:" + std::to_string(tilingData_->wOutDim);
244+ str += " kSizeD:" + std::to_string(tilingData_->kSizeD);
245+ str += " kSizeH:" + std::to_string(tilingData_->kSizeH);
246+ str += " kSizeW:" + std::to_string(tilingData_->kSizeW);
247+ str += " stridesD:" + std::to_string(tilingData_->stridesD);
248+ str += " stridesH:" + std::to_string(tilingData_->stridesH);
249+ str += " stridesW:" + std::to_string(tilingData_->stridesW);
250+ str += " padD:" + std::to_string(tilingData_->padD);
251+ str += " padH:" + std::to_string(tilingData_->padH);
252+ str += " padW:" + std::to_string(tilingData_->padW);
253+ str += " dilationD:" + std::to_string(tilingData_->dilationD);
254+ str += " dilationH:" + std::to_string(tilingData_->dilationH);
255+ str += " dilationW:" + std::to_string(tilingData_->dilationW);
256+ str += " ceilMode:" + std::to_string(tilingData_->ceilMode);
257+ OP_LOGI(context_, "%s", str.c_str());
258+}
259+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2TilingSIMT, 2);
260+} // namespace optiling
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_simt_tiling.h+102-0
@@ -0,0 +1,102 @@
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 max_pool3d_with_argmax_v2_simt_tiling.h
13+ * \brief simt imply for max_pool3d_with_argmax
14+ */
15+ 
16+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL3D_WITH_ARGMAX_V2_SIMT_TILING_H
17+#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL3D_WITH_ARGMAX_V2_SIMT_TILING_H
18+ 
19+#include "max_pool3d_with_argmax_v2_tiling_base.h"
20+#include "../op_kernel/arch35/max_pool3d_with_argmax_v2_tiling_struct.h"
21+ 
22+namespace optiling
23+{
24+const int64_t NCDHW_DIMS = 5;
25+const int64_t KERNEL_POS = 0;
26+const int64_t STRIDE_POS = 1;
27+const int64_t PADDING_POS = 2;
28+const int64_t DILATION_POS = 3;
29+const int64_t CEIL_POS = 4;
30+const int64_t FORMAT_POS = 5;
31+const int64_t N_DIM_ = 0;
32+const int64_t C_DIM_ = 1;
33+const int64_t D_DIM_ = 2;
34+const int64_t H_DIM_ = 3;
35+const int64_t W_DIM_ = 4;
36+const int64_t D_IDX_ = 0;
37+const int64_t H_IDX_ = 1;
38+const int64_t W_IDX_ = 2;
39+const int64_t DOUB = 2;
40+const int64_t FIRPOS = 0;
41+const int64_t SECPOS = 1;
42+constexpr int64_t MAX_INT32 = 2147483647;
43+constexpr uint64_t SIMT_NCDHW_TILING_KEY_INT32 = 600001;
44+constexpr uint64_t SIMT_NDHWC_TILING_KEY_INT32 = 600002;
45+constexpr uint64_t SIMT_NCDHW_TILING_KEY_INT64 = 600011;
46+constexpr uint64_t SIMT_NDHWC_TILING_KEY_INT64 = 600012;
47+constexpr int64_t MAX_THREAD_NUM = 256;
48+constexpr size_t SYS_WORKSPACE_SIZE = 16 * 1024 * 1024;
49+ 
50+struct InputSIMTInfo {
51+ array<uint64_t, NCDHW_DIMS> inputShape;
52+ array<uint64_t, NCDHW_DIMS> outShape;
53+ array<uint64_t, DHW_DIMS> kernelSize;
54+ array<uint64_t, DHW_DIMS> stride;
55+ array<uint64_t, DHW_DIMS> pad;
56+ array<uint64_t, DHW_DIMS> dilation;
57+ bool ceilMode;
58+ std::string data_format;
59+};
60+ 
61+class MaxPool3DWithArgmaxV2TilingSIMT : public MaxPool3DWithArgmaxV2BaseTiling
62+{
63+public:
64+ explicit MaxPool3DWithArgmaxV2TilingSIMT(gert::TilingContext* context) : MaxPool3DWithArgmaxV2BaseTiling(context)
65+ {
66+ }
67+ 
68+ ~MaxPool3DWithArgmaxV2TilingSIMT() override
69+ {
70+ }
71+ 
72+protected:
73+ bool IsCapable() override;
74+ ge::graphStatus GetPlatformInfo() override;
75+ // 获取INPUT/OUTPUT/ATTR信息
76+ ge::graphStatus GetShapeAttrsInfo() override;
77+ // 计算数据切分TilingData
78+ ge::graphStatus DoOpTiling() override;
79+ // 计算TilingKey
80+ uint64_t GetTilingKey() const override;
81+ // 分配workspace
82+ ge::graphStatus GetWorkspaceSize() override;
83+ // 保存Tiling数据
84+ ge::graphStatus PostTiling() override;
85+ // tiling信息打屏
86+ void DumpTilingInfo() override;
87+ 
88+private:
89+ uint64_t GenerateTilingKey(uint64_t innerKey);
90+ MaxPool3DWithArgmaxV2Tiling::MaxPool3DWithArgmaxV2SimtTilingData* tilingData_ =
91+ context_->GetTilingData<MaxPool3DWithArgmaxV2Tiling::MaxPool3DWithArgmaxV2SimtTilingData>();
92+ InputSIMTInfo inputData;
93+ int nDimPos = 0;
94+ int cDimPos = 1;
95+ int dDimPos = 2;
96+ int hDimPos = 3;
97+ int wDimPos = 4;
98+ int64_t outputDataCount = 0;
99+};
100+ 
101+} // namespace optiling
102+#endif // MAX_POOL3D_WITH_ARGMAX_V2_SIMT_TILING_H
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling.cpp+1-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_base.cpp+1-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_base.h+5-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -356,6 +356,10 @@ struct InputInfo {
356 array<uint64_t, DHW_DIMS> pad;356 array<uint64_t, DHW_DIMS> pad;
357 array<uint64_t, DHW_DIMS> dilation;357 array<uint64_t, DHW_DIMS> dilation;
358 bool ceilMode;358 bool ceilMode;
359+ ge::DataType indexDtype;
360+ ge::Format inputFormat;
361+ uint64_t nInput;
362+ uint64_t cInput;
359};363};
360 364 
361struct PadInputInfo {365struct PadInputInfo {
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_basesplit.cpp+1-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_big_kernel_regbase.cpp+330-0
@@ -0,0 +1,330 @@
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 max_pool3d_with_argmax_v2_tiling_big_kernel_regbase.cpp
13+ * \brief big kernel imply for max_pool3d_with_argmax_v2
14+ */
15+ 
16+#include "tiling_base/tiling_util.h"
17+#include "max_pool3d_with_argmax_v2_tiling_big_kernel_regbase.h"
18+#include "error_util.h"
19+ 
20+using namespace std;
21+static const gert::Shape g_vec_1_shape = {1};
22+ 
23+namespace optiling
24+{
25+static const gert::Shape &EnsureNotScalar(const gert::Shape &inShape) {
26+ if (inShape.IsScalar()) {
27+ return g_vec_1_shape;
28+ }
29+ return inShape;
30+}
31+ 
32+ge::graphStatus MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::GetShapeAttrsInfo()
33+{
34+ auto platformInfo = context_->GetPlatformInfo();
35+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
36+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
37+ if (!Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
38+ return ge::GRAPH_PARAM_INVALID;
39+ }
40+ auto inputX = context_->GetInputShape(0);
41+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputX);
42+ auto inputShape = EnsureNotScalar(inputX->GetStorageShape());
43+ 
44+ OP_CHECK_IF(inputShape.GetDimNum() != NCDHW_DIMS,
45+ OP_LOGE(context_->GetNodeName(),
46+ "MaxPool3DWithArgmaxV2: input shape dim = %zu, should be equal 5",
47+ inputShape.GetDimNum()),
48+ return ge::GRAPH_FAILED);
49+ OP_CHECK_IF(inputShape.GetShapeSize() <= 0,
50+ OP_LOGE(context_->GetNodeName(),
51+ "MaxPool3DWithArgmaxV2: input shape size %ld less than zero failed",
52+ inputShape.GetShapeSize()),
53+ return ge::GRAPH_FAILED);
54+ auto inputDesc = context_->GetInputDesc(0);
55+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputDesc);
56+ dtype = inputDesc->GetDataType();
57+ if (dtype != ge::DataType::DT_BF16 && dtype != ge::DataType::DT_FLOAT16 && dtype != ge::DataType::DT_FLOAT) {
58+ OP_LOGE(context_->GetNodeName(), "MaxPool3DWithArgmaxV2: invalid dtype");
59+ return ge::GRAPH_FAILED;
60+ }
61+ 
62+ auto outX = context_->GetOutputShape(0);
63+ OP_CHECK_NULL_WITH_CONTEXT(context_, outX);
64+ auto outShape = EnsureNotScalar(outX->GetStorageShape());
65+ auto indicesX = context_->GetOutputShape(1);
66+ OP_CHECK_NULL_WITH_CONTEXT(context_, indicesX);
67+ auto indicesShape = EnsureNotScalar(indicesX->GetStorageShape());
68+ if (indicesShape != outShape) {
69+ OP_LOGE(context_->GetNodeName(),
70+ "MaxPool3DWithArgmaxV2: indices shape and values shape is different");
71+ return ge::GRAPH_FAILED;
72+ }
73+ auto runtimeAttrs = context_->GetAttrs();
74+ OP_CHECK_NULL_WITH_CONTEXT(context_, runtimeAttrs);
75+ 
76+ std::string inputFormatStr("NCDHW");
77+ const char* inputFormat = runtimeAttrs->GetAttrPointer<char>(FORMAT_POS);
78+ if (inputFormat != nullptr) {
79+ inputFormatStr = inputFormat;
80+ }
81+ int d_dim = MP_MAX_3D_DIM_TWO, h_dim = MP_MAX_3D_DIM_THREE, w_dim = MP_MAX_3D_DIM_FOUR;
82+ if (inputFormatStr == "NCDHW") {
83+ inputData.inputFormat = ge::Format::FORMAT_NCDHW;
84+ inputData.batches = inputShape.GetDim(MP_MAX_3D_DIM_ZERO) * inputShape.GetDim(MP_MAX_3D_DIM_ONE);
85+ inputData.nInput = inputShape.GetDim(MP_MAX_3D_DIM_ZERO);
86+ inputData.cInput = inputShape.GetDim(MP_MAX_3D_DIM_ONE);
87+ } else {
88+ return ge::GRAPH_PARAM_INVALID;
89+ }
90+ 
91+ OP_CHECK_IF(outShape.GetDim(d_dim) < 1 || outShape.GetDim(h_dim) < 1 || outShape.GetDim(w_dim) < 1,
92+ OP_LOGE(context_->GetNodeName(),
93+ "MaxPool3DWithArgmaxV2: output shape [%ld, %ld, %ld] not support",
94+ outShape.GetDim(d_dim), outShape.GetDim(h_dim), outShape.GetDim(w_dim)),
95+ return ge::GRAPH_FAILED);
96+ 
97+ inputData.inputShape =
98+ array<uint64_t, DHW_DIMS>{uint64_t(inputShape.GetDim(d_dim)), uint64_t(inputShape.GetDim(h_dim)), uint64_t(inputShape.GetDim(w_dim))};
99+ inputData.outShape = array<uint64_t, DHW_DIMS>{uint64_t(outShape.GetDim(d_dim)), uint64_t(outShape.GetDim(h_dim)), uint64_t(outShape.GetDim(w_dim))};
100+ 
101+ int32_t dValue = 0;
102+ int32_t hValue = 0;
103+ int32_t wValue = 0;
104+ const gert::TypedContinuousVector<int64_t>* kernelSize = runtimeAttrs->GetListInt(KERNEL_POS);
105+ OP_CHECK_NULL_WITH_CONTEXT(context_, kernelSize);
106+ dValue = *(kernelSize->GetData());
107+ hValue = *(kernelSize->GetData() + 1);
108+ wValue = *(kernelSize->GetData() + 2);
109+ inputData.kernelSize = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
110+ OP_CHECK_IF(
111+ dValue <= 0 || hValue <= 0 || wValue <= 0,
112+ OP_LOGE(context_->GetNodeName(),
113+ "MaxPool3DWithArgmaxV2: not support kernel shape [%d, %d, %d]", dValue, hValue, wValue),
114+ return ge::GRAPH_FAILED);
115+ 
116+ int32_t kdValue = dValue;
117+ int32_t khValue = hValue;
118+ int32_t kwValue = wValue;
119+ const gert::TypedContinuousVector<int64_t>* stride = runtimeAttrs->GetListInt(STRIDE_POS);
120+ OP_CHECK_NULL_WITH_CONTEXT(context_, stride);
121+ dValue = *(stride->GetData());
122+ hValue = *(stride->GetData() + 1);
123+ wValue = *(stride->GetData() + 2);
124+ inputData.stride = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
125+ OP_CHECK_IF(
126+ dValue <= 0 || hValue <= 0 || wValue <= 0,
127+ OP_LOGE(context_->GetNodeName(),
128+ "MaxPool3DWithArgmaxV2: not support stride shape [%d, %d, %d]", dValue, hValue, wValue),
129+ return ge::GRAPH_FAILED);
130+ 
131+ const gert::TypedContinuousVector<int64_t>* padding = runtimeAttrs->GetListInt(PADDING_POS);
132+ OP_CHECK_NULL_WITH_CONTEXT(context_, padding);
133+ dValue = *(padding->GetData());
134+ hValue = *(padding->GetData() + 1);
135+ wValue = *(padding->GetData() + 2);
136+ inputData.pad = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
137+ OP_CHECK_IF(
138+ dValue > kdValue / 2 || hValue > khValue / 2 || wValue > kwValue / 2,
139+ OP_LOGE(context_->GetNodeName(),
140+ "MaxPool3DWithArgmaxV2: not support pad shape [%d, %d, %d] kernel shape [%d, %d, %d]",
141+ dValue, hValue, wValue, kdValue, khValue, kwValue),
142+ return ge::GRAPH_FAILED);
143+ 
144+ inputData.dilation = array<uint64_t, DHW_DIMS>{1, 1, 1};
145+ dValue = 1;
146+ hValue = 1;
147+ wValue = 1;
148+ const gert::TypedContinuousVector<int64_t>* dilation = runtimeAttrs->GetListInt(DILATION_POS);
149+ if (dilation != nullptr) {
150+ dValue = *(dilation->GetData());
151+ hValue = *(dilation->GetData() + 1);
152+ wValue = *(dilation->GetData() + 2);
153+ inputData.dilation = array<uint64_t, DHW_DIMS>{uint64_t(dValue), uint64_t(hValue), uint64_t(wValue)};
154+ OP_CHECK_IF(
155+ dValue <= 0 || hValue <= 0 || wValue <= 0,
156+ OP_LOGE(context_->GetNodeName(),
157+ "MaxPool3DWithArgmaxV2: not support dilation shape [%d, %d, %d]", dValue, hValue, wValue),
158+ return ge::GRAPH_FAILED);
159+ }
160+ 
161+ inputData.ceilMode = false;
162+ const bool* ceilModePtr = runtimeAttrs->GetAttrPointer<bool>(CEIL_POS);
163+ if (ceilModePtr != nullptr) {
164+ inputData.ceilMode = *ceilModePtr;
165+ }
166+ 
167+ int indexDtype = THREE;
168+ const int* indexDtypePtr = runtimeAttrs->GetAttrPointer<int>(DTYPE_POS);
169+ if (indexDtypePtr != nullptr) {
170+ indexDtype = *indexDtypePtr;
171+ }
172+ switch (indexDtype) {
173+ case MP_MAX_3D_TYPE_INT32:
174+ inputData.indexDtype = ge::DataType::DT_INT32;
175+ break;
176+ case MP_MAX_3D_TYPE_INT64:
177+ inputData.indexDtype = ge::DataType::DT_INT64;
178+ break;
179+ default:
180+ inputData.indexDtype = ge::DataType::DT_INT32;
181+ break;
182+ }
183+ return ge::GRAPH_SUCCESS;
184+}
185+ 
186+ge::graphStatus MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::GetPlatformInfo()
187+{
188+ auto platformPtr = context_->GetPlatformInfo();
189+ if (platformPtr == nullptr) {
190+ auto compileInfoPtr = reinterpret_cast<const MaxPool3DWithArgmaxV2CompileInfo*>(context_->GetCompileInfo());
191+ OP_CHECK_IF(compileInfoPtr == nullptr, CUBE_INNER_ERR_REPORT(context_, "compile info is null"),
192+ return ge::GRAPH_FAILED);
193+ coreNum = compileInfoPtr->coreNum;
194+ ubSize = compileInfoPtr->ubSize;
195+ } else {
196+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformPtr);
197+ coreNum = ascendcPlatform.GetCoreNumAiv();
198+ uint64_t ubSizePlatform;
199+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
200+ ubSize = static_cast<int64_t>(ubSizePlatform);
201+ }
202+ OP_CHECK_IF(coreNum == 0, CUBE_INNER_ERR_REPORT(context_, "coreNum is 0"), return ge::GRAPH_FAILED);
203+ return ge::GRAPH_SUCCESS;
204+}
205+ 
206+bool MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::IsCapable()
207+{
208+ int64_t ubAvailable = ubSize - (BYTES_FOUR + BYTES_EIGHT) * OUT_BUFFER_LEN;
209+ maxCount_ = ubAvailable / BUFFER_NUM;
210+ int64_t vRegSize = Ops::Base::GetVRegSize(context_);
211+ maxCount_ = Ops::Base::FloorAlign(maxCount_, vRegSize);
212+ int64_t dtypeSize = ge::GetSizeByDataType(dtype);
213+ OP_CHECK_IF(
214+ dtypeSize <= 0,
215+ OP_LOGE(context_, "dtypeSize must be greater than 0, dtypeSize: %ld", dtypeSize),
216+ return false);
217+ if (dtypeSize != 0) {
218+ maxCount_ = maxCount_ / dtypeSize;
219+ }
220+ if (inputData.dilation[D_DIM] == 1 && inputData.dilation[H_DIM] == 1 && inputData.dilation[W_DIM] == 1 && maxCount_ > MIN_COUNT &&
221+ inputData.inputFormat == ge::Format::FORMAT_NCDHW && inputData.kernelSize[W_DIM] * dtypeSize >= KW_THRESHOLD) {
222+ return true;
223+ }
224+ return false;
225+}
226+ 
227+void MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::DoUBTiling()
228+{
229+ totalIdx_ = inputData.batches * inputData.outShape[D_DIM] * inputData.outShape[H_DIM] * inputData.outShape[W_DIM];
230+ blockFactor_ = totalIdx_ / coreNum;
231+ blockTail_ = totalIdx_ % coreNum;
232+ if (blockFactor_ == 0) {
233+ coreNums_ = totalIdx_;
234+ } else {
235+ coreNums_ = coreNum;
236+ }
237+ isSigOut_ = (inputData.outShape[D_DIM] == 1 && inputData.outShape[H_DIM] == 1 && inputData.outShape[W_DIM] == 1) ? 1 : 0;
238+}
239+ 
240+void MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::SetTilingData()
241+{
242+ tilingData_->dOutDim = inputData.outShape[D_DIM];
243+ tilingData_->hOutDim = inputData.outShape[H_DIM];
244+ tilingData_->wOutDim = inputData.outShape[W_DIM];
245+ tilingData_->dInDim = inputData.inputShape[D_DIM];
246+ tilingData_->hInDim = inputData.inputShape[H_DIM];
247+ tilingData_->wInDim = inputData.inputShape[W_DIM];
248+ tilingData_->kD = inputData.kernelSize[D_DIM];
249+ tilingData_->kH = inputData.kernelSize[H_DIM];
250+ tilingData_->kW = inputData.kernelSize[W_DIM];
251+ tilingData_->sD = inputData.stride[D_DIM];
252+ tilingData_->sH = inputData.stride[H_DIM];
253+ tilingData_->sW = inputData.stride[W_DIM];
254+ tilingData_->pD = inputData.pad[D_DIM];
255+ tilingData_->pH = inputData.pad[H_DIM];
256+ tilingData_->pW = inputData.pad[W_DIM];
257+ tilingData_->dD = inputData.dilation[D_DIM];
258+ tilingData_->dH = inputData.dilation[H_DIM];
259+ tilingData_->dW = inputData.dilation[W_DIM];
260+ tilingData_->blockFactor = blockFactor_;
261+ tilingData_->blockTail = blockTail_;
262+ tilingData_->totalIdx = totalIdx_;
263+ tilingData_->coreNums = coreNums_;
264+ tilingData_->maxCount = maxCount_;
265+ tilingData_->isSigOut = isSigOut_;
266+}
267+ 
268+ge::graphStatus MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::DoOpTiling()
269+{
270+ DoUBTiling();
271+ SetTilingData();
272+ return ge::GRAPH_SUCCESS;
273+}
274+ 
275+ge::graphStatus MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::DoLibApiTiling()
276+{
277+ return ge::GRAPH_SUCCESS;
278+}
279+ 
280+ge::graphStatus MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::GetWorkspaceSize()
281+{
282+ auto sys_workspace = WS_SYS_SIZE;
283+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
284+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
285+ currentWorkspace[0] = sys_workspace;
286+ 
287+ return ge::GRAPH_SUCCESS;
288+}
289+ 
290+ge::graphStatus MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::PostTiling()
291+{
292+ context_->SetBlockDim(coreNums_);
293+ return ge::GRAPH_SUCCESS;
294+}
295+ 
296+uint64_t MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::GetTilingKey() const
297+{
298+ return MAX_POOL_WITH_ARGMAX_V2_TILING_KEY_BIG_KERNEL_REGBASE_NCDHW;
299+}
300+ 
301+void MaxPool3DWithArgmaxV2BigKernelRegbaseTiling::DumpTilingInfo()
302+{
303+ std::string str;
304+ str += " dInDim:" + std::to_string(tilingData_->dInDim);
305+ str += " hInDim:" + std::to_string(tilingData_->hInDim);
306+ str += " wInDim:" + std::to_string(tilingData_->wInDim);
307+ str += " dOutDim:" + std::to_string(tilingData_->dOutDim);
308+ str += " hOutDim:" + std::to_string(tilingData_->hOutDim);
309+ str += " wOutDim:" + std::to_string(tilingData_->wOutDim);
310+ str += " kD:" + std::to_string(tilingData_->kD);
311+ str += " kH:" + std::to_string(tilingData_->kH);
312+ str += " kW:" + std::to_string(tilingData_->kW);
313+ str += " sD:" + std::to_string(tilingData_->sD);
314+ str += " sH:" + std::to_string(tilingData_->sH);
315+ str += " sW:" + std::to_string(tilingData_->sW);
316+ str += " pD:" + std::to_string(tilingData_->pD);
317+ str += " pH:" + std::to_string(tilingData_->pH);
318+ str += " pW:" + std::to_string(tilingData_->pW);
319+ str += " blockFactor:" + std::to_string(tilingData_->blockFactor);
320+ str += " blockTail:" + std::to_string(tilingData_->blockTail);
321+ str += " totalIdx:" + std::to_string(tilingData_->totalIdx);
322+ str += " coreNums:" + std::to_string(tilingData_->coreNums);
323+ str += " maxCount:" + std::to_string(tilingData_->maxCount);
324+ str += " isSigOut:" + std::to_string(tilingData_->isSigOut);
325+ OP_LOGI(context_, "%s", str.c_str());
326+}
327+ 
328+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2BigKernelRegbaseTiling, 1);
329+ 
330+}
Apooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_big_kernel_regbase.h+109-0
@@ -0,0 +1,109 @@
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 max_pool3d_with_argmax_v2_tiling_big_kernel_regbase.h
13+ * \brief big kernel imply for max_pool3d_with_argmax_v2
14+ */
15+ 
16+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL3D_WITH_AGRMAX_V2_TILING_BIG_KERNEL_REGBASE_H_
17+#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL3D_WITH_AGRMAX_V2_TILING_BIG_KERNEL_REGBASE_H_
18+ 
19+#include "max_pool3d_with_argmax_v2_tiling_base.h"
20+#include "../op_kernel/arch35/max_pool3d_with_argmax_v2_tiling_struct.h"
21+#include "error_util.h"
22+#include "op_common/op_host/util/platform_util.h"
23+ 
24+namespace optiling
25+{
26+using Ops::NN::Optiling::TilingBaseClass;
27+ 
28+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V2_TILING_KEY_BIG_KERNEL_REGBASE_NCDHW = 611110;
29+const int64_t INPUT_IDX_X = 0;
30+const int64_t NCDHW_DIMS = 5;
31+const int64_t KERNEL_POS = 0;
32+const int64_t STRIDE_POS = 1;
33+const int64_t PADDING_POS = 2;
34+const int64_t DTYPE_POS = 6;
35+const int64_t DILATION_POS = 3;
36+const int64_t CEIL_POS = 4;
37+const int64_t FORMAT_POS = 5;
38+const int64_t WS_SYS_SIZE = 16 * 1024 * 1024;
39+static const int64_t MP_MAX_3D_DIM_ZERO = 0;
40+static const int64_t MP_MAX_3D_DIM_ONE = 1;
41+static const int64_t MP_MAX_3D_DIM_TWO = 2;
42+static const int64_t MP_MAX_3D_DIM_THREE = 3;
43+static const int64_t MP_MAX_3D_DIM_FOUR = 4;
44+static const int64_t MP_MAX_3D_TYPE_INT32 = 3;
45+static const int64_t MP_MAX_3D_TYPE_INT64 = 9;
46+static constexpr int64_t OUT_BUFFER_LEN = 1024;
47+static constexpr int64_t BUFFER_NUM = 2;
48+static constexpr int64_t MIN_COUNT = 1024;
49+static constexpr int64_t BYTES_FOUR = 4;
50+static constexpr int64_t BYTES_EIGHT = 8;
51+static constexpr int64_t KW_THRESHOLD = 128;
52+static constexpr int64_t THREE = 3;
53+ 
54+struct BigKernelInputInfo {
55+ uint64_t batches;
56+ std::array<uint64_t, DHW_DIMS> inputShape;
57+ std::array<uint64_t, DHW_DIMS> outShape;
58+ std::array<uint64_t, DHW_DIMS> kernelSize;
59+ std::array<uint64_t, DHW_DIMS> stride;
60+ std::array<uint64_t, DHW_DIMS> pad;
61+ std::array<uint64_t, DHW_DIMS> dilation;
62+ bool ceilMode;
63+ ge::DataType indexDtype;
64+ ge::Format inputFormat;
65+ uint64_t nInput;
66+ uint64_t cInput;
67+};
68+ 
69+class MaxPool3DWithArgmaxV2BigKernelRegbaseTiling : public MaxPool3DWithArgmaxV2BaseTiling {
70+public:
71+ explicit MaxPool3DWithArgmaxV2BigKernelRegbaseTiling(gert::TilingContext* context) : MaxPool3DWithArgmaxV2BaseTiling(context)
72+ {
73+ }
74+ ~MaxPool3DWithArgmaxV2BigKernelRegbaseTiling() override
75+ {
76+ }
77+ 
78+private:
79+ ge::graphStatus GetPlatformInfo() override;
80+ ge::graphStatus GetShapeAttrsInfo() override;
81+ ge::graphStatus DoLibApiTiling() override;
82+ ge::graphStatus GetWorkspaceSize() override;
83+
84+ void DoUBTiling();
85+ void SetTilingData();
86+ uint64_t GetTilingKey() const override;
87+ bool IsCapable() override;
88+ ge::graphStatus DoOpTiling() override;
89+ ge::graphStatus PostTiling() override;
90+ void DumpTilingInfo() override;
91+ 
92+ MaxPool3DWithArgmaxV2Tiling::MaxPool3DWithArgmaxV2BigKernelRegbaseTilingData* tilingData_ =
93+ context_->GetTilingData<MaxPool3DWithArgmaxV2Tiling::MaxPool3DWithArgmaxV2BigKernelRegbaseTilingData>();
94+ int64_t totalIdx_{0};
95+ int64_t blockFactor_{0};
96+ int64_t blockTail_{0};
97+ int64_t maxCount_{0};
98+ int64_t isSigOut_{0};
99+ int64_t coreNums_{0};
100+ 
101+public:
102+ BigKernelInputInfo inputData;
103+ ge::DataType dtype = ge::DataType::DT_FLOAT;
104+ uint32_t coreNum = 1;
105+ uint32_t ubSize = 0;
106+};
107+ 
108+}
109+#endif // CANN_MAX_POOL_3D_WITH_ARGMAX_V2_TILING_BIG_KERNEL_REGBASE
CANN-robot
CANN-robotCANN-robot2月4日

代码结构与可维护性: 头文件保护宏的注释与第16-17行的宏定义不匹配。第16-17行使用的是AIR_CXX_RUNTIME_V2_OP_IMPL前缀，而此处注释使用的是CANN前缀，存在不一致。

问题类型: 代码结构与可维护性 文件路径: pooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_big_kernel_regbase.h 行号: 109 问题代码:

#endif  // CANN_MAX_POOL_3D_WITH_ARGMAX_V2_TILING_BIG_KERNEL_REGBASE

修改建议:

将注释更正为与宏定义一致：// AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL3D_WITH_ARGMAX_V2_TILING_BIG_KERNEL_REGBASE_H_（注意修正ARGMAX拼写）

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

likedislike
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_bigkernel.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18// 3, splitD=1, splitH=1, splitW=1, splitKernel = 1, dtype=float19// 3, splitD=1, splitH=1, splitW=1, splitKernel = 1, dtype=float
@@ -27,6 +28,12 @@ namespace optiling {
27 28 
28bool MaxPool3DWithArgmaxV2BigKernelTiling::IsCapable()29bool MaxPool3DWithArgmaxV2BigKernelTiling::IsCapable()
29{30{
31+ auto platformInfo = context_->GetPlatformInfo();
32+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
33+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
34+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
35+ return false;
36+ }
30 if (inputData.dilation[D_DIM] == 1 && inputData.dilation[H_DIM] == 1 && inputData.dilation[W_DIM] == 1) {37 if (inputData.dilation[D_DIM] == 1 && inputData.dilation[H_DIM] == 1 && inputData.dilation[W_DIM] == 1) {
31 return true;38 return true;
32 }39 }
@@ -100,6 +107,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2BigKernelTiling::PostTiling()
100 return ge::GRAPH_SUCCESS;107 return ge::GRAPH_SUCCESS;
101}108}
102 109 
103-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2BigKernelTiling, 1);110+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2BigKernelTiling, 11);
104 111 
105} // namespace optiling112} // namespace optiling
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_hugekernel.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18// 1, splitD=1, splitH=1, splitW=1, splitKernel = 1, dtype=float19// 1, splitD=1, splitH=1, splitW=1, splitKernel = 1, dtype=float
@@ -27,6 +28,12 @@ namespace optiling {
27 28 
28bool MaxPool3DWithArgmaxV2HugeKernelTiling::IsCapable()29bool MaxPool3DWithArgmaxV2HugeKernelTiling::IsCapable()
29{30{
31+ auto platformInfo = context_->GetPlatformInfo();
32+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
33+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
34+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
35+ return false;
36+ }
30 array<uint64_t, DHW_DIMS> parts{37 array<uint64_t, DHW_DIMS> parts{
31 inputData.kernelSize[D_DIM], inputData.kernelSize[H_DIM], inputData.kernelSize[W_DIM]};38 inputData.kernelSize[D_DIM], inputData.kernelSize[H_DIM], inputData.kernelSize[W_DIM]};
32 LargeKernelCalcParts(parts, parts.size(), bufSizes, D_DIM);39 LargeKernelCalcParts(parts, parts.size(), bufSizes, D_DIM);
@@ -224,6 +231,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2HugeKernelTiling::PostTiling()
224 return ge::GRAPH_SUCCESS;231 return ge::GRAPH_SUCCESS;
225}232}
226 233 
227-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2HugeKernelTiling, 6);234+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2HugeKernelTiling, 16);
228 235 
229} // namespace optiling236} // namespace optiling
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_nosplit.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18// 1, splitD=0, splitH=0, splitW=0, splitKernel = 0, dtype=float19// 1, splitD=0, splitH=0, splitW=0, splitKernel = 0, dtype=float
@@ -27,6 +28,12 @@ namespace optiling {
27 28 
28bool MaxPool3DWithArgmaxV2NoSplitTiling::IsCapable()29bool MaxPool3DWithArgmaxV2NoSplitTiling::IsCapable()
29{30{
31+ auto platformInfo = context_->GetPlatformInfo();
32+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
33+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
34+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
35+ return false;
36+ }
30 array<uint64_t, DHW_DIMS> tmpOutShape{37 array<uint64_t, DHW_DIMS> tmpOutShape{
31 inputData.outShape[D_DIM], inputData.outShape[H_DIM], inputData.outShape[W_DIM]};38 inputData.outShape[D_DIM], inputData.outShape[H_DIM], inputData.outShape[W_DIM]};
32 auto summaryMemory = CalcBufferSizes(39 auto summaryMemory = CalcBufferSizes(
@@ -103,6 +110,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2NoSplitTiling::PostTiling()
103 return ge::GRAPH_SUCCESS;110 return ge::GRAPH_SUCCESS;
104}111}
105 112 
106-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2NoSplitTiling, 2);113+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2NoSplitTiling, 12);
107 114 
108} // namespace optiling115} // namespace optiling
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_splitd.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18// 1, splitD=1, splitH=0, splitW=0, splitKernel = 0, dtype=float19// 1, splitD=1, splitH=0, splitW=0, splitKernel = 0, dtype=float
@@ -27,6 +28,12 @@ namespace optiling {
27 28 
28bool MaxPool3DWithArgmaxV2SplitDTiling::IsCapable()29bool MaxPool3DWithArgmaxV2SplitDTiling::IsCapable()
29{30{
31+ auto platformInfo = context_->GetPlatformInfo();
32+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
33+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
34+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
35+ return false;
36+ }
30 array<uint64_t, DHW_DIMS> parts{37 array<uint64_t, DHW_DIMS> parts{
31 padInputData.padInputShape[D_DIM], padInputData.padInputShape[H_DIM], padInputData.padInputShape[W_DIM]};38 padInputData.padInputShape[D_DIM], padInputData.padInputShape[H_DIM], padInputData.padInputShape[W_DIM]};
32 array<uint64_t, DHW_DIMS> partOuts{inputData.outShape[D_DIM], inputData.outShape[H_DIM], inputData.outShape[W_DIM]};39 array<uint64_t, DHW_DIMS> partOuts{inputData.outShape[D_DIM], inputData.outShape[H_DIM], inputData.outShape[W_DIM]};
@@ -161,6 +168,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2SplitDTiling::PostTiling()
161 return ge::GRAPH_SUCCESS;168 return ge::GRAPH_SUCCESS;
162}169}
163 170 
164-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2SplitDTiling, 3);171+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2SplitDTiling, 13);
165 172 
166} // namespace optiling173} // namespace optiling
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_splith.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18// 1, splitD=1, splitH=1, splitW=0, splitKernel = 0, dtype=float19// 1, splitD=1, splitH=1, splitW=0, splitKernel = 0, dtype=float
@@ -26,6 +27,12 @@ namespace optiling {
26 27 
27bool MaxPool3DWithArgmaxV2SplitHTiling::IsCapable()28bool MaxPool3DWithArgmaxV2SplitHTiling::IsCapable()
28{29{
30+ auto platformInfo = context_->GetPlatformInfo();
31+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
32+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
33+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
34+ return false;
35+ }
29 splitData.partD = inputData.dilation[D_DIM] * (inputData.kernelSize[D_DIM] - 1) + 1;36 splitData.partD = inputData.dilation[D_DIM] * (inputData.kernelSize[D_DIM] - 1) + 1;
30 splitData.partOutD = 1UL;37 splitData.partOutD = 1UL;
31 array<uint64_t, DHW_DIMS> parts{38 array<uint64_t, DHW_DIMS> parts{
@@ -167,6 +174,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2SplitHTiling::PostTiling()
167 return ge::GRAPH_SUCCESS;174 return ge::GRAPH_SUCCESS;
168}175}
169 176 
170-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2SplitHTiling, 4);177+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2SplitHTiling, 14);
171 178 
172} // namespace optiling179} // namespace optiling
Mpooling/max_pool3d_with_argmax_v2/op_host/max_pool3d_with_argmax_v2_tiling_splitw.cpp+9-2
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include "tiling_base/tiling_util.h"
16#include "max_pool3d_with_argmax_v2_tiling_base.h"17#include "max_pool3d_with_argmax_v2_tiling_base.h"
17 18 
18// 1, splitD=1, splitH=1, splitW=1, splitKernel = 0, dtype=float19// 1, splitD=1, splitH=1, splitW=1, splitKernel = 0, dtype=float
@@ -26,6 +27,12 @@ namespace optiling {
26 27 
27bool MaxPool3DWithArgmaxV2SplitWTiling::IsCapable()28bool MaxPool3DWithArgmaxV2SplitWTiling::IsCapable()
28{29{
30+ auto platformInfo = context_->GetPlatformInfo();
31+ OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
32+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
33+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context_)){
34+ return false;
35+ }
29 splitData.partD = inputData.dilation[D_DIM] * (inputData.kernelSize[D_DIM] - 1) + 1;36 splitData.partD = inputData.dilation[D_DIM] * (inputData.kernelSize[D_DIM] - 1) + 1;
30 splitData.partOutD = 1UL;37 splitData.partOutD = 1UL;
31 splitData.partH = inputData.dilation[H_DIM] * (inputData.kernelSize[H_DIM] - 1) + 1;38 splitData.partH = inputData.dilation[H_DIM] * (inputData.kernelSize[H_DIM] - 1) + 1;
@@ -176,6 +183,6 @@ ge::graphStatus MaxPool3DWithArgmaxV2SplitWTiling::PostTiling()
176 return ge::GRAPH_SUCCESS;183 return ge::GRAPH_SUCCESS;
177}184}
178 185 
179-REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2SplitWTiling, 5);186+REGISTER_TILING_TEMPLATE("MaxPool3DWithArgmaxV2", MaxPool3DWithArgmaxV2SplitWTiling, 15);
180 187 
181} // namespace optiling188} // namespace optiling
Apooling/max_pool3d_with_argmax_v2/op_kernel/arch35/max_pool3d_with_argmax_v2_base.h+120-0
@@ -0,0 +1,120 @@
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 max_pool3d_with_argmax_v2_base.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL3D_WITH_ARGMAX_V2_BASE_H_
17+#define MAX_POOL3D_WITH_ARGMAX_V2_BASE_H_
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+ 
23+using namespace AscendC;
24+ 
25+template <typename T>
26+__aicore__ inline void CalGatterIndex2D(MicroAPI::RegTensor<T>& indexReg, T rate2D, T num1D, T rate1D = 1)
27+{
28+ AscendC::MicroAPI::Arange(indexReg, 0);
29+ AscendC::MicroAPI::RegTensor<T> segScalarReg;
30+ AscendC::MicroAPI::RegTensor<T> tmpReg;
31+ AscendC::MicroAPI::RegTensor<T> constReg;
32+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
33+ AscendC::MicroAPI::Duplicate(constReg, T(num1D));
34+ AscendC::MicroAPI::Div(segScalarReg, indexReg, constReg, preg);
35+ AscendC::MicroAPI::Muls(tmpReg, segScalarReg, T(num1D), preg);
36+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
37+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rate1D), preg);
38+ AscendC::MicroAPI::Muls(segScalarReg, segScalarReg, T(rate2D), preg);
39+ AscendC::MicroAPI::Add(indexReg, indexReg, segScalarReg, preg);
40+}
41+ 
42+template <typename T>
43+__aicore__ inline void CalGatterIndex3D(MicroAPI::RegTensor<T>& indexReg, T rate3D, T num2D, T rate2D, T num1D, T rate1D = 1)
44+{
45+ AscendC::MicroAPI::Arange(indexReg, 0);
46+ AscendC::MicroAPI::RegTensor<T> segScalarReg;
47+ AscendC::MicroAPI::RegTensor<T> segScalarReg2;
48+ AscendC::MicroAPI::RegTensor<T> tmpReg;
49+ AscendC::MicroAPI::RegTensor<T> constReg;
50+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
51+ AscendC::MicroAPI::Duplicate(constReg, T(num2D));
52+ AscendC::MicroAPI::Div(segScalarReg2, indexReg, constReg, preg);
53+ AscendC::MicroAPI::Muls(tmpReg, segScalarReg2, T(num2D), preg);
54+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
55+ AscendC::MicroAPI::Muls(segScalarReg2, segScalarReg2, T(rate3D), preg);
56+ 
57+ AscendC::MicroAPI::Duplicate(constReg, T(num1D));
58+ AscendC::MicroAPI::Div(segScalarReg, indexReg, constReg, preg);
59+ AscendC::MicroAPI::Muls(tmpReg, segScalarReg, T(num1D), preg);
60+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
61+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rate1D), preg);
62+ AscendC::MicroAPI::Muls(segScalarReg, segScalarReg, T(rate2D), preg);
63+ 
64+ AscendC::MicroAPI::Add(indexReg, indexReg, segScalarReg, preg);
65+ AscendC::MicroAPI::Add(indexReg, indexReg, segScalarReg2, preg);
66+}
67+ 
68+template <typename T>
69+__aicore__ inline void CalGatterIndex4D(MicroAPI::RegTensor<T>& indexReg, T rate4D, T num3D, T rate3D, T num2D, T rate2D, T num1D, T rate1D = 1)
70+{
71+ AscendC::MicroAPI::Arange(indexReg, 0);
72+ AscendC::MicroAPI::RegTensor<T> segScalarReg;
73+ AscendC::MicroAPI::RegTensor<T> segScalarReg2;
74+ AscendC::MicroAPI::RegTensor<T> segScalarReg3;
75+ AscendC::MicroAPI::RegTensor<T> tmpReg;
76+ AscendC::MicroAPI::RegTensor<T> constReg;
77+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
78+ AscendC::MicroAPI::Duplicate(constReg, T(num3D));
79+ AscendC::MicroAPI::Div(segScalarReg3, indexReg, constReg, preg);
80+ AscendC::MicroAPI::Muls(tmpReg, segScalarReg3, T(num3D), preg);
81+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
82+ AscendC::MicroAPI::Muls(segScalarReg3, segScalarReg3, T(rate4D), preg);
83+ 
84+ AscendC::MicroAPI::Duplicate(constReg, T(num2D));
85+ AscendC::MicroAPI::Div(segScalarReg2, indexReg, constReg, preg);
86+ AscendC::MicroAPI::Muls(tmpReg, segScalarReg2, T(num2D), preg);
87+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
88+ AscendC::MicroAPI::Muls(segScalarReg2, segScalarReg2, T(rate3D), preg);
89+ 
90+ AscendC::MicroAPI::Duplicate(constReg, T(num1D));
91+ AscendC::MicroAPI::Div(segScalarReg, indexReg, constReg, preg);
92+ AscendC::MicroAPI::Muls(tmpReg, segScalarReg, T(num1D), preg);
93+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
94+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rate1D), preg);
95+ AscendC::MicroAPI::Muls(segScalarReg, segScalarReg, T(rate2D), preg);
96+ 
97+ AscendC::MicroAPI::Add(indexReg, indexReg, segScalarReg, preg);
98+ AscendC::MicroAPI::Add(indexReg, indexReg, segScalarReg2, preg);
99+ AscendC::MicroAPI::Add(indexReg, indexReg, segScalarReg3, preg);
100+}
101+ 
102+template <typename T>
103+__aicore__ inline void SetNegInfReg(MicroAPI::RegTensor<T>& negInfReg)
104+{
105+ // -inf
106+ constexpr uint32_t FLOAT32_NEG_INF = 0xFF800000;
107+ constexpr uint16_t FLOAT16_NEG_INF = 0xFC00;
108+ constexpr uint16_t BFLOAT16_NEG_INF = 0xFF80;
109+ using computeType = std::conditional_t<std::is_same<T, float>::value, uint32_t, uint16_t>;
110+ 
111+ if constexpr(std::is_same<T, float>::value) {
112+ AscendC::MicroAPI::Duplicate((AscendC::MicroAPI::RegTensor<computeType>&)negInfReg, (FLOAT32_NEG_INF));
113+ } else if constexpr(std::is_same<T, half>::value) {
114+ AscendC::MicroAPI::Duplicate((AscendC::MicroAPI::RegTensor<computeType>&)negInfReg, (FLOAT16_NEG_INF));
115+ } else {
116+ AscendC::MicroAPI::Duplicate((AscendC::MicroAPI::RegTensor<computeType>&)negInfReg, (BFLOAT16_NEG_INF));
117+ }
118+}
119+ 
120+#endif // MAX_POOL_WITH_ARGMAX_V3_BASE_H_