Star914
1830
代码
代码
Issues230
Pull Requests316
Actions
项目讨论
Wiki
分析
Star914
1830
已合并
maxPoolGrad simd模板 #4252
liuchuangdev创建于 4月27日
maxPoolGrad simd模板 #4252
已合并
liuchuangdev创建于 4月27日
从liuchuangdev/ops-nnmax_pool_grad_merge合入到cann/ops-nnmaster
共 23 个文件变更+4605-1817
Mclassify_rule.yaml+2-0
@@ -1934,6 +1934,8 @@ VC7@ops-nn:
1934 - ops/ops-nn/pooling/adaptive_avg_pool2d/op_host/1934 - ops/ops-nn/pooling/adaptive_avg_pool2d/op_host/
1935 - ops/ops-nn/pooling/max_pool3d_grad/op_host/1935 - ops/ops-nn/pooling/max_pool3d_grad/op_host/
1936 - ops/ops-nn/pooling/max_pool3d_grad/op_kernel/1936 - ops/ops-nn/pooling/max_pool3d_grad/op_kernel/
1937+ - ops/ops-nn/pooling/max_pool_grad/op_host/
1938+ - ops/ops-nn/pooling/max_pool_grad/op_kernel/
1937 opensource_style: null1939 opensource_style: null
1938 kernel_style: null1940 kernel_style: null
1939 unrelease:1941 unrelease:
Mpooling/max_pool_grad/CMakeLists.txt+1-1
@@ -13,4 +13,4 @@ set(SUPPORT_COMPUTE_UNIT "ascend950")
13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14set(SUPPORT_TILING_DIR "arch35")14set(SUPPORT_TILING_DIR "arch35")
15add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE max_pool_grad 15add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE max_pool_grad
16- ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE DEPENDENCIES pool_grad_common)16+ ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE DEPENDENCIES pool_grad_common pool_3d_common max_pool_with_argmax_v3)
Apooling/max_pool_grad/op_host/arch35/max_pool_grad_nchw_tiling.cpp+177-0
@@ -0,0 +1,177 @@
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_grad_nchw_tiling.cpp
13+ * \brief NCHW format MaxPoolGrad unified tiling
14+ */
15+#include "platform/platform_info.h"
16+#include "op_host/tiling_templates_registry.h"
17+#include "max_pool_grad_tiling.h"
18+ 
19+using namespace PoolGradNameSpace;
20+ 
21+namespace optiling {
22+static constexpr int64_t KSIZE_THRESHOLD = 128;
23+static constexpr int64_t FLOAT16_SIZE = 2;
24+static constexpr int64_t FLOAT32_SIZE = 4;
25+static constexpr int64_t INT32_SIZE = 4;
26+static constexpr int64_t INT64_SIZE = 8;
27+static constexpr int64_t DOUBLE_BUFFER = 2;
28+static constexpr int64_t DOUBLE = 2;
29+ 
30+void MaxPoolGradNCHWTilingHelper::DoBufferCalculate()
31+{
32+ if (inputData->hKernel * inputData->wKernel < KSIZE_THRESHOLD) {
33+ int64_t hInputInner = Ops::Base::CeilDiv(splitData.hOutputInner + inputData->hKernel - 1, inputData->hStride);
34+ int64_t wInputInner = Ops::Base::CeilDiv(splitData.wOutputInner + inputData->wKernel - 1, inputData->wStride);
35+ int64_t wInputInnerAligned = Ops::Base::CeilAlign(wInputInner, baseData.maxDataNumInOneBlock);
36+ int64_t wOutputInnerAligned = Ops::Base::CeilAlign(splitData.wOutputInner, baseData.maxDataNumInOneBlock);
37+ int64_t wInputAligned = Ops::Base::CeilAlign(
38+ splitData.wOutputInner + ((inputData->wKernel - KERNEL_OFFSET) * inputData->wDilation) * DOUBLE,
39+ baseData.maxDataNumInOneBlock);
40+ 
41+ int64_t inputPlaneSizeHW = hInputInner * wInputInnerAligned;
42+ int64_t outputPlaneSizeHW = splitData.hOutputInner * wOutputInnerAligned;
43+ 
44+ splitData.inputBufferSize =
45+ splitData.highAxisInner *
46+ (splitData.hOutputInner + ((inputData->hKernel - KERNEL_OFFSET) * inputData->hDilation) * DOUBLE) *
47+ wInputAligned * baseData.inputBytes;
48+ splitData.gradBufferSize = splitData.highAxisInner * inputPlaneSizeHW * baseData.inputBytes;
49+ splitData.argmaxBufferSize =
50+ splitData.highAxisInner * inputPlaneSizeHW * (inputData->isInt32Meet ? INT64_SIZE : INT32_SIZE);
51+ splitData.outputBufferSize = splitData.highAxisInner * outputPlaneSizeHW * FLOAT32_SIZE;
52+ 
53+ int64_t tmpTotalBufferSize = splitData.inputBufferSize + splitData.outputBufferSize + splitData.gradBufferSize;
54+ splitData.totalBufferSize = tmpTotalBufferSize * DOUBLE_BUFFER + splitData.argmaxBufferSize;
55+ if (baseData.isPad == 1) {
56+ splitData.totalBufferSize += splitData.inputBufferSize;
57+ }
58+ } else {
59+ const int64_t hArgmaxInner = std::min<int64_t>(
60+ inputData->hGrad,
61+ Ops::Base::CeilDiv(splitData.hOutputInner + inputData->hKernel - 1, inputData->hStride));
62+ const int64_t wArgmaxInner = std::min<int64_t>(
63+ inputData->wGrad,
64+ Ops::Base::CeilDiv(splitData.wOutputInner + inputData->wKernel - 1, inputData->wStride));
65+ 
66+ const int64_t wArgmaxInnerAligned = Ops::Base::CeilAlign(wArgmaxInner, baseData.maxDataNumInOneBlock);
67+ const int64_t wOutputInnerAligned = Ops::Base::CeilAlign(splitData.wOutputInner, baseData.maxDataNumInOneBlock);
68+ 
69+ const int64_t gradPlaneSizeHW = hArgmaxInner * wArgmaxInnerAligned;
70+ 
71+ const int64_t argmaxPlaneSizeHW = hArgmaxInner * wArgmaxInner;
72+ 
73+ const int64_t outputPlaneSizeHW = splitData.hOutputInner * wOutputInnerAligned;
74+ 
75+ splitData.gradBufferSize = splitData.highAxisInner * gradPlaneSizeHW * baseData.inputBytes;
76+ splitData.argmaxBufferSize = splitData.highAxisInner * argmaxPlaneSizeHW * baseData.indexBytes;
77+ splitData.outputBufferSize = splitData.highAxisInner * outputPlaneSizeHW * BIG_FLOAT32_SIZE;
78+ 
79+ const int64_t fullKernelCount = inputData->hKernel * inputData->wKernel;
80+ const int64_t fullKernelBytes = fullKernelCount * baseData.inputBytes;
81+ const int64_t forwardInputAvailableBytes =
82+ std::max<int64_t>(0, baseData.availableUb - splitData.argmaxBufferSize - BIG_MERGE_BUF_ALIGN);
83+ const int64_t maxLoadCount = std::max<int64_t>(1, forwardInputAvailableBytes / baseData.inputBytes);
84+ 
85+ if (fullKernelBytes <= forwardInputAvailableBytes) {
86+ splitData.inputBufferSize = fullKernelBytes;
87+ } else {
88+ splitData.inputBufferSize = maxLoadCount * baseData.inputBytes;
89+ }
90+ 
91+ const int64_t forwardStageBufferSize = splitData.inputBufferSize * BIG_DOUBLE_BUFFER + BIG_MERGE_BUF_ALIGN;
92+ const int64_t backwardStageBufferSize =
93+ splitData.gradBufferSize * BIG_DOUBLE_BUFFER + splitData.outputBufferSize * BIG_DOUBLE_BUFFER;
94+ 
95+ splitData.totalBufferSize =
96+ splitData.argmaxBufferSize + std::max<int64_t>(forwardStageBufferSize, backwardStageBufferSize);
97+ }
98+}
99+ 
100+bool MaxPoolGradNCHWTiling::IsCapable()
101+{
102+ if (inputData.inputFormat != ge::Format::FORMAT_NCHW) {
103+ OP_LOGI("IsCapable", "inputFormat error, expected NCHW");
104+ return false;
105+ }
106+ if (inputData.hDilation != 1 || inputData.wDilation != 1) {
107+ OP_LOGI("IsCapable", "hDilation:%ld, wDilation:%ld", inputData.hDilation, inputData.wDilation);
108+ return false;
109+ }
110+ commonTiling.InitializationVars(context_, &hwInfo);
111+ return commonTiling.CheckUBSize() && (inputData.isInt32Meet == 0);
112+}
113+ 
114+uint64_t MaxPoolGradNCHWTiling::GetTilingKey() const
115+{
116+ uint32_t indicesDtype = (inputData.hX * inputData.wX <= static_cast<int64_t>(MAX_INT32)) ? TPL_INT32 : TPL_INT64;
117+ uint32_t format = TPL_NCHW_FORMAT;
118+ uint32_t isCheckRange = commonTiling.GetSplitData().isCheckRange;
119+ uint32_t kernelMode =
120+ (inputData.hKernel * inputData.wKernel >= KSIZE_THRESHOLD) ? TPL_NCHW_BIG_KERNEL : TPL_NCHW_SMALL_KERNEL;
121+ return GET_TPL_TILING_KEY(kernelMode, format, indicesDtype, isCheckRange);
122+}
123+ 
124+ge::graphStatus MaxPoolGradNCHWTiling::DoOpTiling()
125+{
126+ commonTiling.InitializationVars(context_, &hwInfo);
127+ return commonTiling.DoOpTiling(context_, GetTilingKey());
128+}
129+ 
130+ge::graphStatus MaxPoolGradNCHWTiling::PostTiling()
131+{
132+ return commonTiling.PostTiling(context_);
133+}
134+ 
135+ge::graphStatus MaxPoolGradNCHWTiling::GetShapeAttrsInfo()
136+{
137+ auto ret = MaxPoolGradTilingBase::GetShapeAttrsInfo();
138+ if (ret != ge::GRAPH_SUCCESS) {
139+ return ret;
140+ }
141+ 
142+ if (inputData.inputFormat != ge::Format::FORMAT_NCHW) {
143+ OP_LOGI("GetShapeAttrsInfo", "inputFormat error, expected NCHW");
144+ return ge::GRAPH_PARAM_INVALID;
145+ }
146+ 
147+ return ge::GRAPH_SUCCESS;
148+}
149+ 
150+ge::graphStatus MaxPoolGradNCHWTiling::GetPlatformInfo()
151+{
152+ auto platformPtr = context_->GetPlatformInfo();
153+ if (platformPtr == nullptr) {
154+ auto compileInfoPtr = static_cast<const MaxPoolGradWithArgmaxCompileInfo*>(context_->GetCompileInfo());
155+ OP_CHECK_IF(
156+ compileInfoPtr == nullptr, OP_LOGE(context_->GetNodeName(), "compile info is null"),
157+ return ge::GRAPH_FAILED);
158+ hwInfo.coreNum = compileInfoPtr->coreNum;
159+ hwInfo.ubSize = compileInfoPtr->ubSize;
160+ } else {
161+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformPtr);
162+ hwInfo.coreNum = ascendcPlatform.GetCoreNumAiv();
163+ 
164+ uint64_t ubSizePlatform;
165+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
166+ hwInfo.ubSize = static_cast<int64_t>(ubSizePlatform);
167+ }
168+ 
169+ OP_CHECK_IF(hwInfo.coreNum == 0, OP_LOGE(context_->GetNodeName(), "coreNum is 0"), return ge::GRAPH_FAILED);
170+ coreNum_ = hwInfo.coreNum;
171+ ubSize_ = hwInfo.ubSize;
172+ return ge::GRAPH_SUCCESS;
173+}
174+ 
175+REGISTER_TILING_TEMPLATE("MaxPoolGrad", MaxPoolGradNCHWTiling, 0);
176+ 
177+} // namespace optiling
Mpooling/max_pool_grad/op_host/arch35/max_pool_grad_tiling.h+102-61
@@ -1,62 +1,103 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.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 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.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.8+ * See LICENSE in the root of the software repository for the full text of the License.
9- */9+ */
10- 10+ 
11-/*!11+/*!
12- * \file max_pool_grad_tiling.h12+ * \file max_pool_grad_tiling.h
13- * \brief13+ * \brief
14- */14+ */
15- 15+ 
16-#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_TILING_H_16+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_TILING_H_
17-#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_TILING_H_17+#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_TILING_H_
18- 18+ 
19-#include "../../op_kernel/arch35/max_pool_grad_struct.h"19+#include "../../op_kernel/arch35/max_pool_grad_struct.h"
20-#include "../../../pool_grad_common/op_host/arch35/max_pool_grad_with_argmax_tiling_common.h"20+#include "../../../pool_grad_common/op_host/arch35/max_pool_grad_with_argmax_tiling_common.h"
21-#include "../../../pool_grad_common/op_kernel/arch35/max_pool_grad_with_argmax_struct_common.h"21+#include "../../../pool_grad_common/op_kernel/arch35/max_pool_grad_with_argmax_struct_common.h"
22-#include "../../../pool_grad_common/op_host/arch35/util.h"22+#include "../../../pool_grad_common/op_host/arch35/max_pool_grad_nchw_tiling_common.h"
23-#include "platform/platform_info.h"23+#include "../../../pool_grad_common/op_host/arch35/util.h"
24-#include "atvoss/broadcast/broadcast_tiling.h"24+#include "platform/platform_info.h"
25-#include "op_common/op_host/util/platform_util.h"25+#include "atvoss/broadcast/broadcast_tiling.h"
26-#include "register/op_def_registry.h"26+#include "op_common/op_host/util/platform_util.h"
27- 27+#include "register/op_def_registry.h"
28-namespace optiling {28+ 
29-using Ops::NN::Optiling::TilingBaseClass;29+namespace optiling {
30-using namespace MaxPoolGradWithArgmaxNHWCNameSpace;30+using Ops::NN::Optiling::TilingBaseClass;
31- 31+using namespace MaxPoolGradWithArgmaxNHWCNameSpace;
32-static constexpr int64_t KERNEL_OFFSET = 1;32+ 
33- 33+static constexpr int64_t KERNEL_OFFSET = 1;
34-class MaxPoolGradTilingBase : public MaxPoolGradWithArgmaxTilingCommon {34+ 
35-public:35+constexpr int64_t BIG_FLOAT16_SIZE = 2;
36- explicit MaxPoolGradTilingBase(gert::TilingContext* context) : MaxPoolGradWithArgmaxTilingCommon(context)36+constexpr int64_t BIG_FLOAT32_SIZE = 4;
37- {}37+constexpr int64_t BIG_INT32_SIZE = 4;
38- ~MaxPoolGradTilingBase() override38+constexpr int64_t BIG_INT64_SIZE = 8;
39- {}39+constexpr int64_t BIG_UB_RESERVED_SIZE = 1024;
40- 40+constexpr int64_t BIG_DOUBLE_BUFFER = 2;
41- const std::string nodeName = "MaxPoolGrad";41+ 
42- MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxSimtTilingCommonData* tilingData_ =42+constexpr int64_t BIG_KERNEL_THRESHOLD = 128;
43- context_->GetTilingData<MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxSimtTilingCommonData>();43+constexpr int64_t BIG_MAX_KERNEL_COUNT = 512;
44- MaxPoolGradWithArgmaxInputInfoCommon inputData;44+constexpr int64_t BIG_MERGE_BUF_ALIGN = 32;
45- int64_t coreNum_{0};45+ 
46- int64_t ubSize_{0};46+class MaxPoolGradTilingBase : public MaxPoolGradWithArgmaxTilingCommon {
47- 47+public:
48- bool CheckInputShape();48+ explicit MaxPoolGradTilingBase(gert::TilingContext* context) : MaxPoolGradWithArgmaxTilingCommon(context)
49- ge::graphStatus CheckInputDtype();49+ {}
50- ge::graphStatus CheckAttrShape();50+ ~MaxPoolGradTilingBase() override
51- ge::graphStatus CheckAttrVal();51+ {}
52- ge::graphStatus CheckInputValid();52+ 
53- ge::graphStatus SetInputParams();53+ const std::string nodeName = "MaxPoolGrad";
54- ge::graphStatus SetAttrParams();54+ MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxSimtTilingCommonData* tilingData_ =
55- void SetOtherInputParams();55+ context_->GetTilingData<MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxSimtTilingCommonData>();
56- 56+ MaxPoolGradWithArgmaxInputInfoCommon inputData;
57-protected:57+ int64_t coreNum_{0};
58- ge::graphStatus GetShapeAttrsInfo() override;58+ int64_t ubSize_{0};
59-};59+ 
60-} // namespace optiling60+ bool CheckInputShape();
61- 61+ ge::graphStatus CheckInputDtype();
62+ ge::graphStatus CheckAttrShape();
63+ ge::graphStatus CheckAttrVal();
64+ ge::graphStatus CheckInputValid();
65+ ge::graphStatus SetInputParams();
66+ ge::graphStatus SetAttrParams();
67+ void SetOtherInputParams();
68+ 
69+protected:
70+ ge::graphStatus GetShapeAttrsInfo() override;
71+};
72+ 
73+class MaxPoolGradNCHWTilingHelper : public MaxPoolGradNCHWTilingCommon {
74+public:
75+ MaxPoolGradNCHWTilingHelper(MaxPoolGradWithArgmaxInputInfoCommon* input) : MaxPoolGradNCHWTilingCommon(input)
76+ {}
77+ 
78+protected:
79+ void DoBufferCalculate() override;
80+};
81+ 
82+class MaxPoolGradNCHWTiling : public MaxPoolGradTilingBase {
83+public:
84+ explicit MaxPoolGradNCHWTiling(gert::TilingContext* context) : MaxPoolGradTilingBase(context)
85+ {}
86+ 
87+ ~MaxPoolGradNCHWTiling() override
88+ {}
89+ 
90+private:
91+ MaxPoolGradNCHWTilingHelper commonTiling{&inputData};
92+ MaxPoolGradWithArgmaxHardwareInfo hwInfo;
93+ uint64_t GetTilingKey() const override;
94+ bool IsCapable() override;
95+ ge::graphStatus DoOpTiling() override;
96+ ge::graphStatus PostTiling() override;
97+ ge::graphStatus GetShapeAttrsInfo() override;
98+ ge::graphStatus GetPlatformInfo() override;
99+};
100+ 
101+} // namespace optiling
102+ 
62#endif103#endif
Mpooling/max_pool_grad/op_host/arch35/max_pool_grad_tiling_base.cpp+2-3
@@ -41,7 +41,6 @@ constexpr uint64_t NUM_TWO = 2;
41constexpr size_t HW_DIM_NUM = 3;41constexpr size_t HW_DIM_NUM = 3;
42constexpr uint32_t MAX_BLOCK_COUNT = 4095;42constexpr uint32_t MAX_BLOCK_COUNT = 4095;
43 43 
44-// 参数常量
45constexpr size_t NC_DIM_NUM = 2;44constexpr size_t NC_DIM_NUM = 2;
46constexpr size_t NCHW_DIM_NUM = 4;45constexpr size_t NCHW_DIM_NUM = 4;
47constexpr size_t PADS_ATTR_INDEX = 3U;46constexpr size_t PADS_ATTR_INDEX = 3U;
@@ -245,11 +244,11 @@ ge::graphStatus MaxPoolGradTilingBase::CheckAttrShape()
245 // Check attr dim num244 // Check attr dim num
246 OP_CHECK_IF(245 OP_CHECK_IF(
247 (kSizeDimNum != NCHW_DIM_NUM),246 (kSizeDimNum != NCHW_DIM_NUM),
248- OP_LOGE(context_->GetNodeName(), "Attr kSize dim num invalid, dim num should equal 5."),247+ OP_LOGE(context_->GetNodeName(), "Attr kSize dim num invalid, dim num should equal 4."),
249 return ge::GRAPH_FAILED);248 return ge::GRAPH_FAILED);
250 OP_CHECK_IF(249 OP_CHECK_IF(
251 (stridesDimNum != NCHW_DIM_NUM),250 (stridesDimNum != NCHW_DIM_NUM),
252- OP_LOGE(context_->GetNodeName(), "Attr strides dim num invalid, dim num should equal 5."),251+ OP_LOGE(context_->GetNodeName(), "Attr strides dim num invalid, dim num should equal 4."),
253 return ge::GRAPH_FAILED);252 return ge::GRAPH_FAILED);
254 253 
255 return ge::GRAPH_SUCCESS;254 return ge::GRAPH_SUCCESS;
Apooling/max_pool_grad/op_kernel/arch35/max_pool_grad_nchw_backward_base.h+880-0
@@ -0,0 +1,880 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef MAX_POOL_GRAD_NCHW_BACKWARD_BASE_H_
12+#define MAX_POOL_GRAD_NCHW_BACKWARD_BASE_H_
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "max_pool_grad_struct.h"
17+#include "../inc/platform.h"
18+#include "../pool_grad_common/arch35/max_pool_grad_with_argmax_base_common.h"
19+#include "../pool_grad_common/arch35/max_pool_grad_nchw_scatter_common.h"
20+ 
21+namespace MaxPoolGradNCHWBackwardBaseNameSpace {
22+using namespace AscendC;
23+using MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData;
24+ 
25+using ::PEnd;
26+using ::PStart;
27+using MaxPoolGradNCHWNameSpace::DoMulNCNchw;
28+using MaxPoolGradNCHWNameSpace::DoSingleNCNchw;
29+using MaxPoolGradNCHWNameSpace::Gen2DIndexOne;
30+using MaxPoolGradNCHWNameSpace::Gen3DIndexOne;
31+using MaxPoolGradNCHWNameSpace::GenInitial1DIndices;
32+using MaxPoolGradNCHWNameSpace::GenInitial2DIndices;
33+using MaxPoolGradNCHWNameSpace::GenInitial3DIndices;
34+ 
35+constexpr int32_t BK_BUFFER_NUM = 2;
36+constexpr int64_t RATIO = 2;
37+constexpr int32_t HELP_BUFFER = 4096;
38+constexpr uint32_t DOUBLE = 2;
39+ 
40+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
41+class MaxPoolGradNCHWBackwardBase {
42+public:
43+ __aicore__ inline MaxPoolGradNCHWBackwardBase()
44+ {}
45+ 
46+protected:
47+ __aicore__ inline void ParseTilingData(const MaxPoolGradWithArgmaxNCHWTilingCommonData& tilingData);
48+ __aicore__ inline void ScalarCompute(int64_t loopNum);
49+ __aicore__ inline void CopyInGrad();
50+ __aicore__ inline void CopyOut();
51+ __aicore__ inline void ProcessNoArgmaxBlock();
52+ __aicore__ inline void BackwardCompute(
53+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr);
54+ __aicore__ inline void singleLineProcessVF(
55+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
56+ __local_mem__ uint32_t* helpAddr);
57+ __aicore__ inline void multipleLineHwProcessVF(
58+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
59+ __local_mem__ uint32_t* helpAddr);
60+ __aicore__ inline void multipleLineProcessVF2(
61+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
62+ __local_mem__ uint32_t* helpAddr);
63+ 
64+ TQue<QuePosition::VECIN, BK_BUFFER_NUM> gradQue_;
65+ TQue<QuePosition::VECOUT, BK_BUFFER_NUM> outputQue_;
66+ TBuf<QuePosition::VECCALC> argmaxBuf_;
67+ TBuf<QuePosition::VECCALC> helpBuf_;
68+ 
69+ GlobalTensor<T1> gradGm_;
70+ GlobalTensor<T1> yGm_;
71+ 
72+ uint32_t blockIdx_ = 0;
73+ 
74+ int64_t hArgmax_ = 1;
75+ int64_t wArgmax_ = 1;
76+ int64_t hOutput_ = 1;
77+ int64_t wOutput_ = 1;
78+ int64_t kernelH_ = 1;
79+ int64_t kernelW_ = 1;
80+ int64_t strideH_ = 1;
81+ int64_t strideW_ = 1;
82+ int64_t padH_ = 0;
83+ int64_t padW_ = 0;
84+ int64_t dilationH_ = 1;
85+ int64_t dilationW_ = 1;
86+ int64_t highAxisInner_ = 1;
87+ int64_t highAxisTail_ = 1;
88+ int64_t highAxisOuter_ = 1;
89+ int64_t highAxisActual_ = 1;
90+ int64_t hOutputInner_ = 1;
91+ int64_t hOutputTail_ = 1;
92+ int64_t hOutputOuter_ = 1;
93+ int64_t hOutputActual_ = 1;
94+ int64_t wOutputInner_ = 1;
95+ int64_t wOutputTail_ = 1;
96+ int64_t wOutputOuter_ = 1;
97+ int64_t wOutputActual_ = 1;
98+ int64_t wOutputAligned_ = 1;
99+ int64_t normalCoreProcessNum_ = 1;
100+ int64_t tailCoreProcessNum_ = 1;
101+ int64_t curCoreProcessNum_ = 1;
102+ int64_t usedCoreNum_ = 1;
103+ int64_t inputBufferSize_ = 1;
104+ int64_t outputBufferSize_ = 1;
105+ int64_t gradBufferSize_ = 1;
106+ int64_t argmaxBufferSize_ = 1;
107+ int64_t highAxisIndex_ = 0;
108+ int64_t hAxisIndex_ = 0;
109+ int64_t wAxisIndex_ = 0;
110+ int64_t hArgmaxActual_ = 0;
111+ int64_t wArgmaxActual_ = 0;
112+ int64_t wArgmaxAligned_ = 0;
113+ int64_t hArgmaxActualStart_ = 0;
114+ int64_t wArgmaxActualStart_ = 0;
115+ int64_t highAxisArgmaxOffset_ = 0;
116+ int64_t hAxisArgmaxOffset_ = 0;
117+ int64_t wAxisArgmaxOffset_ = 0;
118+ int64_t argmaxPlaneSize_ = 1;
119+ int64_t hProBatchSize_ = 1;
120+ int64_t wProBatchSize_ = 1;
121+ int64_t curHProBatchSize_ = 1;
122+ int64_t curWProBatchSize_ = 1;
123+ 
124+ constexpr static int32_t BLOCK_SIZE = platform::GetUbBlockSize();
125+ constexpr static int32_t V_REG_SIZE = platform::GetVRegSize();
126+ constexpr static int64_t MAX_DATA_NUM_IN_ONE_BLOCK = BLOCK_SIZE / sizeof(T1);
127+ constexpr static uint16_t vlT2_ = platform::GetVRegSize() / sizeof(T2);
128+};
129+ 
130+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
131+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::ParseTilingData(
132+ const MaxPoolGradWithArgmaxNCHWTilingCommonData& tilingData)
133+{
134+ hArgmax_ = tilingData.hArgmax;
135+ wArgmax_ = tilingData.wArgmax;
136+ hOutput_ = tilingData.hOutput;
137+ wOutput_ = tilingData.wOutput;
138+ kernelH_ = tilingData.hKernel;
139+ kernelW_ = tilingData.wKernel;
140+ strideH_ = tilingData.hStride;
141+ strideW_ = tilingData.wStride;
142+ padH_ = tilingData.padH;
143+ padW_ = tilingData.padW;
144+ dilationH_ = tilingData.dilationH;
145+ dilationW_ = tilingData.dilationW;
146+ highAxisInner_ = tilingData.highAxisInner;
147+ highAxisTail_ = tilingData.highAxisTail;
148+ highAxisOuter_ = tilingData.highAxisOuter;
149+ hOutputInner_ = tilingData.hOutputInner;
150+ hOutputTail_ = tilingData.hOutputTail;
151+ hOutputOuter_ = tilingData.hOutputOuter;
152+ wOutputInner_ = tilingData.wOutputInner;
153+ wOutputTail_ = tilingData.wOutputTail;
154+ wOutputOuter_ = tilingData.wOutputOuter;
155+ normalCoreProcessNum_ = tilingData.normalCoreProcessNum;
156+ tailCoreProcessNum_ = tilingData.tailCoreProcessNum;
157+ usedCoreNum_ = tilingData.usedCoreNum;
158+ inputBufferSize_ = tilingData.inputBufferSize;
159+ outputBufferSize_ = tilingData.outputBufferSize;
160+ gradBufferSize_ = tilingData.gradBufferSize;
161+ argmaxBufferSize_ = tilingData.argmaxBufferSize;
162+ hProBatchSize_ = tilingData.hProBatchSize;
163+ wProBatchSize_ = tilingData.wProBatchSize;
164+ argmaxPlaneSize_ = hArgmax_ * wArgmax_;
165+}
166+ 
167+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
168+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::ScalarCompute(int64_t loopNum)
169+{
170+ int64_t baseBlockIdx = blockIdx_ * normalCoreProcessNum_ + loopNum;
171+ highAxisIndex_ = baseBlockIdx / (hOutputOuter_ * wOutputOuter_);
172+ highAxisActual_ = highAxisIndex_ == (highAxisOuter_ - 1) ? highAxisTail_ : highAxisInner_;
173+ int64_t tempTail = baseBlockIdx % (hOutputOuter_ * wOutputOuter_);
174+ hAxisIndex_ = tempTail / wOutputOuter_;
175+ hOutputActual_ = hAxisIndex_ == (hOutputOuter_ - 1) ? hOutputTail_ : hOutputInner_;
176+ wAxisIndex_ = tempTail % wOutputOuter_;
177+ wOutputActual_ = wAxisIndex_ == (wOutputOuter_ - 1) ? wOutputTail_ : wOutputInner_;
178+ wOutputAligned_ =
179+ (wOutputActual_ + MAX_DATA_NUM_IN_ONE_BLOCK - 1) / MAX_DATA_NUM_IN_ONE_BLOCK * MAX_DATA_NUM_IN_ONE_BLOCK;
180+ 
181+ int64_t hArgmaxActualStart = PStart(hAxisIndex_ * hOutputInner_, padH_, kernelH_, dilationH_, strideH_);
182+ int64_t hArgmaxActualEnd = PEnd(hAxisIndex_ * hOutputInner_ + hOutputActual_ - 1, padH_, strideH_, hArgmax_);
183+ int64_t wArgmaxActualStart = PStart(wAxisIndex_ * wOutputInner_, padW_, kernelW_, dilationW_, strideW_);
184+ int64_t wArgmaxActualEnd = PEnd(wAxisIndex_ * wOutputInner_ + wOutputActual_ - 1, padW_, strideW_, wArgmax_);
185+ wArgmaxActual_ = wArgmaxActualEnd - wArgmaxActualStart;
186+ wArgmaxAligned_ =
187+ (wArgmaxActual_ + MAX_DATA_NUM_IN_ONE_BLOCK - 1) / MAX_DATA_NUM_IN_ONE_BLOCK * MAX_DATA_NUM_IN_ONE_BLOCK;
188+ hArgmaxActual_ = hArgmaxActualEnd - hArgmaxActualStart;
189+ hArgmaxActualStart_ = hArgmaxActualStart;
190+ wArgmaxActualStart_ = wArgmaxActualStart;
191+ curHProBatchSize_ = hProBatchSize_ > hArgmaxActual_ ? hArgmaxActual_ : hProBatchSize_;
192+ curWProBatchSize_ = wProBatchSize_ > wArgmaxActual_ ? wArgmaxActual_ : wProBatchSize_;
193+ highAxisArgmaxOffset_ = highAxisIndex_ * highAxisInner_ * argmaxPlaneSize_;
194+ hAxisArgmaxOffset_ = hArgmaxActualStart * wArgmax_;
195+ wAxisArgmaxOffset_ = wArgmaxActualStart;
196+}
197+ 
198+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
199+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::CopyInGrad()
200+{
201+ LocalTensor<T1> gradLocal = gradQue_.AllocTensor<T1>();
202+ int64_t argmaxGmOffset = highAxisArgmaxOffset_ + hAxisArgmaxOffset_ + wAxisArgmaxOffset_;
203+ DataCopyPadExtParams<T1> paramsT1 = {false, 0, 0, 0};
204+ LoopModeParams loopModeParamsT1;
205+ loopModeParamsT1.loop1Size = highAxisActual_;
206+ loopModeParamsT1.loop1SrcStride = argmaxPlaneSize_ * sizeof(T1);
207+ loopModeParamsT1.loop1DstStride = hArgmaxActual_ * wArgmaxAligned_ * sizeof(T1);
208+ loopModeParamsT1.loop2Size = 1;
209+ loopModeParamsT1.loop2SrcStride = 0;
210+ loopModeParamsT1.loop2DstStride = 0;
211+ SetLoopModePara(loopModeParamsT1, DataCopyMVType::OUT_TO_UB);
212+ DataCopyExtParams copyOutParamT1 = {
213+ static_cast<uint16_t>(hArgmaxActual_), static_cast<uint32_t>(wArgmaxActual_ * sizeof(T1)),
214+ static_cast<uint32_t>((wArgmax_ - wArgmaxActual_) * sizeof(T1)), 0, 0};
215+ DataCopyPad(gradLocal, gradGm_[argmaxGmOffset], copyOutParamT1, paramsT1);
216+ ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
217+ gradQue_.EnQue(gradLocal);
218+}
219+ 
220+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
221+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::CopyOut()
222+{
223+ LocalTensor<T1> yLocal = outputQue_.DeQue<T1>();
224+ int64_t outputPlaneSize = hOutput_ * wOutput_;
225+ int64_t highOutputAxisOffset = highAxisIndex_ * highAxisInner_ * outputPlaneSize;
226+ int64_t hOutputAxisOffset = hAxisIndex_ * hOutputInner_ * wOutput_;
227+ int64_t wOutputAxisOffset = wAxisIndex_ * wOutputInner_;
228+ int64_t outputGmOffset = highOutputAxisOffset + hOutputAxisOffset + wOutputAxisOffset;
229+ 
230+ LoopModeParams loopModeParamsT1;
231+ loopModeParamsT1.loop1Size = highAxisActual_;
232+ loopModeParamsT1.loop1SrcStride = hOutputActual_ * wOutputAligned_ * sizeof(T1);
233+ loopModeParamsT1.loop1DstStride = outputPlaneSize * sizeof(T1);
234+ loopModeParamsT1.loop2Size = 1;
235+ loopModeParamsT1.loop2SrcStride = 0;
236+ loopModeParamsT1.loop2DstStride = 0;
237+ SetLoopModePara(loopModeParamsT1, DataCopyMVType::UB_TO_OUT);
238+ DataCopyExtParams copyOutParamT1 = {
239+ static_cast<uint16_t>(hOutputActual_), static_cast<uint32_t>(wOutputActual_ * sizeof(T1)), 0,
240+ static_cast<uint32_t>((wOutput_ - wOutputActual_) * sizeof(T1)), 0};
241+ DataCopyPad(yGm_[outputGmOffset], yLocal, copyOutParamT1);
242+ ResetLoopModePara(DataCopyMVType::UB_TO_OUT);
243+ outputQue_.FreeTensor(yLocal);
244+}
245+ 
246+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
247+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::ProcessNoArgmaxBlock()
248+{
249+ uint32_t calcCount = static_cast<uint32_t>(outputBufferSize_) / sizeof(T1);
250+ LocalTensor<T1> yLocal = outputQue_.AllocTensor<T1>();
251+ Duplicate(yLocal, T1(0), calcCount);
252+ outputQue_.EnQue(yLocal);
253+ CopyOut();
254+}
255+ 
256+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
257+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::BackwardCompute(
258+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr)
259+{
260+ uint32_t wConcurrentCount = wArgmaxActual_ / curWProBatchSize_;
261+ uint32_t hConcurrentCount = hArgmaxActual_ / curHProBatchSize_;
262+ LocalTensor<uint32_t> helpTensor = helpBuf_.Get<uint32_t>();
263+ __local_mem__ uint32_t* helpAddr = (__local_mem__ uint32_t*)helpTensor.GetPhyAddr();
264+ if (wConcurrentCount * DOUBLE * sizeof(T2) > V_REG_SIZE) {
265+ singleLineProcessVF(yAddr, gradAddr, argmaxAddr, helpAddr);
266+ } else if (wConcurrentCount * hConcurrentCount * DOUBLE * sizeof(T2) > V_REG_SIZE) {
267+ multipleLineHwProcessVF(yAddr, gradAddr, argmaxAddr, helpAddr);
268+ } else {
269+ multipleLineProcessVF2(yAddr, gradAddr, argmaxAddr, helpAddr);
270+ }
271+}
272+ 
273+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
274+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::singleLineProcessVF(
275+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
276+ __local_mem__ uint32_t* helpAddr)
277+{
278+ int64_t wOutput = wOutput_;
279+ int64_t wOutputActual = wOutputActual_;
280+ int64_t wOutputAligned = wOutputAligned_;
281+ int64_t hOutputActual = hOutputActual_;
282+ uint16_t highAxisActual = static_cast<uint16_t>(highAxisActual_);
283+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
284+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
285+ int64_t wArgmaxActual = wArgmaxActual_;
286+ int64_t wArgmaxAligned = wArgmaxAligned_;
287+ uint16_t hArgmaxActual = hArgmaxActual_;
288+ uint16_t wProBatchSize = curWProBatchSize_;
289+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
290+ uint16_t computeSizeT2 = V_REG_SIZE / sizeof(T2);
291+ uint16_t repeatimes = wFullBatchCount / computeSizeT2;
292+ uint16_t wRemain = wArgmaxActual - repeatimes * wProBatchSize * computeSizeT2;
293+ uint32_t wRemainBatchCount = wRemain / wProBatchSize;
294+ uint16_t wRemainTail = wRemain % wProBatchSize;
295+ uint32_t one = 1;
296+ uint32_t all = computeSizeT2;
297+ 
298+ __VEC_SCOPE__
299+ {
300+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
301+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
302+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
303+ if constexpr (IS_CHECK_RANGE == 1) {
304+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
305+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
306+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
307+ }
308+ 
309+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
310+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
311+ 
312+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegGrad;
313+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegIndex;
314+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegGrad;
315+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
316+ 
317+ AscendC::MicroAPI::MaskReg allMaskU32 =
318+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
319+ 
320+ GenInitial1DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegGrad, wProBatchSize);
321+ GenInitial1DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegIndex, wProBatchSize);
322+ 
323+ for (uint16_t highIdx = 0; highIdx < highAxisActual; ++highIdx) {
324+ uint32_t highArgmaxOffset = highIdx * hArgmaxActual * wArgmaxAligned;
325+ uint32_t highIndexOffset = highIdx * hArgmaxActual * wArgmaxActual;
326+ uint32_t highOutputOffset = highIdx * hOutputActual * wOutputAligned;
327+ for (uint16_t hIdx = 0; hIdx < hArgmaxActual; hIdx++) {
328+ for (uint16_t wRepeatIdx = 0; wRepeatIdx < repeatimes; wRepeatIdx++) {
329+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
330+ uint32_t offset =
331+ (wBatchIdx + wRepeatIdx * computeSizeT2 * wProBatchSize + hIdx * wArgmaxAligned +
332+ highArgmaxOffset);
333+ uint32_t indexOffset =
334+ (wBatchIdx + wRepeatIdx * computeSizeT2 * wProBatchSize + hIdx * wArgmaxActual +
335+ highIndexOffset);
336+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGrad, offset, allMaskU32);
337+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, indexOffset, allMaskU32);
338+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
339+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, all, wOutputConstReg,
340+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
341+ }
342+ }
343+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
344+ uint32_t offset =
345+ (wBatchIdx + repeatimes * computeSizeT2 * wProBatchSize + hIdx * wArgmaxAligned +
346+ highArgmaxOffset);
347+ uint32_t indexOffset =
348+ (wBatchIdx + repeatimes * computeSizeT2 * wProBatchSize + hIdx * wArgmaxActual +
349+ highIndexOffset);
350+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGrad, offset, allMaskU32);
351+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, indexOffset, allMaskU32);
352+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
353+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, wRemainBatchCount,
354+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg,
355+ hMaxReg);
356+ }
357+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
358+ uint32_t offset =
359+ (wBatchIdx + wRemainBatchCount * wProBatchSize + repeatimes * computeSizeT2 * wProBatchSize +
360+ hIdx * wArgmaxAligned + highArgmaxOffset);
361+ uint32_t indexOffset =
362+ (wBatchIdx + wRemainBatchCount * wProBatchSize + repeatimes * computeSizeT2 * wProBatchSize +
363+ hIdx * wArgmaxActual + highIndexOffset);
364+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGrad, offset, allMaskU32);
365+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, indexOffset, allMaskU32);
366+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
367+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, one, wOutputConstReg, curHIndex,
368+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
369+ }
370+ }
371+ }
372+ }
373+}
374+ 
375+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
376+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::multipleLineHwProcessVF(
377+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
378+ __local_mem__ uint32_t* helpAddr)
379+{
380+ int64_t wOutput = wOutput_;
381+ int64_t wOutputActual = wOutputActual_;
382+ int64_t wOutputAligned = wOutputAligned_;
383+ int64_t hOutputActual = hOutputActual_;
384+ uint16_t highAxisActual = static_cast<uint16_t>(highAxisActual_);
385+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
386+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
387+ int64_t wArgmaxAligned = wArgmaxAligned_;
388+ int64_t wArgmaxActual = wArgmaxActual_;
389+ uint16_t hArgmaxActual = hArgmaxActual_;
390+ 
391+ uint16_t hProBatchSize = curHProBatchSize_;
392+ uint16_t wProBatchSize = curWProBatchSize_;
393+ 
394+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
395+ uint16_t hFullBatchCount = hArgmaxActual / hProBatchSize;
396+ uint16_t wRemainTail = wArgmaxActual % wProBatchSize;
397+ 
398+ uint16_t hConcurrentCount = V_REG_SIZE / (wFullBatchCount * sizeof(T2));
399+ 
400+ uint16_t blockConcurrentCount = hFullBatchCount / hConcurrentCount;
401+ uint16_t hRemain = hArgmaxActual - blockConcurrentCount * hConcurrentCount * hProBatchSize;
402+ 
403+ uint16_t hRemainBatchCount = hRemain / hProBatchSize;
404+ uint16_t hRemainTail = hRemain - hRemainBatchCount * hProBatchSize;
405+ 
406+ uint32_t blockOne = 1 * hConcurrentCount;
407+ uint32_t remainBatchOne = 1 * hRemainBatchCount;
408+ uint32_t remainTailOne = 1;
409+ uint32_t maskBlock = wFullBatchCount * hConcurrentCount;
410+ uint32_t maskRemainBatch = wFullBatchCount * hRemainBatchCount;
411+ uint32_t maskRemainTail = wFullBatchCount;
412+ 
413+ __VEC_SCOPE__
414+ {
415+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
416+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
417+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
418+ if constexpr (IS_CHECK_RANGE == 1) {
419+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
420+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
421+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
422+ }
423+ 
424+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
425+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
426+ 
427+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegGrad;
428+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegGradOne;
429+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegIndex;
430+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegIndexOne;
431+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegGrad;
432+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
433+ 
434+ AscendC::MicroAPI::MaskReg allMaskU32 =
435+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
436+ GenInitial2DIndices(
437+ (AscendC::MicroAPI::RegTensor<int32_t>&)initialRegGrad, wProBatchSize, hProBatchSize, wArgmaxAligned,
438+ wFullBatchCount);
439+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegGradOne, hProBatchSize, wArgmaxAligned);
440+ GenInitial2DIndices(
441+ (AscendC::MicroAPI::RegTensor<int32_t>&)initialRegIndex, wProBatchSize, hProBatchSize, wArgmaxActual,
442+ wFullBatchCount);
443+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegIndexOne, hProBatchSize, wArgmaxActual);
444+ 
445+ for (uint16_t highIdx = 0; highIdx < highAxisActual; ++highIdx) {
446+ uint32_t highArgmaxOffset = highIdx * hArgmaxActual * wArgmaxAligned;
447+ uint32_t highIndexOffset = highIdx * hArgmaxActual * wArgmaxActual;
448+ uint32_t highOutputOffset = highIdx * hOutputActual * wOutputAligned;
449+ for (uint16_t hIdx = 0; hIdx < blockConcurrentCount; hIdx++) {
450+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
451+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
452+ T2 offset =
453+ (wBatchIdx + hProBatchIdx * wArgmaxAligned +
454+ hIdx * wArgmaxAligned * hProBatchSize * hConcurrentCount + highArgmaxOffset);
455+ T2 indexOffset =
456+ (wBatchIdx + hProBatchIdx * wArgmaxActual +
457+ hIdx * wArgmaxActual * hProBatchSize * hConcurrentCount + highIndexOffset);
458+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGrad, offset, allMaskU32);
459+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, indexOffset, allMaskU32);
460+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
461+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, maskBlock, wOutputConstReg,
462+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
463+ }
464+ 
465+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
466+ T2 offset =
467+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
468+ hIdx * wArgmaxAligned * hProBatchSize * hConcurrentCount + highArgmaxOffset);
469+ T2 indexOffset =
470+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxActual +
471+ hIdx * wArgmaxActual * hProBatchSize * hConcurrentCount + highIndexOffset);
472+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGradOne, offset, allMaskU32);
473+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndexOne, indexOffset, allMaskU32);
474+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
475+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, blockOne, wOutputConstReg,
476+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
477+ }
478+ }
479+ }
480+ 
481+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
482+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
483+ T2 offset =
484+ (wBatchIdx + hProBatchIdx * wArgmaxAligned +
485+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
486+ T2 indexOffset =
487+ (wBatchIdx + hProBatchIdx * wArgmaxActual +
488+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxActual + highIndexOffset);
489+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGrad, offset, allMaskU32);
490+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, indexOffset, allMaskU32);
491+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
492+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, maskRemainBatch,
493+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg,
494+ hMaxReg);
495+ }
496+ 
497+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
498+ T2 offset =
499+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
500+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
501+ T2 indexOffset =
502+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxActual +
503+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxActual + highIndexOffset);
504+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGradOne, offset, allMaskU32);
505+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndexOne, indexOffset, allMaskU32);
506+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
507+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, remainBatchOne, wOutputConstReg,
508+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
509+ }
510+ }
511+ 
512+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
513+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
514+ T2 offset =
515+ (wBatchIdx + hProBatchIdx * wArgmaxAligned +
516+ hRemainBatchCount * hProBatchSize * wArgmaxAligned +
517+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
518+ T2 indexOffset =
519+ (wBatchIdx + hProBatchIdx * wArgmaxActual + hRemainBatchCount * hProBatchSize * wArgmaxActual +
520+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxActual + highIndexOffset);
521+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGrad, offset, allMaskU32);
522+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, indexOffset, allMaskU32);
523+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
524+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, maskRemainTail, wOutputConstReg,
525+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
526+ }
527+ 
528+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
529+ T2 offset =
530+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
531+ hRemainBatchCount * hProBatchSize * wArgmaxAligned +
532+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
533+ T2 indexOffset =
534+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxActual +
535+ hRemainBatchCount * hProBatchSize * wArgmaxActual +
536+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxActual + highIndexOffset);
537+ AscendC::MicroAPI::Adds(parallelRegGrad, initialRegGradOne, offset, allMaskU32);
538+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndexOne, indexOffset, allMaskU32);
539+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
540+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, remainTailOne, wOutputConstReg,
541+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
542+ }
543+ }
544+ }
545+ }
546+}
547+ 
548+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
549+__aicore__ inline void MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>::multipleLineProcessVF2(
550+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
551+ __local_mem__ uint32_t* helpAddr)
552+{
553+ int64_t wOutput = wOutput_;
554+ int64_t wOutputActual = wOutputActual_;
555+ int64_t wOutputAligned = wOutputAligned_;
556+ int64_t hOutputActual = hOutputActual_;
557+ int32_t highOutputPlaneActual = wOutputAligned * hOutputActual;
558+ int64_t highAxisActual = highAxisActual_;
559+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
560+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
561+ int64_t wArgmaxAligned = wArgmaxAligned_;
562+ int64_t wArgmaxActual = wArgmaxActual_;
563+ uint16_t hArgmaxActual = hArgmaxActual_;
564+ 
565+ uint16_t hProBatchSize = curHProBatchSize_;
566+ uint16_t wProBatchSize = curWProBatchSize_;
567+ 
568+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
569+ uint16_t hFullBatchCount = hArgmaxActual / hProBatchSize;
570+ uint16_t wRemainTail = wArgmaxActual % wProBatchSize;
571+ uint32_t whFullBatchCount = wFullBatchCount * hFullBatchCount;
572+ 
573+ uint16_t highConcurrentCount = V_REG_SIZE / (whFullBatchCount * sizeof(T2));
574+ 
575+ uint16_t highBlockConcurrentCount = highAxisActual / highConcurrentCount;
576+ uint16_t highBlockRemainTail = highAxisActual - highBlockConcurrentCount * highConcurrentCount;
577+ 
578+ uint16_t hRemainTail = hArgmaxActual - hFullBatchCount * hProBatchSize;
579+ 
580+ uint32_t mask0 = highConcurrentCount * whFullBatchCount;
581+ uint32_t mask1 = highConcurrentCount * hFullBatchCount * 1;
582+ uint32_t mask2 = highConcurrentCount * 1 * wFullBatchCount;
583+ uint32_t mask3 = highConcurrentCount * 1 * 1;
584+ uint32_t mask4 = highBlockRemainTail * whFullBatchCount;
585+ uint32_t mask5 = highBlockRemainTail * hFullBatchCount * 1;
586+ uint32_t mask6 = highBlockRemainTail * 1 * wFullBatchCount;
587+ uint32_t mask7 = highBlockRemainTail * 1 * 1;
588+ 
589+ __VEC_SCOPE__
590+ {
591+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegGrad;
592+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegGradOne;
593+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegGrad;
594+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegGradOne;
595+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
596+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
597+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
598+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
599+ GenInitial3DIndices(
600+ (AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegGrad, wProBatchSize, hProBatchSize, wArgmaxAligned,
601+ wFullBatchCount, hFullBatchCount, hArgmaxActual);
602+ Gen3DIndexOne(
603+ (AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegGradOne, hProBatchSize, wArgmaxAligned, hFullBatchCount,
604+ hArgmaxActual);
605+ GenInitial2DIndices(
606+ (AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegGrad, wProBatchSize, hArgmaxActual, wArgmaxAligned,
607+ wFullBatchCount);
608+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegGradOne, hArgmaxActual, wArgmaxAligned);
609+ GenInitial3DIndices(
610+ (AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegIndex, wProBatchSize, hProBatchSize, wArgmaxActual,
611+ wFullBatchCount, hFullBatchCount, hArgmaxActual);
612+ Gen3DIndexOne(
613+ (AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegIndexOne, hProBatchSize, wArgmaxActual, hFullBatchCount,
614+ hArgmaxActual);
615+ GenInitial2DIndices(
616+ (AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegIndex, wProBatchSize, hArgmaxActual, wArgmaxActual,
617+ wFullBatchCount);
618+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegIndexOne, hArgmaxActual, wArgmaxActual);
619+ AscendC::MicroAPI::MaskReg allMaskStore =
620+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
621+ AscendC::MicroAPI::DataCopy(helpAddr, initial3DRegGrad, allMaskStore);
622+ AscendC::MicroAPI::DataCopy(helpAddr + V_REG_SIZE / sizeof(uint32_t), initial3DRegGradOne, allMaskStore);
623+ AscendC::MicroAPI::DataCopy(helpAddr + 2 * V_REG_SIZE / sizeof(uint32_t), initial2DRegGrad, allMaskStore);
624+ AscendC::MicroAPI::DataCopy(helpAddr + 3 * V_REG_SIZE / sizeof(uint32_t), initial2DRegGradOne, allMaskStore);
625+ AscendC::MicroAPI::DataCopy(helpAddr + 4 * V_REG_SIZE / sizeof(uint32_t), initial3DRegIndex, allMaskStore);
626+ AscendC::MicroAPI::DataCopy(helpAddr + 5 * V_REG_SIZE / sizeof(uint32_t), initial3DRegIndexOne, allMaskStore);
627+ AscendC::MicroAPI::DataCopy(helpAddr + 6 * V_REG_SIZE / sizeof(uint32_t), initial2DRegIndex, allMaskStore);
628+ AscendC::MicroAPI::DataCopy(helpAddr + 7 * V_REG_SIZE / sizeof(uint32_t), initial2DRegIndexOne, allMaskStore);
629+ }
630+ 
631+ __VEC_SCOPE__
632+ {
633+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
634+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
635+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
636+ if constexpr (IS_CHECK_RANGE == 1) {
637+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
638+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
639+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
640+ }
641+ 
642+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
643+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
644+ 
645+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegGrad;
646+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegGradOne;
647+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
648+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
649+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegGrad;
650+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
651+ 
652+ AscendC::MicroAPI::MaskReg allMaskU32 =
653+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
654+ AscendC::MicroAPI::DataCopy(initial3DRegGrad, helpAddr);
655+ AscendC::MicroAPI::DataCopy(initial3DRegGradOne, helpAddr + V_REG_SIZE / sizeof(uint32_t));
656+ AscendC::MicroAPI::DataCopy(initial3DRegIndex, helpAddr + 4 * V_REG_SIZE / sizeof(uint32_t));
657+ AscendC::MicroAPI::DataCopy(initial3DRegIndexOne, helpAddr + 5 * V_REG_SIZE / sizeof(uint32_t));
658+ 
659+ for (uint16_t highBlockIdx = 0; highBlockIdx < highBlockConcurrentCount; ++highBlockIdx) {
660+ uint32_t highArgmaxOffset = highBlockIdx * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
661+ uint32_t highIndexOffset = highBlockIdx * highConcurrentCount * hArgmaxActual * wArgmaxActual;
662+ uint32_t highOutputOffset = highBlockIdx * highConcurrentCount * hOutputActual * wOutputAligned;
663+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
664+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
665+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
666+ T2 indexOffset = (wBatchIdx + hProBatchIdx * wArgmaxActual + highIndexOffset);
667+ AscendC::MicroAPI::Adds(parallelRegGrad, initial3DRegGrad, offset, allMaskU32);
668+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndex, indexOffset, allMaskU32);
669+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
670+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask0, wOutputConstReg,
671+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
672+ highOutputPlaneActual, whFullBatchCount);
673+ }
674+ 
675+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
676+ T2 offset =
677+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
678+ highArgmaxOffset);
679+ T2 indexOffset =
680+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxActual + highIndexOffset);
681+ AscendC::MicroAPI::Adds(parallelRegGrad, initial3DRegGradOne, offset, allMaskU32);
682+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndexOne, indexOffset, allMaskU32);
683+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
684+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask1, wOutputConstReg,
685+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
686+ highOutputPlaneActual, hFullBatchCount);
687+ }
688+ }
689+ }
690+ }
691+ 
692+ __VEC_SCOPE__
693+ {
694+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
695+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
696+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
697+ if constexpr (IS_CHECK_RANGE == 1) {
698+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
699+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
700+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
701+ }
702+ 
703+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
704+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
705+ 
706+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegGrad;
707+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegGradOne;
708+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
709+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
710+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegGrad;
711+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
712+ 
713+ AscendC::MicroAPI::MaskReg allMaskU32 =
714+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
715+ AscendC::MicroAPI::DataCopy(initial2DRegGrad, helpAddr + 2 * V_REG_SIZE / sizeof(uint32_t));
716+ AscendC::MicroAPI::DataCopy(initial2DRegGradOne, helpAddr + 3 * V_REG_SIZE / sizeof(uint32_t));
717+ AscendC::MicroAPI::DataCopy(initial2DRegIndex, helpAddr + 6 * V_REG_SIZE / sizeof(uint32_t));
718+ AscendC::MicroAPI::DataCopy(initial2DRegIndexOne, helpAddr + 7 * V_REG_SIZE / sizeof(uint32_t));
719+ 
720+ for (uint16_t highBlockIdx = 0; highBlockIdx < highBlockConcurrentCount; ++highBlockIdx) {
721+ uint32_t highArgmaxOffset = highBlockIdx * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
722+ uint32_t highIndexOffset = highBlockIdx * highConcurrentCount * hArgmaxActual * wArgmaxActual;
723+ uint32_t highOutputOffset = highBlockIdx * highConcurrentCount * hOutputActual * wOutputAligned;
724+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
725+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
726+ T2 offset =
727+ (wBatchIdx + (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxAligned +
728+ highArgmaxOffset);
729+ T2 indexOffset =
730+ (wBatchIdx + (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxActual +
731+ highIndexOffset);
732+ AscendC::MicroAPI::Adds(parallelRegGrad, initial2DRegGrad, offset, allMaskU32);
733+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndex, indexOffset, allMaskU32);
734+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
735+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask2, wOutputConstReg,
736+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
737+ highOutputPlaneActual, wFullBatchCount);
738+ }
739+ 
740+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
741+ T2 offset =
742+ (wBatchIdx + wProBatchSize * wFullBatchCount +
743+ (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
744+ T2 indexOffset =
745+ (wBatchIdx + wProBatchSize * wFullBatchCount +
746+ (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxActual + highIndexOffset);
747+ AscendC::MicroAPI::Adds(parallelRegGrad, initial2DRegGradOne, offset, allMaskU32);
748+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndexOne, indexOffset, allMaskU32);
749+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
750+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask3, wOutputConstReg,
751+ curHIndex, curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
752+ highOutputPlaneActual, 1);
753+ }
754+ }
755+ }
756+ }
757+ 
758+ __VEC_SCOPE__
759+ {
760+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
761+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
762+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
763+ if constexpr (IS_CHECK_RANGE == 1) {
764+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
765+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
766+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
767+ }
768+ 
769+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
770+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
771+ 
772+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegGrad;
773+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegGradOne;
774+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
775+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
776+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegGrad;
777+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
778+ 
779+ AscendC::MicroAPI::MaskReg allMaskU32 =
780+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
781+ AscendC::MicroAPI::DataCopy(initial3DRegGrad, helpAddr);
782+ AscendC::MicroAPI::DataCopy(initial3DRegGradOne, helpAddr + V_REG_SIZE / sizeof(uint32_t));
783+ AscendC::MicroAPI::DataCopy(initial3DRegIndex, helpAddr + 4 * V_REG_SIZE / sizeof(uint32_t));
784+ AscendC::MicroAPI::DataCopy(initial3DRegIndexOne, helpAddr + 5 * V_REG_SIZE / sizeof(uint32_t));
785+ 
786+ uint32_t highArgmaxOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
787+ uint32_t highIndexOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxActual;
788+ uint32_t highOutputOffset = highBlockConcurrentCount * highConcurrentCount * hOutputActual * wOutputAligned;
789+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
790+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
791+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
792+ T2 indexOffset = (wBatchIdx + hProBatchIdx * wArgmaxActual + highIndexOffset);
793+ AscendC::MicroAPI::Adds(parallelRegGrad, initial3DRegGrad, offset, allMaskU32);
794+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndex, indexOffset, allMaskU32);
795+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
796+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask4, wOutputConstReg, curHIndex,
797+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual,
798+ whFullBatchCount);
799+ }
800+ 
801+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
802+ T2 offset =
803+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
804+ T2 indexOffset =
805+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxActual + highIndexOffset);
806+ AscendC::MicroAPI::Adds(parallelRegGrad, initial3DRegGradOne, offset, allMaskU32);
807+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndexOne, indexOffset, allMaskU32);
808+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
809+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask5, wOutputConstReg, curHIndex,
810+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual,
811+ hFullBatchCount);
812+ }
813+ }
814+ }
815+ 
816+ __VEC_SCOPE__
817+ {
818+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
819+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
820+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
821+ if constexpr (IS_CHECK_RANGE == 1) {
822+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
823+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
824+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
825+ }
826+ 
827+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
828+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
829+ 
830+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegGrad;
831+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegGradOne;
832+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
833+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
834+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegGrad;
835+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
836+ 
837+ AscendC::MicroAPI::MaskReg allMaskU32 =
838+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
839+ AscendC::MicroAPI::DataCopy(initial2DRegGrad, helpAddr + 2 * V_REG_SIZE / sizeof(uint32_t));
840+ AscendC::MicroAPI::DataCopy(initial2DRegGradOne, helpAddr + 3 * V_REG_SIZE / sizeof(uint32_t));
841+ AscendC::MicroAPI::DataCopy(initial2DRegIndex, helpAddr + 6 * V_REG_SIZE / sizeof(uint32_t));
842+ AscendC::MicroAPI::DataCopy(initial2DRegIndexOne, helpAddr + 7 * V_REG_SIZE / sizeof(uint32_t));
843+ 
844+ uint32_t highArgmaxOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
845+ uint32_t highIndexOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxActual;
846+ uint32_t highOutputOffset = highBlockConcurrentCount * highConcurrentCount * hOutputActual * wOutputAligned;
847+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
848+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
849+ T2 offset =
850+ (wBatchIdx + (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
851+ T2 indexOffset =
852+ (wBatchIdx + (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxActual + highIndexOffset);
853+ AscendC::MicroAPI::Adds(parallelRegGrad, initial2DRegGrad, offset, allMaskU32);
854+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndex, indexOffset, allMaskU32);
855+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
856+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask6, wOutputConstReg, curHIndex,
857+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual,
858+ wFullBatchCount);
859+ }
860+ 
861+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
862+ T2 offset =
863+ (wBatchIdx + wProBatchSize * wFullBatchCount +
864+ (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
865+ T2 indexOffset =
866+ (wBatchIdx + wProBatchSize * wFullBatchCount +
867+ (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxActual + highIndexOffset);
868+ AscendC::MicroAPI::Adds(parallelRegGrad, initial2DRegGradOne, offset, allMaskU32);
869+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndexOne, indexOffset, allMaskU32);
870+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
871+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, parallelRegGrad, mask7, wOutputConstReg, curHIndex,
872+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual,
873+ 1);
874+ }
875+ }
876+ }
877+}
878+ 
879+} // namespace MaxPoolGradNCHWBackwardBaseNameSpace
880+#endif // MAX_POOL_GRAD_NCHW_BACKWARD_BASE_H_
Apooling/max_pool_grad/op_kernel/arch35/max_pool_grad_nchw_big_kernel.h+426-0
@@ -0,0 +1,426 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef MAX_POOL_GRAD_BIG_KERNEL_H_
12+#define MAX_POOL_GRAD_BIG_KERNEL_H_
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "max_pool_grad_struct.h"
17+#include "../inc/platform.h"
18+#include "max_pool_grad_nchw_backward_base.h"
19+#include "../pool_3d_common/arch35/pool_big_kernel_utils.h"
20+#include "../../max_pool_with_argmax_v3/arch35/max_pool_with_argmax_v3_base.h"
21+#include <algorithm>
22+#include <type_traits>
23+ 
24+namespace MaxPoolGradNCHWBigKernelNameSpace {
25+using namespace AscendC;
26+using MaxPoolGradNCHWTilingCommonData = MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData;
27+using PoolBigKernelUtils::CalcRealIndex;
28+using PoolBigKernelUtils::DuplicateNegInf;
29+using PoolBigKernelUtils::LoadOneElement;
30+using PoolBigKernelUtils::LoadOneTensor;
31+using PoolBigKernelUtils::ReduceMaxWithIndex;
32+using PoolBigKernelUtils::StoreOneElement;
33+ 
34+constexpr int32_t BK_BUFFER_NUM = 2;
35+constexpr int32_t MERGE_BUF_ALIGN = 32;
36+ 
37+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
38+class MaxPoolGradNCHWBigKernel
39+ : public MaxPoolGradNCHWBackwardBaseNameSpace::MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE> {
40+ using Base = MaxPoolGradNCHWBackwardBaseNameSpace::MaxPoolGradNCHWBackwardBase<T1, T2, T3, IS_CHECK_RANGE>;
41+ 
42+public:
43+ __aicore__ inline MaxPoolGradNCHWBigKernel()
44+ {}
45+ 
46+ __aicore__ inline void Init(
47+ GM_ADDR x, GM_ADDR origY, GM_ADDR grad, GM_ADDR y, TPipe& pipeIn,
48+ const MaxPoolGradNCHWTilingCommonData& tilingData)
49+ {
50+ Base::ParseTilingData(tilingData);
51+ Base::blockIdx_ = GetBlockIdx();
52+ if (Base::blockIdx_ >= Base::usedCoreNum_) {
53+ return;
54+ }
55+ 
56+ Base::curCoreProcessNum_ =
57+ (Base::blockIdx_ + 1 == Base::usedCoreNum_) ? Base::tailCoreProcessNum_ : Base::normalCoreProcessNum_;
58+ 
59+ xGm_.SetGlobalBuffer((__gm__ T1*)x);
60+ origYGm_.SetGlobalBuffer((__gm__ T1*)origY);
61+ Base::gradGm_.SetGlobalBuffer((__gm__ T1*)grad);
62+ Base::yGm_.SetGlobalBuffer((__gm__ T1*)y);
63+ 
64+ inHW_ = Base::hOutput_ * Base::wOutput_;
65+ int64_t maxCountBySize = Base::inputBufferSize_ / static_cast<int64_t>(sizeof(T1));
66+ maxCount_ = (maxCountBySize < 1) ? 1 : maxCountBySize;
67+ 
68+ const int64_t forwardStageBufferSize = Base::inputBufferSize_ * BK_BUFFER_NUM + MERGE_BUF_ALIGN;
69+ const int64_t backwardStageBufferSize =
70+ Base::gradBufferSize_ * BK_BUFFER_NUM + Base::outputBufferSize_ * BK_BUFFER_NUM;
71+ totalStageBufferSize_ =
72+ (forwardStageBufferSize > backwardStageBufferSize) ? forwardStageBufferSize : backwardStageBufferSize;
73+ 
74+ pipeIn.InitBuffer(Base::argmaxBuf_, Base::argmaxBufferSize_);
75+ pipeIn.InitBuffer(Base::helpBuf_, MaxPoolGradNCHWBackwardBaseNameSpace::HELP_BUFFER);
76+ pipeIn.InitBufPool(forwardBufPool_, static_cast<uint32_t>(totalStageBufferSize_));
77+ pipeIn.InitBufPool(backwardBufPool_, static_cast<uint32_t>(totalStageBufferSize_), forwardBufPool_);
78+ }
79+ 
80+ __aicore__ inline void Process()
81+ {
82+ if (Base::blockIdx_ >= Base::usedCoreNum_) {
83+ return;
84+ }
85+ 
86+ for (int64_t loopNum = 0; loopNum < Base::curCoreProcessNum_; ++loopNum) {
87+ Base::ScalarCompute(loopNum);
88+ 
89+ if (Base::hArgmaxActual_ <= 0 || Base::wArgmaxActual_ <= 0) {
90+ InitBackwardStageBuffers();
91+ Base::ProcessNoArgmaxBlock();
92+ backwardBufPool_.Reset();
93+ continue;
94+ }
95+ 
96+ InitForwardStageBuffers();
97+ ForwardComputeTile();
98+ forwardBufPool_.Reset();
99+ 
100+ InitBackwardStageBuffers();
101+ Base::CopyInGrad();
102+ BackwardScatter();
103+ Base::CopyOut();
104+ backwardBufPool_.Reset();
105+ }
106+ }
107+ 
108+private:
109+ TBufPool<TPosition::VECCALC> forwardBufPool_;
110+ TBufPool<TPosition::VECCALC> backwardBufPool_;
111+ TQue<QuePosition::VECIN, BK_BUFFER_NUM> inputQue_;
112+ TBuf<TPosition::VECCALC> maxValBuf_;
113+ GlobalTensor<T1> xGm_;
114+ GlobalTensor<T1> origYGm_;
115+ int64_t inHW_ = 1;
116+ int64_t maxCount_ = 1;
117+ int64_t totalStageBufferSize_ = 0;
118+ 
119+private:
120+ __aicore__ inline void InitForwardStageBuffers()
121+ {
122+ forwardBufPool_.InitBuffer(inputQue_, BK_BUFFER_NUM, static_cast<uint32_t>(Base::inputBufferSize_));
123+ forwardBufPool_.InitBuffer(maxValBuf_, static_cast<uint32_t>(MERGE_BUF_ALIGN));
124+ }
125+ 
126+ __aicore__ inline void InitBackwardStageBuffers()
127+ {
128+ backwardBufPool_.InitBuffer(Base::outputQue_, BK_BUFFER_NUM, static_cast<uint32_t>(Base::outputBufferSize_));
129+ backwardBufPool_.InitBuffer(Base::gradQue_, BK_BUFFER_NUM, static_cast<uint32_t>(Base::gradBufferSize_));
130+ }
131+ 
132+ __aicore__ inline void ForwardComputeTile()
133+ {
134+ LocalTensor<T2> argmaxLocal = Base::argmaxBuf_.template Get<T2>();
135+ Duplicate(argmaxLocal, T2(-1), Base::argmaxBufferSize_ / sizeof(T2));
136+ PipeBarrier<PIPE_ALL>();
137+ 
138+ const int64_t tileHW = Base::hArgmaxActual_ * Base::wArgmaxActual_;
139+ for (int64_t highIdx = 0; highIdx < Base::highAxisActual_; ++highIdx) {
140+ for (int64_t hIdx = 0; hIdx < Base::hArgmaxActual_; ++hIdx) {
141+ for (int64_t wIdx = 0; wIdx < Base::wArgmaxActual_; ++wIdx) {
142+ int64_t curkH = 0;
143+ int64_t curkW = 0;
144+ int64_t curInOffset = 0;
145+ int64_t curOriginIndex = 0;
146+ CalcKernelSize(highIdx, hIdx, wIdx, curkH, curkW, curInOffset, curOriginIndex);
147+ 
148+ if (curkH <= 0 || curkW <= 0) {
149+ continue;
150+ }
151+ 
152+ const int64_t bufferOffset = highIdx * tileHW + hIdx * Base::wArgmaxActual_ + wIdx;
153+ if (curkH * curkW <= maxCount_) {
154+ NoSplitKernelProcess(curkH, curkW, curInOffset, curOriginIndex, bufferOffset);
155+ } else {
156+ SplitKernelProcess(curkH, curkW, curInOffset, curOriginIndex, bufferOffset);
157+ }
158+ PipeBarrier<PIPE_ALL>();
159+ }
160+ }
161+ }
162+ }
163+ 
164+ __aicore__ inline void BackwardScatter()
165+ {
166+ const uint32_t calcCount = Base::outputBufferSize_ / sizeof(computeType);
167+ 
168+ LocalTensor<computeType> yLocal = Base::outputQue_.template AllocTensor<computeType>();
169+ Duplicate(yLocal, computeType(0), calcCount);
170+ PipeBarrier<PIPE_ALL>();
171+ 
172+ LocalTensor<T1> gradLocal = Base::gradQue_.template DeQue<T1>();
173+ LocalTensor<T2> argmaxLocal = Base::argmaxBuf_.template Get<T2>();
174+ __local_mem__ computeType* yAddr = (__local_mem__ computeType*)yLocal.GetPhyAddr();
175+ __local_mem__ T1* gradAddr = (__local_mem__ T1*)gradLocal.GetPhyAddr();
176+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
177+ 
178+ Base::BackwardCompute(yAddr, gradAddr, argmaxAddr);
179+ PipeBarrier<PIPE_ALL>();
180+ 
181+ if constexpr (!std::is_same<T1, float>::value) {
182+ Cast(yLocal.ReinterpretCast<T1>(), yLocal, RoundMode::CAST_RINT, calcCount);
183+ PipeBarrier<PIPE_ALL>();
184+ }
185+ 
186+ Base::outputQue_.EnQue(yLocal);
187+ Base::gradQue_.FreeTensor(gradLocal);
188+ }
189+ 
190+ __aicore__ inline void CalcKernelSize(
191+ int64_t highIdx, int64_t hIdx, int64_t wIdx, int64_t& curkH, int64_t& curkW, int64_t& curInOffset,
192+ int64_t& curOriginIndex)
193+ {
194+ const int64_t ncOffset = Base::highAxisIndex_ * Base::highAxisInner_ + highIdx;
195+ const int64_t ho = Base::hArgmaxActualStart_ + hIdx;
196+ const int64_t wo = Base::wArgmaxActualStart_ + wIdx;
197+ 
198+ int64_t curOriginH = ho * Base::strideH_ - Base::padH_;
199+ int64_t curOriginW = wo * Base::strideW_ - Base::padW_;
200+ curkH = Base::kernelH_;
201+ curkW = Base::kernelW_;
202+ 
203+ if (curOriginH < 0) {
204+ curkH += curOriginH;
205+ curOriginH = 0;
206+ }
207+ if (curOriginH + curkH > Base::hOutput_) {
208+ curkH = Base::hOutput_ - curOriginH;
209+ }
210+ 
211+ if (curOriginW < 0) {
212+ curkW += curOriginW;
213+ curOriginW = 0;
214+ }
215+ if (curOriginW + curkW > Base::wOutput_) {
216+ curkW = Base::wOutput_ - curOriginW;
217+ }
218+ 
219+ curOriginIndex = curOriginH * Base::wOutput_ + curOriginW;
220+ curInOffset = ncOffset * inHW_ + curOriginIndex;
221+ }
222+ 
223+ __aicore__ inline void CopyInMultiRows(int64_t offset, int64_t blockLen, int64_t blockCount)
224+ {
225+ LocalTensor<T1> xLocal = inputQue_.AllocTensor<T1>();
226+ 
227+ DataCopyPadExtParams<T1> padExtParams = {false, 0, 0, 0};
228+ DataCopyExtParams extParams;
229+ extParams.blockCount = blockCount;
230+ extParams.blockLen = blockLen * sizeof(T1);
231+ extParams.srcStride = (Base::wOutput_ - blockLen) * sizeof(T1);
232+ extParams.dstStride = 0;
233+ 
234+ DataCopyPad<T1, PaddingMode::Compact>(xLocal, xGm_[offset], extParams, padExtParams);
235+ PipeBarrier<PIPE_ALL>();
236+ inputQue_.EnQue(xLocal);
237+ }
238+ 
239+ __aicore__ inline void CopyInSingleRow(int64_t offset, int64_t blockLen)
240+ {
241+ LocalTensor<T1> xLocal = inputQue_.AllocTensor<T1>();
242+ 
243+ DataCopyPadExtParams<T1> padExtParams = {false, 0, 0, 0};
244+ DataCopyExtParams extParams;
245+ extParams.blockCount = 1;
246+ extParams.blockLen = blockLen * sizeof(T1);
247+ extParams.srcStride = 0;
248+ extParams.dstStride = 0;
249+ 
250+ DataCopyPad(xLocal, xGm_[offset], extParams, padExtParams);
251+ PipeBarrier<PIPE_ALL>();
252+ inputQue_.EnQue(xLocal);
253+ }
254+ 
255+ __aicore__ inline void NoSplitKernelProcess(
256+ int64_t curkH, int64_t curkW, int64_t curInOffset, int64_t curOriginIndex, int64_t bufferOffset)
257+ {
258+ CopyInMultiRows(curInOffset, curkW, curkH);
259+ ComputeSingleArgmax<false, false>(curkW * curkH, curkW, curOriginIndex, bufferOffset);
260+ }
261+ 
262+ __aicore__ inline void SplitKernelProcess(
263+ int64_t curkH, int64_t curkW, int64_t curInOffset, int64_t curOriginIndex, int64_t bufferOffset)
264+ {
265+ InitMergeBuffer(bufferOffset, curOriginIndex);
266+ 
267+ if (curkW <= 0 || curkH <= 0 || maxCount_ <= 0) {
268+ return;
269+ }
270+ 
271+ if (curkW <= maxCount_) {
272+ const int64_t hFactor = maxCount_ / ((curkW > 0) ? curkW : 1);
273+ const int64_t hLoops = (curkH + hFactor - 1) / hFactor;
274+ const int64_t hTail = curkH - (hLoops - 1) * hFactor;
275+ 
276+ int64_t inputOffset = curInOffset;
277+ int64_t kernelOffset = curOriginIndex;
278+ for (int64_t hLoop = 0; hLoop < hLoops; ++hLoop) {
279+ const int64_t curhFactor = (hLoop == hLoops - 1) ? hTail : hFactor;
280+ CopyInMultiRows(inputOffset, curkW, curhFactor);
281+ ComputeSingleArgmax<true, false>(
282+ curkW * curhFactor, ((curkW > 0) ? curkW : 1), kernelOffset, bufferOffset);
283+ PipeBarrier<PIPE_ALL>();
284+ inputOffset += curhFactor * Base::wOutput_;
285+ kernelOffset += curhFactor * Base::wOutput_;
286+ }
287+ } else {
288+ const int64_t hLoops = curkH;
289+ const int64_t wFactor = maxCount_;
290+ const int64_t wLoops = (curkW + wFactor - 1) / wFactor;
291+ const int64_t wTail = curkW - (wLoops - 1) * wFactor;
292+ 
293+ for (int64_t hLoop = 0; hLoop < hLoops; ++hLoop) {
294+ int64_t inputOffset = curInOffset + hLoop * Base::wOutput_;
295+ int64_t kernelOffset = curOriginIndex + hLoop * Base::wOutput_;
296+ for (int64_t wLoop = 0; wLoop < wLoops; ++wLoop) {
297+ const int64_t curFactor = (wLoop == wLoops - 1) ? wTail : wFactor;
298+ CopyInSingleRow(inputOffset, curFactor);
299+ ComputeSingleArgmax<true, true>(
300+ curFactor, ((curkW > 0) ? curkW : 1), kernelOffset, bufferOffset);
301+ PipeBarrier<PIPE_ALL>();
302+ inputOffset += curFactor;
303+ kernelOffset += curFactor;
304+ }
305+ }
306+ }
307+ }
308+ 
309+ __aicore__ inline void InitMergeBuffer(int64_t bufferOffset, int64_t initIndex)
310+ {
311+ LocalTensor<T2> argmaxLocal = Base::argmaxBuf_.template Get<T2>();
312+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
313+ 
314+ LocalTensor<float> maxValLocal = maxValBuf_.template Get<float>();
315+ __local_mem__ float* maxValAddr = (__local_mem__ float*)maxValLocal.GetPhyAddr();
316+ 
317+ const float negInf = AscendC::NumericLimits<float>::NegativeInfinity();
318+ 
319+ __VEC_SCOPE__
320+ {
321+ MicroAPI::RegTensor<float> negInfVal;
322+ MicroAPI::Duplicate(negInfVal, negInf);
323+ MicroAPI::MaskReg pregOne = MicroAPI::CreateMask<float, MicroAPI::MaskPattern::VL1>();
324+ StoreOneElement<float, float>(maxValAddr, negInfVal, pregOne, 0);
325+ 
326+ MicroAPI::RegTensor<T2> initResIndex;
327+ MicroAPI::Duplicate(initResIndex, static_cast<T2>(initIndex));
328+ StoreOneElement<T2, T2>(argmaxAddr, initResIndex, pregOne, bufferOffset);
329+ }
330+ PipeBarrier<PIPE_ALL>();
331+ }
332+ 
333+ template <bool MERGE, bool SPLITKW>
334+ __aicore__ inline void ComputeSingleArgmax(
335+ int64_t dataCount, int64_t curKw, int64_t curOriginIndex, int64_t bufferOffset)
336+ {
337+ LocalTensor<T1> xLocal = inputQue_.DeQue<T1>();
338+ 
339+ __local_mem__ T1* xLocalAddr = (__local_mem__ T1*)xLocal.GetPhyAddr();
340+ 
341+ LocalTensor<T2> argmaxLocal = Base::argmaxBuf_.template Get<T2>();
342+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
343+ 
344+ LocalTensor<float> maxValLocal = maxValBuf_.template Get<float>();
345+ __local_mem__ float* maxValAddr = (__local_mem__ float*)maxValLocal.GetPhyAddr();
346+ 
347+ const float negInf = AscendC::NumericLimits<float>::NegativeInfinity();
348+ 
349+ constexpr uint32_t repeatElm = platform::GetVRegSize() / sizeof(float);
350+ const uint16_t repeatTimes =
351+ static_cast<uint16_t>((dataCount + repeatElm - 1) / repeatElm);
352+ uint32_t num = repeatTimes * repeatElm;
353+ const uint32_t padNum = num - dataCount;
354+ constexpr int32_t padIndex = -1;
355+ 
356+ __VEC_SCOPE__
357+ {
358+ DuplicateNegInf<T1>(xLocalAddr, padNum, dataCount);
359+ 
360+ MicroAPI::RegTensor<float> res;
361+ MicroAPI::RegTensor<int32_t> resIndex;
362+ MicroAPI::RegTensor<int32_t> index;
363+ MicroAPI::RegTensor<float> vd0;
364+ 
365+ MicroAPI::MaskReg nanMaskReg;
366+ MicroAPI::MaskReg cmpMaskReg;
367+ MicroAPI::MaskReg maskAll = MicroAPI::CreateMask<float, MicroAPI::MaskPattern::ALL>();
368+ 
369+ MicroAPI::Duplicate(resIndex, padIndex);
370+ MicroAPI::Duplicate(res, negInf);
371+ MicroAPI::Arange(index, 0);
372+ 
373+ for (uint16_t i = 0; i < repeatTimes; ++i) {
374+ uint32_t maskNum = num;
375+ MicroAPI::MaskReg p0 = MicroAPI::UpdateMask<float>(maskNum);
376+ MicroAPI::AddrReg offset = MicroAPI::CreateAddrReg<T1>(i, repeatElm);
377+ LoadOneTensor<T1>(xLocalAddr, vd0, p0, offset);
378+ 
379+ MicroAPI::Compare<float, CMPMODE::NE>(nanMaskReg, vd0, vd0, maskAll);
380+ MicroAPI::Compare<float, CMPMODE::GT>(cmpMaskReg, vd0, res, maskAll);
381+ MicroAPI::MaskXor(cmpMaskReg, cmpMaskReg, nanMaskReg, maskAll);
382+ MicroAPI::Select(res, vd0, res, cmpMaskReg);
383+ MicroAPI::Select(resIndex, index, resIndex, cmpMaskReg);
384+ MicroAPI::Adds(index, index, repeatElm, maskAll);
385+ }
386+ 
387+ ReduceMaxWithIndex<float>(res, index, res, resIndex, padIndex);
388+ 
389+ MicroAPI::MaskReg pregOne = MicroAPI::CreateMask<float, MicroAPI::MaskPattern::VL1>();
390+ MicroAPI::RegTensor<T2> realResIndex;
391+ 
392+ CalcRealIndex<T2, SPLITKW>(realResIndex, index, curKw, Base::wOutput_, curOriginIndex);
393+ 
394+ if constexpr (MERGE) {
395+ MicroAPI::RegTensor<T2> lastResIndex;
396+ LoadOneElement<T2, T2>(argmaxAddr, lastResIndex, pregOne, bufferOffset);
397+ 
398+ MicroAPI::RegTensor<float> lastRes;
399+ LoadOneElement<float, float>(maxValAddr, lastRes, pregOne, 0);
400+ 
401+ MicroAPI::MaskReg curNanMaskReg;
402+ MicroAPI::MaskReg selReg;
403+ 
404+ MicroAPI::Compare<float, CMPMODE::NE>(curNanMaskReg, res, res, maskAll);
405+ MicroAPI::Compare<float, CMPMODE::GT>(selReg, res, lastRes, maskAll);
406+ MicroAPI::MaskXor(selReg, selReg, curNanMaskReg, maskAll);
407+ 
408+ MicroAPI::Select(res, res, lastRes, selReg);
409+ MicroAPI::Select(realResIndex, realResIndex, lastResIndex, selReg);
410+ 
411+ MicroAPI::LocalMemBar<MicroAPI::MemType::VEC_LOAD, MicroAPI::MemType::VEC_STORE>();
412+ }
413+ 
414+ StoreOneElement<T2, T2>(argmaxAddr, realResIndex, pregOne, bufferOffset);
415+ if constexpr (MERGE) {
416+ StoreOneElement<float, float>(maxValAddr, res, pregOne, 0);
417+ }
418+ }
419+ 
420+ PipeBarrier<PIPE_ALL>();
421+ inputQue_.FreeTensor(xLocal);
422+ }
423+};
424+ 
425+} // namespace MaxPoolGradNCHWBigKernelNameSpace
426+#endif // MAX_POOL_GRAD_BIG_KERNEL_H_
Apooling/max_pool_grad/op_kernel/arch35/max_pool_grad_nchw_small_kernel.h+624-0
@@ -0,0 +1,624 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef MAX_POOL_GRAD_SMALL_KERNEL_H
12+#define MAX_POOL_GRAD_SMALL_KERNEL_H
13+ 
14+#include "kernel_operator.h"
15+#include "kernel_tiling/kernel_tiling.h"
16+#include "../inc/platform.h"
17+#include "max_pool_grad_struct.h"
18+#include "max_pool_grad_nchw_backward_base.h"
19+#include "../../max_pool_with_argmax_v3/arch35/max_pool_with_argmax_v3_base.h"
20+#include "../pool_3d_common/arch35/pool_big_kernel_utils.h"
21+ 
22+namespace MaxPoolGradNCHWSmallKernelNameSpace {
23+using namespace AscendC;
24+using MaxPoolGradNCHWNameSpace::GenInitial1DIndices;
25+using MaxPoolGradNCHWNameSpace::GenInitial2DIndices;
26+using MaxPoolGradNCHWNameSpace::GenInitial3DIndices;
27+using MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData;
28+ 
29+constexpr uint32_t BUFFER_NUM = 2;
30+constexpr int64_t RATIO = 2;
31+ 
32+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
33+class PoolGradNCHWSmallKernel : public MaxPoolGradNCHWBackwardBaseNameSpace::MaxPoolGradNCHWBackwardBase<
34+ TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE> {
35+ using Base =
36+ MaxPoolGradNCHWBackwardBaseNameSpace::MaxPoolGradNCHWBackwardBase<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>;
37+ 
38+public:
39+ __aicore__ inline PoolGradNCHWSmallKernel(
40+ TPipe& pipeIn, const MaxPoolGradWithArgmaxNCHWTilingCommonData& tilingData)
41+ : pipe_(pipeIn), tilingData_(tilingData)
42+ {}
43+ 
44+ __aicore__ inline void Init(GM_ADDR orig_x, GM_ADDR orig_y, GM_ADDR grads, GM_ADDR y);
45+ __aicore__ inline void Process();
46+ __aicore__ inline void ForwardScalarCompute();
47+ __aicore__ inline void ConvertIndexWithoutPadAlign(
48+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, TYPE_ARGMAX left, TYPE_ARGMAX wInput,
49+ TYPE_ARGMAX hIndexBase, MicroAPI::RegTensor<TYPE_ARGMAX>& dstReg, int32_t ncInputOffset);
50+ __aicore__ inline void ProcessW(
51+ __local_mem__ TYPE_ORIG_X* computeAddr, int32_t hOffset, uint16_t wStrideOffset,
52+ MicroAPI::RegTensor<int32_t>& indexReg, uint16_t hKernel, uint16_t wKernel, uint16_t repeatElem,
53+ MicroAPI::RegTensor<int32_t>& maxIndexReg, uint32_t hDilation, uint32_t wDilation);
54+ __aicore__ inline void ConvertIndexWithoutPadAlignNc(
55+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, TYPE_ARGMAX left, TYPE_ARGMAX wInput,
56+ TYPE_ARGMAX hIndexBase, MicroAPI::RegTensor<TYPE_ARGMAX>& dstReg, int32_t ncInputOffset, int32_t ncOutputCount,
57+ int32_t inputNcSize);
58+ __aicore__ inline void MultiRowGather(
59+ __local_mem__ TYPE_ORIG_X* computeAddr, __local_mem__ TYPE_ARGMAX* argmaxAddr);
60+ __aicore__ inline void SingleRowGather(
61+ __local_mem__ TYPE_ORIG_X* computeAddr, __local_mem__ TYPE_ARGMAX* argmaxAddr);
62+ __aicore__ inline void MultiNcGather(__local_mem__ TYPE_ORIG_X* computeAddr, __local_mem__ TYPE_ARGMAX* argmaxAddr);
63+ __aicore__ inline void DupBufferNegInf(
64+ __local_mem__ TYPE_ORIG_X* dstAddr, uint32_t repeatElm, uint16_t loop, uint32_t tail);
65+ __aicore__ inline void CopyToCalcBuffer(
66+ __local_mem__ TYPE_ORIG_X* dstAddr, __local_mem__ TYPE_ORIG_X* srcAddr, uint16_t batch, uint16_t rows,
67+ uint16_t loopCols, uint16_t tailCols, uint32_t repeatElm, uint32_t srcBatchStride, uint32_t srcRowStride,
68+ uint32_t dstBatchStride, uint32_t dstRowStride, uint32_t dstRowOffset, uint32_t dstColOffset);
69+ __aicore__ inline void DupAndCopyToCalcBuffer(
70+ __local_mem__ TYPE_ORIG_X* dstAddr, __local_mem__ TYPE_ORIG_X* srcAddr);
71+ __aicore__ inline void ForwardCopyIn();
72+ __aicore__ inline void Forward();
73+ __aicore__ inline void Backward();
74+ 
75+ TPipe& pipe_;
76+ const MaxPoolGradWithArgmaxNCHWTilingCommonData& tilingData_;
77+ TQue<QuePosition::VECIN, BUFFER_NUM> inputQue_;
78+ TBuf<TPosition::VECCALC> inputCalcBuff_;
79+ TBuf<TPosition::VECCALC> argmaxBuff_;
80+ 
81+ GlobalTensor<TYPE_ORIG_X> xGm_;
82+ 
83+ int64_t hInputActualPad_ = 0;
84+ int64_t wInputActualPad_ = 0;
85+ int64_t wInputActualAlignedPad_ = 0;
86+ int64_t leftOffsetToInputLeft_ = 0;
87+ int64_t rightOffsetToInputRight_ = 0;
88+ int64_t topOffsetToInputTop_ = 0;
89+ int64_t downOffsetToInputDown_ = 0;
90+ int64_t highInputOffset_ = 0;
91+ int64_t forwardhInputOffset_ = 0;
92+ int64_t forwardwInputOffset_ = 0;
93+ int64_t hInputActualNoPad_ = 0;
94+ int64_t wInputActualNoPad_ = 0;
95+ int64_t forwardHighAxisIndex_ = 0;
96+ int64_t forwardhighAxisActual_ = 0;
97+ int64_t forwardHAxisIndex_ = 0;
98+ int64_t hOutputReal_ = 0;
99+ int64_t wOutputReal_ = 0;
100+ bool isPad_ = false;
101+ 
102+ constexpr static int32_t BLOCK_SIZE = platform::GetUbBlockSize();
103+ constexpr static int32_t V_REG_SIZE = platform::GetVRegSize();
104+ constexpr static int64_t MAX_DATA_NUM_IN_ONE_BLOCK =
105+ BLOCK_SIZE / sizeof(TYPE_ORIG_X) >= BLOCK_SIZE / sizeof(TYPE_ARGMAX) ? BLOCK_SIZE / sizeof(TYPE_ORIG_X) :
106+ BLOCK_SIZE / sizeof(TYPE_ARGMAX);
107+ constexpr static uint16_t vlT1_ = platform::GetVRegSize() / sizeof(TYPE_ORIG_X);
108+};
109+ 
110+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
111+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::Init(
112+ GM_ADDR orig_x, GM_ADDR orig_y, GM_ADDR grads, GM_ADDR y)
113+{
114+ Base::ParseTilingData(tilingData_);
115+ Base::blockIdx_ = GetBlockIdx();
116+ if (Base::blockIdx_ >= Base::usedCoreNum_) {
117+ return;
118+ }
119+ Base::curCoreProcessNum_ =
120+ (Base::blockIdx_ + 1 == Base::usedCoreNum_) ? Base::tailCoreProcessNum_ : Base::normalCoreProcessNum_;
121+ xGm_.SetGlobalBuffer((__gm__ TYPE_ORIG_X*)orig_x);
122+ Base::gradGm_.SetGlobalBuffer((__gm__ TYPE_ORIG_X*)grads);
123+ Base::yGm_.SetGlobalBuffer((__gm__ TYPE_ORIG_X*)y);
124+ 
125+ isPad_ = tilingData_.isPad != 0;
126+ 
127+ pipe_.InitBuffer(inputQue_, BUFFER_NUM, tilingData_.inputBufferSize);
128+ if (isPad_) {
129+ pipe_.InitBuffer(inputCalcBuff_, tilingData_.inputBufferSize);
130+ }
131+ 
132+ pipe_.InitBuffer(argmaxBuff_, Base::argmaxBufferSize_);
133+ pipe_.InitBuffer(Base::outputQue_, BUFFER_NUM, Base::outputBufferSize_);
134+ pipe_.InitBuffer(Base::gradQue_, BUFFER_NUM, Base::gradBufferSize_);
135+ pipe_.InitBuffer(Base::helpBuf_, HELP_BUFFER);
136+}
137+ 
138+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
139+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::Forward()
140+{
141+ LocalTensor<TYPE_ORIG_X> inputLocal = inputQue_.template DeQue<TYPE_ORIG_X>();
142+ __local_mem__ TYPE_ORIG_X* inputQueAddr = (__local_mem__ TYPE_ORIG_X*)inputLocal.GetPhyAddr();
143+ __local_mem__ TYPE_ORIG_X* computeAddr = inputQueAddr;
144+ if (isPad_) {
145+ LocalTensor<TYPE_ORIG_X> caclBuffLocal = inputCalcBuff_.template Get<TYPE_ORIG_X>();
146+ __local_mem__ TYPE_ORIG_X* inputBuffAddr = (__local_mem__ TYPE_ORIG_X*)caclBuffLocal.GetPhyAddr();
147+ DupAndCopyToCalcBuffer(inputBuffAddr, inputQueAddr);
148+ computeAddr = inputBuffAddr;
149+ }
150+ LocalTensor<TYPE_ARGMAX> argmaxLocal = argmaxBuff_.template Get<TYPE_ARGMAX>();
151+ __local_mem__ TYPE_ARGMAX* argmaxAddr = (__local_mem__ TYPE_ARGMAX*)argmaxLocal.GetPhyAddr();
152+ 
153+ if (Base::wArgmaxActual_ * RATIO > Base::vlT2_) {
154+ SingleRowGather(computeAddr, argmaxAddr);
155+ } else if (Base::hArgmaxActual_ * Base::wArgmaxActual_ * RATIO > Base::vlT2_) {
156+ MultiRowGather(computeAddr, argmaxAddr);
157+ } else {
158+ MultiNcGather(computeAddr, argmaxAddr);
159+ }
160+ 
161+ inputQue_.FreeTensor(inputLocal);
162+}
163+ 
164+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
165+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::Backward()
166+{
167+ uint32_t calCount = Base::outputBufferSize_ / sizeof(computeType);
168+ LocalTensor<computeType> yLocal = Base::outputQue_.template AllocTensor<computeType>();
169+ Duplicate(yLocal, computeType(0), calCount);
170+ LocalTensor<TYPE_ORIG_X> gradLocal = Base::gradQue_.template DeQue<TYPE_ORIG_X>();
171+ LocalTensor<TYPE_ARGMAX> argmaxLocal = argmaxBuff_.template Get<TYPE_ARGMAX>();
172+ __local_mem__ computeType* yAddr = (__local_mem__ computeType*)yLocal.GetPhyAddr();
173+ __local_mem__ TYPE_ORIG_X* gradAddr = (__local_mem__ TYPE_ORIG_X*)gradLocal.GetPhyAddr();
174+ __local_mem__ TYPE_ARGMAX* argmaxAddr = (__local_mem__ TYPE_ARGMAX*)argmaxLocal.GetPhyAddr();
175+ 
176+ Base::BackwardCompute(yAddr, gradAddr, argmaxAddr);
177+ 
178+ if constexpr (std::negation<std::is_same<TYPE_ORIG_X, float>>::value) {
179+ Cast(yLocal.ReinterpretCast<TYPE_ORIG_X>(), yLocal, RoundMode::CAST_RINT, calCount);
180+ }
181+ Base::outputQue_.EnQue(yLocal);
182+ Base::gradQue_.FreeTensor(gradLocal);
183+}
184+ 
185+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
186+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::ForwardScalarCompute()
187+{
188+ forwardHighAxisIndex_ = Base::highAxisIndex_;
189+ forwardhighAxisActual_ = Base::highAxisActual_;
190+ 
191+ forwardHAxisIndex_ = Base::hAxisIndex_;
192+ hOutputReal_ = Base::hArgmaxActual_;
193+ 
194+ wOutputReal_ = Base::wArgmaxActual_;
195+ 
196+ hInputActualPad_ = (hOutputReal_ - 1) * tilingData_.hStride + (tilingData_.hKernel - 1) * tilingData_.dilationH + 1;
197+ wInputActualPad_ = (wOutputReal_ - 1) * tilingData_.wStride + (tilingData_.wKernel - 1) * tilingData_.dilationW + 1;
198+ 
199+ wInputActualAlignedPad_ =
200+ CeilDivision(wInputActualPad_, BLOCK_SIZE / sizeof(TYPE_ORIG_X)) * (BLOCK_SIZE / sizeof(TYPE_ORIG_X));
201+ int64_t inputPlaneSize = tilingData_.hOutput * tilingData_.wOutput;
202+ highInputOffset_ = Base::highAxisIndex_ * tilingData_.highAxisInner * inputPlaneSize;
203+ forwardhInputOffset_ = Base::hArgmaxActualStart_ * tilingData_.hStride * tilingData_.wOutput;
204+ forwardwInputOffset_ = Base::wArgmaxActualStart_ * tilingData_.wStride;
205+ 
206+ if (isPad_) {
207+ int64_t tRelBoundDistance = Base::hArgmaxActualStart_ * tilingData_.hStride - tilingData_.padH;
208+ int64_t bRelBoundDistance = Base::hArgmaxActualStart_ * tilingData_.hStride +
209+ (hOutputReal_ - 1) * tilingData_.hStride + tilingData_.hKernel -
210+ tilingData_.hOutput - tilingData_.padH;
211+ int64_t lRelBoundDistance = Base::wArgmaxActualStart_ * tilingData_.wStride - tilingData_.padW;
212+ int64_t rRelBoundDistance = Base::wArgmaxActualStart_ * tilingData_.wStride +
213+ (wOutputReal_ - 1) * tilingData_.wStride + tilingData_.wKernel -
214+ tilingData_.wOutput - tilingData_.padW;
215+ leftOffsetToInputLeft_ = lRelBoundDistance >= 0 ? 0 : -lRelBoundDistance;
216+ rightOffsetToInputRight_ = rRelBoundDistance >= 0 ? rRelBoundDistance : 0;
217+ topOffsetToInputTop_ = tRelBoundDistance >= 0 ? 0 : -tRelBoundDistance;
218+ downOffsetToInputDown_ = bRelBoundDistance >= 0 ? bRelBoundDistance : 0;
219+ hInputActualNoPad_ = hInputActualPad_ - topOffsetToInputTop_ - downOffsetToInputDown_;
220+ wInputActualNoPad_ = wInputActualPad_ - leftOffsetToInputLeft_ - rightOffsetToInputRight_;
221+ forwardhInputOffset_ =
222+ topOffsetToInputTop_ == 0 ? forwardhInputOffset_ - tilingData_.padH * tilingData_.wOutput : 0;
223+ forwardwInputOffset_ = leftOffsetToInputLeft_ == 0 ? forwardwInputOffset_ - tilingData_.padW : 0;
224+ }
225+}
226+ 
227+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
228+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::ForwardCopyIn()
229+{
230+ LocalTensor<TYPE_ORIG_X> xLocal = inputQue_.template AllocTensor<TYPE_ORIG_X>();
231+ int64_t xGmOffset = highInputOffset_ + forwardhInputOffset_ + forwardwInputOffset_;
232+ 
233+ LoopModeParams loopModeParams;
234+ if (isPad_) {
235+ int64_t wInputActualAlignedNoPad =
236+ CeilDivision(wInputActualNoPad_, BLOCK_SIZE / sizeof(TYPE_ORIG_X)) * (BLOCK_SIZE / sizeof(TYPE_ORIG_X));
237+ loopModeParams.loop1Size = Base::highAxisActual_;
238+ loopModeParams.loop1SrcStride = Base::hOutput_ * Base::wOutput_ * sizeof(TYPE_ORIG_X);
239+ loopModeParams.loop1DstStride = hInputActualNoPad_ * wInputActualAlignedNoPad * sizeof(TYPE_ORIG_X);
240+ loopModeParams.loop2Size = 1;
241+ loopModeParams.loop2SrcStride = 0;
242+ loopModeParams.loop2DstStride = 0;
243+ } else {
244+ loopModeParams.loop1Size = Base::highAxisActual_;
245+ loopModeParams.loop1SrcStride = Base::hOutput_ * Base::wOutput_ * sizeof(TYPE_ORIG_X);
246+ loopModeParams.loop1DstStride = hInputActualPad_ * wInputActualAlignedPad_ * sizeof(TYPE_ORIG_X);
247+ loopModeParams.loop2Size = 1;
248+ loopModeParams.loop2SrcStride = 0;
249+ loopModeParams.loop2DstStride = 0;
250+ }
251+ SetLoopModePara(loopModeParams, DataCopyMVType::OUT_TO_UB);
252+ DataCopyPadExtParams<TYPE_ORIG_X> padParams = {false, 0, 0, 0};
253+ DataCopyExtParams copyParams;
254+ if (isPad_) {
255+ copyParams.blockCount = static_cast<uint16_t>(hInputActualNoPad_);
256+ copyParams.blockLen = static_cast<uint32_t>(wInputActualNoPad_ * sizeof(TYPE_ORIG_X));
257+ copyParams.srcStride = static_cast<uint32_t>((Base::wOutput_ - wInputActualNoPad_) * sizeof(TYPE_ORIG_X));
258+ copyParams.dstStride = 0;
259+ copyParams.rsv = 0;
260+ } else {
261+ copyParams.blockCount = static_cast<uint16_t>(hInputActualPad_);
262+ copyParams.blockLen = static_cast<uint32_t>(wInputActualPad_ * sizeof(TYPE_ORIG_X));
263+ copyParams.srcStride = static_cast<uint32_t>((Base::wOutput_ - wInputActualPad_) * sizeof(TYPE_ORIG_X));
264+ copyParams.dstStride = 0;
265+ copyParams.rsv = 0;
266+ }
267+ DataCopyPad(xLocal, xGm_[xGmOffset], copyParams, padParams);
268+ ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
269+ inputQue_.EnQue(xLocal);
270+}
271+ 
272+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
273+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::Process()
274+{
275+ if (Base::blockIdx_ >= Base::usedCoreNum_) {
276+ return;
277+ }
278+ 
279+ for (int64_t loopNum = 0; loopNum < Base::curCoreProcessNum_; loopNum++) {
280+ Base::ScalarCompute(loopNum);
281+ PipeBarrier<PIPE_ALL>();
282+ if (Base::hArgmaxActual_ <= 0 || Base::wArgmaxActual_ <= 0) {
283+ Base::ProcessNoArgmaxBlock();
284+ continue;
285+ }
286+ ForwardScalarCompute();
287+ ForwardCopyIn();
288+ Forward();
289+ Base::CopyInGrad();
290+ Backward();
291+ Base::CopyOut();
292+ }
293+}
294+ 
295+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
296+__aicore__ inline void
297+PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::ConvertIndexWithoutPadAlign(
298+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, TYPE_ARGMAX left, TYPE_ARGMAX wInput,
299+ TYPE_ARGMAX hIndexBase, MicroAPI::RegTensor<TYPE_ARGMAX>& dstReg, int32_t ncInputOffset)
300+{
301+ if (isPad_) {
302+ ConvertIndexWithoutPadAlignCommon<TYPE_ARGMAX, 1>(
303+ srcReg, wStrideOffset, left, wInput, hIndexBase, dstReg, ncInputOffset);
304+ } else {
305+ ConvertIndexWithoutPadAlignCommon<TYPE_ARGMAX, 0>(
306+ srcReg, wStrideOffset, left, wInput, hIndexBase, dstReg, ncInputOffset);
307+ }
308+}
309+ 
310+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
311+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::ProcessW(
312+ __local_mem__ TYPE_ORIG_X* computeAddr, int32_t hOffset, uint16_t wStrideOffset,
313+ MicroAPI::RegTensor<int32_t>& indexReg, uint16_t hKernel, uint16_t wKernel, uint16_t repeatElem,
314+ MicroAPI::RegTensor<int32_t>& maxIndexReg, uint32_t hDilation, uint32_t wDilation)
315+{
316+ MicroAPI::RegTensor<int32_t> indexWithOffset;
317+ MicroAPI::RegTensor<TYPE_ORIG_X> calcReg;
318+ MicroAPI::RegTensor<int32_t> calcMaxIndexReg;
319+ uint32_t maskCount = repeatElem;
320+ MicroAPI::MaskReg allMaskU32 = MicroAPI::CreateMask<int32_t, MicroAPI::MaskPattern::ALL>();
321+ MicroAPI::MaskReg gatherMask = MicroAPI::UpdateMask<TYPE_ORIG_X>(maskCount);
322+ MicroAPI::RegTensor<TYPE_ORIG_X> maxReg;
323+ MicroAPI::MaskReg neMask;
324+ MicroAPI::MaskReg gtMask;
325+ MicroAPI::MaskReg tmpMask;
326+ MicroAPI::UnalignReg u0;
327+ DuplicateNegInfReg<TYPE_ORIG_X>(maxReg);
328+ 
329+ MicroAPI::Adds(maxIndexReg, indexReg, hOffset, allMaskU32);
330+ for (int32_t hIndex = 0; hIndex < hKernel; hIndex++) {
331+ for (int32_t wIndex = 0; wIndex < wKernel; wIndex++) {
332+ int32_t relIndex = hIndex * wStrideOffset * hDilation + wIndex * wDilation;
333+ int32_t offset = static_cast<int32_t>(hOffset + relIndex);
334+ MicroAPI::Adds(indexWithOffset, indexReg, offset, allMaskU32);
335+ if constexpr (std::is_same<TYPE_ORIG_X, float>::value) {
336+ MicroAPI::DataCopyGather(
337+ calcReg, computeAddr, (MicroAPI::RegTensor<uint32_t>&)indexWithOffset, gatherMask);
338+ } else {
339+ MicroAPI::RegTensor<uint16_t> indexConvert;
340+ MicroAPI::Pack(indexConvert, indexWithOffset);
341+ MicroAPI::DataCopyGather(calcReg, computeAddr, indexConvert, gatherMask);
342+ }
343+ 
344+ MicroAPI::Compare<TYPE_ORIG_X, CMPMODE::GT>(gtMask, calcReg, maxReg, gatherMask);
345+ MicroAPI::Compare<TYPE_ORIG_X, CMPMODE::NE>(neMask, calcReg, calcReg, gatherMask);
346+ MicroAPI::MaskOr(gtMask, gtMask, neMask, gatherMask);
347+ 
348+ if constexpr (sizeof(int32_t) / sizeof(TYPE_ORIG_X) == 1) {
349+ MicroAPI::Select(maxIndexReg, indexWithOffset, maxIndexReg, gtMask);
350+ } else {
351+ MicroAPI::MaskUnPack(tmpMask, gtMask);
352+ MicroAPI::Select(maxIndexReg, indexWithOffset, maxIndexReg, tmpMask);
353+ }
354+ MicroAPI::Max(maxReg, maxReg, calcReg, gatherMask);
355+ }
356+ }
357+}
358+ 
359+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
360+__aicore__ inline void
361+PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::ConvertIndexWithoutPadAlignNc(
362+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, TYPE_ARGMAX left, TYPE_ARGMAX wInput,
363+ TYPE_ARGMAX hIndexBase, MicroAPI::RegTensor<TYPE_ARGMAX>& dstReg, int32_t ncInputOffset, int32_t ncOutputCount,
364+ int32_t inputNcSize)
365+{
366+ if (isPad_) {
367+ ConvertIndexWithoutPadAlignNcCommon<TYPE_ARGMAX, 1>(
368+ srcReg, wStrideOffset, left, wInput, hIndexBase, dstReg, ncInputOffset, ncOutputCount, inputNcSize);
369+ } else {
370+ ConvertIndexWithoutPadAlignNcCommon<TYPE_ARGMAX, 0>(
371+ srcReg, wStrideOffset, left, wInput, hIndexBase, dstReg, ncInputOffset, ncOutputCount, inputNcSize);
372+ }
373+}
374+ 
375+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
376+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::SingleRowGather(
377+ __local_mem__ TYPE_ORIG_X* computeAddr, __local_mem__ TYPE_ARGMAX* argmaxAddr)
378+{
379+ uint16_t loopW = static_cast<uint16_t>(Base::wArgmaxActual_) / Base::vlT2_;
380+ uint16_t repeatsElem = Base::vlT2_;
381+ uint16_t tailRepeatsElem = static_cast<uint16_t>(Base::wArgmaxActual_) - loopW * Base::vlT2_;
382+ if (tailRepeatsElem == 0) {
383+ loopW = loopW - 1;
384+ tailRepeatsElem = repeatsElem;
385+ }
386+ uint16_t hKernel = Base::kernelH_;
387+ uint16_t wKernel = Base::kernelW_;
388+ uint32_t wStride = Base::strideW_;
389+ TYPE_ARGMAX left = static_cast<TYPE_ARGMAX>(Base::wArgmaxActualStart_ * wStride - Base::padW_);
390+ TYPE_ARGMAX hIndexBase = static_cast<TYPE_ARGMAX>(Base::hArgmaxActualStart_ * Base::strideH_ - Base::padH_);
391+ TYPE_ARGMAX wInput = static_cast<TYPE_ARGMAX>(Base::wOutput_);
392+ uint32_t highAxisActual = static_cast<uint32_t>(forwardhighAxisActual_);
393+ uint32_t hOutputActual = static_cast<uint32_t>(Base::hArgmaxActual_);
394+ uint32_t wOutputActual = static_cast<uint32_t>(Base::wArgmaxActual_);
395+ uint32_t hInputActualPad = static_cast<uint32_t>(hInputActualPad_);
396+ uint32_t wInputActualAlignedPad = static_cast<uint32_t>(wInputActualAlignedPad_);
397+ uint32_t hDilation = Base::dilationH_;
398+ uint32_t wDilation = Base::dilationW_;
399+ 
400+ for (uint16_t nc = 0; nc < static_cast<uint16_t>(highAxisActual); nc++) {
401+ for (uint16_t hLoop = 0; hLoop < static_cast<uint16_t>(hOutputActual); hLoop++) {
402+ __VEC_SCOPE__
403+ {
404+ MicroAPI::RegTensor<int32_t> indexReg;
405+ MicroAPI::RegTensor<int32_t> maxIndexReg;
406+ MicroAPI::RegTensor<TYPE_ARGMAX> maxIndexConvertReg;
407+ MicroAPI::UnalignReg u1;
408+ MicroAPI::Arange(indexReg, static_cast<int32_t>(0));
409+ AscendC::MicroAPI::MaskReg allMaskI32 =
410+ AscendC::MicroAPI::CreateMask<int32_t, AscendC::MicroAPI::MaskPattern::ALL>();
411+ MicroAPI::Muls(indexReg, indexReg, static_cast<int32_t>(wStride), allMaskI32);
412+ int32_t ncInputOffset = nc * hInputActualPad * wInputActualAlignedPad;
413+ int32_t ncOutOffset = nc * hOutputActual * wOutputActual;
414+ int32_t vfMaxAddrOffset = ncOutOffset + hLoop * wOutputActual;
415+ __local_mem__ TYPE_ARGMAX* argmaxAddrLocal = argmaxAddr + vfMaxAddrOffset;
416+ for (uint16_t wLoop = 0; wLoop < loopW; wLoop++) {
417+ int32_t wOffset =
418+ ncInputOffset + hLoop * wInputActualAlignedPad * Base::strideH_ + wLoop * repeatsElem * wStride;
419+ ProcessW(
420+ computeAddr, wOffset, static_cast<uint16_t>(wInputActualAlignedPad), indexReg, hKernel, wKernel,
421+ repeatsElem, maxIndexReg, hDilation, wDilation);
422+ ConvertIndexWithoutPadAlign(
423+ maxIndexReg, static_cast<uint32_t>(wInputActualAlignedPad), left, wInput, hIndexBase,
424+ maxIndexConvertReg, ncInputOffset);
425+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, repeatsElem);
426+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
427+ }
428+ int32_t wOffsetTail =
429+ ncInputOffset + hLoop * wInputActualAlignedPad * Base::strideH_ + loopW * repeatsElem * wStride;
430+ ProcessW(
431+ computeAddr, wOffsetTail, static_cast<uint16_t>(wInputActualAlignedPad), indexReg, hKernel, wKernel,
432+ tailRepeatsElem, maxIndexReg, hDilation, wDilation);
433+ ConvertIndexWithoutPadAlign(
434+ maxIndexReg, static_cast<uint32_t>(wInputActualAlignedPad), left, wInput, hIndexBase,
435+ maxIndexConvertReg, ncInputOffset);
436+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, tailRepeatsElem);
437+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
438+ }
439+ }
440+ }
441+ return;
442+}
443+ 
444+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
445+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::MultiRowGather(
446+ __local_mem__ TYPE_ORIG_X* computeAddr, __local_mem__ TYPE_ARGMAX* argmaxAddr)
447+{
448+ uint32_t wOutputActual = static_cast<uint32_t>(Base::wArgmaxActual_);
449+ uint16_t hKernel = Base::kernelH_;
450+ uint16_t wKernel = Base::kernelW_;
451+ uint32_t wStride = Base::strideW_;
452+ uint32_t rate2D = wInputActualAlignedPad_ * Base::strideH_;
453+ uint16_t hBatchCount = Base::vlT2_ / wOutputActual;
454+ uint16_t hLoopTimes = static_cast<uint16_t>(Base::hArgmaxActual_) / hBatchCount;
455+ uint16_t hTail = static_cast<uint16_t>(Base::hArgmaxActual_) - hLoopTimes * hBatchCount;
456+ if (hTail == 0) {
457+ hLoopTimes = hLoopTimes - 1;
458+ hTail = hBatchCount;
459+ }
460+ uint16_t repeatsElem = hBatchCount * wOutputActual;
461+ uint16_t tailRepeatsElem = hTail * wOutputActual;
462+ TYPE_ARGMAX left = static_cast<TYPE_ARGMAX>(Base::wArgmaxActualStart_ * wStride - Base::padW_);
463+ TYPE_ARGMAX hIndexBase = static_cast<TYPE_ARGMAX>(Base::hArgmaxActualStart_ * Base::strideH_ - Base::padH_);
464+ TYPE_ARGMAX wInput = static_cast<TYPE_ARGMAX>(Base::wOutput_);
465+ uint32_t highAxisActual = static_cast<uint32_t>(Base::highAxisActual_);
466+ uint32_t hInputActualPad = static_cast<uint32_t>(hInputActualPad_);
467+ uint32_t wInputActualAlignedPad = static_cast<uint32_t>(wInputActualAlignedPad_);
468+ uint32_t hOutputActual = static_cast<uint32_t>(Base::hArgmaxActual_);
469+ uint32_t hStride = Base::strideH_;
470+ uint32_t hDilation = Base::dilationH_;
471+ uint32_t wDilation = Base::dilationW_;
472+ 
473+ __VEC_SCOPE__
474+ {
475+ MicroAPI::RegTensor<int32_t> indexReg;
476+ MicroAPI::RegTensor<int32_t> maxIndexReg;
477+ MicroAPI::RegTensor<TYPE_ARGMAX> maxIndexConvertReg;
478+ MicroAPI::UnalignReg u1;
479+ __local_mem__ TYPE_ARGMAX* argmaxAddrLocal = argmaxAddr;
480+ GenGatterIndex2D<int32_t>(
481+ indexReg, static_cast<int32_t>(rate2D), static_cast<int32_t>(wOutputActual), static_cast<int32_t>(wStride));
482+ for (uint16_t nc = 0; nc < static_cast<uint16_t>(highAxisActual); nc++) {
483+ int32_t ncInputOffset = nc * hInputActualPad * wInputActualAlignedPad;
484+ for (uint16_t hLoop = 0; hLoop < hLoopTimes; hLoop++) {
485+ int32_t wOffset = ncInputOffset + hLoop * hBatchCount * hStride * wInputActualAlignedPad;
486+ ProcessW(
487+ computeAddr, wOffset, static_cast<uint16_t>(wInputActualAlignedPad), indexReg, hKernel, wKernel,
488+ repeatsElem, maxIndexReg, hDilation, wDilation);
489+ ConvertIndexWithoutPadAlign(
490+ maxIndexReg, static_cast<uint32_t>(wInputActualAlignedPad), left, wInput, hIndexBase,
491+ maxIndexConvertReg, ncInputOffset);
492+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, repeatsElem);
493+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
494+ }
495+ int32_t wOffsetTail = ncInputOffset + hLoopTimes * hBatchCount * hStride * wInputActualAlignedPad;
496+ ProcessW(
497+ computeAddr, wOffsetTail, static_cast<uint16_t>(wInputActualAlignedPad), indexReg, hKernel, wKernel,
498+ tailRepeatsElem, maxIndexReg, hDilation, wDilation);
499+ ConvertIndexWithoutPadAlign(
500+ maxIndexReg, static_cast<uint32_t>(wInputActualAlignedPad), left, wInput, hIndexBase,
501+ maxIndexConvertReg, ncInputOffset);
502+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, tailRepeatsElem);
503+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
504+ }
505+ }
506+ return;
507+}
508+ 
509+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
510+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::MultiNcGather(
511+ __local_mem__ TYPE_ORIG_X* computeAddr, __local_mem__ TYPE_ARGMAX* argmaxAddr)
512+{
513+ uint16_t wKernel = Base::kernelW_;
514+ uint16_t hKernel = Base::kernelH_;
515+ uint32_t wStride = Base::strideW_;
516+ uint16_t rate3D = hInputActualPad_ * wInputActualAlignedPad_;
517+ uint16_t num2D = static_cast<uint16_t>(Base::hArgmaxActual_ * Base::wArgmaxActual_);
518+ uint16_t rate2D = Base::strideH_ * wInputActualAlignedPad_;
519+ uint16_t wOutputActual = static_cast<uint16_t>(Base::wArgmaxActual_);
520+ uint16_t eachBatchCount = static_cast<uint16_t>(Base::hArgmaxActual_ * Base::wArgmaxActual_);
521+ uint16_t ncBatchCount = Base::vlT2_ / eachBatchCount;
522+ uint16_t ncLoopTimes = static_cast<uint16_t>(forwardhighAxisActual_) / ncBatchCount;
523+ uint16_t ncTail = static_cast<uint16_t>(forwardhighAxisActual_) - ncLoopTimes * ncBatchCount;
524+ if (ncTail == 0) {
525+ ncLoopTimes = ncLoopTimes - 1;
526+ ncTail = ncBatchCount;
527+ }
528+ uint16_t repeatsElem = ncBatchCount * eachBatchCount;
529+ uint16_t tailRepeatsElem = ncTail * eachBatchCount;
530+ TYPE_ARGMAX left = static_cast<TYPE_ARGMAX>(Base::wArgmaxActualStart_ * wStride - Base::padW_);
531+ TYPE_ARGMAX hIndexBase = static_cast<TYPE_ARGMAX>(Base::hArgmaxActualStart_ * Base::strideH_ - Base::padH_);
532+ TYPE_ARGMAX wInput = static_cast<TYPE_ARGMAX>(Base::wOutput_);
533+ uint32_t hInputActualPad = static_cast<uint32_t>(hInputActualPad_);
534+ uint32_t wInputActualAlignedPad = static_cast<uint32_t>(wInputActualAlignedPad_);
535+ uint32_t hOutputActual = static_cast<uint32_t>(Base::hArgmaxActual_);
536+ uint32_t hDilation = Base::dilationH_;
537+ uint32_t wDilation = Base::dilationW_;
538+ 
539+ __VEC_SCOPE__
540+ {
541+ MicroAPI::RegTensor<int32_t> indexReg;
542+ MicroAPI::RegTensor<int32_t> maxIndexReg;
543+ MicroAPI::RegTensor<TYPE_ARGMAX> maxIndexConvertReg;
544+ MicroAPI::UnalignReg u1;
545+ __local_mem__ TYPE_ARGMAX* argmaxAddrLocal = argmaxAddr;
546+ GenGatterIndex3D<int32_t>(
547+ indexReg, static_cast<int32_t>(rate3D), static_cast<int32_t>(num2D), static_cast<int32_t>(rate2D),
548+ static_cast<int32_t>(wOutputActual), static_cast<int32_t>(wStride));
549+ for (uint16_t nc = 0; nc < ncLoopTimes; nc++) {
550+ uint32_t ncInputOffset = nc * ncBatchCount * hInputActualPad * wInputActualAlignedPad;
551+ int32_t hOffset = static_cast<int32_t>(nc) * ncBatchCount * rate3D;
552+ ProcessW(
553+ computeAddr, hOffset, static_cast<uint16_t>(wInputActualAlignedPad), indexReg, hKernel, wKernel,
554+ repeatsElem, maxIndexReg, hDilation, wDilation);
555+ ConvertIndexWithoutPadAlignNc(
556+ maxIndexReg, static_cast<uint32_t>(wInputActualAlignedPad), left, wInput, hIndexBase,
557+ maxIndexConvertReg, ncInputOffset, num2D, rate3D);
558+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, repeatsElem);
559+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
560+ }
561+ uint32_t ncInputOffsetTail = ncLoopTimes * ncBatchCount * hInputActualPad * wInputActualAlignedPad;
562+ int32_t hOffset = static_cast<int32_t>(ncLoopTimes) * ncBatchCount * rate3D;
563+ ProcessW(
564+ computeAddr, hOffset, static_cast<uint16_t>(wInputActualAlignedPad), indexReg, hKernel, wKernel,
565+ tailRepeatsElem, maxIndexReg, hDilation, wDilation);
566+ ConvertIndexWithoutPadAlignNc(
567+ maxIndexReg, static_cast<uint32_t>(wInputActualAlignedPad), left, wInput, hIndexBase, maxIndexConvertReg,
568+ ncInputOffsetTail, num2D, rate3D);
569+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, tailRepeatsElem);
570+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
571+ }
572+ return;
573+}
574+ 
575+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
576+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::DupBufferNegInf(
577+ __local_mem__ TYPE_ORIG_X* dstAddr, uint32_t repeatElm, uint16_t loop, uint32_t tail)
578+{
579+ DupBufferNegInfCommon<TYPE_ORIG_X>(dstAddr, repeatElm, loop, tail);
580+}
581+ 
582+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
583+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::CopyToCalcBuffer(
584+ __local_mem__ TYPE_ORIG_X* dstAddr, __local_mem__ TYPE_ORIG_X* srcAddr, uint16_t batch, uint16_t rows,
585+ uint16_t loopCols, uint16_t tailCols, uint32_t repeatElm, uint32_t srcBatchStride, uint32_t srcRowStride,
586+ uint32_t dstBatchStride, uint32_t dstRowStride, uint32_t dstRowOffset, uint32_t dstColOffset)
587+{
588+ CopyToCalcBuffer2DCommon<TYPE_ORIG_X>(
589+ dstAddr, srcAddr, batch, rows, loopCols, tailCols, repeatElm, srcBatchStride, srcRowStride, dstBatchStride,
590+ dstRowStride, dstRowOffset, dstColOffset);
591+}
592+ 
593+template <typename TYPE_ORIG_X, typename TYPE_ARGMAX, typename T3, const uint32_t IS_CHECK_RANGE>
594+__aicore__ inline void PoolGradNCHWSmallKernel<TYPE_ORIG_X, TYPE_ARGMAX, T3, IS_CHECK_RANGE>::DupAndCopyToCalcBuffer(
595+ __local_mem__ TYPE_ORIG_X* dstAddr, __local_mem__ TYPE_ORIG_X* srcAddr)
596+{
597+ uint32_t wInputPadActualAlign =
598+ CeilDivision(wInputActualPad_, MAX_DATA_NUM_IN_ONE_BLOCK) * MAX_DATA_NUM_IN_ONE_BLOCK;
599+ uint16_t hPad = topOffsetToInputTop_;
600+ uint16_t wPad = leftOffsetToInputLeft_;
601+ uint16_t hRows = hInputActualNoPad_;
602+ uint16_t wCols = wInputActualNoPad_;
603+ uint16_t highAxis = Base::highAxisActual_;
604+ uint32_t srcBatchStride =
605+ hInputActualNoPad_ * CeilDivision(wInputActualNoPad_, MAX_DATA_NUM_IN_ONE_BLOCK) * MAX_DATA_NUM_IN_ONE_BLOCK;
606+ uint32_t dstBatchStride = (hRows + hPad + downOffsetToInputDown_) * wInputPadActualAlign;
607+ uint32_t srcRowStride = CeilDivision(wInputActualNoPad_, MAX_DATA_NUM_IN_ONE_BLOCK) * MAX_DATA_NUM_IN_ONE_BLOCK;
608+ uint32_t dstRowStride = wInputPadActualAlign;
609+ uint32_t repeatElm = MAX_DATA_NUM_IN_ONE_BLOCK;
610+ uint16_t loopCols = wCols / repeatElm;
611+ uint16_t tailCols = wCols % repeatElm;
612+ __VEC_SCOPE__
613+ {
614+ DupBufferNegInf(
615+ dstAddr, repeatElm, highAxis * (hRows + hPad + downOffsetToInputDown_) * (wInputPadActualAlign / repeatElm),
616+ repeatElm);
617+ CopyToCalcBuffer(
618+ dstAddr, srcAddr, highAxis, hRows, loopCols, tailCols, repeatElm, srcBatchStride, srcRowStride,
619+ dstBatchStride, dstRowStride, hPad, wPad);
620+ }
621+};
622+ 
623+} // namespace MaxPoolGradNCHWSmallKernelNameSpace
624+#endif // MAX_POOL_GRAD_SMALL_KERNEL_H
Mpooling/max_pool_grad/op_kernel/arch35/max_pool_grad_struct.h+2-11
@@ -51,19 +51,10 @@ ASCENDC_TPL_SEL(
51 ASCENDC_TPL_UINT_SEL(isCheckRange, ASCENDC_TPL_UI_LIST, TPL_NO_CHECK_RANGE),51 ASCENDC_TPL_UINT_SEL(isCheckRange, ASCENDC_TPL_UI_LIST, TPL_NO_CHECK_RANGE),
52 ASCENDC_TPL_TILING_STRUCT_SEL(MaxPoolGradWithArgmaxSimtTilingCommonData)),52 ASCENDC_TPL_TILING_STRUCT_SEL(MaxPoolGradWithArgmaxSimtTilingCommonData)),
53 53 
54- // NCHW BIG kernel - format must be NCHW54+ // SIMD kernel - format must be NCHW
55 ASCENDC_TPL_ARGS_SEL(55 ASCENDC_TPL_ARGS_SEL(
56 ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),56 ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
57- ASCENDC_TPL_UINT_SEL(kernelMode, ASCENDC_TPL_UI_LIST, TPL_NCHW_BIG_KERNEL),57+ ASCENDC_TPL_UINT_SEL(kernelMode, ASCENDC_TPL_UI_LIST, TPL_NCHW_BIG_KERNEL, TPL_NCHW_SMALL_KERNEL),
58- ASCENDC_TPL_UINT_SEL(format, ASCENDC_TPL_UI_LIST, TPL_NCHW_FORMAT),
59- ASCENDC_TPL_UINT_SEL(indicesDtype, ASCENDC_TPL_UI_LIST, TPL_INT64, TPL_INT32),
60- ASCENDC_TPL_UINT_SEL(isCheckRange, ASCENDC_TPL_UI_LIST, TPL_NO_CHECK_RANGE, TPL_CHECK_RANGE),
61- ASCENDC_TPL_TILING_STRUCT_SEL(MaxPoolGradWithArgmaxNCHWTilingCommonData)),
62- 
63- // NCHW SMALL kernel - format must be NCHW
64- ASCENDC_TPL_ARGS_SEL(
65- ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
66- ASCENDC_TPL_UINT_SEL(kernelMode, ASCENDC_TPL_UI_LIST, TPL_NCHW_SMALL_KERNEL),
67 ASCENDC_TPL_UINT_SEL(format, ASCENDC_TPL_UI_LIST, TPL_NCHW_FORMAT),58 ASCENDC_TPL_UINT_SEL(format, ASCENDC_TPL_UI_LIST, TPL_NCHW_FORMAT),
68 ASCENDC_TPL_UINT_SEL(indicesDtype, ASCENDC_TPL_UI_LIST, TPL_INT64, TPL_INT32),59 ASCENDC_TPL_UINT_SEL(indicesDtype, ASCENDC_TPL_UI_LIST, TPL_INT64, TPL_INT32),
69 ASCENDC_TPL_UINT_SEL(isCheckRange, ASCENDC_TPL_UI_LIST, TPL_NO_CHECK_RANGE, TPL_CHECK_RANGE),60 ASCENDC_TPL_UINT_SEL(isCheckRange, ASCENDC_TPL_UI_LIST, TPL_NO_CHECK_RANGE, TPL_CHECK_RANGE),
Mpooling/max_pool_grad/op_kernel/max_pool_grad.cpp+43-0
@@ -12,11 +12,15 @@
12#include "kernel_tiling/kernel_tiling.h"12#include "kernel_tiling/kernel_tiling.h"
13#include "arch35/max_pool_grad_struct.h"13#include "arch35/max_pool_grad_struct.h"
14#include "arch35/max_pool_grad_simt.h"14#include "arch35/max_pool_grad_simt.h"
15+#include "arch35/max_pool_grad_nchw_big_kernel.h"
16+#include "arch35/max_pool_grad_nchw_small_kernel.h"
15 17 
16using namespace AscendC;18using namespace AscendC;
17using MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData;19using MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData;
18using MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxSimtTilingCommonData;20using MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxSimtTilingCommonData;
19using namespace PoolGradNameSpace;21using namespace PoolGradNameSpace;
22+using namespace MaxPoolGradNCHWBigKernelNameSpace;
23+using namespace MaxPoolGradNCHWSmallKernelNameSpace;
20 24 
21template <25template <
22 uint64_t KERNEL_MODE = TPL_SIMT_KERNEL, uint64_t FORMAT = TPL_NCHW_FORMAT, uint64_t INDICES_DTYPE = TPL_INT32,26 uint64_t KERNEL_MODE = TPL_SIMT_KERNEL, uint64_t FORMAT = TPL_NCHW_FORMAT, uint64_t INDICES_DTYPE = TPL_INT32,
@@ -28,6 +32,7 @@ __global__ __aicore__ void max_pool_grad(
28 return;32 return;
29 }33 }
30 TPipe pipe;34 TPipe pipe;
35+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
31 REGISTER_TILING_DEFAULT(MaxPoolGradWithArgmaxSimtTilingCommonData);36 REGISTER_TILING_DEFAULT(MaxPoolGradWithArgmaxSimtTilingCommonData);
32 37 
33 if constexpr (KERNEL_MODE == TPL_SIMT_KERNEL) {38 if constexpr (KERNEL_MODE == TPL_SIMT_KERNEL) {
@@ -41,5 +46,43 @@ __global__ __aicore__ void max_pool_grad(
41 op.Init(orig_x, orig_y, grads, y, workspace);46 op.Init(orig_x, orig_y, grads, y, workspace);
42 op.Process();47 op.Process();
43 }48 }
49+ } else if constexpr (KERNEL_MODE == TPL_NCHW_BIG_KERNEL) {
50+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxNCHWTilingCommonData, tilingData, tiling);
51+ if constexpr (INDICES_DTYPE == TPL_INT32 && IS_CHECK_RANGE == TPL_CHECK_RANGE) {
52+ MaxPoolGradNCHWBigKernel<DTYPE_X1, int32_t, int32_t, true> op;
53+ op.Init(orig_x, orig_y, grads, y, pipe, tilingData);
54+ op.Process();
55+ } else if constexpr (INDICES_DTYPE == TPL_INT32 && IS_CHECK_RANGE == TPL_NO_CHECK_RANGE) {
56+ MaxPoolGradNCHWBigKernel<DTYPE_X1, int32_t, int32_t , false> op;
57+ op.Init(orig_x, orig_y, grads, y, pipe, tilingData);
58+ op.Process();
59+ } else if constexpr (INDICES_DTYPE == TPL_INT64 && IS_CHECK_RANGE == TPL_CHECK_RANGE) {
60+ MaxPoolGradNCHWBigKernel<DTYPE_X1, int64_t, int64_t , true> op;
61+ op.Init(orig_x, orig_y, grads, y, pipe, tilingData);
62+ op.Process();
63+ } else if constexpr (INDICES_DTYPE == TPL_INT64 && IS_CHECK_RANGE == TPL_NO_CHECK_RANGE) {
64+ MaxPoolGradNCHWBigKernel<DTYPE_X1, int64_t, int64_t , false> op;
65+ op.Init(orig_x, orig_y, grads, y, pipe, tilingData);
66+ op.Process();
67+ }
68+ } else if constexpr (KERNEL_MODE == TPL_NCHW_SMALL_KERNEL) {
69+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxNCHWTilingCommonData, tilingData, tiling);
70+ if constexpr (INDICES_DTYPE == TPL_INT32 && IS_CHECK_RANGE == TPL_CHECK_RANGE) {
71+ PoolGradNCHWSmallKernel<DTYPE_X1, int32_t, int32_t, true> op(pipe, tilingData);
72+ op.Init(orig_x, orig_y, grads, y);
73+ op.Process();
74+ } else if constexpr (INDICES_DTYPE == TPL_INT64 && IS_CHECK_RANGE == TPL_CHECK_RANGE) {
75+ PoolGradNCHWSmallKernel<DTYPE_X1, int64_t, int64_t, true> op(pipe, tilingData);
76+ op.Init(orig_x, orig_y, grads, y);
77+ op.Process();
78+ } else if constexpr (INDICES_DTYPE == TPL_INT32 && IS_CHECK_RANGE == TPL_NO_CHECK_RANGE) {
79+ PoolGradNCHWSmallKernel<DTYPE_X1, int32_t, int32_t, false> op(pipe, tilingData);
80+ op.Init(orig_x, orig_y, grads, y);
81+ op.Process();
82+ } else if constexpr (INDICES_DTYPE == TPL_INT64 && IS_CHECK_RANGE == TPL_NO_CHECK_RANGE) {
83+ PoolGradNCHWSmallKernel<DTYPE_X1, int64_t, int64_t, false> op(pipe, tilingData);
84+ op.Init(orig_x, orig_y, grads, y);
85+ op.Process();
86+ }
44 }87 }
45}88}
Mpooling/max_pool_grad/tests/ut/op_host/test_max_pool_grad_tiling.cpp+110-0
@@ -254,3 +254,113 @@ TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_fp16_valid_nchw_k3)
254 x1Shape, x2Shape, gradShape, yShape, {1, 3, 3, 1}, {1, 3, 3, 1}, "VALID", "NCHW", ge::DT_FLOAT16,254 x1Shape, x2Shape, gradShape, yShape, {1, 3, 3, 1}, {1, 3, 3, 1}, "VALID", "NCHW", ge::DT_FLOAT16,
255 ge::GRAPH_SUCCESS);255 ge::GRAPH_SUCCESS);
256}256}
257+ 
258+// ============================================================
259+// Big kernel cases
260+// ============================================================
261+ 
262+TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_big_kernel_fp16_valid_nchw_k16_s16)
263+{
264+ gert::StorageShape x1Shape = {{1, 4, 64, 64}, {1, 4, 64, 64}};
265+ gert::StorageShape x2Shape = {{1, 4, 4, 4}, {1, 4, 4, 4}};
266+ gert::StorageShape gradShape = {{1, 4, 4, 4}, {1, 4, 4, 4}};
267+ gert::StorageShape yShape = {{1, 4, 64, 64}, {1, 4, 64, 64}};
268+ 
269+ ExecuteTilingTestCase(
270+ x1Shape, x2Shape, gradShape, yShape,
271+ {1, 1, 16, 16},
272+ {1, 1, 16, 16},
273+ "VALID",
274+ "NCHW",
275+ ge::DT_FLOAT16,
276+ ge::GRAPH_SUCCESS);
277+}
278+ 
279+TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_big_kernel_bf16_valid_nchw_k20_s10)
280+{
281+ gert::StorageShape x1Shape = {{1, 4, 80, 80}, {1, 4, 80, 80}};
282+ gert::StorageShape x2Shape = {{1, 4, 7, 7}, {1, 4, 7, 7}};
283+ gert::StorageShape gradShape = {{1, 4, 7, 7}, {1, 4, 7, 7}};
284+ gert::StorageShape yShape = {{1, 4, 80, 80}, {1, 4, 80, 80}};
285+ 
286+ ExecuteTilingTestCase(
287+ x1Shape, x2Shape, gradShape, yShape,
288+ {1, 1, 20, 20},
289+ {1, 1, 10, 10},
290+ "VALID",
291+ "NCHW",
292+ ge::DT_BF16,
293+ ge::GRAPH_SUCCESS);
294+}
295+ 
296+TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_big_kernel_bf16_same_nchw_k18_10_large_w)
297+{
298+ gert::StorageShape x1Shape = {{4, 4, 150, 3417}, {4, 4, 150, 3417}};
299+ gert::StorageShape x2Shape = {{4, 4, 15, 285}, {4, 4, 15, 285}};
300+ gert::StorageShape gradShape = {{4, 4, 15, 285}, {4, 4, 15, 285}};
301+ gert::StorageShape yShape = {{4, 4, 150, 3417}, {4, 4, 150, 3417}};
302+ 
303+ ExecuteTilingTestCase(
304+ x1Shape, x2Shape, gradShape, yShape,
305+ {1, 1, 18, 10},
306+ {1, 1, 10, 12},
307+ "SAME",
308+ "NCHW",
309+ ge::DT_BF16,
310+ ge::GRAPH_SUCCESS);
311+}
312+ 
313+// ============================================================
314+// Small kernel cases
315+// ============================================================
316+ 
317+TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_small_kernel_fp32_valid_nchw_k2_s2)
318+{
319+ gert::StorageShape x1Shape = {{2, 16, 64, 128}, {2, 16, 64, 128}};
320+ gert::StorageShape x2Shape = {{2, 16, 32, 64}, {2, 16, 32, 64}};
321+ gert::StorageShape gradShape = {{2, 16, 32, 64}, {2, 16, 32, 64}};
322+ gert::StorageShape yShape = {{2, 16, 64, 128}, {2, 16, 64, 128}};
323+ 
324+ ExecuteTilingTestCase(
325+ x1Shape, x2Shape, gradShape, yShape,
326+ {1, 1, 2, 2},
327+ {1, 1, 2, 2},
328+ "VALID",
329+ "NCHW",
330+ ge::DT_FLOAT,
331+ ge::GRAPH_SUCCESS);
332+}
333+ 
334+TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_small_kernel_fp32_valid_nchw_k3_5_s2_4)
335+{
336+ gert::StorageShape x1Shape = {{1, 32, 35, 137}, {1, 32, 35, 137}};
337+ gert::StorageShape x2Shape = {{1, 32, 17, 34}, {1, 32, 17, 34}};
338+ gert::StorageShape gradShape = {{1, 32, 17, 34}, {1, 32, 17, 34}};
339+ gert::StorageShape yShape = {{1, 32, 35, 137}, {1, 32, 35, 137}};
340+ 
341+ ExecuteTilingTestCase(
342+ x1Shape, x2Shape, gradShape, yShape,
343+ {1, 1, 3, 5},
344+ {1, 1, 2, 4},
345+ "VALID",
346+ "NCHW",
347+ ge::DT_FLOAT,
348+ ge::GRAPH_SUCCESS);
349+}
350+ 
351+TEST_F(MaxPoolGradTiling, MaxPoolGrad_tiling_success_small_kernel_fp32_same_nchw_k2_s2_odd_shape)
352+{
353+ gert::StorageShape x1Shape = {{1, 24, 33, 65}, {1, 24, 33, 65}};
354+ gert::StorageShape x2Shape = {{1, 24, 17, 33}, {1, 24, 17, 33}};
355+ gert::StorageShape gradShape = {{1, 24, 17, 33}, {1, 24, 17, 33}};
356+ gert::StorageShape yShape = {{1, 24, 33, 65}, {1, 24, 33, 65}};
357+ 
358+ ExecuteTilingTestCase(
359+ x1Shape, x2Shape, gradShape, yShape,
360+ {1, 1, 2, 2},
361+ {1, 1, 2, 2},
362+ "SAME",
363+ "NCHW",
364+ ge::DT_FLOAT,
365+ ge::GRAPH_SUCCESS);
366+}
Mpooling/max_pool_grad_with_argmax/op_host/arch35/max_pool_grad_with_argmax_tiling_base.cpp+0-1
@@ -41,7 +41,6 @@ static constexpr int64_t DTYPE_INT64 = 9;
41static constexpr int64_t INPUT_X = 0;41static constexpr int64_t INPUT_X = 0;
42static constexpr int64_t INPUT_GRAD = 1;42static constexpr int64_t INPUT_GRAD = 1;
43static constexpr int64_t INPUT_ARGMAX = 2;43static constexpr int64_t INPUT_ARGMAX = 2;
44-static constexpr int64_t DIGIT_TWO = 2;
45 44 
46 ge::graphStatus MaxPoolGradWithArgmaxBaseTiling::GetShapeAttrsInfo() {45 ge::graphStatus MaxPoolGradWithArgmaxBaseTiling::GetShapeAttrsInfo() {
47 OP_LOGD("MaxPoolGradWithArgmax", "MaxPoolGradWithArgmaxBaseTiling::GetShapeAttrsInfo()");46 OP_LOGD("MaxPoolGradWithArgmax", "MaxPoolGradWithArgmaxBaseTiling::GetShapeAttrsInfo()");
Mpooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nchw_tiling.cpp+7-375
@@ -10,61 +10,16 @@
10 10 
11/*!11/*!
12 * \file max_pool_grad_with_argmax_v3_nchw_tiling.cpp12 * \file max_pool_grad_with_argmax_v3_nchw_tiling.cpp
13- * \brief13+ * \brief MaxPoolGradWithArgmaxV3 NCHW格式Tiling实现，使用公共Tiling类
14 */14 */
15#include "platform/platform_info.h"15#include "platform/platform_info.h"
16#include "op_host/tiling_templates_registry.h"16#include "op_host/tiling_templates_registry.h"
17#include "max_pool_grad_with_argmax_v3_nchw_tiling.h"17#include "max_pool_grad_with_argmax_v3_nchw_tiling.h"
18 18 
19namespace optiling {19namespace optiling {
20-static constexpr int64_t FLOAT16_SIZE = 2;
21-static constexpr int64_t FLOAT32_SIZE = 4;
22-static constexpr int64_t INT32_SIZE = 4;
23-static constexpr int64_t INT64_SIZE = 8;
24-static constexpr int64_t UB_RESVERVED_SIZE = 1024;
25static constexpr int64_t NO_CHECK_RANGE_TILING_KEY_NCHW = 100;20static constexpr int64_t NO_CHECK_RANGE_TILING_KEY_NCHW = 100;
26static constexpr int64_t CHECK_RANGE_TILING_KEY_NCHW = 101;21static constexpr int64_t CHECK_RANGE_TILING_KEY_NCHW = 101;
27static constexpr int64_t T3_INT64 = 10;22static constexpr int64_t T3_INT64 = 10;
28-static constexpr int64_t DOUBLE_BUFFER = 2;
29- 
30-void MaxPoolGradWithArgmaxV3NCHWTiling::InitializationVars()
31-{
32- baseData.vRegSize = Ops::Base::GetVRegSize(context_);
33- baseData.ubBlockSize = Ops::Base::GetUbBlockSize(context_);
34- baseData.inputBytes = inputData.inputDtype == ge::DT_FLOAT ? FLOAT32_SIZE : FLOAT16_SIZE;
35- baseData.indexBytes = inputData.indexDtype == ge::DT_INT32 ? INT32_SIZE : INT64_SIZE;
36- baseData.availableUb = hardwareData.ubSize - UB_RESVERVED_SIZE;
37- baseData.totalCoreNum = hardwareData.coreNum;
38- baseData.coreUsedForBestPerformance = baseData.totalCoreNum;
39- 
40- int64_t oneBlockNumT1 = baseData.ubBlockSize / baseData.inputBytes;
41- int64_t oneBlockNumT2 = baseData.ubBlockSize / baseData.indexBytes;
42- 
43- baseData.maxDataNumInOneBlock = std::max(oneBlockNumT1, oneBlockNumT2);
44- 
45- baseData.proDataNumInOneBeatT2 = baseData.vRegSize / baseData.ubBlockSize * oneBlockNumT2;
46- baseData.inputNCSize = inputData.nX * inputData.cX;
47- 
48- baseData.isPad = 0;
49- if (inputData.hPad != 0 || inputData.wPad != 0) {
50- baseData.isPad = 1;
51- }
52- 
53- baseData.hProBatchSize = 1;
54- if (inputData.hKernel > inputData.hStride) {
55- baseData.hProBatchSize = Ops::Base::CeilDiv(inputData.hKernel, inputData.hStride);
56- }
57- 
58- baseData.wProBatchSize = 1;
59- if (inputData.wKernel > inputData.wStride) {
60- baseData.wProBatchSize = Ops::Base::CeilDiv(inputData.wKernel, inputData.wStride);
61- }
62- 
63- baseData.isOverlap = 0;
64- if (baseData.wProBatchSize != 1 || baseData.hProBatchSize != 1) {
65- baseData.isOverlap = 1;
66- }
67-}
68 23 
69bool MaxPoolGradWithArgmaxV3NCHWTiling::IsCapable()24bool MaxPoolGradWithArgmaxV3NCHWTiling::IsCapable()
70{25{
@@ -72,22 +27,14 @@ bool MaxPoolGradWithArgmaxV3NCHWTiling::IsCapable()
72 return false;27 return false;
73 }28 }
74 29 
75- InitializationVars();30+ nchwTilingCommon.InitializationVars(context_, &hardwareData);
76- // all the h and w is overlapped.31+ return nchwTilingCommon.CheckUBSize();
77- if (baseData.hProBatchSize >= inputData.hGrad && baseData.wProBatchSize >= inputData.wGrad) {
78- return false;
79- }
80- // ub is not enough
81- splitData.highAxisInner = 1;
82- splitData.hOutputInner = 1;
83- splitData.wOutputInner = std::min(inputData.wX, baseData.proDataNumInOneBeatT2);
84- DoBufferCalculate();
85- return splitData.totalBufferSize <= baseData.availableUb;
86}32}
87 33 
88uint64_t MaxPoolGradWithArgmaxV3NCHWTiling::GetTilingKey() const34uint64_t MaxPoolGradWithArgmaxV3NCHWTiling::GetTilingKey() const
89{35{
90 uint64_t tilingKey = NO_CHECK_RANGE_TILING_KEY_NCHW;36 uint64_t tilingKey = NO_CHECK_RANGE_TILING_KEY_NCHW;
37+ auto splitData = nchwTilingCommon.GetSplitData();
91 if (splitData.isCheckRange == 1) {38 if (splitData.isCheckRange == 1) {
92 tilingKey = CHECK_RANGE_TILING_KEY_NCHW;39 tilingKey = CHECK_RANGE_TILING_KEY_NCHW;
93 }40 }
@@ -98,331 +45,16 @@ uint64_t MaxPoolGradWithArgmaxV3NCHWTiling::GetTilingKey() const
98 return tilingKey;45 return tilingKey;
99}46}
100 47 
101-void MaxPoolGradWithArgmaxV3NCHWTiling::DoBufferCalculate()
102-{
103- // The calculation only involves inner.
104- int64_t hInputInner = Ops::Base::CeilDiv(splitData.hOutputInner + inputData.hKernel - 1, inputData.hStride);
105- int64_t wInputInner = Ops::Base::CeilDiv(splitData.wOutputInner + inputData.wKernel - 1, inputData.wStride);
106- int64_t wInputInnerAligned = Ops::Base::CeilAlign(wInputInner, baseData.maxDataNumInOneBlock);
107- int64_t wOutputInnerAligned = Ops::Base::CeilAlign(splitData.wOutputInner, baseData.maxDataNumInOneBlock);
108- 
109- int64_t inputPlaneSizeHW = hInputInner * wInputInnerAligned;
110- int64_t outputPlaneSizeHW = splitData.hOutputInner * wOutputInnerAligned;
111- 
112- splitData.gradBufferSize = splitData.highAxisInner * inputPlaneSizeHW * baseData.inputBytes;
113- splitData.argmaxBufferSize = splitData.highAxisInner * inputPlaneSizeHW * baseData.indexBytes;
114- splitData.outputBufferSize = splitData.highAxisInner * outputPlaneSizeHW * FLOAT32_SIZE; // 累加需要提高精度
115- 
116- int64_t tmpTotalBufferSize = splitData.outputBufferSize + splitData.gradBufferSize + splitData.argmaxBufferSize;
117- splitData.totalBufferSize = tmpTotalBufferSize * DOUBLE_BUFFER;
118-}
119- 
120-bool MaxPoolGradWithArgmaxV3NCHWTiling::IsMeetTargetCoreNum() const
121-{
122- // The calculation only involves inner.
123- int64_t tmpWOutputOuter = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputInner);
124- int64_t tmpHOutputOuter = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputInner);
125- int64_t tmpHighAxisOutputOuter = Ops::Base::CeilDiv(baseData.inputNCSize, splitData.highAxisInner);
126- 
127- return tmpWOutputOuter * tmpHOutputOuter * tmpHighAxisOutputOuter >= baseData.coreUsedForBestPerformance;
128-}
129- 
130-bool MaxPoolGradWithArgmaxV3NCHWTiling::IsMeetUBSize()
131-{
132- DoBufferCalculate();
133- return splitData.totalBufferSize <= baseData.availableUb;
134-}
135- 
136-bool MaxPoolGradWithArgmaxV3NCHWTiling::TrySplitNC()
137-{
138- splitData.wOutputInner = inputData.wX;
139- splitData.hOutputInner = inputData.hX;
140- 
141- splitData.highAxisInner = Ops::Base::CeilDiv(baseData.inputNCSize, baseData.coreUsedForBestPerformance);
142- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
143- return true;
144- }
145- 
146- splitData.highAxisInner = 1;
147- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
148- int64_t left = 1;
149- int64_t right = baseData.inputNCSize;
150- int64_t bestSplit = 1;
151- 
152- while (left <= right) {
153- int64_t mid = left + (right - left) / 2;
154- splitData.highAxisInner = mid;
155- 
156- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
157- bestSplit = mid;
158- left = mid + 1;
159- } else {
160- right = mid - 1;
161- }
162- }
163- 
164- splitData.highAxisInner = bestSplit;
165- return true;
166- } else {
167- return false;
168- }
169-}
170- 
171-bool MaxPoolGradWithArgmaxV3NCHWTiling::TrySplitAlignH()
172-{
173- splitData.highAxisInner = 1;
174- splitData.wOutputInner = inputData.wX;
175- 
176- splitData.hOutputInner = inputData.hStride;
177- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
178- int64_t left = 1;
179- int64_t right = Ops::Base::CeilDiv(inputData.hX / 2, inputData.hStride);
180- int64_t bestSplit = 1;
181- 
182- while (left <= right) {
183- int64_t mid = left + (right - left) / 2;
184- splitData.hOutputInner = mid * inputData.hStride;
185- 
186- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
187- bestSplit = mid;
188- left = mid + 1;
189- } else {
190- right = mid - 1;
191- }
192- }
193- 
194- splitData.hOutputInner = bestSplit * inputData.hStride;
195- return true;
196- } else {
197- return false;
198- }
199-}
200- 
201-bool MaxPoolGradWithArgmaxV3NCHWTiling::TrySplitAlignW()
202-{
203- splitData.highAxisInner = 1;
204- splitData.hOutputInner = inputData.hStride;
205- 
206- splitData.wOutputInner = inputData.wStride;
207- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
208- int64_t left = 1;
209- int64_t right = Ops::Base::CeilDiv(inputData.wX / 2, inputData.wStride);
210- int64_t bestSplit = 1;
211- 
212- while (left <= right) {
213- int64_t mid = left + (right - left) / 2;
214- splitData.wOutputInner = mid * inputData.wStride;
215- 
216- if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
217- bestSplit = mid;
218- left = mid + 1;
219- } else {
220- right = mid - 1;
221- }
222- }
223- 
224- splitData.wOutputInner = bestSplit * inputData.wStride;
225- return true;
226- } else {
227- return false;
228- }
229-}
230- 
231-void MaxPoolGradWithArgmaxV3NCHWTiling::SplitUnalignHW()
232-{
233- splitData.highAxisInner = 1;
234- if (baseData.isPad == 0 && baseData.isOverlap == 0) {
235- splitData.hOutputInner = inputData.hStride;
236- splitData.wOutputInner = inputData.wStride;
237- } else {
238- splitData.hOutputInner = inputData.hX;
239- splitData.wOutputInner = inputData.wX;
240- }
241- 
242- splitData.wOutputOuter = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputInner);
243- splitData.hOutputOuter = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputInner);
244- 
245- while (splitData.hOutputInner != 1 || splitData.wOutputInner > baseData.proDataNumInOneBeatT2) {
246- if (!IsMeetTargetCoreNum() || !IsMeetUBSize()) {
247- DynamicAdjustmentWH();
248- } else {
249- return;
250- }
251- }
252- 
253- splitData.wOutputInner = std::min(inputData.wX, baseData.proDataNumInOneBeatT2);
254- return;
255-}
256- 
257-void MaxPoolGradWithArgmaxV3NCHWTiling::DynamicAdjustmentWH()
258-{
259- if (splitData.hOutputInner == 1) {
260- splitData.wOutputOuter++;
261- splitData.wOutputInner = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputOuter);
262- } else {
263- splitData.hOutputOuter++;
264- splitData.hOutputInner = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputOuter);
265- }
266-}
267- 
268-void MaxPoolGradWithArgmaxV3NCHWTiling::SearchBestTiling()
269-{
270- splitData.isCheckRange = 0;
271- if (TrySplitNC()) {
272- return;
273- }
274- 
275- if (baseData.isPad == 0 && baseData.isOverlap == 0) {
276- if (TrySplitAlignH()) {
277- return;
278- }
279- 
280- if (TrySplitAlignW()) {
281- return;
282- }
283- }
284- 
285- // 带pad 或者overlap 或者 最小整切仍然不满足条件需要更细粒度切分HW
286- splitData.isCheckRange = 1;
287- SplitUnalignHW();
288- return;
289-}
290- 
291-void MaxPoolGradWithArgmaxV3NCHWTiling::DoUBTiling()
292-{
293- SearchBestTiling();
294- DoBufferCalculate();
295- splitData.wOutputOuter = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputInner);
296- int64_t tempWOutputTail = inputData.wX % splitData.wOutputInner;
297- splitData.wOutputTail = tempWOutputTail == 0 ? splitData.wOutputInner : tempWOutputTail;
298- 
299- splitData.hOutputOuter = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputInner);
300- int64_t tempHOutputTail = inputData.hX % splitData.hOutputInner;
301- splitData.hOutputTail = tempHOutputTail == 0 ? splitData.hOutputInner : tempHOutputTail;
302- 
303- splitData.highAxisOuter = Ops::Base::CeilDiv(baseData.inputNCSize, splitData.highAxisInner);
304- int64_t tempHighAxisTail = baseData.inputNCSize % splitData.highAxisInner;
305- splitData.highAxisTail = tempHighAxisTail == 0 ? splitData.highAxisInner : tempHighAxisTail;
306-}
307- 
308-void MaxPoolGradWithArgmaxV3NCHWTiling::DoBlockTiling()
309-{
310- splitData.totalBaseBlockNum = splitData.highAxisOuter * splitData.hOutputOuter * splitData.wOutputOuter;
311- splitData.normalCoreProcessNum = Ops::Base::CeilDiv(splitData.totalBaseBlockNum, baseData.totalCoreNum);
312- splitData.usedCoreNum = Ops::Base::CeilDiv(splitData.totalBaseBlockNum, splitData.normalCoreProcessNum);
313- splitData.tailCoreProcessNum =
314- splitData.totalBaseBlockNum - splitData.normalCoreProcessNum * (splitData.usedCoreNum - 1);
315-}
316- 
317-void MaxPoolGradWithArgmaxV3NCHWTiling::PrintBaseData() const
318-{
319- OP_LOGD("MaxPoolGradWithArgmaxV3NCHW", "[MaxPoolGradWithArgmaxV3NCHW] PrintBaseData start running");
320- 
321- std::ostringstream info;
322- info << "baseData.vRegSize: " << baseData.vRegSize << std::endl;
323- info << "baseData.ubBlockSize: " << baseData.ubBlockSize << std::endl;
324- info << "baseData.inputBytes: " << baseData.inputBytes << std::endl;
325- info << "baseData.indexBytes: " << baseData.indexBytes << std::endl;
326- info << "baseData.availableUb: " << baseData.availableUb << std::endl;
327- info << "baseData.maxDataNumInOneBlock: " << baseData.maxDataNumInOneBlock << std::endl;
328- info << "baseData.proDataNumInOneBeatT2: " << baseData.proDataNumInOneBeatT2 << std::endl;
329- info << "baseData.totalCoreNum: " << baseData.totalCoreNum << std::endl;
330- info << "baseData.coreUsedForBestPerformance: " << baseData.coreUsedForBestPerformance << std::endl;
331- info << "baseData.isPad: " << baseData.isPad << std::endl;
332- info << "baseData.isOverlap: " << baseData.isOverlap << std::endl;
333- info << "baseData.hProBatchSize: " << baseData.hProBatchSize << std::endl;
334- info << "baseData.wProBatchSize: " << baseData.wProBatchSize << std::endl;
335- info << "baseData.inputNCSize: " << baseData.inputNCSize << std::endl;
336- 
337- OP_LOGI("MaxPoolGradWithArgmaxV3NCHW", "%s", info.str().c_str());
338-}
339- 
340-void MaxPoolGradWithArgmaxV3NCHWTiling::PrintSplitData() const
341-{
342- OP_LOGD("MaxPoolGradWithArgmaxV3NCHW", "[MaxPoolGradWithArgmaxV3NCHW] PrintSplitData start running");
343- 
344- std::ostringstream info;
345- info << "splitData.isCheckRange: " << splitData.isCheckRange << std::endl;
346- 
347- info << "splitData.highAxisInner: " << splitData.highAxisInner << std::endl;
348- info << "splitData.highAxisTail: " << splitData.highAxisTail << std::endl;
349- info << "splitData.highAxisOuter: " << splitData.highAxisOuter << std::endl;
350- 
351- info << "splitData.hOutputInner: " << splitData.hOutputInner << std::endl;
352- info << "splitData.hOutputTail: " << splitData.hOutputTail << std::endl;
353- info << "splitData.hOutputOuter: " << splitData.hOutputOuter << std::endl;
354- 
355- info << "splitData.wOutputInner: " << splitData.wOutputInner << std::endl;
356- info << "splitData.wOutputTail: " << splitData.wOutputTail << std::endl;
357- info << "splitData.wOutputOuter: " << splitData.wOutputOuter << std::endl;
358- 
359- info << "splitData.normalCoreProcessNum: " << splitData.normalCoreProcessNum << std::endl;
360- info << "splitData.tailCoreProcessNum: " << splitData.tailCoreProcessNum << std::endl;
361- info << "splitData.usedCoreNum: " << splitData.usedCoreNum << std::endl;
362- info << "splitData.totalBaseBlockNum: " << splitData.totalBaseBlockNum << std::endl;
363- 
364- info << "splitData.outputBufferSize: " << splitData.outputBufferSize << std::endl;
365- info << "splitData.gradBufferSize: " << splitData.gradBufferSize << std::endl;
366- info << "splitData.argmaxBufferSize: " << splitData.argmaxBufferSize << std::endl;
367- info << "splitData.totalBufferSize: " << splitData.totalBufferSize << std::endl;
368- 
369- OP_LOGI("MaxPoolGradWithArgmaxV3NCHW", "%s", info.str().c_str());
370-}
371- 
372-void MaxPoolGradWithArgmaxV3NCHWTiling::SetTilingData()
373-{
374- MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData* tilingData =
375- context_->GetTilingData<MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData>();
376- tilingData->hArgmax = inputData.hGrad;
377- tilingData->wArgmax = inputData.wGrad;
378- tilingData->hOutput = inputData.hX;
379- tilingData->wOutput = inputData.wX;
380- tilingData->hKernel = inputData.hKernel;
381- tilingData->wKernel = inputData.wKernel;
382- tilingData->hStride = inputData.hStride;
383- tilingData->wStride = inputData.wStride;
384- tilingData->padH = inputData.hPad;
385- tilingData->padW = inputData.wPad;
386- tilingData->dilationH = inputData.hDilation;
387- tilingData->dilationW = inputData.wDilation;
388- tilingData->highAxisInner = splitData.highAxisInner;
389- tilingData->highAxisTail = splitData.highAxisTail;
390- tilingData->highAxisOuter = splitData.highAxisOuter;
391- tilingData->hOutputInner = splitData.hOutputInner;
392- tilingData->hOutputTail = splitData.hOutputTail;
393- tilingData->hOutputOuter = splitData.hOutputOuter;
394- tilingData->wOutputInner = splitData.wOutputInner;
395- tilingData->wOutputTail = splitData.wOutputTail;
396- tilingData->wOutputOuter = splitData.wOutputOuter;
397- tilingData->normalCoreProcessNum = splitData.normalCoreProcessNum;
398- tilingData->tailCoreProcessNum = splitData.tailCoreProcessNum;
399- tilingData->usedCoreNum = splitData.usedCoreNum;
400- tilingData->outputBufferSize = splitData.outputBufferSize;
401- tilingData->gradBufferSize = splitData.gradBufferSize;
402- tilingData->argmaxBufferSize = splitData.argmaxBufferSize;
403- tilingData->hProBatchSize = baseData.hProBatchSize;
404- tilingData->wProBatchSize = baseData.wProBatchSize;
405- tilingData->tilingKey = GetTilingKey();
406-}
407- 
408ge::graphStatus MaxPoolGradWithArgmaxV3NCHWTiling::DoOpTiling()48ge::graphStatus MaxPoolGradWithArgmaxV3NCHWTiling::DoOpTiling()
409{49{
410- DoUBTiling();50+ return nchwTilingCommon.DoOpTiling(context_, GetTilingKey());
411- DoBlockTiling();
412- SetTilingData();
413- PrintBaseData();
414- PrintSplitData();
415- return ge::GRAPH_SUCCESS;
416}51}
417 52 
418ge::graphStatus MaxPoolGradWithArgmaxV3NCHWTiling::PostTiling()53ge::graphStatus MaxPoolGradWithArgmaxV3NCHWTiling::PostTiling()
419{54{
420- MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData* tilingData =55+ return nchwTilingCommon.PostTiling(context_);
421- context_->GetTilingData<MaxPoolGradWithArgmaxNHWCNameSpace::MaxPoolGradWithArgmaxNCHWTilingCommonData>();
422- context_->SetBlockDim(tilingData->usedCoreNum);
423- return ge::GRAPH_SUCCESS;
424}56}
425 57 
426REGISTER_OPS_TILING_TEMPLATE(MaxPoolGradWithArgmaxV3, MaxPoolGradWithArgmaxV3NCHWTiling, 2);58REGISTER_OPS_TILING_TEMPLATE(MaxPoolGradWithArgmaxV3, MaxPoolGradWithArgmaxV3NCHWTiling, 2);
427 59 
428-} // namespace optiling60+} // namespace optiling
Mpooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nchw_tiling.h+4-65
@@ -10,94 +10,33 @@
10 10 
11/*!11/*!
12 * \file max_pool_grad_with_argmax_v3_nchw_tiling.h12 * \file max_pool_grad_with_argmax_v3_nchw_tiling.h
13- * \brief13+ * \brief MaxPoolGradWithArgmaxV3 NCHW格式Tiling实现，继承自公共基类
14 */14 */
15 15 
16#ifndef MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_H_16#ifndef MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_H_
17#define MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_H_17#define MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_H_
18 18 
19#include "max_pool_grad_with_argmax_v3_tiling_base.h"19#include "max_pool_grad_with_argmax_v3_tiling_base.h"
20+#include "../../../pool_grad_common/op_host/arch35/max_pool_grad_nchw_tiling_common.h"
20 21 
21namespace optiling {22namespace optiling {
22 23 
23-struct MaxPoolGradWithArgmaxV3NCHWBaseInfo {
24- int64_t vRegSize{0};
25- int64_t ubBlockSize{0};
26- int64_t inputBytes{0};
27- int64_t indexBytes{0};
28- int64_t availableUb{0};
29- int64_t totalCoreNum{0};
30- int64_t coreUsedForBestPerformance{0};
31- int64_t hProBatchSize{0};
32- int64_t wProBatchSize{0};
33- int64_t inputNCSize{0};
34- int64_t maxDataNumInOneBlock{0};
35- int64_t proDataNumInOneBeatT2{0};
36- int64_t isPad{0};
37- int64_t isOverlap{0};
38-};
39- 
40-struct MaxPoolGradWithArgmaxV3NCHWSplitInfo {
41- // DoUBTiling
42- int64_t isCheckRange{0};
43- 
44- int64_t highAxisInner{0};
45- int64_t highAxisTail{0};
46- int64_t highAxisOuter{0};
47- 
48- int64_t hOutputInner{0};
49- int64_t hOutputTail{0};
50- int64_t hOutputOuter{0};
51- 
52- int64_t wOutputInner{0};
53- int64_t wOutputTail{0};
54- int64_t wOutputOuter{0};
55- 
56- // DoBlockTiling
57- int64_t normalCoreProcessNum{0};
58- int64_t tailCoreProcessNum{0};
59- int64_t usedCoreNum{0};
60- int64_t totalBaseBlockNum{0};
61- 
62- // DoBufferCalculate
63- int64_t outputBufferSize{0};
64- int64_t gradBufferSize{0};
65- int64_t argmaxBufferSize{0};
66- int64_t totalBufferSize{0};
67-};
68- 
69class MaxPoolGradWithArgmaxV3NCHWTiling : public MaxPoolGradWithArgmaxV3BaseTiling {24class MaxPoolGradWithArgmaxV3NCHWTiling : public MaxPoolGradWithArgmaxV3BaseTiling {
70public:25public:
71 explicit MaxPoolGradWithArgmaxV3NCHWTiling(gert::TilingContext* context)26 explicit MaxPoolGradWithArgmaxV3NCHWTiling(gert::TilingContext* context)
72- : MaxPoolGradWithArgmaxV3BaseTiling(context)27+ : MaxPoolGradWithArgmaxV3BaseTiling(context), nchwTilingCommon(&inputData)
73 {}28 {}
74 29 
75 ~MaxPoolGradWithArgmaxV3NCHWTiling() override30 ~MaxPoolGradWithArgmaxV3NCHWTiling() override
76 {}31 {}
77 32 
78private:33private:
79- void DoUBTiling();
80- void InitializationVars();
81- bool TrySplitNC();
82- bool TrySplitAlignH();
83- bool TrySplitAlignW();
84- void SplitUnalignHW();
85- bool IsMeetTargetCoreNum() const;
86- bool IsMeetUBSize();
87- void SearchBestTiling();
88- void DynamicAdjustmentWH();
89- void SetTilingData();
90 uint64_t GetTilingKey() const override;34 uint64_t GetTilingKey() const override;
91- void PrintBaseData() const;
92- void PrintSplitData() const;
93- void DoBlockTiling();
94- void DoBufferCalculate();
95 bool IsCapable() override;35 bool IsCapable() override;
96 ge::graphStatus DoOpTiling() override;36 ge::graphStatus DoOpTiling() override;
97 ge::graphStatus PostTiling() override;37 ge::graphStatus PostTiling() override;
98 38 
99- MaxPoolGradWithArgmaxV3NCHWBaseInfo baseData;39+ MaxPoolGradNCHWTilingCommon nchwTilingCommon;
100- MaxPoolGradWithArgmaxV3NCHWSplitInfo splitData;
101};40};
102 41 
103} // namespace optiling42} // namespace optiling