Star872
1766
代码
代码
Issues241
Pull Requests323
Actions
项目讨论
Wiki
分析
Star872
1766
已合并
补充AvgPoolV2Grad、MaxPoolWithArgmax、MaskedScatterWithPosition算子资料 #4624
ideal_创建于 5月8日
补充AvgPoolV2Grad、MaxPoolWithArgmax、MaskedScatterWithPosition算子资料 #4624
已合并
ideal_创建于 5月8日
从wang-shilong32/ops-nnfix_materials_avg_max_mswp合入到cann/ops-nnmaster
共 6 个文件变更+823-1
Mdocs/zh/op_list.md+20-0
@@ -2996,6 +2996,16 @@
2996 <td>AI Core</td>2996 <td>AI Core</td>
2997 <td>对输入Tensor进行窗口为kH * kW、步长为sH * sW的二维平均池化操作，其中k为kernelSize，表示池化窗口的大小，s为stride，表示池化操作的步长。</td>2997 <td>对输入Tensor进行窗口为kH * kW、步长为sH * sW的二维平均池化操作，其中k为kernelSize，表示池化窗口的大小，s为stride，表示池化操作的步长。</td>
2998 </tr>2998 </tr>
2999+ <tr>
3000+ <td>pooling</td>
3001+ <td><a href="../../pooling/avg_pool_v2_grad/README.md">avg_pool_v2_grad</a></td>
3002+ <td>✓</td>
3003+ <td>✓</td>
3004+ <td>✓</td>
3005+ <td>✓</td>
3006+ <td>AI Core</td>
3007+ <td>二维平均池化的反向传播，计算二维平均池化正向传播的输入梯度。</td>
3008+ </tr>
2999 <tr>3009 <tr>
3000 <td>pooling</td>3010 <td>pooling</td>
3001 <td><a href="../../pooling/avg_pool3_d/README.md">avg_pool3_d</a></td>3011 <td><a href="../../pooling/avg_pool3_d/README.md">avg_pool3_d</a></td>
@@ -3036,6 +3046,16 @@
3036 <td>AI Core</td>3046 <td>AI Core</td>
3037 <td>对于3维或4维的输入张量，进行最大池化（max pooling）操作。</td>3047 <td>对于3维或4维的输入张量，进行最大池化（max pooling）操作。</td>
3038 </tr>3048 </tr>
3049+ <tr>
3050+ <td>pooling</td>
3051+ <td><a href="../../pooling/max_pool_with_argmax/README.md">max_pool_with_argmax</a></td>
3052+ <td>✓</td>
3053+ <td>✓</td>
3054+ <td>✗</td>
3055+ <td>✓</td>
3056+ <td>AI Core</td>
3057+ <td>对于输入数据计算2维最大池化操作，同时输出池化后的最大值和对应位置的索引。</td>
3058+ </tr>
3039 <tr>3059 <tr>
3040 <td>pooling</td>3060 <td>pooling</td>
3041 <td><a href="../../pooling/max_pool_with_argmax_v3/README.md">max_pool_with_argmax_v3</a></td>3061 <td><a href="../../pooling/max_pool_with_argmax_v3/README.md">max_pool_with_argmax_v3</a></td>
Mindex/masked_scatter_with_position/README.md+1-1
@@ -1,4 +1,4 @@
1-# Maskedscatterwithposition1+# MaskedScatterWithPosition
2 2 
3## 产品支持情况3## 产品支持情况
4 4 
Apooling/avg_pool_v2_grad/README.md+151-0
@@ -0,0 +1,151 @@
1+# AvgPoolV2Grad
2+ 
3+## 产品支持情况
4+ 
5+|产品 | 是否支持 |
6+|:-------------------------|:----------:|
7+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8+| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | x |
9+| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | x |
10+| <term>Atlas 200I/500 A2 推理产品</term> | x |
11+| <term>Atlas 推理系列产品</term> | x |
12+| <term>Atlas 训练系列产品</term> | x |
13+ 
14+## 功能说明
15+ 
16+- 接口功能：
17+二维平均池化的反向传播，计算二维平均池化正向传播的输入梯度。
18+ 
19+- 计算公式：
20+ - 对于前向平均池化输出梯度，计算原始输入的梯度：
21+
22+ $$
23+ out\_grad(i,j) = \frac{1}{divisor} \sum_{k \in pooling\_window} input\_grad(k)
24+ $$
25+
26+ 其中divisor为池化窗口大小（若exclusive为true，则不包含padding区域；若指定divisor_override则使用该值）。
27+ 
28+## 参数说明
29+ 
30+<table style="undefined;table-layout: fixed; width: 1005px"><colgroup>
31+ <col style="width: 170px">
32+ <col style="width: 170px">
33+ <col style="width: 352px">
34+ <col style="width: 213px">
35+ <col style="width: 100px">
36+ </colgroup>
37+ <thead>
38+ <tr>
39+ <th>参数名</th>
40+ <th>输入/输出/属性</th>
41+ <th>描述</th>
42+ <th>数据类型</th>
43+ <th>数据格式</th>
44+ </tr></thead>
45+ <tbody>
46+ <tr>
47+ <td>orig_input_shape</td>
48+ <td>输入</td>
49+ <td>原始输入的形状，描述前向AvgPoolV2的输入维度[N,C,H,W]或[N,H,W,C]。</td>
50+ <td>INT32</td>
51+ <td>ND</td>
52+ </tr>
53+ <tr>
54+ <td>input_grad</td>
55+ <td>输入</td>
56+ <td>平均池化输出的梯度张量。</td>
57+ <td>FLOAT16、FLOAT、BFLOAT16</td>
58+ <td>ND</td>
59+ </tr>
60+ <tr>
61+ <td>ksize</td>
62+ <td>属性</td>
63+ <td>池化窗口的大小，长度为4，数值大于0。</td>
64+ <td>LIST_INT</td>
65+ <td>-</td>
66+ </tr>
67+ <tr>
68+ <td>strides</td>
69+ <td>属性</td>
70+ <td>滑动窗口的步长，长度为4，数值大于0。</td>
71+ <td>LIST_INT</td>
72+ <td>-</td>
73+ </tr>
74+ <tr>
75+ <td>padding_mode</td>
76+ <td>属性</td>
77+ <td>填充算法类型，支持"VALID"、"SAME"或"CALCULATED"。默认"CALCULATED"。</td>
78+ <td>STRING</td>
79+ <td>-</td>
80+ </tr>
81+ <tr>
82+ <td>pads</td>
83+ <td>属性</td>
84+ <td>输入特征图的填充大小，默认{0,0,0,0}。</td>
85+ <td>LIST_INT</td>
86+ <td>-</td>
87+ </tr>
88+ <tr>
89+ <td>data_format</td>
90+ <td>属性</td>
91+ <td>数据格式，支持"NCHW"或"NHWC"。默认"NCHW"。</td>
92+ <td>STRING</td>
93+ <td>-</td>
94+ </tr>
95+ <tr>
96+ <td>global_pooling</td>
97+ <td>属性</td>
98+ <td>是否使用全局池化，若为true则忽略ksize和pads。默认false。</td>
99+ <td>BOOL</td>
100+ <td>-</td>
101+ </tr>
102+ <tr>
103+ <td>ceil_mode</td>
104+ <td>属性</td>
105+ <td>推导的输出out的shape是否向上取整。默认false。</td>
106+ <td>BOOL</td>
107+ <td>-</td>
108+ </tr>
109+ <tr>
110+ <td>exclusive</td>
111+ <td>属性</td>
112+ <td>计算平均池化时是否排除padding点。默认true。</td>
113+ <td>BOOL</td>
114+ <td>-</td>
115+ </tr>
116+ <tr>
117+ <td>divisor_override</td>
118+ <td>属性</td>
119+ <td>指定除数覆盖值，若为0则使用池化区域大小。默认0。</td>
120+ <td>INT</td>
121+ <td>-</td>
122+ </tr>
123+ <tr>
124+ <td>out_grad</td>
125+ <td>输出</td>
126+ <td>输出的梯度张量，形状与orig_input_shape相同。</td>
127+ <td>FLOAT16、FLOAT、BFLOAT16</td>
128+ <td>ND</td>
129+ </tr>
130+ </tbody></table>
131+ 
132+## 约束说明
133+ 
134+- **值域限制说明：**
135+ - orig_input_shape：必须是一维张量，长度为4，描述[N,C,H,W]或[N,H,W,C]。
136+ - input_grad：4D张量，支持FLOAT16、FLOAT、BFLOAT16类型。
137+ - ksize：长度为 4，必须大于0，指定池化窗口大小。
138+ - strides：长度为 4，必须大于0，指定滑动步长。
139+ - padding_mode：只支持"VALID"、"SAME"、"CALCULATED"三种模式。
140+ - pads：长度为 4，必须大于0，仅在padding_mode为"CALCULATED"时生效。不超过ksize对应位置的1/2。
141+ - data_format：支持"NCHW"或"NHWC"。
142+ - global_pooling：若为true，ksize 和 pads 将被忽略。
143+ - ceil_mode：若为true，使用ceil计算输出形状；否则使用floor。
144+ - exclusive：若为true，计算均值时排除padding区域。
145+ - divisor_override：若非0，将作为除数使用；否则使用池化区域大小。
146+ 
147+## 调用说明
148+ 
149+| 调用方式 | 样例代码 | 说明 |
150+| ---------------- | --------------------------- | --------------------------------------------------- |
151+| aclnn接口 | [test_aclnn_avg_pool_v2_grad](examples/test_aclnn_avg_pool_v2_grad.cpp) | 通过[aclnnAvgPool2dBackward](../../pooling/avg_pool3_d_grad/docs/aclnnAvgPool2dBackward.md)接口方式调用AvgPoolV2Grad算子。 |
Apooling/avg_pool_v2_grad/examples/arch35/test_aclnn_avg_pool_v2_grad.cpp+184-0
@@ -0,0 +1,184 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <cstdio>
12+#include <iostream>
13+#include <vector>
14+#include "acl/acl.h"
15+#include "aclnnop/aclnn_avgpool2d_backward.h"
16+ 
17+#define CHECK_RET(cond, return_expr) \
18+ do { \
19+ if (!(cond)) { \
20+ return_expr; \
21+ } \
22+ } while (0)
23+ 
24+#define LOG_PRINT(message, ...) \
25+ do { \
26+ printf(message, ##__VA_ARGS__); \
27+ } while (0)
28+ 
29+int64_t GetShapeSize(const std::vector<int64_t>& shape) {
30+ int64_t shapeSize = 1;
31+ for (auto i : shape) {
32+ shapeSize *= i;
33+ }
34+ return shapeSize;
35+}
36+ 
37+int Init(int32_t deviceId, aclrtStream* stream) {
38+ // 固定写法，资源初始化
39+ auto ret = aclInit(nullptr);
40+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
41+ ret = aclrtSetDevice(deviceId);
42+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
43+ ret = aclrtCreateStream(stream);
44+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
45+ return 0;
46+}
47+ 
48+template <typename T>
49+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
50+ aclDataType dataType, aclTensor** tensor, aclFormat Format = aclFormat::ACL_FORMAT_ND) {
51+ auto size = GetShapeSize(shape) * sizeof(T);
52+ // 调用aclrtMalloc申请device侧内存
53+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
54+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
55+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
56+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
57+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
58+ 
59+ // 计算连续tensor的strides
60+ std::vector<int64_t> strides(shape.size(), 1);
61+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
62+ strides[i] = shape[i + 1] * strides[i + 1];
63+ }
64+ 
65+ // 调用aclCreateTensor接口创建aclTensor
66+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, Format, shape.data(), shape.size(), *deviceAddr);
67+ return 0;
68+}
69+ 
70+int main() {
71+ // 1. （固定写法）device/stream初始化，参考acl对外接口列表
72+ // 根据自己的实际device填写deviceId
73+ int32_t deviceId = 0;
74+ aclrtStream stream;
75+ auto ret = Init(deviceId, &stream);
76+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
77+ 
78+ // 2. 构造输入与输出，需要根据API的接口自定义构造
79+ std::vector<int64_t> gradOutputShape = {1, 16, 1, 1};
80+ std::vector<int64_t> selfShape = {1, 16, 4, 4};
81+ std::vector<int64_t> kernelDims = {4, 4};
82+ std::vector<int64_t> strideDims = {1, 1};
83+ std::vector<int64_t> paddingDims = {0, 0};
84+ bool ceilMode = false;
85+ int64_t divisorOverride = 0;
86+ bool countIncludePad = true;
87+ int8_t cubeMathType = 1;
88+ std::vector<int64_t> gradInputShape = {1, 16, 4, 4};
89+ 
90+ void* gradOutputDeviceAddr = nullptr;
91+ void* selfDeviceAddr = nullptr;
92+ void* gradInputDeviceAddr = nullptr;
93+ 
94+ aclTensor* gradOutput = nullptr;
95+ aclTensor* self = nullptr;
96+ aclTensor* gradInput = nullptr;
97+ 
98+ std::vector<float> gradOutputHostData(GetShapeSize(gradOutputShape), 1);
99+ std::vector<float> selfHostData(GetShapeSize(selfShape), 1);
100+ std::vector<float> gradInputHostData(GetShapeSize(gradInputShape), 1);
101+ 
102+ // 创建gradOutput aclTensor
103+ ret = CreateAclTensor(gradOutputHostData, gradOutputShape, &gradOutputDeviceAddr, aclDataType::ACL_FLOAT, &gradOutput, aclFormat::ACL_FORMAT_NCHW);
104+ CHECK_RET(ret == ACL_SUCCESS, return ret);
105+ 
106+ // 创建self aclTensor
107+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self, aclFormat::ACL_FORMAT_NCHW);
108+ CHECK_RET(ret == ACL_SUCCESS, return ret);
109+ 
110+ // 创建gradInput aclTensor
111+ ret = CreateAclTensor(gradInputHostData, gradInputShape, &gradInputDeviceAddr, aclDataType::ACL_FLOAT, &gradInput, aclFormat::ACL_FORMAT_NCHW);
112+ CHECK_RET(ret == ACL_SUCCESS, return ret);
113+ 
114+ // 创建kernel aclIntArray
115+ aclIntArray *kernelSize = aclCreateIntArray(kernelDims.data(), 2);
116+ 
117+ // 创建stride aclIntArray
118+ aclIntArray *stride = aclCreateIntArray(strideDims.data(), 2);
119+ 
120+ // 创建paddings aclIntArray
121+ aclIntArray *padding = aclCreateIntArray(paddingDims.data(), 2);
122+ 
123+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
124+ uint64_t workspaceSize = 0;
125+ aclOpExecutor* executor;
126+ // 调用aclnnAvgPool2dBackward第一段接口
127+ ret = aclnnAvgPool2dBackwardGetWorkspaceSize(gradOutput,
128+ self,
129+ kernelSize,
130+ stride,
131+ padding,
132+ ceilMode,
133+ countIncludePad,
134+ divisorOverride,
135+ cubeMathType,
136+ gradInput,
137+ &workspaceSize,
138+ &executor);
139+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAvgPool2dBackwardGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
140+ // 根据第一段接口计算出的workspaceSize申请device内存
141+ void* workspaceAddr = nullptr;
142+ if (workspaceSize > 0) {
143+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
144+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
145+ }
146+ // 调用aclnnAvgPool2dBackward第二段接口
147+ ret = aclnnAvgPool2dBackward(workspaceAddr, workspaceSize, executor, stream);
148+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAvgPool2dBackward failed. ERROR: %d\n", ret); return ret);
149+ 
150+ // 4. （固定写法）同步等待任务执行结束
151+ ret = aclrtSynchronizeStream(stream);
152+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
153+ 
154+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
155+ auto size = GetShapeSize(gradInputShape);
156+ std::vector<float> outData(size, 0);
157+ ret = aclrtMemcpy(outData.data(), outData.size() * sizeof(outData[0]), gradInputDeviceAddr,
158+ size * sizeof(outData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
159+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
160+ for (int64_t i = 0; i < size; i++) {
161+ LOG_PRINT("out result[%ld] is: %f\n", i, outData[i]);
162+ }
163+ 
164+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
165+ aclDestroyTensor(gradOutput);
166+ aclDestroyTensor(self);
167+ aclDestroyTensor(gradInput);
168+ aclDestroyIntArray(kernelSize);
169+ aclDestroyIntArray(stride);
170+ aclDestroyIntArray(padding);
171+ 
172+ // 7. 释放device资源，需要根据具体API的接口定义修改
173+ aclrtFree(gradOutputDeviceAddr);
174+ aclrtFree(selfDeviceAddr);
175+ aclrtFree(gradInputDeviceAddr);
176+ if (workspaceSize > 0) {
177+ aclrtFree(workspaceAddr);
178+ }
179+ aclrtDestroyStream(stream);
180+ aclrtResetDevice(deviceId);
181+ aclFinalize();
182+ 
183+ return 0;
184+}
Apooling/max_pool_with_argmax/README.md+136-0
@@ -0,0 +1,136 @@
1+# MaxPoolWithArgmax
2+ 
3+## 产品支持情况
4+ 
5+|产品 | 是否支持 |
6+|:-------------------------|:----------:|
7+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8+| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | x |
9+| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | x |
10+| <term>Atlas 200I/500 A2 推理产品</term> | x |
11+| <term>Atlas 推理系列产品</term> | x |
12+| <term>Atlas 训练系列产品</term> | x |
13+ 
14+## 功能说明
15+ 
16+- 接口功能：
17+执行最大池化操作，同时输出池化后的最大值和对应位置的索引。
18+ 
19+- 计算公式：
20+ - output tensor中每个元素的计算公式，以 NCHW 为例：
21+
22+ $$
23+ y(N_i, C_j, h, w) = \max\limits_{{k\in[0,k_{H}-1],m\in[0,k_{W}-1]}}x(N_i,C_j,stride[2]\times h + k, stride[3]\times w + m)
24+ $$
25+
26+ - argmax输出最大值在池化窗口中的索引位置。
27+ 
28+## 参数说明
29+ 
30+<table style="undefined;table-layout: fixed; width: 1005px"><colgroup>
31+ <col style="width: 170px">
32+ <col style="width: 170px">
33+ <col style="width: 352px">
34+ <col style="width: 213px">
35+ <col style="width: 100px">
36+ </colgroup>
37+ <thead>
38+ <tr>
39+ <th>参数名</th>
40+ <th>输入/输出/属性</th>
41+ <th>描述</th>
42+ <th>数据类型</th>
43+ <th>数据格式</th>
44+ </tr></thead>
45+ <tbody>
46+ <tr>
47+ <td>x</td>
48+ <td>输入</td>
49+ <td>输入的4D张量。</td>
50+ <td>FLOAT16、FLOAT、BFLOAT16</td>
51+ <td>ND</td>
52+ </tr>
53+ <tr>
54+ <td>ksize</td>
55+ <td>属性</td>
56+ <td>池化窗口的大小，长度为4，且数组元素必须都大于0。</td>
57+ <td>LIST_INT</td>
58+ <td>-</td>
59+ </tr>
60+ <tr>
61+ <td>strides</td>
62+ <td>属性</td>
63+ <td>滑动窗口的步长，长度为4，且数组元素必须都大于0。</td>
64+ <td>LIST_INT</td>
65+ <td>-</td>
66+ </tr>
67+ <tr>
68+ <td>padding</td>
69+ <td>属性</td>
70+ <td>指定padding的模式，支持"SAME"或"VALID"。</td>
71+ <td>STRING</td>
72+ <td>-</td>
73+ </tr>
74+ <tr>
75+ <td>Targmax</td>
76+ <td>属性</td>
77+ <td>指定argmax输出的数据类型，支持int32(3)或int64(9)。默认int64(9)。</td>
78+ <td>INT</td>
79+ <td>-</td>
80+ </tr>
81+ <tr>
82+ <td>include_batch_in_index</td>
83+ <td>属性</td>
84+ <td>计算argmax索引时是否包含batch维度。目前仅支持false。</td>
85+ <td>BOOL</td>
86+ <td>-</td>
87+ </tr>
88+ <tr>
89+ <td>data_format</td>
90+ <td>属性</td>
91+ <td>数据格式，支持"NCHW"或"NHWC"。默认"NHWC"。</td>
92+ <td>STRING</td>
93+ <td>-</td>
94+ </tr>
95+ <tr>
96+ <td>nan_prop</td>
97+ <td>属性</td>
98+ <td>是否处理NAN值。true时NAN参与比较，false时忽略NAN。默认false。</td>
99+ <td>BOOL</td>
100+ <td>-</td>
101+ </tr>
102+ <tr>
103+ <td>y</td>
104+ <td>输出</td>
105+ <td>池化后的最大值张量，与输入x类型相同。</td>
106+ <td>FLOAT16、FLOAT、BFLOAT16</td>
107+ <td>ND</td>
108+ </tr>
109+ <tr>
110+ <td>argmax</td>
111+ <td>输出</td>
112+ <td>最大值对应的索引位置。</td>
113+ <td>INT32、INT64</td>
114+ <td>ND</td>
115+ </tr>
116+ </tbody></table>
117+ 
118+## 约束说明
119+ 
120+- **值域限制说明：**
121+ - x：4D张量，支持FLOAT16、FLOAT、BFLOAT16类型，支持NCHW和NHWC格式。
122+ - ksize：长度为4的列表，[NCHW]格式要求ksize[0]=1和ksize[1]=1；[NHWC]格式要求ksize[0]=1和ksize[3]=1。
123+ - strides：长度为4的列表，[NCHW]格式要求strides[0]=1和strides[1]=1；[NHWC]格式要求strides[0]=1和strides[3]=1。
124+ - padding：只支持"SAME"或"VALID"模式。
125+ - SAME：填充使输出形状等于ceil(输入形状/步长)，当步长为1时输出等于输入。
126+ - VALID：不填充，仅在有有效区域滑动，输出较小。
127+ - Targmax：只支持3(int32)或9(int64)。
128+ - include_batch_in_index：目前仅支持false，表示索引计算不包含batch维度。
129+ - data_format：支持"NCHW"或"NHWC"。
130+ - nan_prop：true时NAN值参与比较，最大值可以是NAN；false时忽略NAN值。
131+ 
132+## 调用说明
133+ 
134+| 调用方式 | 样例代码 | 说明 |
135+| ---------------- | --------------------------- | --------------------------------------------------- |
136+| aclnn模式接口 | [test_aclnn_max_pool_with_argmax](examples/test_aclnn_max_pool_with_argmax.cpp) | 通过aclnn[aclnnMaxPool2dWithIndices](../max_pool3d_with_argmax_v2/docs/aclnnMaxPool2dWithIndices.md)方式调用MaxPoolWithArgmax算子。 |
Apooling/max_pool_with_argmax/examples/arch35/test_geir_max_pool_with_argmax.cpp+331-0
@@ -0,0 +1,331 @@
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+#include <iostream>
12+#include <fstream>
13+#include <string.h>
14+#include <stdint.h>
15+#include <vector>
16+#include <string>
17+#include <map>
18+#include "assert.h"
19+ 
20+#include "graph.h"
21+#include "types.h"
22+#include "tensor.h"
23+#include "ge_error_codes.h"
24+#include "ge_api_types.h"
25+#include "ge_api.h"
26+#include "array_ops.h"
27+#include "ge_ir_build.h"
28+ 
29+#include "experiment_ops.h"
30+#include "nn_other.h"
31+#include "../../op_graph/max_pool_with_argmax_proto.h"
32+ 
33+#define FAILED -1
34+#define SUCCESS 0
35+ 
36+using namespace ge;
37+using std::map;
38+using std::string;
39+using std::vector;
40+#define ADD_INPUT(inputIndex, inputName, inputDtype, inputShape) \
41+ vector<int64_t> placeholder##inputIndex##_shape = inputShape; \
42+ auto placeholder##inputIndex = op::Data("placeholder" + inputIndex).set_attr_index(0); \
43+ TensorDesc placeholder##inputIndex##_desc = \
44+ TensorDesc(ge::Shape(placeholder##inputIndex##_shape), FORMAT_ND, inputDtype); \
45+ placeholder##inputIndex##_desc.SetPlacement(ge::kPlacementHost); \
46+ placeholder##inputIndex##_desc.SetFormat(FORMAT_ND); \
47+ Tensor tensor_placeholder##inputIndex; \
48+ ret = GenOnesData( \
49+ placeholder##inputIndex##_shape, tensor_placeholder##inputIndex, placeholder##inputIndex##_desc, inputDtype, \
50+ 2); \
51+ if (ret != SUCCESS) { \
52+ printf("%s - ERROR - [XIR]: Generate input data failed\n", GetTime().c_str()); \
53+ return FAILED; \
54+ } \
55+ placeholder##inputIndex.update_input_desc_x(placeholder##inputIndex##_desc); \
56+ input.push_back(tensor_placeholder##inputIndex); \
57+ graph.AddOp(placeholder##inputIndex); \
58+ add1.set_input_##inputName(placeholder##inputIndex); \
59+ inputs.push_back(placeholder##inputIndex)
60+ 
61+#define ADD_OUTPUT(outputIndex, outputName, outputDtype, outputShape) \
62+ TensorDesc outputName##outputIndex##_desc = TensorDesc(ge::Shape(outputShape), FORMAT_ND, outputDtype); \
63+ add1.update_output_desc_##outputName(outputName##outputIndex##_desc)
64+ 
65+#define LOG_PRINT(message, ...) \
66+ do { \
67+ printf(message, ##__VA_ARGS__); \
68+ } while (0)
69+ 
70+#define ADD_INPUT_ATTR(attrName, attrValue) add1.set_attr_##attrName(attrValue)
71+ 
72+string GetTime()
73+{
74+ time_t timep;
75+ time(&timep);
76+ char tmp[64];
77+ strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime(&timep));
78+ return tmp;
79+}
80+ 
81+uint32_t GetDataTypeSize(DataType dt)
82+{
83+ uint32_t dilation = 1;
84+ uint32_t oneByte = 1;
85+ uint32_t twoByte = 2;
86+ uint32_t fourByte = 4;
87+ uint32_t eightByte = 8;
88+ 
89+ if (dt == ge::DT_FLOAT) {
90+ dilation = fourByte;
91+ } else if (dt == ge::DT_FLOAT16) {
92+ dilation = twoByte;
93+ } else if (dt == ge::DT_BF16) {
94+ dilation = twoByte;
95+ } else if (dt == ge::DT_INT16) {
96+ dilation = twoByte;
97+ } else if (dt == ge::DT_UINT16) {
98+ dilation = twoByte;
99+ } else if (dt == ge::DT_INT32) {
100+ dilation = fourByte;
101+ } else if (dt == ge::DT_UINT32) {
102+ dilation = fourByte;
103+ } else if (dt == ge::DT_INT64) {
104+ dilation = eightByte;
105+ } else if (dt == ge::DT_UINT64) {
106+ dilation = eightByte;
107+ } else if (dt == ge::DT_INT8) {
108+ dilation = oneByte;
109+ }
110+ return dilation;
111+}
112+ 
113+int32_t GenOnesData(
114+ vector<int64_t> shapes, Tensor& input_tensor, TensorDesc& input_tensor_desc, DataType data_type, int value)
115+{
116+ input_tensor_desc.SetRealDimCnt(shapes.size());
117+ size_t size = 1;
118+ for (uint32_t i = 0; i < shapes.size(); i++) {
119+ size *= shapes[i];
120+ }
121+ uint32_t data_len = size * GetDataTypeSize(data_type);
122+ int32_t* pData = new (std::nothrow) int32_t[data_len];
123+ for (uint32_t i = 0; i < size; ++i) {
124+ *(pData + i) = value;
125+ }
126+ input_tensor = Tensor(input_tensor_desc, reinterpret_cast<uint8_t*>(pData), data_len);
127+ return SUCCESS;
128+}
129+ 
130+int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData)
131+{
132+ FILE* fp = fopen(bin_file.c_str(), "w");
133+ fwrite(inputData, sizeof(uint8_t), data_size, fp);
134+ fclose(fp);
135+ return SUCCESS;
136+}
137+ 
138+int CreateOppInGraph(
139+ DataType inDtype, std::vector<ge::Tensor>& input, std::vector<Operator>& inputs, std::vector<Operator>& outputs,
140+ Graph& graph)
141+{
142+ Status ret = SUCCESS;
143+ // 自定义代码：添加单算子定义到图中
144+ auto add1 = op::MaxPoolWithArgmax("max_pool_with_argmax");
145+ vector<vector<int64_t>> shapes = {
146+ {4, 1, 4, 4}, // x shape
147+ {4, 1, 2, 2}, // y shape (output)
148+ {4, 1, 2, 2} // argmax shape
149+ };
150+ vector<vector<int64_t>> attrs = {
151+ {1, 1, 2, 2}, // ksize
152+ {1, 1, 2, 2} // strides
153+ };
154+ 
155+ ADD_INPUT(1, x, inDtype, shapes[0]);
156+ ADD_OUTPUT(2, y, inDtype, shapes[1]);
157+ ADD_OUTPUT(3, argmax, DT_INT32, shapes[2]);
158+ ADD_INPUT_ATTR(ksize, attrs[0]);
159+ ADD_INPUT_ATTR(strides, attrs[1]);
160+ ADD_INPUT_ATTR(padding, "VALID");
161+ ADD_INPUT_ATTR(Targmax, 3); // int32
162+ ADD_INPUT_ATTR(include_batch_in_index, false);
163+ ADD_INPUT_ATTR(data_format, "NCHW");
164+ ADD_INPUT_ATTR(nan_prop, false);
165+ outputs.push_back(add1);
166+ // 添加完毕
167+ return SUCCESS;
168+}
169+ 
170+bool InitEnv()
171+{
172+ printf("%s - INFO - [XIR]: Start to initialize ge using ge global options\n", GetTime().c_str());
173+ std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}};
174+ Status ret = ge::GEInitialize(global_options);
175+ if (ret != SUCCESS) {
176+ printf("%s - INFO - [XIR]: Initialize ge using ge global options failed\n", GetTime().c_str());
177+ return false;
178+ }
179+ printf("%s - INFO - [XIR]: Initialize ge using ge global options success\n", GetTime().c_str());
180+ return true;
181+}
182+ 
183+bool CreateAndConfigGraph(Graph& graph, std::vector<ge::Tensor>& input)
184+{
185+ printf("%s - INFO - [XIR]: Start to CreateAndConfigGraph\n", GetTime().c_str());
186+ std::vector<Operator> inputs{};
187+ std::vector<Operator> outputs{};
188+ 
189+ DataType inDtype = DT_FLOAT16;
190+ std::cout << inDtype << std::endl;
191+ 
192+ Status ret = CreateOppInGraph(inDtype, input, inputs, outputs, graph);
193+ if (ret != SUCCESS) {
194+ printf("%s - ERROR - [XIR]: Create ir session using build options failed\n", GetTime().c_str());
195+ return false;
196+ }
197+ 
198+ if (!inputs.empty() && !outputs.empty()) {
199+ graph.SetInputs(inputs).SetOutputs(outputs);
200+ }
201+ return true;
202+}
203+ 
204+bool AddGraphToSession(ge::Session* session, Graph& graph, uint32_t graph_id)
205+{
206+ printf("%s - INFO - [XIR]: Create ir session using build options success\n", GetTime().c_str());
207+ 
208+ printf("%s - INFO - [XIR]: Start to add compute graph to ir session\n", GetTime().c_str());
209+ 
210+ std::map<AscendString, AscendString> graph_options = {
211+ 
212+ };
213+ 
214+ Status ret = session->AddGraph(graph_id, graph, graph_options);
215+ if (ret != SUCCESS) {
216+ printf("%s - ERROR - [XIR]: Add graph to session failed\n", GetTime().c_str());
217+ return false;
218+ }
219+ return true;
220+}
221+ 
222+bool RunGraph(ge::Session* session, uint32_t graph_id, std::vector<ge::Tensor>& input, std::vector<ge::Tensor>& output)
223+{
224+ printf("%s - INFO - [XIR]: Start to run graph\n", GetTime().c_str());
225+ 
226+ Status ret = session->RunGraph(graph_id, input, output);
227+ if (ret != SUCCESS) {
228+ printf("%s - ERROR - [XIR]: Run graph failed\n", GetTime().c_str());
229+ return false;
230+ }
231+ return true;
232+}
233+ 
234+bool FinalizeEnv()
235+{
236+ printf("%s - INFO - [XIR]: Start to finalize ge\n", GetTime().c_str());
237+ Status ret = ge::GEFinalize();
238+ if (ret != SUCCESS) {
239+ printf("%s - INFO - [XIR]: Finalize ge failed\n", GetTime().c_str());
240+ return false;
241+ }
242+ printf("%s - INFO - [XIR]: Finalize ge success\n", GetTime().c_str());
243+ return true;
244+}
245+ 
246+int main(int argc, char* argv[])
247+{
248+ std::vector<ge::Tensor> input{};
249+ std::vector<ge::Tensor> output{};
250+ Graph graph("max_pool_with_argmax_graph");
251+ 
252+ if (!InitEnv()) {
253+ return FAILED;
254+ }
255+ 
256+ if (!CreateAndConfigGraph(graph, input)) {
257+ return FAILED;
258+ }
259+ 
260+ std::map<ge::AscendString, ge::AscendString> session_options = {};
261+ ge::Session* session = new ge::Session(session_options);
262+ if (session == nullptr) {
263+ return FAILED;
264+ }
265+ 
266+ if (!AddGraphToSession(session, graph, 0)) {
267+ return FAILED;
268+ }
269+ 
270+ if (!RunGraph(session, 0, input, output)) {
271+ return FAILED;
272+ }
273+ printf("%s - INFO - [XIR]: Session run ir compute graph success\n", GetTime().c_str());
274+ 
275+ int input_num = input.size();
276+ for (int i = 0; i < input_num; i++) {
277+ std::cout << "input " << i << " dtype : " << input[i].GetTensorDesc().GetDataType() << std::endl;
278+ string input_file = "./max_pool_with_argmax_npu_input_" + std::to_string(i) + ".bin";
279+ uint8_t* input_data_i = input[i].GetData();
280+ int64_t input_shape = input[i].GetTensorDesc().GetShape().GetShapeSize();
281+ std::cout << "this is " << i << "th input, input shape size =" << input_shape << std::endl;
282+ uint32_t data_size = input_shape * GetDataTypeSize(input[i].GetTensorDesc().GetDataType());
283+ WriteDataToFile((const char*)input_file.c_str(), data_size, input_data_i);
284+ }
285+ 
286+ int output_num = output.size();
287+ for (int i = 0; i < output_num; i++) {
288+ std::cout << "output " << i << " dtype : " << output[i].GetTensorDesc().GetDataType() << std::endl;
289+ string output_file = "./max_pool_with_argmax_npu_output_" + std::to_string(i) + ".bin";
290+ uint8_t* output_data_i = output[i].GetData();
291+ int64_t output_shape = output[i].GetTensorDesc().GetShape().GetShapeSize();
292+ std::cout << "this is " << i << "th output, output shape size =" << output_shape << std::endl;
293+ uint32_t data_size = output_shape * GetDataTypeSize(output[i].GetTensorDesc().GetDataType());
294+ WriteDataToFile((const char*)output_file.c_str(), data_size, output_data_i);
295+ 
296+ if (i == 0) {
297+ std::cout << "output y (pooling result):" << std::endl;
298+ if (output[i].GetTensorDesc().GetDataType() == DT_FLOAT16) {
299+ uint16_t* data = reinterpret_cast<uint16_t*>(output_data_i);
300+ for (int j = 0; j < std::min(output_shape, (int64_t)16); j++) {
301+ printf(" y[%d] = %u (float16)\n", j, data[j]);
302+ }
303+ }
304+ } else if (i == 1) {
305+ std::cout << "output argmax (indices):" << std::endl;
306+ int32_t* data = reinterpret_cast<int32_t*>(output_data_i);
307+ for (int j = 0; j < std::min(output_shape, (int64_t)16); j++) {
308+ printf(" argmax[%d] = %d\n", j, data[j]);
309+ }
310+ }
311+ }
312+ 
313+ ge::AscendString error_msg = ge::GEGetErrorMsgV2();
314+ std::string error_str(error_msg.GetString());
315+ std::cout << "Error message: " << error_str << std::endl;
316+ ge::AscendString warning_msg = ge::GEGetWarningMsgV2();
317+ std::string warning_str(warning_msg.GetString());
318+ std::cout << "Warning message: " << warning_str << std::endl;
319+ printf("%s - INFO - [XIR]: Precision is ok\n", GetTime().c_str());
320+ 
321+ if (!FinalizeEnv()) {
322+ return FAILED;
323+ }
324+ 
325+ if (session != nullptr) {
326+ delete session;
327+ }
328+ 
329+ printf("%s - INFO - [XIR]: Test case passed successfully\n", GetTime().c_str());
330+ return SUCCESS;
331+}