Star869
1749
代码
代码
Issues214
Pull Requests312
Actions
项目讨论
Wiki
分析
Star869
1749
已合并
A5 support MseLoss MaxPoolWithArgmaxV3 MaxPoolGradWithArgmaxV3 #533
TangPC创建于 2025年12月29日
A5 support MseLoss MaxPoolWithArgmaxV3 MaxPoolGradWithArgmaxV3 #533
已合并
TangPC创建于 2025年12月29日
从pingchuantang/ops-nn-tpcopen合入到cann/ops-nnmaster
共 98 个文件变更+17103-467
Mloss/mse_loss/CMakeLists.txt+7-12
@@ -1,20 +1,15 @@
1# ----------------------------------------------------------------------------1# ----------------------------------------------------------------------------
2# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 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").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-message(STATUS "=== Debug: start ops.loss.mse_loss.CMakeLists.txt ")11+# 设置算子定义时支持的芯片类型
12-file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)12+set(SUPPORT_COMPUTE_UNIT "ascend910_95")
13-if(NOT ENABLE_TEST AND NOT BENCHMARK)13+# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14- list(REMOVE_ITEM CURRENT_DIRS tests)14+set(SUPPORT_TILING_DIR "arch35")
15-endif()15+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE mse_loss ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
16-foreach(SUB_DIR ${CURRENT_DIRS})
17- if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
18- add_subdirectory(${SUB_DIR})
19- endif()
20-endforeach()
Mloss/mse_loss/README.md+73-1
@@ -1,3 +1,75 @@
1# MseLoss1# MseLoss
2 2 
3-本目录仅包含MseLoss算子对应的aclnn接口；如您想要贡献该算子的AscendC实现，请参考[贡献流程](../../CONTRIBUTING.md)。3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| ---- | :----:|
7+|昇腾910_95 AI处理器|√|
8+|Atlas A3 训练系列产品/Atlas A3 推理系列产品|√|
9+|Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件|√|
10+|Atlas 200I/500 A2推理产品|×|
11+|Atlas 推理系列产品|√|
12+|Atlas 训练系列产品|√|
13+|Atlas 200/300/500 推理产品|×|
14+ 
15+## 功能说明
16+ 
17+- 算子功能：计算输入x和目标y中每个元素之间的均方误差。
18+ 
19+## 参数说明
20+ 
21+<table style="undefined;table-layout: fixed; width: 980px"><colgroup>
22+ <col style="width: 100px">
23+ <col style="width: 150px">
24+ <col style="width: 280px">
25+ <col style="width: 330px">
26+ <col style="width: 120px">
27+ </colgroup>
28+ <thead>
29+ <tr>
30+ <th>参数名</th>
31+ <th>输入/输出/属性</th>
32+ <th>描述</th>
33+ <th>数据类型</th>
34+ <th>数据格式</th>
35+ </tr></thead>
36+ <tbody>
37+ <tr>
38+ <td>predict</td>
39+ <td>输入</td>
40+ <td>输入的概率</td>
41+ <td>FLOAT16、BFLOAT16、FLOAT</td>
42+ <td>ND</td>
43+ </tr>
44+ <tr>
45+ <td>label</td>
46+ <td>输入</td>
47+ <td>输入的目标标签</td>
48+ <td>FLOAT16、BFLOAT16、FLOAT</td>
49+ <td>ND</td>
50+ </tr>
51+ <tr>
52+ <td>reduction</td>
53+ <td>输入属性</td>
54+ <td>指定要应用到输出的缩减</td>
55+ <td>String</td>
56+ <td>ND</td>
57+ </tr>
58+ </tr>
59+ <td>y</td>
60+ <td>输出</td>
61+ <td>输出的损失tensor</td>
62+ <td>FLOAT16、BFLOAT16、FLOAT</td>
63+ <td>ND</td>
64+ </tr>
65+ </tbody></table>
66+ 
67+## 约束说明
68+ 
69+无
70+ 
71+## 调用说明
72+ 
73+| 调用方式 | 调用样例 | 说明 |
74+|--------------|------------------------------------------------------------------------|--------------------------------------------------------------|
75+| aclnn调用 | [test_aclnn_mse_loss](./examples/test_aclnn_mes_loss.cpp) | 通过[aclnnMseLoss](./docs/aclnnMseLoss.md)接口方式调用MseLoss算子。 |
Mloss/mse_loss/docs/aclnnMseLoss.md+16-3
@@ -1,11 +1,18 @@
1# aclnnMseLoss1# aclnnMseLoss
2 2 
3+[📄 查看源码](https://gitcode.com/cann/ops-nn/tree/master/loss/mse_loss)
4+ 
3## 产品支持情况5## 产品支持情况
4 6 
5| 产品 | 是否支持 |7| 产品 | 是否支持 |
6| :----------------------------------------------------------- | :------: |8| :----------------------------------------------------------- | :------: |
9+| <term>昇腾910_95 AI处理器</term> | √ |
7| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |10| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
8| <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ |11| <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ |
12+| <term>Atlas 200I/500 A2 推理产品</term> | × |
13+| <term>Atlas 推理系列产品 </term> | √ |
14+| <term>Atlas 训练系列产品</term> | √ |
15+| <term>Atlas 200/300/500 推理产品</term> | × |
9 16 
10## 功能说明17## 功能说明
11 18 
@@ -13,6 +20,7 @@
13 20 
14- 计算公式：21- 计算公式：
15 22 
23+ 
16 当`reduction`为`none`时：24 当`reduction`为`none`时：
17 25 
18$$26$$
@@ -42,22 +50,26 @@ $$
42- **参数说明：**50- **参数说明：**
43 51 
44 - self(aclTensor*, 计算输入)：公式中的输入`x`，Device侧的aclTensor，self与target的数据类型满足数据类型推导规则（参见[互推导关系](../../../docs/zh/context/互推导关系.md)）。self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。52 - self(aclTensor*, 计算输入)：公式中的输入`x`，Device侧的aclTensor，self与target的数据类型满足数据类型推导规则（参见[互推导关系](../../../docs/zh/context/互推导关系.md)）。self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。
45- - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。53+ - <term>Atlas 推理系列产品</term>、<term>Atlas 训练系列产品</term>：数据类型支持FLOAT16、FLOAT。
54+ - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>昇腾910_95 AI处理器</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。
46 55 
47 - target(aclTensor*, 计算输入)：公式中的输入`y`，Device侧的aclTensor，self与target的数据类型满足数据类型推导规则（参见[互推导关系](../../../docs/zh/context/互推导关系.md)）。self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。56 - target(aclTensor*, 计算输入)：公式中的输入`y`，Device侧的aclTensor，self与target的数据类型满足数据类型推导规则（参见[互推导关系](../../../docs/zh/context/互推导关系.md)）。self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。
48- - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。57+ - <term>Atlas 推理系列产品</term>、<term>Atlas 训练系列产品</term>：数据类型支持FLOAT16、FLOAT。
58+ - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>昇腾910_95 AI处理器</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。
49 59 
50 - reduction(int64_t, 计算输入)：公式中的参数`reduction`，指定要应用到输出的缩减，支持 0('none') | 1('mean') | 2('sum')。60 - reduction(int64_t, 计算输入)：公式中的参数`reduction`，指定要应用到输出的缩减，支持 0('none') | 1('mean') | 2('sum')。
51 61 
52 'none' 表示不应用缩减，'mean' 表示输出的总和将除以输出中的元素数，'sum' 表示输出将被求和。62 'none' 表示不应用缩减，'mean' 表示输出的总和将除以输出中的元素数，'sum' 表示输出将被求和。
53 63 
54 - out(aclTensor*, 计算输出)：公式中的输出$\ell(x, y)$，Device侧的aclTensor，数据类型需要是self与target推导之后可转换的数据类型（参见[互转换关系](../../../docs/zh/context/互转换关系.md)）。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND。当reduction的值为0时，out与self、target做broadcast后的tensor的shape一致；当reduction的值为1或2时，out是0维tensor。64 - out(aclTensor*, 计算输出)：公式中的输出$\ell(x, y)$，Device侧的aclTensor，数据类型需要是self与target推导之后可转换的数据类型（参见[互转换关系](../../../docs/zh/context/互转换关系.md)）。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND。当reduction的值为0时，out与self、target做broadcast后的tensor的shape一致；当reduction的值为1或2时，out是0维tensor。
55- - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。65+ - <term>Atlas 推理系列产品</term>、<term>Atlas 训练系列产品</term>：数据类型支持FLOAT16、FLOAT。
66+ - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>昇腾910_95 AI处理器</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。
56 67 
57 - workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。68 - workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
58 69 
59 - executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。70 - executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。
60 71 
72+ 
61- **返回值：**73- **返回值：**
62 74 
63 aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。75 aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。
@@ -85,6 +97,7 @@ $$
85 97 
86 - stream(aclrtStream, 入参)：指定执行任务的Stream。98 - stream(aclrtStream, 入参)：指定执行任务的Stream。
87 99 
100+ 
88- **返回值：**101- **返回值：**
89 102 
90 aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。103 aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。
Mloss/mse_loss/docs/aclnnMseLossOut.md+11-3
@@ -4,8 +4,13 @@
4 4 
5| 产品 | 是否支持 |5| 产品 | 是否支持 |
6| :----------------------------------------------------------- | :------: |6| :----------------------------------------------------------- | :------: |
7+| <term>昇腾910_95 AI处理器</term> | √ |
7| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
8| <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ |9| <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ |
10+| <term>Atlas 200I/500 A2 推理产品</term> | × |
11+| <term>Atlas 推理系列产品 </term> | × |
12+| <term>Atlas 训练系列产品</term> | √ |
13+| <term>Atlas 200/300/500 推理产品</term> | × |
9 14 
10## 功能说明15## 功能说明
11 16 
@@ -42,17 +47,20 @@
42- **参数说明：**47- **参数说明：**
43 48 
44 - self(aclTensor*, 计算输入)：公式中的输入`x`，Device侧的aclTensor，self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。49 - self(aclTensor*, 计算输入)：公式中的输入`x`，Device侧的aclTensor，self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。
45- - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。50+ - <term>Atlas 训练系列产品</term>：数据类型支持FLOAT16、FLOAT。
51+ - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>昇腾910_95 AI处理器</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。
46 52 
47 - target(aclTensor*, 计算输入)：公式中的输入`y`，Device侧的aclTensor，self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。53 - target(aclTensor*, 计算输入)：公式中的输入`y`，Device侧的aclTensor，self与target的shape满足[broadcast关系](../../../docs/zh/context/broadcast关系.md)。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND，shape支持0到8维。
48- - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。54+ - <term>Atlas 训练系列产品</term>：数据类型支持FLOAT16、FLOAT。
55+ - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>昇腾910_95 AI处理器</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。
49 56 
50 - reduction(int64_t, 计算输入)：公式中的参数`reduction`，指定要应用到输出的缩减，支持 0('none') | 1('mean') | 2('sum')。57 - reduction(int64_t, 计算输入)：公式中的参数`reduction`，指定要应用到输出的缩减，支持 0('none') | 1('mean') | 2('sum')。
51 58 
52 'none' 表示不应用减少，'mean' 表示输出将被reduce 0轴求均值，'sum' 表示输出将被reduce 0轴求和。59 'none' 表示不应用减少，'mean' 表示输出将被reduce 0轴求均值，'sum' 表示输出将被reduce 0轴求和。
53 60 
54 - out(aclTensor*, 计算输出)：公式中的输出$\ell(x, y)$，Device侧的aclTensor。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND。当reduction的值为0时，out与self、target做broadcast后的tensor的shape一致；当reduction的值为1或2时，out与self、target做broadcast后，再reduce 0轴得到的tensor的shape一致。61 - out(aclTensor*, 计算输出)：公式中的输出$\ell(x, y)$，Device侧的aclTensor。支持[非连续的Tensor](../../../docs/zh/context/非连续的Tensor.md)，[数据格式](../../../docs/zh/context/数据格式.md)支持ND。当reduction的值为0时，out与self、target做broadcast后的tensor的shape一致；当reduction的值为1或2时，out与self、target做broadcast后，再reduce 0轴得到的tensor的shape一致。
55- - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。62+ - <term>Atlas 训练系列产品</term>：数据类型支持FLOAT16、FLOAT。
63+ - <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>昇腾910_95 AI处理器</term>：数据类型支持BFLOAT16、FLOAT16、FLOAT。
56 64 
57 - workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。65 - workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
58 66 
Rloss/mse_loss/examples/test_aclnn_mse_loss.cpp→loss/mse_loss/examples/arch35/test_aclnn_mse_loss.cpp+4-4
@@ -1,10 +1,10 @@
1/**1/**
2+ * This program is free software, you can redistribute it and/or modify.
2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of4+ * This file is a part of the CANN Open Software.
4- * CANN Open Software License Agreement Version 2.0 (the "License").5+ * Licensed under 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 * 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+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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#include <iostream>10#include <iostream>
Rloss/mse_loss/examples/test_aclnn_mse_loss_out.cpp→loss/mse_loss/examples/arch35/test_aclnn_mse_loss_out.cpp+4-4
@@ -1,10 +1,10 @@
1/**1/**
2+ * This program is free software, you can redistribute it and/or modify.
2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of4+ * This file is a part of the CANN Open Software.
4- * CANN Open Software License Agreement Version 2.0 (the "License").5+ * Licensed under 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 * 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+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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#include <iostream>10#include <iostream>
Rloss/mse_loss/op_host/op_api/aclnn_mse_loss.cpp→loss/mse_loss/op_api/aclnn_mse_loss.cpp+0-0
文件重命名但无更改。
Rloss/mse_loss/op_host/op_api/aclnn_mse_loss.h→loss/mse_loss/op_api/aclnn_mse_loss.h+1-0
@@ -7,6 +7,7 @@
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#ifndef OP_API_INC_MSE_LOSS_H_11#ifndef OP_API_INC_MSE_LOSS_H_
11#define OP_API_INC_MSE_LOSS_H_12#define OP_API_INC_MSE_LOSS_H_
12 13 
Rloss/mse_loss/op_host/op_api/aclnn_mse_loss_out.cpp→loss/mse_loss/op_api/aclnn_mse_loss_out.cpp+0-0
文件重命名但无更改。
Rloss/mse_loss/op_host/op_api/aclnn_mse_loss_out.h→loss/mse_loss/op_api/aclnn_mse_loss_out.h+1-0
@@ -7,6 +7,7 @@
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#ifndef OP_API_INC_MSE_LOSS_OUT_H_11#ifndef OP_API_INC_MSE_LOSS_OUT_H_
11#define OP_API_INC_MSE_LOSS_OUT_H_12#define OP_API_INC_MSE_LOSS_OUT_H_
12 13 
Rloss/mse_loss/op_host/op_api/mse_loss.cpp→loss/mse_loss/op_api/mse_loss.cpp+0-0
文件重命名但无更改。
Rloss/mse_loss/op_host/op_api/mse_loss.h→loss/mse_loss/op_api/mse_loss.h+0-0
文件重命名但无更改。
Aloss/mse_loss/op_graph/mse_loss_proto.h+43-0
@@ -0,0 +1,43 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
CANN-robot
CANN-robotCANN-robot2025年12月29日

文档/注释不一致: 文件头注释中的文件名与实际文件名不一致。注释中描述的文件名为 'nn_norm_ops.h'，但实际文件名是 'mse_loss_proto.h'。这种不一致可能导致开发者混淆，影响代码的可读性和维护性。

问题类型: 文档/注释不一致 文件路径: loss/mse_loss/op_graph/mse_loss_proto.h 行号: 11 问题代码:

/*!
 * \file nn_norm_ops.h
 * \brief
 */

修改建议:

将文件头注释中的文件名更新为实际文件名 'mse_loss_proto.h'，以保持一致性。

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

likedislike
12+ * \file nn_norm_ops.h
13+ * \brief
14+ */
15+#ifndef OPS_BUILT_IN_OP_PROTO_INC_NN_NORM_OPS_H_
CANN-robot
CANN-robotCANN-robot2025年12月29日

宏定义命名冲突风险: 头文件保护宏的命名 'OPS_BUILT_IN_OP_PROTO_INC_NN_NORM_OPS_H_' 与注释中提到的文件名 'nn_norm_ops.h' 相关，但与当前实际文件名 'mse_loss_proto.h' 不符。这可能导致不同头文件使用相同宏定义的风险，引发编译错误或包含错误版本的头文件。

问题类型: 宏定义命名冲突风险 文件路径: loss/mse_loss/op_graph/mse_loss_proto.h 行号: 15 问题代码:

#ifndef OPS_BUILT_IN_OP_PROTO_INC_NN_NORM_OPS_H_
#define OPS_BUILT_IN_OP_PROTO_INC_NN_NORM_OPS_H_

修改建议:

将头文件保护宏的命名更新为与实际文件名一致，例如改为 'OPS_BUILT_IN_OP_PROTO_INC_MSE_LOSS_PROTO_H_'。

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

likedislike
16+#define OPS_BUILT_IN_OP_PROTO_INC_NN_NORM_OPS_H_
17+ 
18+#include "graph/operator_reg.h"
19+namespace ge {
20+ 
21+/**
22+ * @brief Computes mse loss.
23+ * @par Inputs:
24+ * two inputs, including:
25+ * @li predict: An ND Tensor of dtype float16, float32 or bfloat16.
26+ * @li label: An ND Tensor of dtype float16, float32 or bfloat16.\n
27+ *
28+ * @par Attributes:
29+ * reduction:An optional str from sum, none, mean, Defaults to "mean".\n
CANN-robot
CANN-robotCANN-robot2025年12月29日

代码结构与可维护性: 属性'reduction'的注释描述中缺少空格，格式不规范。注释中写为'reduction:An optional str'，冒号后应有一个空格。虽然不影响功能，但影响代码的可读性和规范性。

问题类型: 代码结构与可维护性 文件路径: loss/mse_loss/op_graph/mse_loss_proto.h 行号: 29 问题代码:

 * reduction:An optional str from sum, none, mean, Defaults to "mean".\n

修改建议:

在冒号后添加空格，修改为'reduction: An optional str from sum, none, mean, Defaults to "mean".'，保持代码注释格式的一致性。

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

likedislike
30+ *
31+ * @par Outputs:
32+ * y: when reduction=sum/mean, y is scale. when reduction=none, y has
33+ * same type and shape as "predict".\n
34+ */
35+REG_OP(MseLoss)
36+ .INPUT(predict, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
37+ .INPUT(label, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
38+ .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
39+ .ATTR(reduction, String, "mean")
40+ .OP_END_FACTORY_REG(MseLoss)
41+ 
42+} // namespace ge
43+#endif // OPS_BUILT_IN_OP_PROTO_INC_NN_NORM_OPS_H_
Aloss/mse_loss/op_host/arch35/mse_loss_tiling.cpp+57-0文件内容审核中，请稍后刷新重试
Aloss/mse_loss/op_host/arch35/mse_loss_tiling.h+32-0文件内容审核中，请稍后刷新重试
Aloss/mse_loss/op_host/arch35/mse_loss_tiling_arch35.cpp+206-0文件内容审核中，请稍后刷新重试
Aloss/mse_loss/op_host/arch35/mse_loss_tiling_arch35.h+56-0文件内容审核中，请稍后刷新重试
Aloss/mse_loss/op_host/mse_loss_def.cpp+52-0
@@ -0,0 +1,52 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file mse_loss_def.cpp
13+ * \brief mse_loss def
14+ */
15+ 
16+#include <cstdint>
17+#include "register/op_def_registry.h"
18+ 
19+namespace ops {
20+class MseLoss : public OpDef {
21+public:
22+ explicit MseLoss(const char* name) : OpDef(name)
23+ {
24+ this->Input("predict")
25+ .ParamType(REQUIRED)
26+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
27+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
28+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
29+ this->Input("label")
30+ .ParamType(REQUIRED)
31+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
32+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
33+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
34+ this->Output("y")
35+ .ParamType(REQUIRED)
36+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
37+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
38+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
39+ this->Attr("reduction").AttrType(OPTIONAL).String("mean");
40+ 
41+ OpAICoreConfig aicoreConfig;
42+ aicoreConfig.DynamicCompileStaticFlag(true)
43+ .DynamicRankSupportFlag(true)
44+ .DynamicShapeSupportFlag(true)
45+ .PrecisionReduceFlag(false)
46+ .ExtendCfgInfo("opFile.value", "mse_loss_apt");
47+ this->AICore().AddConfig("ascend910_95", aicoreConfig);
48+ }
49+};
50+ 
51+OP_ADD(MseLoss);
52+} // namespace ops
Aloss/mse_loss/op_host/mse_loss_infershape.cpp+60-0
@@ -0,0 +1,60 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file mse_loss_infershape.cpp
13+ * \brief
14+ */
15+ 
16+#include "graph/utils/type_utils.h"
17+#include "runtime/infer_shape_context.h"
18+#include "register/op_impl_registry.h"
19+#include "log/log.h"
20+#include "util/shape_util.h"
21+#include "util/math_util.h"
22+ 
23+using namespace ge;
24+namespace ops {
25+static ge::graphStatus InferShapeTwoInOneOutWithReduction(gert::InferShapeContext* context)
26+{
27+ auto input_x_shape = context->GetInputShape(0);
28+ OP_CHECK_NULL_WITH_CONTEXT(context, input_x_shape);
29+ auto input_y_shape = context->GetInputShape(1);
30+ OP_CHECK_NULL_WITH_CONTEXT(context, input_y_shape);
31+ 
32+ OP_CHECK_IF(
33+ *input_x_shape != *input_y_shape,
34+ OP_LOGE(
35+ context->GetNodeName(), "input_x shape %s must be same as input_y shape %s",
36+ Ops::Base::ToString(*input_x_shape).c_str(), Ops::Base::ToString(*input_y_shape).c_str()),
37+ return GRAPH_FAILED);
38+ 
39+ auto out_shape = context->GetOutputShape(0);
40+ OP_CHECK_NULL_WITH_CONTEXT(context, out_shape);
41+ auto attrs = context->GetAttrs();
42+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
43+ 
44+ const char* reduction = attrs->GetAttrPointer<char>(0);
45+ OP_CHECK_NULL_WITH_CONTEXT(context, reduction);
46+ if (strcmp(reduction, "none") == 0) {
47+ auto in_shape = context->GetInputShape(0);
48+ OP_CHECK_NULL_WITH_CONTEXT(context, in_shape);
49+ *out_shape = *in_shape;
50+ } else {
51+ // if reduction == "mean" or reduction == "sum" , output a scalar
52+ out_shape->SetDimNum(0);
53+ }
54+ 
55+ return GRAPH_SUCCESS;
56+}
57+ 
58+IMPL_OP_INFERSHAPE(MseLoss).InferShape(InferShapeTwoInOneOutWithReduction);
59+ 
60+} // namespace ops
Aloss/mse_loss/op_kernel/arch35/mse_loss_dag.h+82-0
@@ -0,0 +1,82 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file mse_loss_dag.h
13+ * \brief
14+ */
15+ 
16+#ifndef CANN_CUSTOM_OPS_MSE_LOSS_DAG_H
17+#define CANN_CUSTOM_OPS_MSE_LOSS_DAG_H
18+ 
19+#include "atvoss/util/dag.h"
20+#include "atvoss/util/vec.h"
21+#include "atvoss/util/placeholder.h"
22+#include "atvoss/reduce/reduce_operator.h"
23+ 
24+namespace MseLoss {
25+using namespace Ops::Base;
26+using namespace AscendC;
27+ 
28+template <typename T, typename PromteT = float>
29+struct MseLossOp {
30+ // 通过Compute构造计算图
31+ // (a-b)^2
32+ using OpCopyIn0 = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
33+ using OpCopyIn1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
34+ 
35+ using OpCopyIn0Cast = Bind<Vec::Cast<PromteT, T, 0>, OpCopyIn0>;
36+ using OpCopyIn1Cast = Bind<Vec::Cast<PromteT, T, 0>, OpCopyIn1>;
37+ 
38+ using OpSub = Bind<Vec::Sub<PromteT>, OpCopyIn0Cast, OpCopyIn1Cast>;
39+ using OpMul = Bind<Vec::Mul<PromteT>, OpSub, OpSub>;
40+ 
41+ using OpResultCast = Bind<Vec::Cast<T, PromteT, 1>, OpMul>;
42+ using OpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, OpResultCast>;
43+ // 指定输出节点
44+ using Outputs = Elems<OpCopyOut>; // 设置输出
45+ // 指定计算顺序
46+ using OpDag = DAGSch<Outputs>;
47+};
48+template <typename T, typename PromteT>
49+struct MseLossSumDag {
50+ using OpCopyIn0 = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
51+ using OpCopyIn1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
52+ using OpCopyIn0Cast = Bind<Vec::Cast<PromteT, T, 0>, OpCopyIn0>;
53+ using OpCopyIn1Cast = Bind<Vec::Cast<PromteT, T, 0>, OpCopyIn1>;
54+ using OpSub = Bind<Vec::Sub<PromteT>, OpCopyIn0Cast, OpCopyIn1Cast>;
55+ using OpMul = Bind<Vec::Mul<PromteT>, OpSub, OpSub>;
56+ using ReduceOp0 = Bind<Vec::ReduceSumOp<PromteT>, OpMul>;
57+ using Cast1 = Bind<Vec::Cast<T, PromteT, 1>, ReduceOp0>;
58+ using OpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, Cast1>;
59+ using Outputs = Elems<OpCopyOut>;
60+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
61+ using OpDag = DAGSch<Outputs, void, MemCfg>;
62+};
63+ 
64+template <typename T, typename PromteT>
65+struct MseLossMeanDag {
66+ using OpCopyIn0 = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
67+ using OpCopyIn1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
68+ using OpCopyIn0Cast = Bind<Vec::Cast<PromteT, T, 0>, OpCopyIn0>;
69+ using OpCopyIn1Cast = Bind<Vec::Cast<PromteT, T, 0>, OpCopyIn1>;
70+ using OpSub = Bind<Vec::Sub<PromteT>, OpCopyIn0Cast, OpCopyIn1Cast>;
71+ using OpMul = Bind<Vec::Mul<PromteT>, OpSub, OpSub>;
72+ using ReduceOp0 = Bind<Vec::ReduceSumOp<PromteT>, OpMul>;
73+ using Mul0 = Bind<Vec::Muls<PromteT>, ReduceOp0, Placeholder::Var<PromteT, 0>>;
74+ using Cast1 = Bind<Vec::Cast<T, PromteT, 1>, Mul0>;
75+ using OpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, Cast1>;
76+ using Outputs = Elems<OpCopyOut>;
77+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
78+ using OpDag = DAGSch<Outputs, void, MemCfg>;
79+};
80+} // namespace MseLoss
81+ 
82+#endif // CANN_CUSTOM_OPS_MSE_LOSS_DAG_H
Aloss/mse_loss/op_kernel/arch35/mse_loss_tiling_def.h+46-0文件内容审核中，请稍后刷新重试
Aloss/mse_loss/op_kernel/arch35/mse_loss_tiling_key.h+70-0
@@ -0,0 +1,70 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file mse_loss_tiling_key.h
13+ * \brief mseloss tiling key
14+ */
15+ 
16+#ifndef _MSE_LOSS_TILING_KEY_H_
17+#define _MSE_LOSS_TILING_KEY_H_
18+#define REDUCTION_BIT_WIDTH 2
19+#define MSELOSS_SUM 1
20+#define MSELOSS_MEAN 2
21+#define MSELOSS_TPL_FP16 10
22+#define MSELOSS_TPL_FP32 20
23+#define MSELOSS_TPL_BFP16 30
24+ 
25+#include "atvoss/reduce/reduce_tiling_key_decl.h"
26+ 
27+ASCENDC_TPL_ARGS_DECL(
28+ MSELOSS, REDUCE_TPL_KEY_DECL(),
29+ ASCENDC_TPL_UINT_DECL(Reduction, REDUCTION_BIT_WIDTH, ASCENDC_TPL_UI_LIST, 0, MSELOSS_SUM, MSELOSS_MEAN),
30+ ASCENDC_TPL_UINT_DECL(
31+ Dtype, BIT_WIDTH, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP16, MSELOSS_TPL_FP32, MSELOSS_TPL_BFP16));
32+ 
33+ASCENDC_TPL_SEL(
34+ ASCENDC_TPL_ARGS_SEL(
35+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY), ASCENDC_TPL_UINT_SEL(PatternID, ASCENDC_TPL_UI_LIST, 0),
36+ ASCENDC_TPL_UINT_SEL(LoopARCount, ASCENDC_TPL_UI_LIST, 0),
37+ ASCENDC_TPL_UINT_SEL(LoopInnerARCount, ASCENDC_TPL_UI_LIST, 0),
38+ ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, 0),
39+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP16)),
40+ 
41+ ASCENDC_TPL_ARGS_SEL(
42+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY), ASCENDC_TPL_UINT_SEL(PatternID, ASCENDC_TPL_UI_LIST, 0),
43+ ASCENDC_TPL_UINT_SEL(LoopARCount, ASCENDC_TPL_UI_LIST, 0),
44+ ASCENDC_TPL_UINT_SEL(LoopInnerARCount, ASCENDC_TPL_UI_LIST, 0),
45+ ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, 0),
46+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP32)),
47+ 
48+ ASCENDC_TPL_ARGS_SEL(
49+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY), ASCENDC_TPL_UINT_SEL(PatternID, ASCENDC_TPL_UI_LIST, 0),
50+ ASCENDC_TPL_UINT_SEL(LoopARCount, ASCENDC_TPL_UI_LIST, 0),
51+ ASCENDC_TPL_UINT_SEL(LoopInnerARCount, ASCENDC_TPL_UI_LIST, 0),
52+ ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, 0),
53+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_BFP16)),
54+ 
55+ ASCENDC_TPL_ARGS_SEL(
56+ REDUCE_TPL_KEY_SEL_EMPTY(), ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, MSELOSS_SUM, MSELOSS_MEAN),
57+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP32)),
58+ 
59+ ASCENDC_TPL_ARGS_SEL(
60+ REDUCE_TPL_KEY_SEL_AR_NORMAL(), ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, MSELOSS_SUM, MSELOSS_MEAN),
61+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP32)),
62+ 
63+ ASCENDC_TPL_ARGS_SEL(
64+ REDUCE_TPL_KEY_SEL_AR_GROUP(), ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, MSELOSS_SUM, MSELOSS_MEAN),
65+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP32)),
66+ 
67+ ASCENDC_TPL_ARGS_SEL(
68+ REDUCE_TPL_KEY_SEL_A(), ASCENDC_TPL_UINT_SEL(Reduction, ASCENDC_TPL_UI_LIST, MSELOSS_SUM, MSELOSS_MEAN),
69+ ASCENDC_TPL_UINT_SEL(Dtype, ASCENDC_TPL_UI_LIST, MSELOSS_TPL_FP32)));
70+#endif
Aloss/mse_loss/op_kernel/mse_loss_apt.cpp+54-0
@@ -0,0 +1,54 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file mse_loss_apt.cpp
13+ * \brief
14+ */
15+#include <cmath>
16+#include "kernel_operator.h"
17+#include "kernel_tiling/kernel_tiling.h"
18+#include "./arch35/mse_loss_dag.h"
19+#include "./arch35/mse_loss_tiling_key.h"
20+#include "atvoss/elewise/elewise_sch.h"
21+#include "atvoss/util/dfx.h"
22+#include "atvoss/reduce/reduce_sch.h"
23+#include "./arch35/mse_loss_tiling_def.h"
24+#include "atvoss/reduce/reduce_sch.h"
25+ 
26+using namespace ReduceOpTmpl;
27+using namespace AscendC;
28+ 
29+template <REDUCE_TPL_PARAM, uint32_t Reduction, uint32_t Dtype>
30+__global__ __aicore__ void mse_loss(GM_ADDR predict, GM_ADDR label, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
31+{
32+ REGISTER_TILING_DEFAULT(MseLossTilingData);
33+ GET_TILING_DATA_WITH_STRUCT(MseLossTilingData, tilingData, tiling);
34+ TPipe pipe;
35+ using PromoteType = __reduceType::GetPromoteType<DTYPE_PREDICT>::T;
36+ if constexpr (Reduction == 0) {
37+ ElementwiseSch<0UL, MseLoss::MseLossOp<DTYPE_PREDICT>::OpDag> sch(
38+ &(tilingData.baseTiling),
39+ &pipe); // 获取Schedule
40+ sch.Init(predict, label, y);
41+ sch.Process();
42+ } else if constexpr (Reduction == 1) {
43+ using Op = ReduceSch<REDUCE_TPL_VALUE, MseLoss::MseLossSumDag<DTYPE_PREDICT, PromoteType>::OpDag>;
44+ Op op((ReduceOpTilingData*)&tilingData.reduceTiling);
45+ op.Init(&pipe, predict, label, y, workspace);
46+ op.Process();
47+ } else if constexpr (Reduction == 2) {
48+ using Op = ReduceSch<REDUCE_TPL_VALUE, MseLoss::MseLossMeanDag<DTYPE_PREDICT, PromoteType>::OpDag>;
49+ Op op((ReduceOpTilingData*)&tilingData.reduceTiling);
50+ op.template SetVar<PromoteType, 0>(tilingData.reduceTiling.meanVar);
51+ op.Init(&pipe, predict, label, y, workspace);
52+ op.Process(static_cast<DTYPE_PREDICT>(NAN));
53+ }
54+}
Mloss/mse_loss/tests/CMakeLists.txt+7-8
@@ -1,18 +1,17 @@
1# ----------------------------------------------------------------------------1# ----------------------------------------------------------------------------
2# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 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").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-message(STATUS "=== Debug: start ops.loss.mse_loss.tests.CMakeLists.txt ")11+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
12-file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)12+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
13-message(STATUS "=== Debug: CURRENT_DIRS =${CURRENT_DIRS} ")13+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
14-foreach(SUB_DIR ${CURRENT_DIRS})14+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
15- if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
16 add_subdirectory(${SUB_DIR})15 add_subdirectory(${SUB_DIR})
17 endif()16 endif()
18-endforeach()17+endforeach()
Mloss/mse_loss/tests/ut/CMakeLists.txt+6-7
@@ -1,18 +1,17 @@
1# ----------------------------------------------------------------------------1# ----------------------------------------------------------------------------
2# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 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").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-# 每个目录下需要生成的可执行文件，具体参考：ops/built-in/test/CMakeLists.txt： 50~12411+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
12-file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)12+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
13-message(STATUS "=== Debug: CURRENT_DIRS =${CURRENT_DIRS} ")13+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
14-foreach(SUB_DIR ${CURRENT_DIRS})14+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
15- if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
16 add_subdirectory(${SUB_DIR})15 add_subdirectory(${SUB_DIR})
17 endif()16 endif()
18endforeach()17endforeach()
Rpooling/max_pool_with_argmax_v3/op_host/CMakeLists.txt→loss/mse_loss/tests/ut/op_api/CMakeLists.txt+14-12
@@ -1,12 +1,14 @@
1-# ----------------------------------------------------------------------------1+# ----------------------------------------------------------------------------
2-# Copyright (c) 2025 Huawei Technologies Co., Ltd.2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 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").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-message(STATUS "=== Debug: start ops.pooling.max_pool3d_with_argmax_v3.op_host.CMakeLists.txt ")10+ 
11- 11+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE max_pool3d_with_argmax_v3 ACLNNTYPE aclnn_exclude)12+if(UT_TEST_ALL OR OP_API_UT)
13+ add_modules_ut_sources(HOSTNAME ${OP_API_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14+endif()
Rloss/mse_loss/tests/ut/op_host/test_aclnn_mse_loss.cpp→loss/mse_loss/tests/ut/op_api/test_aclnn_mse_loss.cpp+2-1
@@ -7,11 +7,12 @@
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#include "gtest/gtest.h"11#include "gtest/gtest.h"
11#include <array>12#include <array>
12#include <vector>13#include <vector>
13 14 
14-#include "../../../op_host/op_api/aclnn_mse_loss.h"15+#include "../../../op_api/aclnn_mse_loss.h"
15 16 
16#include "op_api_ut_common/op_api_ut.h"17#include "op_api_ut_common/op_api_ut.h"
17#include "op_api_ut_common/scalar_desc.h"18#include "op_api_ut_common/scalar_desc.h"
Aloss/mse_loss/tests/ut/op_api/test_aclnn_mse_loss_out.cpp+376-0
@@ -0,0 +1,376 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "gtest/gtest.h"
12+#include <array>
13+#include <vector>
14+ 
15+#include "../../../op_api/aclnn_mse_loss_out.h"
16+ 
17+#include "op_api_ut_common/op_api_ut.h"
18+#include "op_api_ut_common/scalar_desc.h"
19+#include "op_api_ut_common/tensor_desc.h"
20+#include "opdev/platform.h"
21+ 
22+#include <unistd.h>
23+ 
24+using namespace op;
25+using namespace std;
26+ 
27+class l2_mse_loss_out_test : public testing::Test {
28+protected:
29+ static void SetUpTestCase()
30+ {
31+ std::cout << "mse_loss_test SetUp" << std::endl;
32+ }
33+ 
34+ static void TearDownTestCase()
35+ {
36+ std::cout << "mse_loss_test TearDown" << std::endl;
37+ }
38+};
39+ 
40+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_01_float_nd_none)
41+{
42+ auto selfDesc = TensorDesc({1, 2, 3, 2}, ACL_FLOAT, ACL_FORMAT_NCHW).ValueRange(-1, 1);
43+ auto targetDesc = TensorDesc({2, 2, 1, 2}, ACL_FLOAT, ACL_FORMAT_NCHW).ValueRange(-1, 1);
44+ int64_t reduction = 0;
45+ 
46+ auto outDesc = TensorDesc({2, 2, 3, 2}, ACL_FLOAT, ACL_FORMAT_NCHW).Precision(0.0001, 0.0001);
47+ 
48+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
49+ 
50+ // SAMPLE: only test GetWorkspaceSize
51+ uint64_t workspace_size = 0;
52+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
53+ EXPECT_EQ(aclRet, ACL_SUCCESS);
54+ 
55+ // SAMPLE: precision simulate
56+ ut.TestPrecision();
57+}
58+ 
59+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_02_float16_nchw_mean)
60+{
61+ auto selfDesc = TensorDesc({2, 3, 2}, ACL_FLOAT16, ACL_FORMAT_NCHW).ValueRange(-1, 1);
62+ auto targetDesc = TensorDesc({2, 3, 2}, ACL_FLOAT16, ACL_FORMAT_NCHW).ValueRange(-1, 1);
63+ int64_t reduction = 1;
64+ 
65+ auto outDesc = TensorDesc({3, 2}, ACL_FLOAT16, ACL_FORMAT_NCHW).Precision(0.001, 0.001);
66+ 
67+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
68+ 
69+ // SAMPLE: only test GetWorkspaceSize
70+ uint64_t workspace_size = 0;
71+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
72+ EXPECT_EQ(aclRet, ACL_SUCCESS);
73+ 
74+ // SAMPLE: precision simulate
75+ // ut.TestPrecision();
76+}
77+ 
78+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_03_float_float16_nhwc_sum)
79+{
80+ auto selfDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_NHWC).ValueRange(-1, 1);
81+ auto targetDesc = TensorDesc({2, 2}, ACL_FLOAT16, ACL_FORMAT_NHWC).ValueRange(-1, 1);
82+ int64_t reduction = 2;
83+ 
84+ auto outDesc = TensorDesc({2}, ACL_FLOAT, ACL_FORMAT_NHWC).Precision(0.0001, 0.0001);
85+ 
86+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
87+ 
88+ // SAMPLE: only test GetWorkspaceSize
89+ uint64_t workspace_size = 0;
90+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
91+ EXPECT_EQ(aclRet, ACL_SUCCESS);
92+ 
93+ // SAMPLE: precision simulate
94+ // ut.TestPrecision();
95+}
96+ 
97+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_04_float16_float_nhwc_mean)
98+{
99+ auto selfDesc = TensorDesc({2, 2}, ACL_FLOAT16, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
100+ auto targetDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
101+ int64_t reduction = 1;
102+ 
103+ auto outDesc = TensorDesc({2}, ACL_FLOAT, ACL_FORMAT_NDHWC).Precision(0.0001, 0.0001);
104+ 
105+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
106+ 
107+ // SAMPLE: only test GetWorkspaceSize
108+ uint64_t workspace_size = 0;
109+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
110+ EXPECT_EQ(aclRet, ACL_SUCCESS);
111+ 
112+ // SAMPLE: precision simulate
113+ // ut.TestPrecision();
114+}
115+ 
116+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_05_float_nd_empty_tensor_none)
117+{
118+ auto selfDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
119+ auto targetDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
120+ int64_t reduction = 0;
121+ 
122+ auto outDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
123+ 
124+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
125+ 
126+ // SAMPLE: only test GetWorkspaceSize
127+ uint64_t workspace_size = 0;
128+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
129+ EXPECT_EQ(aclRet, ACL_SUCCESS);
130+ 
131+ // SAMPLE: precision simulate
132+ ut.TestPrecision();
133+}
134+ 
135+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_06_float_nd_empty_tensor_mean)
136+{
137+ auto selfDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
138+ auto targetDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
139+ int64_t reduction = 1;
140+ 
141+ auto outDesc = TensorDesc({0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
142+ 
143+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
144+ 
145+ // SAMPLE: only test GetWorkspaceSize
146+ uint64_t workspace_size = 0;
147+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
148+ EXPECT_EQ(aclRet, ACL_SUCCESS);
149+ 
150+ // SAMPLE: precision simulate
151+ ut.TestPrecision();
152+}
153+ 
154+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_07_float_nd_empty_tensor_sum)
155+{
156+ auto selfDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
157+ auto targetDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
158+ int64_t reduction = 2;
159+ 
160+ auto outDesc = TensorDesc({0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
161+ 
162+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
163+ 
164+ // SAMPLE: only test GetWorkspaceSize
165+ uint64_t workspace_size = 0;
166+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
167+ EXPECT_EQ(aclRet, ACL_SUCCESS);
168+ 
169+ // SAMPLE: precision simulate
170+ ut.TestPrecision();
171+}
172+ 
173+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_11_float_hwcn_input_not_contiguous)
174+{
175+ auto selfDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_HWCN, {1, 5}, 0, {4, 5}).ValueRange(-2, 2);
176+ auto targetDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_HWCN, {1, 5}, 0, {4, 5}).ValueRange(-2, 2);
177+ int64_t reduction = 0;
178+ 
179+ auto outDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_HWCN).Precision(0.0001, 0.0001);
180+ 
181+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
182+ 
183+ // SAMPLE: only test GetWorkspaceSize
184+ uint64_t workspace_size = 0;
185+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
186+ EXPECT_EQ(aclRet, ACL_SUCCESS);
187+ 
188+ // SAMPLE: precision simulate
189+ ut.TestPrecision();
190+}
191+ 
192+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_12_float_ncdhw_out_not_contiguous)
193+{
194+ auto selfDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_NCDHW).ValueRange(-2, 2);
195+ auto targetDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_NCDHW, {1, 5}, 0, {4, 5}).ValueRange(-2, 2);
196+ int64_t reduction = 0;
197+ 
198+ auto outDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_NCDHW, {1, 5}, 0, {4, 5}).Precision(0.0001, 0.0001);
199+ 
200+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
201+ 
202+ // SAMPLE: only test GetWorkspaceSize
203+ uint64_t workspace_size = 0;
204+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
205+ EXPECT_EQ(aclRet, ACL_SUCCESS);
206+ 
207+ // SAMPLE: precision simulate
208+ ut.TestPrecision();
209+}
210+ 
211+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_13_reduction_error)
212+{
213+ auto selfDesc = TensorDesc({8}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
214+ auto targetDesc = TensorDesc({8}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
215+ int64_t reduction = 3;
216+ 
217+ auto outDesc = TensorDesc({8}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
218+ 
219+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
220+ 
221+ // SAMPLE: only test GetWorkspaceSize
222+ uint64_t workspace_size = 0;
223+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
224+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
225+}
226+ 
227+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_14_input_out_nullptr)
228+{
229+ auto tensor_desc = TensorDesc({10, 5}, ACL_FLOAT, ACL_FORMAT_ND);
230+ int64_t reduction = 0;
231+ 
232+ auto ut_l = OP_API_UT(aclnnMseLossOut, INPUT(nullptr, tensor_desc, reduction), OUTPUT(tensor_desc));
233+ // SAMPLE: only test GetWorkspaceSize
234+ uint64_t workspace_size = 0;
235+ aclnnStatus aclRet = ut_l.TestGetWorkspaceSize(&workspace_size);
236+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
237+ 
238+ auto ut_r = OP_API_UT(aclnnMseLossOut, INPUT(tensor_desc, nullptr, reduction), OUTPUT(tensor_desc));
239+ // SAMPLE: only test GetWorkspaceSize
240+ aclRet = ut_r.TestGetWorkspaceSize(&workspace_size);
241+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
242+ 
243+ auto ut_o = OP_API_UT(aclnnMseLossOut, INPUT(tensor_desc, tensor_desc, reduction), OUTPUT(nullptr));
244+ // SAMPLE: only test GetWorkspaceSize
245+ aclRet = ut_o.TestGetWorkspaceSize(&workspace_size);
246+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
247+}
248+ 
249+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_15_aclnnMseLossOut_input_error_shape)
250+{
251+ auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
252+ auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
253+ int64_t reduction = 0;
254+ 
255+ auto outDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
256+ 
257+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
258+ // SAMPLE: only test GetWorkspaceSize
259+ uint64_t workspace_size = 0;
260+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
261+ // EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
262+}
263+ 
264+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_16_aclnnMseLossOut_error_output_dtype)
265+{
266+ auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
267+ auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
268+ int64_t reduction = 0;
269+ 
270+ auto outDesc = TensorDesc({8, 2}, ACL_INT32, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
271+ 
272+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
273+ // SAMPLE: only test GetWorkspaceSize
274+ uint64_t workspace_size = 0;
275+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
276+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
277+}
278+ 
279+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_17_aclnnMseLossOut_error_input_dtype)
280+{
281+ auto selfDesc = TensorDesc({6, 2, 1, 2}, ACL_INT32, ACL_FORMAT_ND).ValueRange(-2, 2);
282+ auto targetDesc = TensorDesc({6, 2, 1, 2}, ACL_INT16, ACL_FORMAT_ND).ValueRange(-2, 2);
283+ int64_t reduction = 0;
284+ 
285+ auto outDesc = TensorDesc({6, 2, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
286+ 
287+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
288+ // SAMPLE: only test GetWorkspaceSize
289+ uint64_t workspace_size = 0;
290+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
291+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
292+}
293+ 
294+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_18_aclnnMseLossOut_input_error_shape_len)
295+{
296+ auto tensorDesc9 = TensorDesc({2, 3, 4, 5, 6, 7, 8, 9, 10}, ACL_INT32, ACL_FORMAT_ND).ValueRange(-2, 2);
297+ auto tensorDesc1 = TensorDesc({7, 8, 9, 10}, ACL_INT16, ACL_FORMAT_ND).ValueRange(-2, 2);
298+ auto tensorDesc2 = TensorDesc({7, 8, 9, 10}, ACL_INT16, ACL_FORMAT_ND).ValueRange(-2, 2);
299+ int64_t reduction = 0;
300+ 
301+ auto ut_self = OP_API_UT(aclnnMseLossOut, INPUT(tensorDesc9, tensorDesc1, reduction), OUTPUT(tensorDesc2));
302+ // SAMPLE: only test GetWorkspaceSize
303+ uint64_t workspace_size = 0;
304+ aclnnStatus aclRet = ut_self.TestGetWorkspaceSize(&workspace_size);
305+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
306+ 
307+ auto ut_tar = OP_API_UT(aclnnMseLossOut, INPUT(tensorDesc1, tensorDesc9, reduction), OUTPUT(tensorDesc2));
308+ // SAMPLE: only test GetWorkspaceSize
309+ aclRet = ut_tar.TestGetWorkspaceSize(&workspace_size);
310+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
311+}
312+ 
313+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_19_float_hwcn_ndhwc_mean)
314+{
315+ auto selfDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_HWCN).ValueRange(-1, 1);
316+ auto targetDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
317+ int64_t reduction = 1;
318+ 
319+ auto outDesc = TensorDesc({2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
320+ 
321+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
322+ 
323+ // SAMPLE: only test GetWorkspaceSize
324+ uint64_t workspace_size = 0;
325+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
326+ EXPECT_EQ(aclRet, ACL_SUCCESS);
327+ 
328+ // SAMPLE: precision simulate
329+ ut.TestPrecision();
330+}
331+ 
332+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_20_aclnnMseLossOut_output_error_shape_none)
333+{
334+ auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
335+ auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
336+ int64_t reduction = 0;
337+ 
338+ auto outDesc = TensorDesc({123, 8, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
339+ 
340+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
341+ // SAMPLE: only test GetWorkspaceSize
342+ uint64_t workspace_size = 0;
343+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
344+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
345+}
346+ 
347+TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_21_aclnnMseLossOut_output_error_shape_mean)
348+{
349+ auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
350+ auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
351+ int64_t reduction = 1;
352+ 
353+ auto outDesc = TensorDesc({}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
354+ 
355+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
356+ // SAMPLE: only test GetWorkspaceSize
357+ uint64_t workspace_size = 0;
358+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
359+ EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
360+}
361+ 
362+TEST_F(l2_mse_loss_out_test, ascend910B2_aclnnMseLossOut_bf16_float_nhwc_mean)
363+{
364+ auto selfDesc = TensorDesc({2, 2}, ACL_BF16, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
365+ auto targetDesc = TensorDesc({2, 2}, ACL_BF16, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
366+ int64_t reduction = 1;
367+ 
368+ auto outDesc = TensorDesc({2}, ACL_BF16, ACL_FORMAT_NDHWC).Precision(0.0001, 0.0001);
369+ 
370+ auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
371+ 
372+ // SAMPLE: only test GetWorkspaceSize
373+ uint64_t workspaceSize = 0;
374+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
375+ EXPECT_EQ(aclRet, ACLNN_SUCCESS);
376+}
Mloss/mse_loss/tests/ut/op_host/CMakeLists.txt+5-7
@@ -1,17 +1,15 @@
1+# ----------------------------------------------------------------------------
1# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
2-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
3# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
4# 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.
5-# 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,
6# 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.
7# 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.
8-# ============================================================================9+# ----------------------------------------------------------------------------
9 10 
10file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
11if(UT_TEST_ALL OR OP_HOST_UT)12if(UT_TEST_ALL OR OP_HOST_UT)
12 add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})13 add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13 add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})14 add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14-endif()15+endif()
15-if(UT_TEST_ALL OR OP_API_UT)
16- add_modules_ut_sources(HOSTNAME ${OP_API_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
17-endif()
Aloss/mse_loss/tests/ut/op_host/test_MseLoss_infershape.cpp+65-0
@@ -0,0 +1,65 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file test_MseLoss_infershape.cpp
13+ * \brief
14+ */
15+ 
16+#include <iostream>
17+#include "exe_graph/runtime/storage_format.h"
18+#include "exe_graph/runtime/storage_shape.h"
19+#include <gtest/gtest.h>
20+#include "kernel_run_context_facker.h"
21+#include "infershape_test_util.h"
22+#include "ut_op_common.h"
23+#include "register/op_impl_registry.h"
24+#include "log/log.h"
25+#include "platform/platform_info.h"
26+#include "../../../op_graph/mse_loss_proto.h"
27+ 
28+class mse_loss : public testing::Test {
29+protected:
30+ static void SetUpTestCase()
31+ {
32+ std::cout << "mse_loss Proto Test SetUp" << std::endl;
33+ }
34+ 
35+ static void TearDownTestCase()
36+ {
37+ std::cout << "mse_loss Proto Test TearDown" << std::endl;
38+ }
39+};
40+ 
41+TEST_F(mse_loss, mse_loss_infershape_diff_test)
42+{
43+ ge::op::MseLoss op;
44+ std::vector<std::pair<int64_t, int64_t>> shape_range = {{15, 16}, {8, 8}, {375, 375}};
45+ auto tensor_desc =
46+ create_desc_shape_range({-1, 8, 375}, ge::DT_FLOAT16, ge::FORMAT_ND, {16, 8, 375}, ge::FORMAT_ND, shape_range);
47+ op.UpdateInputDesc("predict", tensor_desc);
48+ op.UpdateInputDesc("label", tensor_desc);
49+ op.SetAttr("reduction", "mean");
50+ // auto ret = op.InferShapeAndType();
51+ // EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
52+ auto output_y1_desc = op.GetOutputDesc("y");
53+ // EXPECT_EQ(output_y1_desc.GetDataType(), ge::DT_FLOAT16);
54+ std::vector<int64_t> expected_output_shape = {};
55+ EXPECT_EQ(output_y1_desc.GetShape().GetDims(), expected_output_shape);
56+ std::vector<std::pair<int64_t, int64_t>> output_shape_range;
57+ EXPECT_EQ(output_y1_desc.GetShapeRange(output_shape_range), ge::GRAPH_SUCCESS);
58+ std::vector<std::pair<int64_t, int64_t>> expected_shape_range = {};
59+ EXPECT_EQ(output_shape_range, expected_shape_range);
60+ 
61+ Runtime2TestParam param{{"reduction"}};
62+ EXPECT_EQ(InferShapeTest(op, param), ge::GRAPH_SUCCESS);
63+ auto output0_desc = op.GetOutputDesc(0);
64+ EXPECT_EQ(output0_desc.GetShape().GetDims(), expected_output_shape);
65+}
Dloss/mse_loss/tests/ut/op_host/test_aclnn_mse_loss_out.cpp+0-354
@@ -1,354 +0,0 @@
1-/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10-#include "gtest/gtest.h"
11-#include <array>
12-#include <vector>
13- 
14-#include "../../../op_host/op_api/aclnn_mse_loss_out.h"
15- 
16-#include "op_api_ut_common/op_api_ut.h"
17-#include "op_api_ut_common/scalar_desc.h"
18-#include "op_api_ut_common/tensor_desc.h"
19-#include "opdev/platform.h"
20- 
21-#include <unistd.h>
22- 
23-using namespace op;
24-using namespace std;
25- 
26-class l2_mse_loss_out_test : public testing::Test {
27- protected:
28- static void SetUpTestCase() {
29- std::cout << "mse_loss_test SetUp" << std::endl;
30- }
31- 
32- static void TearDownTestCase() {
33- std::cout << "mse_loss_test TearDown" << std::endl;
34- }
35-};
36- 
37-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_01_float_nd_none) {
38- auto selfDesc = TensorDesc({1, 2, 3, 2}, ACL_FLOAT, ACL_FORMAT_NCHW).ValueRange(-1, 1);
39- auto targetDesc = TensorDesc({2, 2, 1, 2}, ACL_FLOAT, ACL_FORMAT_NCHW).ValueRange(-1, 1);
40- int64_t reduction = 0;
41- 
42- auto outDesc = TensorDesc({2, 2, 3, 2}, ACL_FLOAT, ACL_FORMAT_NCHW).Precision(0.0001, 0.0001);
43- 
44- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
45- 
46- // SAMPLE: only test GetWorkspaceSize
47- uint64_t workspace_size = 0;
48- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
49- EXPECT_EQ(aclRet, ACL_SUCCESS);
50- 
51- // SAMPLE: precision simulate
52- ut.TestPrecision();
53-}
54- 
55-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_02_float16_nchw_mean) {
56- auto selfDesc = TensorDesc({2, 3, 2}, ACL_FLOAT16, ACL_FORMAT_NCHW).ValueRange(-1, 1);
57- auto targetDesc = TensorDesc({2, 3, 2}, ACL_FLOAT16, ACL_FORMAT_NCHW).ValueRange(-1, 1);
58- int64_t reduction = 1;
59- 
60- auto outDesc = TensorDesc({3, 2}, ACL_FLOAT16, ACL_FORMAT_NCHW).Precision(0.001, 0.001);
61- 
62- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
63- 
64- // SAMPLE: only test GetWorkspaceSize
65- uint64_t workspace_size = 0;
66- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
67- EXPECT_EQ(aclRet, ACL_SUCCESS);
68- 
69- // SAMPLE: precision simulate
70- // ut.TestPrecision();
71-}
72- 
73-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_03_float_float16_nhwc_sum) {
74- auto selfDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_NHWC).ValueRange(-1, 1);
75- auto targetDesc = TensorDesc({2, 2}, ACL_FLOAT16, ACL_FORMAT_NHWC).ValueRange(-1, 1);
76- int64_t reduction = 2;
77- 
78- auto outDesc = TensorDesc({2}, ACL_FLOAT, ACL_FORMAT_NHWC).Precision(0.0001, 0.0001);
79- 
80- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
81- 
82- // SAMPLE: only test GetWorkspaceSize
83- uint64_t workspace_size = 0;
84- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
85- EXPECT_EQ(aclRet, ACL_SUCCESS);
86- 
87- // SAMPLE: precision simulate
88- // ut.TestPrecision();
89-}
90- 
91-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_04_float16_float_nhwc_mean) {
92- auto selfDesc = TensorDesc({2, 2}, ACL_FLOAT16, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
93- auto targetDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
94- int64_t reduction = 1;
95- 
96- auto outDesc = TensorDesc({2}, ACL_FLOAT, ACL_FORMAT_NDHWC).Precision(0.0001, 0.0001);
97- 
98- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
99- 
100- // SAMPLE: only test GetWorkspaceSize
101- uint64_t workspace_size = 0;
102- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
103- EXPECT_EQ(aclRet, ACL_SUCCESS);
104- 
105- // SAMPLE: precision simulate
106- // ut.TestPrecision();
107-}
108- 
109-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_05_float_nd_empty_tensor_none) {
110- auto selfDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
111- auto targetDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
112- int64_t reduction = 0;
113- 
114- auto outDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
115- 
116- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
117- 
118- // SAMPLE: only test GetWorkspaceSize
119- uint64_t workspace_size = 0;
120- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
121- EXPECT_EQ(aclRet, ACL_SUCCESS);
122- 
123- // SAMPLE: precision simulate
124- ut.TestPrecision();
125-}
126- 
127-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_06_float_nd_empty_tensor_mean) {
128- auto selfDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
129- auto targetDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
130- int64_t reduction = 1;
131- 
132- auto outDesc = TensorDesc({0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
133- 
134- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
135- 
136- // SAMPLE: only test GetWorkspaceSize
137- uint64_t workspace_size = 0;
138- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
139- EXPECT_EQ(aclRet, ACL_SUCCESS);
140- 
141- // SAMPLE: precision simulate
142- ut.TestPrecision();
143-}
144- 
145-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_07_float_nd_empty_tensor_sum) {
146- auto selfDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
147- auto targetDesc = TensorDesc({1, 0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND);
148- int64_t reduction = 2;
149- 
150- auto outDesc = TensorDesc({0, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
151- 
152- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
153- 
154- // SAMPLE: only test GetWorkspaceSize
155- uint64_t workspace_size = 0;
156- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
157- EXPECT_EQ(aclRet, ACL_SUCCESS);
158- 
159- // SAMPLE: precision simulate
160- ut.TestPrecision();
161-}
162- 
163-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_11_float_hwcn_input_not_contiguous) {
164- auto selfDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_HWCN, {1, 5}, 0, {4, 5}).ValueRange(-2, 2);
165- auto targetDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_HWCN, {1, 5}, 0, {4, 5}).ValueRange(-2, 2);
166- int64_t reduction = 0;
167- 
168- auto outDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_HWCN).Precision(0.0001, 0.0001);
169- 
170- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
171- 
172- // SAMPLE: only test GetWorkspaceSize
173- uint64_t workspace_size = 0;
174- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
175- EXPECT_EQ(aclRet, ACL_SUCCESS);
176- 
177- // SAMPLE: precision simulate
178- ut.TestPrecision();
179-}
180- 
181-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_12_float_ncdhw_out_not_contiguous) {
182- auto selfDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_NCDHW).ValueRange(-2, 2);
183- auto targetDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_NCDHW, {1, 5}, 0, {4, 5}).ValueRange(-2, 2);
184- int64_t reduction = 0;
185- 
186- auto outDesc = TensorDesc({5, 4}, ACL_FLOAT, ACL_FORMAT_NCDHW, {1, 5}, 0, {4, 5}).Precision(0.0001, 0.0001);
187- 
188- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
189- 
190- // SAMPLE: only test GetWorkspaceSize
191- uint64_t workspace_size = 0;
192- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
193- EXPECT_EQ(aclRet, ACL_SUCCESS);
194- 
195- // SAMPLE: precision simulate
196- ut.TestPrecision();
197-}
198- 
199-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_13_reduction_error) {
200- auto selfDesc = TensorDesc({8}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
201- auto targetDesc = TensorDesc({8}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
202- int64_t reduction = 3;
203- 
204- auto outDesc = TensorDesc({8}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
205- 
206- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
207- 
208- // SAMPLE: only test GetWorkspaceSize
209- uint64_t workspace_size = 0;
210- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
211- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
212-}
213- 
214-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_14_input_out_nullptr) {
215- auto tensor_desc = TensorDesc({10, 5}, ACL_FLOAT, ACL_FORMAT_ND);
216- int64_t reduction = 0;
217- 
218- auto ut_l = OP_API_UT(aclnnMseLossOut, INPUT(nullptr, tensor_desc, reduction), OUTPUT(tensor_desc));
219- // SAMPLE: only test GetWorkspaceSize
220- uint64_t workspace_size = 0;
221- aclnnStatus aclRet = ut_l.TestGetWorkspaceSize(&workspace_size);
222- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
223- 
224- auto ut_r = OP_API_UT(aclnnMseLossOut, INPUT(tensor_desc, nullptr, reduction), OUTPUT(tensor_desc));
225- // SAMPLE: only test GetWorkspaceSize
226- aclRet = ut_r.TestGetWorkspaceSize(&workspace_size);
227- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
228- 
229- auto ut_o = OP_API_UT(aclnnMseLossOut, INPUT(tensor_desc, tensor_desc, reduction), OUTPUT(nullptr));
230- // SAMPLE: only test GetWorkspaceSize
231- aclRet = ut_o.TestGetWorkspaceSize(&workspace_size);
232- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
233-}
234- 
235-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_15_aclnnMseLossOut_input_error_shape) {
236- auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
237- auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
238- int64_t reduction = 0;
239- 
240- auto outDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
241- 
242- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
243- // SAMPLE: only test GetWorkspaceSize
244- uint64_t workspace_size = 0;
245- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
246- // EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
247-}
248- 
249-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_16_aclnnMseLossOut_error_output_dtype) {
250- auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
251- auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
252- int64_t reduction = 0;
253- 
254- auto outDesc = TensorDesc({8, 2}, ACL_INT32, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
255- 
256- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
257- // SAMPLE: only test GetWorkspaceSize
258- uint64_t workspace_size = 0;
259- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
260- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
261-}
262- 
263-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_17_aclnnMseLossOut_error_input_dtype) {
264- auto selfDesc = TensorDesc({6, 2, 1, 2}, ACL_INT32, ACL_FORMAT_ND).ValueRange(-2, 2);
265- auto targetDesc = TensorDesc({6, 2, 1, 2}, ACL_INT16, ACL_FORMAT_ND).ValueRange(-2, 2);
266- int64_t reduction = 0;
267- 
268- auto outDesc = TensorDesc({6, 2, 1, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
269- 
270- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
271- // SAMPLE: only test GetWorkspaceSize
272- uint64_t workspace_size = 0;
273- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
274- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
275-}
276- 
277-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_18_aclnnMseLossOut_input_error_shape_len) {
278- auto tensorDesc9 = TensorDesc({2, 3, 4, 5, 6, 7, 8, 9, 10}, ACL_INT32, ACL_FORMAT_ND).ValueRange(-2, 2);
279- auto tensorDesc1 = TensorDesc({7, 8, 9, 10}, ACL_INT16, ACL_FORMAT_ND).ValueRange(-2, 2);
280- auto tensorDesc2 = TensorDesc({7, 8, 9, 10}, ACL_INT16, ACL_FORMAT_ND).ValueRange(-2, 2);
281- int64_t reduction = 0;
282- 
283- auto ut_self = OP_API_UT(aclnnMseLossOut, INPUT(tensorDesc9, tensorDesc1, reduction), OUTPUT(tensorDesc2));
284- // SAMPLE: only test GetWorkspaceSize
285- uint64_t workspace_size = 0;
286- aclnnStatus aclRet = ut_self.TestGetWorkspaceSize(&workspace_size);
287- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
288- 
289- auto ut_tar = OP_API_UT(aclnnMseLossOut, INPUT(tensorDesc1, tensorDesc9, reduction), OUTPUT(tensorDesc2));
290- // SAMPLE: only test GetWorkspaceSize
291- aclRet = ut_tar.TestGetWorkspaceSize(&workspace_size);
292- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
293-}
294- 
295-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_19_float_hwcn_ndhwc_mean) {
296- auto selfDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_HWCN).ValueRange(-1, 1);
297- auto targetDesc = TensorDesc({2, 2}, ACL_FLOAT, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
298- int64_t reduction = 1;
299- 
300- auto outDesc = TensorDesc({2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
301- 
302- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
303- 
304- // SAMPLE: only test GetWorkspaceSize
305- uint64_t workspace_size = 0;
306- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
307- EXPECT_EQ(aclRet, ACL_SUCCESS);
308- 
309- // SAMPLE: precision simulate
310- ut.TestPrecision();
311-}
312- 
313-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_20_aclnnMseLossOut_output_error_shape_none) {
314- auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
315- auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
316- int64_t reduction = 0;
317- 
318- auto outDesc = TensorDesc({123, 8, 2}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
319- 
320- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
321- // SAMPLE: only test GetWorkspaceSize
322- uint64_t workspace_size = 0;
323- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
324- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
325-}
326- 
327-TEST_F(l2_mse_loss_out_test, aclnnMseLossOut_21_aclnnMseLossOut_output_error_shape_mean) {
328- auto selfDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
329- auto targetDesc = TensorDesc({8, 2}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
330- int64_t reduction = 1;
331- 
332- auto outDesc = TensorDesc({}, ACL_FLOAT, ACL_FORMAT_ND).Precision(0.0001, 0.0001);
333- 
334- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
335- // SAMPLE: only test GetWorkspaceSize
336- uint64_t workspace_size = 0;
337- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
338- EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
339-}
340- 
341-TEST_F(l2_mse_loss_out_test, ascend910B2_aclnnMseLossOut_bf16_float_nhwc_mean) {
342- auto selfDesc = TensorDesc({2, 2}, ACL_BF16, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
343- auto targetDesc = TensorDesc({2, 2}, ACL_BF16, ACL_FORMAT_NDHWC).ValueRange(-1, 1);
344- int64_t reduction = 1;
345- 
346- auto outDesc = TensorDesc({2}, ACL_BF16, ACL_FORMAT_NDHWC).Precision(0.0001, 0.0001);
347- 
348- auto ut = OP_API_UT(aclnnMseLossOut, INPUT(selfDesc, targetDesc, reduction), OUTPUT(outDesc));
349- 
350- // SAMPLE: only test GetWorkspaceSize
351- uint64_t workspaceSize = 0;
352- aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
353- EXPECT_EQ(aclRet, ACLNN_SUCCESS);
354-}
Aloss/mse_loss/tests/ut/op_host/test_mse_loss_tiling.cpp+304-0文件内容审核中，请稍后刷新重试
Mpooling/max_pool3d_grad_with_argmax/op_host/op_api/aclnn_max_pool2d_with_indices_backward.cpp+1-1
@@ -10,7 +10,7 @@
10 10 
11#include "aclnn_max_pool2d_with_indices_backward.h"11#include "aclnn_max_pool2d_with_indices_backward.h"
12#include "max_pool_grad_with_argmax_v1.h"12#include "max_pool_grad_with_argmax_v1.h"
13-#include "../../../max_pool_grad_with_argmax_v3/op_host/op_api/max_pool_grad_with_argmax_v3.h"13+#include "../../../max_pool_grad_with_argmax_v3/op_api/max_pool_grad_with_argmax_v3.h"
14#include "max_pool3d_grad_with_argmax.h"14#include "max_pool3d_grad_with_argmax.h"
15#include "aclnn_kernels/contiguous.h"15#include "aclnn_kernels/contiguous.h"
16#include "level0/unsqueeze.h"16#include "level0/unsqueeze.h"
Mpooling/max_pool3d_with_argmax_v2/op_host/op_api/aclnn_max_pool2d_with_indices.cpp+1-1文件内容审核中，请稍后刷新重试
Mpooling/max_pool_grad_with_argmax_v3/CMakeLists.txt+6-9
@@ -8,12 +8,9 @@
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-file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11+# 设置算子定义时支持的芯片类型
12-if(NOT ENABLE_TEST AND NOT BENCHMARK)12+set(SUPPORT_COMPUTE_UNIT "ascend910_95")
13- list(REMOVE_ITEM CURRENT_DIRS tests)13+# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14-endif()14+set(SUPPORT_TILING_DIR "arch35")
15-foreach(SUB_DIR ${CURRENT_DIRS})15+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE max_pool_grad_with_argmax_v3
16- if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")16+ ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
17- add_subdirectory(${SUB_DIR})
18- endif()
19-endforeach()
Apooling/max_pool_grad_with_argmax_v3/examples/arch35/test_aclnn_max_pool_grad_with_argmax_v3.cpp+172-0
@@ -0,0 +1,172 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include <vector>
13+#include "acl/acl.h"
14+#include "aclnnop/aclnn_max_pool2d_with_indices_backward.h"
15+ 
16+#define CHECK_RET(cond, return_expr) \
17+ do { \
18+ if (!(cond)) { \
19+ return_expr; \
20+ } \
21+ } while (0)
22+ 
23+#define LOG_PRINT(message, ...) \
24+ do { \
25+ printf(message, ##__VA_ARGS__); \
26+ } while (0)
27+ 
28+int64_t GetShapeSize(const std::vector<int64_t>& shape) {
29+ int64_t shapeSize = 1;
30+ for (auto i : shape) {
31+ shapeSize *= i;
32+ }
33+ return shapeSize;
34+}
35+ 
36+int Init(int32_t deviceId, aclrtStream* stream) {
37+ // 固定写法，资源初始化
38+ auto ret = aclInit(nullptr);
39+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
40+ ret = aclrtSetDevice(deviceId);
41+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
42+ ret = aclrtCreateStream(stream);
43+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
44+ return 0;
45+}
46+ 
47+template <typename T>
48+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
49+ aclDataType dataType, aclTensor** tensor) {
50+ auto size = GetShapeSize(shape) * sizeof(T);
51+ // 调用aclrtMalloc申请Device侧内存
52+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
53+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
54+ 
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, aclFormat::ACL_FORMAT_NCHW,
67+ shape.data(), shape.size(), *deviceAddr);
68+ return 0;
69+}
70+ 
71+int main() {
72+ // 1. （固定写法）device/stream初始化，参考acl API手册
73+ // 根据自己的实际device填写deviceId
74+ int32_t deviceId = 0;
75+ aclrtStream stream;
76+ auto ret = Init(deviceId, &stream);
77+ // check根据自己的需要处理
78+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
79+ 
80+ // 2. 构造输入与输出，需要根据API的接口自定义构造
81+ std::vector<int64_t> gradOutShape = {1, 1, 2, 1};
82+ std::vector<int64_t> selfShape = {1, 1, 4, 3};
83+ std::vector<int64_t> indicesShape = {1, 1, 2, 1};
84+ std::vector<int64_t> gradInShape = {1, 1, 4, 3};
85+ std::vector<int64_t> kernelSizeData = {2, 2};
86+ std::vector<int64_t> strideData = {2, 2};
87+ std::vector<int64_t> paddingData = {0, 0};
88+ std::vector<int64_t> dilationData = {1, 1};
89+ void* gradOutDeviceAddr = nullptr;
90+ void* selfDeviceAddr = nullptr;
91+ void* indicesDeviceAddr = nullptr;
92+ void* gradInDeviceAddr = nullptr;
93+ aclTensor* gradOut = nullptr;
94+ aclTensor* self = nullptr;
95+ aclTensor* indices = nullptr;
96+ aclTensor* gradIn = nullptr;
97+ std::vector<float> gradOutHostData = {0.4757, 0.1726};
98+ std::vector<float> selfHostData = {0.0850, -0.5147, -0.0212, -0.5654, -0.3222, 0.5847, 1.7510, 0.9954, 0.1842, 0.8392, 0.4835, 0.9213};
99+ std::vector<int32_t> indicesHostData = {0, 6};
100+ std::vector<float> gradInHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
101+ 
102+ // 创建gradOut aclTensor
103+ ret = CreateAclTensor(gradOutHostData, gradOutShape, &gradOutDeviceAddr, aclDataType::ACL_FLOAT, &gradOut);
104+ CHECK_RET(ret == ACL_SUCCESS, return ret);
105+ // 创建self aclTensor
106+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
107+ CHECK_RET(ret == ACL_SUCCESS, return ret);
108+ // 创建indices aclTensor
109+ ret = CreateAclTensor(indicesHostData, indicesShape, &indicesDeviceAddr, aclDataType::ACL_INT32, &indices);
110+ CHECK_RET(ret == ACL_SUCCESS, return ret);
111+ // 创建gradIn aclTensor
112+ ret = CreateAclTensor(gradInHostData, gradInShape, &gradInDeviceAddr, aclDataType::ACL_FLOAT, &gradIn);
113+ CHECK_RET(ret == ACL_SUCCESS, return ret);
114+ 
115+ // 创建输入数组
116+ aclIntArray* kernelSize = aclCreateIntArray(kernelSizeData.data(), 2);
117+ aclIntArray* stride = aclCreateIntArray(strideData.data(), 2);
118+ aclIntArray* padding = aclCreateIntArray(paddingData.data(), 2);
119+ aclIntArray* dilation = aclCreateIntArray(dilationData.data(), 2);
120+ const bool ceilMode = false;
121+ 
122+ uint64_t workspaceSize = 0;
123+ aclOpExecutor* executor;
124+ 
125+ // aclnnMaxPool2dWithIndicesBackward接口调用示例
126+ // 3. 调用CANN算子库API，需要修改为具体的API名称
127+ // 调用aclnnMaxPool2dWithIndicesBackward第一段接口
128+ ret = aclnnMaxPool2dWithIndicesBackwardGetWorkspaceSize(gradOut, self, indices, kernelSize, stride, padding, dilation, ceilMode, gradIn, &workspaceSize, &executor);
129+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaxPool2dWithIndicesBackwardGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
130+ // 根据第一段接口计算出的workspaceSize申请device内存
131+ void* workspaceAddr = nullptr;
132+ if (workspaceSize > 0) {
133+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
134+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
135+ }
136+ // 调用aclnnMaxPool2dWithIndicesBackward第二段接口
137+ ret = aclnnMaxPool2dWithIndicesBackward(workspaceAddr, workspaceSize, executor, stream);
138+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaxPool2dWithIndicesBackward failed. ERROR: %d\n", ret); return ret);
139+ 
140+ // 4. （固定写法）同步等待任务执行结束
141+ ret = aclrtSynchronizeStream(stream);
142+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
143+ 
144+ // 5. 获取输出的值，将Device侧内存上的结果拷贝至Host侧，需要根据具体API的接口定义修改
145+ auto size = GetShapeSize(gradInShape);
146+ std::vector<float> resultData(size, 0);
147+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), gradInDeviceAddr,
148+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
149+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy gradIn result from device to host failed. ERROR: %d\n", ret); return ret);
150+ for (int64_t i = 0; i < size; i++) {
151+ LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
152+ }
153+ 
154+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
155+ aclDestroyTensor(gradOut);
156+ aclDestroyTensor(self);
157+ aclDestroyTensor(indices);
158+ aclDestroyTensor(gradIn);
159+ 
160+ // 7. 释放device资源，需要根据具体API的接口定义修改
161+ aclrtFree(gradOutDeviceAddr);
162+ aclrtFree(selfDeviceAddr);
163+ aclrtFree(indicesDeviceAddr);
164+ aclrtFree(gradInDeviceAddr);
165+ if (workspaceSize > 0) {
166+ aclrtFree(workspaceAddr);
167+ }
168+ aclrtDestroyStream(stream);
169+ aclrtResetDevice(deviceId);
170+ aclFinalize();
171+ return 0;
172+}
Rpooling/max_pool_grad_with_argmax_v3/op_host/op_api/max_pool_grad_with_argmax_v3.cpp→pooling/max_pool_grad_with_argmax_v3/op_api/max_pool_grad_with_argmax_v3.cpp+3-3
@@ -1,12 +1,12 @@
1/**1/**
2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions 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 #include "max_pool_grad_with_argmax_v3.h"11 #include "max_pool_grad_with_argmax_v3.h"
12 #include "opdev/data_type_utils.h"12 #include "opdev/data_type_utils.h"
Rpooling/max_pool_grad_with_argmax_v3/op_host/op_api/max_pool_grad_with_argmax_v3.h→pooling/max_pool_grad_with_argmax_v3/op_api/max_pool_grad_with_argmax_v3.h+3-3
@@ -1,12 +1,12 @@
1/**1/**
2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions 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#ifndef OP_API_INC_LEVEL0_MAX_POOL_GRAD_WITH_ARGMAX_V3_H_10#ifndef OP_API_INC_LEVEL0_MAX_POOL_GRAD_WITH_ARGMAX_V3_H_
11#define OP_API_INC_LEVEL0_MAX_POOL_GRAD_WITH_ARGMAX_V3_H_11#define OP_API_INC_LEVEL0_MAX_POOL_GRAD_WITH_ARGMAX_V3_H_
12 12 
Apooling/max_pool_grad_with_argmax_v3/op_graph/max_pool_grad_with_argmax_v3_proto.h+76-0
@@ -0,0 +1,76 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_proto.h
13+ * \brief
14+ */
15+ 
16+#ifndef OPS_BUILT_IN_OP_PROTO_INC_MAX_POOLING_GRAD_WITH_ARGMAX_V3_PROTO_H_
17+#define OPS_BUILT_IN_OP_PROTO_INC_MAX_POOLING_GRAD_WITH_ARGMAX_V3_PROTO_H_
18+ 
19+#include "graph/operator_reg.h"
20+#include "graph/operator.h"
21+ 
22+namespace ge {
23+ 
24+/**
25+* @brief Performs the backpropagation of MaxPoolGradWithArgmaxV3.
26+ 
27+* @par Inputs:
28+* Three inputs, including:
29+* @li x: A tensor of dtype bfloat16, float16, float32, the shape is `[batch, channels, height_in, width_in]` or
30+ `[batch, height_in, width_in, channels]` , the format is `NCHW` or `NHWC`.
31+* @li grad: A tensor has the same dtype and format as input "x", the shape is `[batch, channels, height_out, width_out]`
32+or
33+ `[batch, height_out, width_out, channels]`.
34+* @li argmax: A tensor has the same shape and format as input "grad", the dtype is int32 or int64.
35+ 
36+* @par Attributes:
37+* @li ksize: A required list of int64 values,
38+* specifying the size of the window for each dimension of the input tensor. No default value.
39+* @li strides: A required list of int64 values,
40+* specifying the stride of the sliding window for each dimension of the input tensor. No default value.
41+* @li pads: A required list of int64 values,
42+* specifying the pad of the input feature map. No default value.
43+ 
44+* @par Outputs:
45+* y: A Tensor. Has the same dtype , shape and format as input "x".
46+ 
47+* @attention Constraints:
48+* @li The MaxPoolGradWithArgmaxV3 operator has the same function, and it is recommended to use the V3 operator.
49+* @li ksize: a list that has length 2:
50+* @li strides: a list that has length 2:
51+* @li pads: a list that has length 2:
52+* 1 <= pads[0] <= (ksize[0]//2), 1 <= pads[1] <= (ksize[1]//2).
53+* @li dilation: a list that has length 2. default value is {1,1}.
54+* @li dtype: A optional int. default value is 3.
55+* @li ceil_mode: defaults to False.
56+* @li data_format: defaults to "NCHW".
57+ 
58+* @par Third-party framework compatibility
59+* Compatible with the Pytorch backward operator of max_pool2d_with_indices.
60+*/
61+REG_OP(MaxPoolGradWithArgmaxV3)
62+ .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
63+ .INPUT(grad, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
64+ .INPUT(argmax, TensorType({DT_INT32, DT_INT64}))
65+ .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
66+ .REQUIRED_ATTR(ksize, ListInt)
67+ .REQUIRED_ATTR(strides, ListInt)
68+ .REQUIRED_ATTR(pads, ListInt)
69+ .ATTR(dtype, Int, 3)
70+ .ATTR(dilation, ListInt, {1, 1})
71+ .ATTR(ceil_mode, Bool, false)
72+ .ATTR(data_format, String, "NCHW")
73+ .OP_END_FACTORY_REG(MaxPoolGradWithArgmaxV3)
74+ 
75+} // namespace ge
76+#endif // OPS_BUILT_IN_OP_PROTO_INC_NN_POOLING_OPS_H
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nchw_tiling.cpp+430-0
@@ -0,0 +1,430 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nchw_tiling.cpp
13+ * \brief
14+ */
15+#include "platform/platform_info.h"
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_grad_with_argmax_v3_nchw_tiling.h"
18+ 
19+namespace 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;
25+static constexpr int64_t NO_CHECK_RANGE_TILING_KEY_NCHW = 100;
26+static constexpr int64_t CHECK_RANGE_TILING_KEY_NCHW = 101;
27+static 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+ 
69+bool MaxPoolGradWithArgmaxV3NCHWTiling::IsCapable()
70+{
71+ if (inputData.inputFormat != ge::Format::FORMAT_NCHW) {
72+ return false;
73+ }
74+ 
75+ InitializationVars();
76+ // all the h and w is overlapped.
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+}
87+ 
88+uint64_t MaxPoolGradWithArgmaxV3NCHWTiling::GetTilingKey() const
89+{
90+ uint64_t tilingKey = NO_CHECK_RANGE_TILING_KEY_NCHW;
91+ if (splitData.isCheckRange == 1) {
92+ tilingKey = CHECK_RANGE_TILING_KEY_NCHW;
93+ }
94+ if (inputData.isInt32Meet == 0) {
95+ tilingKey += T3_INT64;
96+ }
97+ 
98+ return tilingKey;
99+}
100+ 
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+ tilingData.set_hArgmax(inputData.hGrad);
375+ tilingData.set_wArgmax(inputData.wGrad);
376+ tilingData.set_hOutput(inputData.hX);
377+ tilingData.set_wOutput(inputData.wX);
378+ tilingData.set_hKernel(inputData.hKernel);
379+ tilingData.set_wKernel(inputData.wKernel);
380+ tilingData.set_hStride(inputData.hStride);
381+ tilingData.set_wStride(inputData.wStride);
382+ tilingData.set_padH(inputData.hPad);
383+ tilingData.set_padW(inputData.wPad);
384+ tilingData.set_dilationH(inputData.hDilation);
385+ tilingData.set_dilationW(inputData.wDilation);
386+ tilingData.set_highAxisInner(splitData.highAxisInner);
387+ tilingData.set_highAxisTail(splitData.highAxisTail);
388+ tilingData.set_highAxisOuter(splitData.highAxisOuter);
389+ tilingData.set_hOutputInner(splitData.hOutputInner);
390+ tilingData.set_hOutputTail(splitData.hOutputTail);
391+ tilingData.set_hOutputOuter(splitData.hOutputOuter);
392+ tilingData.set_wOutputInner(splitData.wOutputInner);
393+ tilingData.set_wOutputTail(splitData.wOutputTail);
394+ tilingData.set_wOutputOuter(splitData.wOutputOuter);
395+ tilingData.set_normalCoreProcessNum(splitData.normalCoreProcessNum);
396+ tilingData.set_tailCoreProcessNum(splitData.tailCoreProcessNum);
397+ tilingData.set_usedCoreNum(splitData.usedCoreNum);
398+ tilingData.set_outputBufferSize(splitData.outputBufferSize);
399+ tilingData.set_gradBufferSize(splitData.gradBufferSize);
400+ tilingData.set_argmaxBufferSize(splitData.argmaxBufferSize);
401+ tilingData.set_hProBatchSize(baseData.hProBatchSize);
402+ tilingData.set_wProBatchSize(baseData.wProBatchSize);
403+ tilingData.set_tilingKey(GetTilingKey());
404+}
405+ 
406+ge::graphStatus MaxPoolGradWithArgmaxV3NCHWTiling::DoOpTiling()
407+{
408+ DoUBTiling();
409+ DoBlockTiling();
410+ SetTilingData();
411+ PrintBaseData();
412+ PrintSplitData();
413+ return ge::GRAPH_SUCCESS;
414+}
415+ 
416+ge::graphStatus MaxPoolGradWithArgmaxV3NCHWTiling::PostTiling()
417+{
418+ context_->SetBlockDim(tilingData.get_usedCoreNum());
419+ if (tilingData.GetDataSize() > context_->GetRawTilingData()->GetCapacity()) {
420+ return ge::GRAPH_FAILED;
421+ }
422+ 
423+ tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
424+ context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
425+ return ge::GRAPH_SUCCESS;
426+}
427+ 
428+REGISTER_OPS_TILING_TEMPLATE(MaxPoolGradWithArgmaxV3, MaxPoolGradWithArgmaxV3NCHWTiling, 0);
429+ 
430+} // namespace optiling
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nchw_tiling.h+146-0
@@ -0,0 +1,146 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nchw_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_H_
17+#define MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_H_
18+ 
19+#include "max_pool_grad_with_argmax_v3_tiling_base.h"
20+ 
21+namespace optiling {
22+ 
23+BEGIN_TILING_DATA_DEF(MaxPoolGradWithArgmaxV3NCHWTilingData)
24+TILING_DATA_FIELD_DEF(int64_t, hArgmax);
25+TILING_DATA_FIELD_DEF(int64_t, wArgmax);
26+TILING_DATA_FIELD_DEF(int64_t, hOutput);
27+TILING_DATA_FIELD_DEF(int64_t, wOutput);
28+TILING_DATA_FIELD_DEF(int64_t, hKernel);
29+TILING_DATA_FIELD_DEF(int64_t, wKernel);
30+TILING_DATA_FIELD_DEF(int64_t, hStride);
31+TILING_DATA_FIELD_DEF(int64_t, wStride);
32+TILING_DATA_FIELD_DEF(int64_t, padH);
33+TILING_DATA_FIELD_DEF(int64_t, padW);
34+TILING_DATA_FIELD_DEF(int64_t, dilationH);
35+TILING_DATA_FIELD_DEF(int64_t, dilationW);
36+TILING_DATA_FIELD_DEF(int64_t, highAxisInner);
37+TILING_DATA_FIELD_DEF(int64_t, highAxisTail);
38+TILING_DATA_FIELD_DEF(int64_t, highAxisOuter);
39+TILING_DATA_FIELD_DEF(int64_t, hOutputInner);
40+TILING_DATA_FIELD_DEF(int64_t, hOutputTail);
41+TILING_DATA_FIELD_DEF(int64_t, hOutputOuter);
42+TILING_DATA_FIELD_DEF(int64_t, wOutputInner);
43+TILING_DATA_FIELD_DEF(int64_t, wOutputTail);
44+TILING_DATA_FIELD_DEF(int64_t, wOutputOuter);
45+TILING_DATA_FIELD_DEF(int64_t, normalCoreProcessNum);
46+TILING_DATA_FIELD_DEF(int64_t, tailCoreProcessNum);
47+TILING_DATA_FIELD_DEF(int64_t, usedCoreNum);
48+TILING_DATA_FIELD_DEF(int64_t, outputBufferSize);
49+TILING_DATA_FIELD_DEF(int64_t, gradBufferSize);
50+TILING_DATA_FIELD_DEF(int64_t, argmaxBufferSize);
51+TILING_DATA_FIELD_DEF(int64_t, hProBatchSize);
52+TILING_DATA_FIELD_DEF(int64_t, wProBatchSize);
53+TILING_DATA_FIELD_DEF(int64_t, tilingKey);
54+END_TILING_DATA_DEF;
55+ 
56+// no check range 100 check range 101
57+// no check range int64 110 check range int64 111
58+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_100, MaxPoolGradWithArgmaxV3NCHWTilingData);
59+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_101, MaxPoolGradWithArgmaxV3NCHWTilingData);
60+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_110, MaxPoolGradWithArgmaxV3NCHWTilingData);
61+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_111, MaxPoolGradWithArgmaxV3NCHWTilingData);
62+ 
63+struct MaxPoolGradWithArgmaxV3NCHWBaseInfo {
64+ int64_t vRegSize{0};
65+ int64_t ubBlockSize{0};
66+ int64_t inputBytes{0};
67+ int64_t indexBytes{0};
68+ int64_t availableUb{0};
69+ int64_t totalCoreNum{0};
70+ int64_t coreUsedForBestPerformance{0};
71+ int64_t hProBatchSize{0};
72+ int64_t wProBatchSize{0};
73+ int64_t inputNCSize{0};
74+ int64_t maxDataNumInOneBlock{0};
75+ int64_t proDataNumInOneBeatT2{0};
76+ int64_t isPad{0};
77+ int64_t isOverlap{0};
78+};
79+ 
80+struct MaxPoolGradWithArgmaxV3NCHWSplitInfo {
81+ // DoUBTiling
82+ int64_t isCheckRange{0};
83+ 
84+ int64_t highAxisInner{0};
85+ int64_t highAxisTail{0};
86+ int64_t highAxisOuter{0};
87+ 
88+ int64_t hOutputInner{0};
89+ int64_t hOutputTail{0};
90+ int64_t hOutputOuter{0};
91+ 
92+ int64_t wOutputInner{0};
93+ int64_t wOutputTail{0};
94+ int64_t wOutputOuter{0};
95+ 
96+ // DoBlockTiling
97+ int64_t normalCoreProcessNum{0};
98+ int64_t tailCoreProcessNum{0};
99+ int64_t usedCoreNum{0};
100+ int64_t totalBaseBlockNum{0};
101+ 
102+ // DoBufferCalculate
103+ int64_t outputBufferSize{0};
104+ int64_t gradBufferSize{0};
105+ int64_t argmaxBufferSize{0};
106+ int64_t totalBufferSize{0};
107+};
108+ 
109+class MaxPoolGradWithArgmaxV3NCHWTiling : public MaxPoolGradWithArgmaxV3BaseTiling {
110+public:
111+ explicit MaxPoolGradWithArgmaxV3NCHWTiling(gert::TilingContext* context)
112+ : MaxPoolGradWithArgmaxV3BaseTiling(context)
113+ {}
114+ 
115+ ~MaxPoolGradWithArgmaxV3NCHWTiling() override
116+ {}
117+ 
118+private:
119+ void DoUBTiling();
120+ void InitializationVars();
121+ bool TrySplitNC();
122+ bool TrySplitAlignH();
123+ bool TrySplitAlignW();
124+ void SplitUnalignHW();
125+ bool IsMeetTargetCoreNum() const;
126+ bool IsMeetUBSize();
127+ void SearchBestTiling();
128+ void DynamicAdjustmentWH();
129+ void SetTilingData();
130+ uint64_t GetTilingKey() const override;
131+ void PrintBaseData() const;
132+ void PrintSplitData() const;
133+ void DoBlockTiling();
134+ void DoBufferCalculate();
135+ bool IsCapable() override;
136+ ge::graphStatus DoOpTiling() override;
137+ ge::graphStatus PostTiling() override;
138+ 
139+ MaxPoolGradWithArgmaxV3NCHWTilingData tilingData;
140+ MaxPoolGradWithArgmaxV3NCHWBaseInfo baseData;
141+ MaxPoolGradWithArgmaxV3NCHWSplitInfo splitData;
142+};
143+ 
144+} // namespace optiling
145+ 
146+#endif
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nchw_tiling_scalar.cpp+216-0
@@ -0,0 +1,216 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nchw_tiling_scalar.cpp
13+ * \brief
14+ */
15+#include "max_pool_grad_with_argmax_v3_nchw_tiling_scalar.h"
16+#include "platform/platform_info.h"
17+#include "tiling_base/tiling_templates_registry.h"
18+namespace optiling {
19+static constexpr int64_t HALF = 2;
20+static constexpr int64_t DOUBLE_BUFFER = 2;
21+static constexpr int64_t CHECK_RANGE_TILING_KEY_NCHW_SCALAR = 301;
22+static constexpr int64_t FLOAT32_SIZE = 4;
23+bool MaxPoolGradWithArgmaxV3NCHWScalarTiling::IsCapable()
24+{
25+ if (inputData.inputFormat != ge::Format::FORMAT_NCHW) {
26+ return false;
27+ }
28+ return true;
29+}
30+ge::graphStatus MaxPoolGradWithArgmaxV3NCHWScalarTiling::DoOpTiling()
31+{
32+ CalcBase();
33+ CalcParamsEachCore();
34+ ge::graphStatus result = CalcGradArgmax();
35+ if (result != ge::GRAPH_SUCCESS) {
36+ return result;
37+ }
38+ SetTilingData();
39+ PrintData();
40+ return ge::GRAPH_SUCCESS;
41+}
42+void MaxPoolGradWithArgmaxV3NCHWScalarTiling::CalcBase()
43+{
44+ int64_t ncTotal = inputData.nX * inputData.cX;
45+ int64_t ncSizePerCore = Ops::Base::CeilDiv(ncTotal, hardwareData.coreNum);
46+ int64_t hwSize = inputData.hX * inputData.wX * FLOAT32_SIZE;
47+ int64_t inputUbSize = hardwareData.ubSize / DOUBLE_BUFFER / HALF;
48+ scalarTilingData_.outputBufferSize = inputUbSize;
49+ if (ncSizePerCore * hwSize <= inputUbSize) {
50+ scalarTilingData_.highAxisInner = ncSizePerCore;
51+ scalarTilingData_.hOutputInner = inputData.hX;
52+ scalarTilingData_.wOutputInner = inputData.wX;
53+ return;
54+ }
55+ 
56+ if (hwSize <= inputUbSize) {
57+ scalarTilingData_.highAxisInner = inputUbSize / hwSize;
58+ scalarTilingData_.hOutputInner = inputData.hX;
59+ scalarTilingData_.wOutputInner = inputData.wX;
60+ return;
61+ }
62+ 
63+ int64_t perHSize = 1 * inputData.wX * FLOAT32_SIZE;
64+ if (perHSize <= inputUbSize) {
65+ scalarTilingData_.highAxisInner = 1;
66+ scalarTilingData_.hOutputInner = inputUbSize / perHSize;
67+ scalarTilingData_.wOutputInner = inputData.wX;
68+ return;
69+ }
70+ 
71+ int64_t perWSize = 1 * FLOAT32_SIZE;
72+ if (perWSize <= inputUbSize) {
73+ scalarTilingData_.highAxisInner = 1;
74+ scalarTilingData_.hOutputInner = 1;
75+ scalarTilingData_.wOutputInner = inputUbSize / perWSize;
76+ return;
77+ }
78+}
79+void MaxPoolGradWithArgmaxV3NCHWScalarTiling::CalcParamsEachCore()
80+{
81+ int64_t ncTotal = inputData.nX * inputData.cX;
82+ scalarTilingData_.highAxisOuter = Ops::Base::CeilDiv(ncTotal, scalarTilingData_.highAxisInner);
83+ scalarTilingData_.hOutputOuter = Ops::Base::CeilDiv(inputData.hX, scalarTilingData_.hOutputInner);
84+ scalarTilingData_.wOutputOuter = Ops::Base::CeilDiv(inputData.wX, scalarTilingData_.wOutputInner);
85+ scalarTilingData_.highAxisTail = ncTotal - (scalarTilingData_.highAxisOuter - 1) * scalarTilingData_.highAxisInner;
86+ scalarTilingData_.hOutputTail =
87+ inputData.hX - (scalarTilingData_.hOutputOuter - 1) * scalarTilingData_.hOutputInner;
88+ scalarTilingData_.wOutputTail =
89+ inputData.wX - (scalarTilingData_.wOutputOuter - 1) * scalarTilingData_.wOutputInner;
90+ int64_t totalCount =
91+ scalarTilingData_.highAxisOuter * scalarTilingData_.hOutputOuter * scalarTilingData_.wOutputOuter;
92+ scalarTilingData_.normalCoreProcessNum = Ops::Base::CeilDiv(totalCount, hardwareData.coreNum);
93+ scalarTilingData_.usedCoreNum = Ops::Base::CeilDiv(totalCount, scalarTilingData_.normalCoreProcessNum);
94+ scalarTilingData_.tailCoreProcessNum =
95+ totalCount - (scalarTilingData_.usedCoreNum - 1) * scalarTilingData_.normalCoreProcessNum;
96+ return;
97+}
98+ge::graphStatus MaxPoolGradWithArgmaxV3NCHWScalarTiling::CalcGradArgmaxInner(int64_t argmaxCountInUB)
99+{
100+ int64_t hInputInner = Ops::Base::CeilDiv(scalarTilingData_.hOutputInner + inputData.hKernel - 1, inputData.hStride);
101+ int64_t wInputInner = Ops::Base::CeilDiv(scalarTilingData_.wOutputInner + inputData.wKernel - 1, inputData.wStride);
102+ hInputInner = std::min(hInputInner, inputData.hGrad);
103+ wInputInner = std::min(wInputInner, inputData.wGrad);
104+ if (hInputInner == 0 || wInputInner == 0) {
105+ return ge::GRAPH_FAILED;
106+ }
107+ int64_t inputPlaneSize = hInputInner * wInputInner;
108+ if (scalarTilingData_.highAxisInner * hInputInner * wInputInner <= argmaxCountInUB) {
109+ scalarTilingData_.argmaxNcInner = scalarTilingData_.highAxisInner;
110+ scalarTilingData_.argmaxHInner = hInputInner;
111+ scalarTilingData_.argmaxWInner = wInputInner;
112+ } else if (inputPlaneSize <= argmaxCountInUB) {
113+ scalarTilingData_.argmaxNcInner = argmaxCountInUB / inputPlaneSize;
114+ scalarTilingData_.argmaxHInner = hInputInner;
115+ scalarTilingData_.argmaxWInner = wInputInner;
116+ } else if (wInputInner <= argmaxCountInUB) {
117+ scalarTilingData_.argmaxNcInner = 1;
118+ scalarTilingData_.argmaxHInner = argmaxCountInUB / wInputInner;
119+ scalarTilingData_.argmaxWInner = wInputInner;
120+ } else {
121+ scalarTilingData_.argmaxNcInner = 1;
122+ scalarTilingData_.argmaxHInner = 1;
123+ scalarTilingData_.argmaxWInner = argmaxCountInUB;
124+ }
125+ 
126+ scalarTilingData_.argmaxNcOuter =
127+ Ops::Base::CeilDiv(scalarTilingData_.highAxisInner, scalarTilingData_.argmaxNcInner);
128+ scalarTilingData_.argmaxHOuter = Ops::Base::CeilDiv(hInputInner, scalarTilingData_.argmaxHInner);
129+ scalarTilingData_.argmaxWOuter = Ops::Base::CeilDiv(wInputInner, scalarTilingData_.argmaxWInner);
130+ scalarTilingData_.argmaxNcTail =
131+ scalarTilingData_.highAxisInner - (scalarTilingData_.argmaxNcOuter - 1) * scalarTilingData_.argmaxNcInner;
132+ scalarTilingData_.argmaxHTail = hInputInner - (scalarTilingData_.argmaxHOuter - 1) * scalarTilingData_.argmaxHInner;
133+ scalarTilingData_.argmaxWTail = wInputInner - (scalarTilingData_.argmaxWOuter - 1) * scalarTilingData_.argmaxWInner;
134+ scalarTilingData_.argmaxInnerLoop =
135+ scalarTilingData_.argmaxNcOuter * scalarTilingData_.argmaxHOuter * scalarTilingData_.argmaxWOuter;
136+ return ge::GRAPH_SUCCESS;
137+}
138+ge::graphStatus MaxPoolGradWithArgmaxV3NCHWScalarTiling::CalcGradArgmax()
139+{
140+ int64_t totalGradArgmaxUBSize = hardwareData.ubSize / DOUBLE_BUFFER / HALF;
141+ int64_t inputDtypeSize = ge::GetSizeByDataType(inputData.inputDtype);
142+ int64_t indexDtypeSize = ge::GetSizeByDataType(inputData.indexDtype);
143+ int64_t alignTypeSize = (inputDtypeSize < indexDtypeSize ? inputDtypeSize : indexDtypeSize);
144+ int64_t ubBlockSize = Ops::Base::GetUbBlockSize(context_);
145+ int64_t baseAlignedCount = ubBlockSize / alignTypeSize;
146+ int64_t argmaxCountInUB = totalGradArgmaxUBSize / (ge::GetSizeByDataType(inputData.inputDtype) +
147+ ge::GetSizeByDataType(inputData.indexDtype));
148+ argmaxCountInUB = argmaxCountInUB / baseAlignedCount * baseAlignedCount;
149+ scalarTilingData_.gradBufferSize = argmaxCountInUB * ge::GetSizeByDataType(inputData.inputDtype);
150+ scalarTilingData_.argmaxBufferSize = argmaxCountInUB * ge::GetSizeByDataType(inputData.indexDtype);
151+ return CalcGradArgmaxInner(argmaxCountInUB);
152+}
153+uint64_t MaxPoolGradWithArgmaxV3NCHWScalarTiling::GetTilingKey() const
154+{
155+ return CHECK_RANGE_TILING_KEY_NCHW_SCALAR;
156+}
157+ 
158+void MaxPoolGradWithArgmaxV3NCHWScalarTiling::SetTilingData()
159+{
160+ tilingData_.set_hArgmax(inputData.hGrad);
161+ tilingData_.set_wArgmax(inputData.wGrad);
162+ tilingData_.set_hOutput(inputData.hX);
163+ tilingData_.set_wOutput(inputData.wX);
164+ tilingData_.set_hKernel(inputData.hKernel);
165+ tilingData_.set_wKernel(inputData.wKernel);
166+ tilingData_.set_hStride(inputData.hStride);
167+ tilingData_.set_wStride(inputData.wStride);
168+ tilingData_.set_padH(inputData.hPad);
169+ tilingData_.set_padW(inputData.wPad);
170+ tilingData_.set_dilationH(inputData.hDilation);
171+ tilingData_.set_dilationW(inputData.wDilation);
172+ tilingData_.set_highAxisInner(scalarTilingData_.highAxisInner);
173+ tilingData_.set_highAxisTail(scalarTilingData_.highAxisTail);
174+ tilingData_.set_highAxisOuter(scalarTilingData_.highAxisOuter);
175+ tilingData_.set_hOutputInner(scalarTilingData_.hOutputInner);
176+ tilingData_.set_hOutputTail(scalarTilingData_.hOutputTail);
177+ tilingData_.set_hOutputOuter(scalarTilingData_.hOutputOuter);
178+ tilingData_.set_wOutputInner(scalarTilingData_.wOutputInner);
179+ tilingData_.set_wOutputTail(scalarTilingData_.wOutputTail);
180+ tilingData_.set_wOutputOuter(scalarTilingData_.wOutputOuter);
181+ tilingData_.set_normalCoreProcessNum(scalarTilingData_.normalCoreProcessNum);
182+ tilingData_.set_tailCoreProcessNum(scalarTilingData_.tailCoreProcessNum);
183+ tilingData_.set_usedCoreNum(scalarTilingData_.usedCoreNum);
184+ tilingData_.set_outputBufferSize(scalarTilingData_.outputBufferSize);
185+ tilingData_.set_gradBufferSize(scalarTilingData_.gradBufferSize);
186+ tilingData_.set_argmaxBufferSize(scalarTilingData_.argmaxBufferSize);
187+ tilingData_.set_argmaxNcInner(scalarTilingData_.argmaxNcInner);
188+ tilingData_.set_argmaxNcOuter(scalarTilingData_.argmaxNcOuter);
189+ tilingData_.set_argmaxNcTail(scalarTilingData_.argmaxNcTail);
190+ tilingData_.set_argmaxHInner(scalarTilingData_.argmaxHInner);
191+ tilingData_.set_argmaxHOuter(scalarTilingData_.argmaxHOuter);
192+ tilingData_.set_argmaxHTail(scalarTilingData_.argmaxHTail);
193+ tilingData_.set_argmaxWInner(scalarTilingData_.argmaxWInner);
194+ tilingData_.set_argmaxWOuter(scalarTilingData_.argmaxWOuter);
195+ tilingData_.set_argmaxWTail(scalarTilingData_.argmaxWTail);
196+ tilingData_.set_argmaxInnerLoop(scalarTilingData_.argmaxInnerLoop);
197+ return;
198+}
199+void MaxPoolGradWithArgmaxV3NCHWScalarTiling::PrintData() const
200+{
201+ OP_LOGI("PrintData", "%s", scalarTilingData_.ToString().c_str());
202+ return;
203+}
204+ge::graphStatus MaxPoolGradWithArgmaxV3NCHWScalarTiling::PostTiling()
205+{
206+ context_->SetBlockDim(tilingData_.get_usedCoreNum());
207+ if (tilingData_.GetDataSize() > context_->GetRawTilingData()->GetCapacity()) {
208+ return ge::GRAPH_FAILED;
209+ }
210+ 
211+ tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
212+ context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
213+ return ge::GRAPH_SUCCESS;
214+}
215+REGISTER_OPS_TILING_TEMPLATE(MaxPoolGradWithArgmaxV3, MaxPoolGradWithArgmaxV3NCHWScalarTiling, 10);
216+} // namespace optiling
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nchw_tiling_scalar.h+139-0
@@ -0,0 +1,139 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nchw_tiling_scalar.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_SCALAR_H_
17+#define MAX_POOL_GRAD_WITH_AGRMAX_V3_NCHW_TILING_SCALAR_H_
18+ 
19+#include "max_pool_grad_with_argmax_v3_tiling_base.h"
20+ 
21+namespace optiling {
22+BEGIN_TILING_DATA_DEF(MaxPoolGradWithArgmaxV3NCHWScalarTilingData)
23+TILING_DATA_FIELD_DEF(int64_t, hArgmax);
24+TILING_DATA_FIELD_DEF(int64_t, wArgmax);
25+TILING_DATA_FIELD_DEF(int64_t, hOutput);
26+TILING_DATA_FIELD_DEF(int64_t, wOutput);
27+TILING_DATA_FIELD_DEF(int64_t, hKernel);
28+TILING_DATA_FIELD_DEF(int64_t, wKernel);
29+TILING_DATA_FIELD_DEF(int64_t, hStride);
30+TILING_DATA_FIELD_DEF(int64_t, wStride);
31+TILING_DATA_FIELD_DEF(int64_t, padH);
32+TILING_DATA_FIELD_DEF(int64_t, padW);
33+TILING_DATA_FIELD_DEF(int64_t, dilationH);
34+TILING_DATA_FIELD_DEF(int64_t, dilationW);
35+TILING_DATA_FIELD_DEF(int64_t, highAxisInner);
36+TILING_DATA_FIELD_DEF(int64_t, highAxisTail);
37+TILING_DATA_FIELD_DEF(int64_t, highAxisOuter);
38+TILING_DATA_FIELD_DEF(int64_t, hOutputInner);
39+TILING_DATA_FIELD_DEF(int64_t, hOutputTail);
40+TILING_DATA_FIELD_DEF(int64_t, hOutputOuter);
41+TILING_DATA_FIELD_DEF(int64_t, wOutputInner);
42+TILING_DATA_FIELD_DEF(int64_t, wOutputTail);
43+TILING_DATA_FIELD_DEF(int64_t, wOutputOuter);
44+TILING_DATA_FIELD_DEF(int64_t, normalCoreProcessNum);
45+TILING_DATA_FIELD_DEF(int64_t, tailCoreProcessNum);
46+TILING_DATA_FIELD_DEF(int64_t, usedCoreNum);
47+TILING_DATA_FIELD_DEF(int64_t, outputBufferSize);
48+TILING_DATA_FIELD_DEF(int64_t, gradBufferSize);
49+TILING_DATA_FIELD_DEF(int64_t, argmaxBufferSize);
50+TILING_DATA_FIELD_DEF(int64_t, argmaxNcInner);
51+TILING_DATA_FIELD_DEF(int64_t, argmaxNcOuter);
52+TILING_DATA_FIELD_DEF(int64_t, argmaxNcTail);
53+TILING_DATA_FIELD_DEF(int64_t, argmaxHInner);
54+TILING_DATA_FIELD_DEF(int64_t, argmaxHOuter);
55+TILING_DATA_FIELD_DEF(int64_t, argmaxHTail);
56+TILING_DATA_FIELD_DEF(int64_t, argmaxWInner);
57+TILING_DATA_FIELD_DEF(int64_t, argmaxWOuter);
58+TILING_DATA_FIELD_DEF(int64_t, argmaxWTail);
59+TILING_DATA_FIELD_DEF(int64_t, argmaxInnerLoop);
60+END_TILING_DATA_DEF;
61+ 
62+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_301, MaxPoolGradWithArgmaxV3NCHWScalarTilingData);
63+ 
64+class MaxPoolGradWithArgmaxV3ScalarTilingInfo {
65+public:
66+ int64_t highAxisInner = 0;
67+ int64_t highAxisTail = 0;
68+ int64_t highAxisOuter = 0;
69+ int64_t hOutputInner = 0;
70+ int64_t hOutputTail = 0;
71+ int64_t hOutputOuter = 0;
72+ int64_t wOutputInner = 0;
73+ int64_t wOutputTail = 0;
74+ int64_t wOutputOuter = 0;
75+ int64_t normalCoreProcessNum = 0;
76+ int64_t tailCoreProcessNum = 0;
77+ int64_t usedCoreNum = 0;
78+ int64_t outputBufferSize = 0;
79+ int64_t gradBufferSize = 0;
80+ int64_t argmaxBufferSize = 0;
81+ int64_t argmaxNcInner = 0;
82+ int64_t argmaxNcOuter = 0;
83+ int64_t argmaxNcTail = 0;
84+ int64_t argmaxHInner = 0;
85+ int64_t argmaxHOuter = 0;
86+ int64_t argmaxHTail = 0;
87+ int64_t argmaxWInner = 0;
88+ int64_t argmaxWOuter = 0;
89+ int64_t argmaxWTail = 0;
90+ int64_t argmaxInnerLoop = 0;
91+ 
92+public:
93+ std::string ToString() const
94+ {
95+ std::stringstream info;
96+ info << "MaxPoolGradWithArgmaxV3ScalarTilingInfo {";
97+ info << "highAxisInner:" << highAxisInner << ",highAxisTail:" << highAxisTail
98+ << ",highAxisOuter:" << highAxisOuter << ",hOutputInner:" << hOutputInner
99+ << ", hOutputTail:" << hOutputTail << ", hOutputOuter:" << hOutputOuter
100+ << ", wOutputInner:" << wOutputInner << ", wOutputTail:" << wOutputTail
101+ << ", wOutputOuter:" << wOutputOuter << ", normalCoreProcessNum:" << normalCoreProcessNum
102+ << ", tailCoreProcessNum:" << tailCoreProcessNum << ", usedCoreNum:" << usedCoreNum
103+ << ", outputBufferSize:" << outputBufferSize << ", gradBufferSize:" << gradBufferSize
104+ << ", argmaxBufferSize:" << argmaxBufferSize << ", argmaxNcInner:" << argmaxNcInner
105+ << ", argmaxNcOuter:" << argmaxNcOuter << ", argmaxNcTail:" << argmaxNcTail
106+ << ", argmaxHInner:" << argmaxHInner << ", argmaxHOuter:" << argmaxHOuter
107+ << ", argmaxHTail:" << argmaxHTail << ", argmaxWInner:" << argmaxWInner
108+ << ", argmaxWOuter:" << argmaxWOuter << ", argmaxWTail:" << argmaxWTail
109+ << ", argmaxInnerLoop:" << argmaxInnerLoop;
110+ info << " }";
111+ return info.str();
112+ }
113+};
114+ 
115+class MaxPoolGradWithArgmaxV3NCHWScalarTiling : public MaxPoolGradWithArgmaxV3BaseTiling {
116+public:
117+ explicit MaxPoolGradWithArgmaxV3NCHWScalarTiling(gert::TilingContext* context)
118+ : MaxPoolGradWithArgmaxV3BaseTiling(context)
119+ {}
120+ ~MaxPoolGradWithArgmaxV3NCHWScalarTiling() override
121+ {}
122+ 
123+protected:
124+ void SetTilingData();
125+ uint64_t GetTilingKey() const override;
126+ bool IsCapable() override;
127+ ge::graphStatus DoOpTiling() override;
128+ ge::graphStatus PostTiling() override;
129+ void PrintData() const;
130+ void CalcBase();
131+ void CalcParamsEachCore();
132+ ge::graphStatus CalcGradArgmaxInner(int64_t argmaxCountInUB);
133+ ge::graphStatus CalcGradArgmax();
134+ MaxPoolGradWithArgmaxV3NCHWScalarTilingData tilingData_;
135+ MaxPoolGradWithArgmaxV3ScalarTilingInfo scalarTilingData_;
136+};
137+ 
138+} // namespace optiling
139+#endif
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nhwc_tiling.cpp+509-0
@@ -0,0 +1,509 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nhwc_tiling.cpp
13+ * \brief
14+ */
15+#include "platform/platform_info.h"
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_grad_with_argmax_v3_nhwc_tiling.h"
18+ 
19+namespace 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;
CANN-robot
CANN-robotCANN-robot2025年12月29日

代码结构与可维护性: 第24行常量UB_RESVERVED_SIZE拼写错误，应该是UB_RESERVED_SIZE。虽然这不影响功能，但影响代码可读性和维护性。

问题类型: 代码结构与可维护性 文件路径: pooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nhwc_tiling.cpp 行号: 24 问题代码:

static constexpr int64_t UB_RESVERVED_SIZE = 1024;

修改建议:

将UB_RESVERVED_SIZE更正为UB_RESERVED_SIZE。

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

likedislike
25+static constexpr int64_t EXTRA_BUFFER_SIZE = 256;
26+static constexpr int64_t NO_CHECK_RANGE_TILING_KEY_NHWC = 200;
27+static constexpr int64_t CHECK_RANGE_TILING_KEY_NHWC = 201;
28+static constexpr int64_t T3_INT64 = 10;
29+static constexpr int64_t DOUBLE_BUFFER = 2;
30+static constexpr int64_t CACHE_LINE_SIZE = 128;
31+ 
32+void MaxPoolGradWithArgmaxV3NHWCTiling::InitializationVars()
33+{
34+ baseData.vRegSize = Ops::Base::GetVRegSize(context_);
35+ baseData.ubBlockSize = Ops::Base::GetUbBlockSize(context_);
36+ baseData.inputBytes = inputData.inputDtype == ge::DT_FLOAT ? FLOAT32_SIZE : FLOAT16_SIZE;
37+ baseData.indexBytes = inputData.indexDtype == ge::DT_INT32 ? INT32_SIZE : INT64_SIZE;
38+ baseData.availableUb = hardwareData.ubSize - UB_RESVERVED_SIZE;
39+ baseData.totalCoreNum = hardwareData.coreNum;
40+ baseData.coreUsedForBestPerformance = baseData.totalCoreNum;
41+ 
42+ int64_t oneBlockNumT1 = baseData.ubBlockSize / baseData.inputBytes;
43+ int64_t oneBlockNumT2 = baseData.ubBlockSize / baseData.indexBytes;
44+ 
45+ baseData.maxDataNumInOneBlock = std::max(oneBlockNumT1, oneBlockNumT2);
46+ 
47+ baseData.proDataNumInOneBeatT2 = baseData.vRegSize / baseData.ubBlockSize * oneBlockNumT2;
48+ baseData.moveDataNumCacheLineT2 = CACHE_LINE_SIZE / baseData.indexBytes;
49+ 
50+ baseData.isPad = 0;
51+ if (inputData.hPad != 0 || inputData.wPad != 0) {
52+ baseData.isPad = 1;
53+ }
54+ 
55+ baseData.hProBatchSize = 1;
56+ if (inputData.hKernel > inputData.hStride) {
57+ baseData.hProBatchSize = Ops::Base::CeilDiv(inputData.hKernel, inputData.hStride);
58+ }
59+ 
60+ baseData.wProBatchSize = 1;
61+ if (inputData.wKernel > inputData.wStride) {
62+ baseData.wProBatchSize = Ops::Base::CeilDiv(inputData.wKernel, inputData.wStride);
63+ }
64+ 
65+ baseData.isOverlap = 0;
66+ if (baseData.wProBatchSize != 1 || baseData.hProBatchSize != 1) {
67+ baseData.isOverlap = 1;
68+ }
69+}
70+ 
71+bool MaxPoolGradWithArgmaxV3NHWCTiling::IsCapable()
72+{
73+ if (inputData.inputFormat != ge::Format::FORMAT_NHWC) {
74+ return false;
75+ }
76+ 
77+ InitializationVars();
78+ return true;
79+}
80+ 
81+uint64_t MaxPoolGradWithArgmaxV3NHWCTiling::GetTilingKey() const
82+{
83+ uint64_t tilingKey = NO_CHECK_RANGE_TILING_KEY_NHWC;
84+ if (splitData.isCheckRange == 1) {
85+ tilingKey = CHECK_RANGE_TILING_KEY_NHWC;
86+ }
87+ 
88+ if (inputData.isInt32Meet == 0) {
89+ tilingKey += T3_INT64;
90+ }
91+ 
92+ return tilingKey;
93+}
94+ 
95+void MaxPoolGradWithArgmaxV3NHWCTiling::DoBufferCalculate()
96+{
97+ // The calculation only involves inner.
98+ int64_t hInputInner = Ops::Base::CeilDiv(splitData.hOutputInner + inputData.hKernel - 1, inputData.hStride);
99+ int64_t wInputInner = Ops::Base::CeilDiv(splitData.wOutputInner + inputData.wKernel - 1, inputData.wStride);
100+ 
101+ int64_t inputPlaneSizeHW = hInputInner * wInputInner;
102+ int64_t outputPlaneSizeHW = splitData.hOutputInner * splitData.wOutputInner;
103+ int64_t cOutputAligned = Ops::Base::CeilAlign(splitData.cOutputInner, baseData.maxDataNumInOneBlock);
104+ int64_t ncPlaneAlignedSize = cOutputAligned * splitData.nOutputInner;
105+ 
106+ splitData.gradBufferSize = ncPlaneAlignedSize * inputPlaneSizeHW * baseData.inputBytes + EXTRA_BUFFER_SIZE;
107+ splitData.argmaxBufferSize = ncPlaneAlignedSize * inputPlaneSizeHW * baseData.indexBytes + EXTRA_BUFFER_SIZE;
108+ 
109+ splitData.outputBufferSize = ncPlaneAlignedSize * outputPlaneSizeHW * FLOAT32_SIZE;
110+ 
111+ int64_t tmpTotalBufferSize = splitData.outputBufferSize + splitData.gradBufferSize + splitData.argmaxBufferSize;
112+ splitData.totalBufferSize = tmpTotalBufferSize * DOUBLE_BUFFER;
113+}
114+ 
115+bool MaxPoolGradWithArgmaxV3NHWCTiling::IsMeetTargetCoreNum() const
116+{
117+ // The calculation only involves inner.
118+ int64_t tmpWOutputOuter = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputInner);
119+ int64_t tmpHOutputOuter = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputInner);
120+ int64_t tmpNOutputOuter = Ops::Base::CeilDiv(inputData.nX, splitData.nOutputInner);
121+ int64_t tmpCOutputOuter = Ops::Base::CeilDiv(inputData.cX, splitData.cOutputInner);
122+ 
123+ return tmpWOutputOuter * tmpHOutputOuter * tmpNOutputOuter * tmpCOutputOuter >= baseData.coreUsedForBestPerformance;
124+}
125+ 
126+bool MaxPoolGradWithArgmaxV3NHWCTiling::IsMeetUBSize()
127+{
128+ DoBufferCalculate();
129+ return splitData.totalBufferSize <= baseData.availableUb;
130+}
131+ 
132+bool MaxPoolGradWithArgmaxV3NHWCTiling::TrySplitN()
133+{
134+ splitData.wOutputInner = inputData.wX;
135+ splitData.hOutputInner = inputData.hX;
136+ splitData.cOutputInner = inputData.cX;
137+ 
138+ splitData.nOutputInner = Ops::Base::CeilDiv(inputData.nX, baseData.coreUsedForBestPerformance);
139+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
140+ return true;
141+ }
142+ 
143+ splitData.nOutputInner = 1;
144+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
145+ int64_t left = 1;
146+ int64_t right = inputData.nX;
147+ int64_t bestSplit = 1;
148+ 
149+ while (left <= right) {
150+ int64_t mid = left + (right - left) / 2;
151+ splitData.nOutputInner = mid;
152+ 
153+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
154+ bestSplit = mid;
155+ left = mid + 1;
156+ } else {
157+ right = mid - 1;
158+ }
159+ }
160+ 
161+ splitData.nOutputInner = bestSplit;
162+ return true;
163+ } else {
164+ return false;
165+ }
166+}
167+ 
168+bool MaxPoolGradWithArgmaxV3NHWCTiling::TrySplitAlignH()
169+{
170+ splitData.nOutputInner = 1;
171+ splitData.wOutputInner = inputData.wX;
172+ splitData.cOutputInner = inputData.cX;
173+ 
174+ splitData.hOutputInner = inputData.hStride;
175+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
176+ int64_t left = 1;
177+ int64_t right = Ops::Base::CeilDiv(inputData.hX / 2, inputData.hStride);
178+ int64_t bestSplit = 1;
179+ 
180+ while (left <= right) {
181+ int64_t mid = left + (right - left) / 2;
182+ splitData.hOutputInner = mid * inputData.hStride;
183+ 
184+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
185+ bestSplit = mid;
186+ left = mid + 1;
187+ } else {
188+ right = mid - 1;
189+ }
190+ }
191+ 
192+ splitData.hOutputInner = bestSplit * inputData.hStride;
193+ return true;
194+ } else {
195+ return false;
196+ }
197+}
198+ 
199+bool MaxPoolGradWithArgmaxV3NHWCTiling::TrySplitAlignW()
200+{
201+ splitData.nOutputInner = 1;
202+ splitData.hOutputInner = inputData.hStride;
203+ splitData.cOutputInner = inputData.cX;
204+ 
205+ splitData.wOutputInner = inputData.wStride;
206+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
207+ int64_t left = 1;
208+ int64_t right = Ops::Base::CeilDiv(inputData.wX / 2, inputData.wStride);
209+ int64_t bestSplit = 1;
210+ 
211+ while (left <= right) {
212+ int64_t mid = left + (right - left) / 2;
213+ splitData.wOutputInner = mid * inputData.wStride;
214+ 
215+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
216+ bestSplit = mid;
217+ left = mid + 1;
218+ } else {
219+ right = mid - 1;
220+ }
221+ }
222+ 
223+ splitData.wOutputInner = bestSplit * inputData.wStride;
224+ return true;
225+ } else {
226+ return false;
227+ }
228+}
229+ 
230+bool MaxPoolGradWithArgmaxV3NHWCTiling::TrySplitAlignC()
231+{
232+ splitData.nOutputInner = 1;
233+ splitData.hOutputInner = inputData.hStride;
234+ splitData.wOutputInner = inputData.wStride;
235+ 
236+ int64_t tmpCAligned =
237+ inputData.cX < baseData.moveDataNumCacheLineT2 ? inputData.cX : baseData.moveDataNumCacheLineT2;
238+ splitData.cOutputInner = tmpCAligned;
239+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
240+ int64_t left = 1;
241+ int64_t right = Ops::Base::CeilDiv(inputData.cX / 2, baseData.moveDataNumCacheLineT2);
242+ int64_t bestSplit = 1;
243+ 
244+ while (left <= right) {
245+ int64_t mid = left + (right - left) / 2;
246+ splitData.cOutputInner = mid * baseData.moveDataNumCacheLineT2;
247+ 
248+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
249+ bestSplit = mid;
250+ left = mid + 1;
251+ } else {
252+ right = mid - 1;
253+ }
254+ }
255+ 
256+ splitData.cOutputInner = bestSplit * baseData.moveDataNumCacheLineT2;
257+ return true;
258+ } else {
259+ // hw stride 较大场景 或者 nhwc超小场景 ---> 应该对hw做更小的切分
260+ return false;
261+ }
262+}
263+ 
264+void MaxPoolGradWithArgmaxV3NHWCTiling::SplitUnalignHWC()
265+{
266+ splitData.nOutputInner = 1;
267+ if (baseData.isPad == 0 && baseData.isOverlap == 0) {
268+ splitData.hOutputInner = inputData.hStride;
269+ splitData.wOutputInner = inputData.wStride;
270+ int64_t tmpCAligned =
271+ inputData.cX < baseData.moveDataNumCacheLineT2 ? inputData.cX : baseData.moveDataNumCacheLineT2;
272+ splitData.cOutputInner = tmpCAligned;
273+ } else {
274+ splitData.wOutputInner = inputData.wX;
275+ splitData.hOutputInner = inputData.hX;
276+ splitData.cOutputInner = inputData.cX;
277+ }
278+ 
279+ splitData.wOutputOuter = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputInner);
280+ splitData.hOutputOuter = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputInner);
281+ 
282+ while (splitData.hOutputInner != 1 || splitData.wOutputInner != 1) {
283+ if (!IsMeetTargetCoreNum() || !IsMeetUBSize()) {
284+ DynamicAdjustmentWH();
285+ } else {
286+ return;
287+ }
288+ }
289+ 
290+ // NHW全切为1 C 超大场景 或者 NHW超小场景
291+ if (inputData.cX <= baseData.proDataNumInOneBeatT2) {
292+ return;
293+ } else if (IsMeetUBSize()) {
294+ splitData.cOutputInner = baseData.proDataNumInOneBeatT2;
295+ return;
296+ } else {
297+ int64_t left = 1;
298+ int64_t right = Ops::Base::CeilDiv(inputData.cX / 2, baseData.proDataNumInOneBeatT2);
299+ int64_t bestSplit = 1;
300+ while (left <= right) {
301+ int64_t mid = left + (right - left) / 2;
302+ splitData.cOutputInner = mid * baseData.proDataNumInOneBeatT2;
303+ 
304+ if (IsMeetUBSize()) {
305+ bestSplit = mid;
306+ left = mid + 1;
307+ } else {
308+ right = mid - 1;
309+ }
310+ }
311+ splitData.cOutputInner = bestSplit * baseData.proDataNumInOneBeatT2;
312+ return;
313+ }
CANN-robotCANN-robot
CANN-robotCANN-robot2025年12月29日

代码结构与可维护性: 在SplitUnalignHWC函数第314行有一个多余的return语句，该函数在第312行已经有一个return，第314行的return永远不会执行，属于死代码。

问题类型: 代码结构与可维护性 文件路径: pooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nhwc_tiling.cpp 行号: 313 问题代码:

    }
    return;
}

修改建议:

删除第314行多余的return语句。

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

likedislike
CANN-robotCANN-robot2025年12月29日

冗余代码: 在 'SplitUnalignHWC' 函数的末尾，第314行有一个多余的 'return;' 语句。因为第312行已经有一个 'return;'，所以第314行的 'return;' 永远不会被执行，是死代码。

问题类型: 冗余代码 文件路径: pooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nhwc_tiling.cpp 行号: 313 问题代码:

    }
    return;
}

修改建议:

删除第314行的冗余 'return;' 语句。

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

likedislike
314+ return;
315+}
316+ 
317+void MaxPoolGradWithArgmaxV3NHWCTiling::DynamicAdjustmentWH()
318+{
319+ if (splitData.hOutputInner == 1) {
320+ splitData.wOutputOuter++;
321+ splitData.wOutputInner = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputOuter);
322+ } else {
323+ splitData.hOutputOuter++;
324+ splitData.hOutputInner = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputOuter);
325+ }
326+}
327+ 
328+void MaxPoolGradWithArgmaxV3NHWCTiling::SearchBestTiling()
329+{
330+ splitData.isCheckRange = 0;
331+ if (TrySplitN()) {
332+ return;
333+ }
334+ 
335+ if (baseData.isPad == 0 && baseData.isOverlap == 0) {
336+ if (TrySplitAlignH()) {
337+ return;
338+ }
339+ 
340+ if (TrySplitAlignW()) {
341+ return;
342+ }
343+ 
344+ if (TrySplitAlignC()) {
345+ return;
346+ }
347+ }
348+ 
349+ // 带pad 或者 最小整切仍然不满足条件需要更细粒度切分HWC
350+ splitData.isCheckRange = 1;
351+ SplitUnalignHWC();
352+ return;
353+}
354+ 
355+void MaxPoolGradWithArgmaxV3NHWCTiling::DoUBTiling()
356+{
357+ SearchBestTiling();
358+ DoBufferCalculate();
359+ splitData.wOutputOuter = Ops::Base::CeilDiv(inputData.wX, splitData.wOutputInner);
360+ int64_t tempWOutputTail = inputData.wX % splitData.wOutputInner;
361+ splitData.wOutputTail = tempWOutputTail == 0 ? splitData.wOutputInner : tempWOutputTail;
362+ 
363+ splitData.hOutputOuter = Ops::Base::CeilDiv(inputData.hX, splitData.hOutputInner);
364+ int64_t tempHOutputTail = inputData.hX % splitData.hOutputInner;
365+ splitData.hOutputTail = tempHOutputTail == 0 ? splitData.hOutputInner : tempHOutputTail;
366+ 
367+ splitData.nOutputOuter = Ops::Base::CeilDiv(inputData.nX, splitData.nOutputInner);
368+ int64_t tempNOutputTail = inputData.nX % splitData.nOutputInner;
369+ splitData.nOutputTail = tempNOutputTail == 0 ? splitData.nOutputInner : tempNOutputTail;
370+ 
371+ splitData.cOutputOuter = Ops::Base::CeilDiv(inputData.cX, splitData.cOutputInner);
372+ int64_t tempCOutputTail = inputData.cX % splitData.cOutputInner;
373+ splitData.cOutputTail = tempCOutputTail == 0 ? splitData.cOutputInner : tempCOutputTail;
374+}
375+ 
376+void MaxPoolGradWithArgmaxV3NHWCTiling::DoBlockTiling()
377+{
378+ splitData.totalBaseBlockNum =
379+ splitData.nOutputOuter * splitData.cOutputOuter * splitData.hOutputOuter * splitData.wOutputOuter;
380+ splitData.normalCoreProcessNum = Ops::Base::CeilDiv(splitData.totalBaseBlockNum, baseData.totalCoreNum);
381+ splitData.usedCoreNum = Ops::Base::CeilDiv(splitData.totalBaseBlockNum, splitData.normalCoreProcessNum);
382+ splitData.tailCoreProcessNum =
383+ splitData.totalBaseBlockNum - splitData.normalCoreProcessNum * (splitData.usedCoreNum - 1);
384+}
385+ 
386+void MaxPoolGradWithArgmaxV3NHWCTiling::PrintBaseData() const
387+{
388+ OP_LOGD("MaxPoolGradWithArgmaxV3NHWC", "[MaxPoolGradWithArgmaxV3NHWC] PrintBaseData start running");
389+ 
390+ std::ostringstream info;
391+ info << "baseData.vRegSize: " << baseData.vRegSize << std::endl;
392+ info << "baseData.ubBlockSize: " << baseData.ubBlockSize << std::endl;
393+ 
394+ info << "baseData.inputBytes: " << baseData.inputBytes << std::endl;
395+ info << "baseData.indexBytes: " << baseData.indexBytes << std::endl;
396+ info << "baseData.availableUb: " << baseData.availableUb << std::endl;
397+ info << "baseData.maxDataNumInOneBlock: " << baseData.maxDataNumInOneBlock << std::endl;
398+ info << "baseData.proDataNumInOneBeatT2: " << baseData.proDataNumInOneBeatT2 << std::endl;
399+ info << "baseData.totalCoreNum: " << baseData.totalCoreNum << std::endl;
400+ info << "baseData.coreUsedForBestPerformance: " << baseData.coreUsedForBestPerformance << std::endl;
401+ 
402+ info << "baseData.isPad: " << baseData.isPad << std::endl;
403+ info << "baseData.isOverlap: " << baseData.isOverlap << std::endl;
404+ info << "baseData.hProBatchSize: " << baseData.hProBatchSize << std::endl;
405+ info << "baseData.wProBatchSize: " << baseData.wProBatchSize << std::endl;
406+ info << "baseData.moveDataNumCacheLineT2: " << baseData.moveDataNumCacheLineT2 << std::endl;
407+ 
408+ OP_LOGI("MaxPoolGradWithArgmaxV3NHWC", "%s", info.str().c_str());
409+}
410+ 
411+void MaxPoolGradWithArgmaxV3NHWCTiling::PrintSplitData() const
412+{
413+ OP_LOGD("MaxPoolGradWithArgmaxV3NHWC", "[MaxPoolGradWithArgmaxV3NHWC] PrintSplitData start running");
414+ 
415+ std::ostringstream info;
416+ info << "splitData.isCheckRange: " << splitData.isCheckRange << std::endl;
417+ 
418+ info << "splitData.nOutputInner: " << splitData.nOutputInner << std::endl;
419+ info << "splitData.nOutputTail: " << splitData.nOutputTail << std::endl;
420+ info << "splitData.nOutputOuter: " << splitData.nOutputOuter << std::endl;
421+ 
422+ info << "splitData.hOutputInner: " << splitData.hOutputInner << std::endl;
423+ info << "splitData.hOutputTail: " << splitData.hOutputTail << std::endl;
424+ info << "splitData.hOutputOuter: " << splitData.hOutputOuter << std::endl;
425+ 
426+ info << "splitData.wOutputInner: " << splitData.wOutputInner << std::endl;
427+ info << "splitData.wOutputTail: " << splitData.wOutputTail << std::endl;
428+ info << "splitData.wOutputOuter: " << splitData.wOutputOuter << std::endl;
429+ 
430+ info << "splitData.cOutputInner: " << splitData.cOutputInner << std::endl;
431+ info << "splitData.cOutputTail: " << splitData.cOutputTail << std::endl;
432+ info << "splitData.cOutputOuter: " << splitData.cOutputOuter << std::endl;
433+ 
434+ info << "splitData.normalCoreProcessNum: " << splitData.normalCoreProcessNum << std::endl;
435+ info << "splitData.tailCoreProcessNum: " << splitData.tailCoreProcessNum << std::endl;
436+ info << "splitData.usedCoreNum: " << splitData.usedCoreNum << std::endl;
437+ info << "splitData.totalBaseBlockNum: " << splitData.totalBaseBlockNum << std::endl;
438+ 
439+ info << "splitData.outputBufferSize: " << splitData.outputBufferSize << std::endl;
440+ info << "splitData.gradBufferSize: " << splitData.gradBufferSize << std::endl;
441+ info << "splitData.argmaxBufferSize: " << splitData.argmaxBufferSize << std::endl;
442+ info << "splitData.totalBufferSize: " << splitData.totalBufferSize << std::endl;
443+ 
444+ OP_LOGI("MaxPoolGradWithArgmaxV3NHWC", "%s", info.str().c_str());
445+}
446+ 
447+void MaxPoolGradWithArgmaxV3NHWCTiling::SetTilingData()
448+{
449+ tilingData.set_hArgmax(inputData.hGrad);
450+ tilingData.set_wArgmax(inputData.wGrad);
451+ tilingData.set_cOutput(inputData.cX);
452+ tilingData.set_hOutput(inputData.hX);
453+ tilingData.set_wOutput(inputData.wX);
454+ tilingData.set_hKernel(inputData.hKernel);
455+ tilingData.set_wKernel(inputData.wKernel);
456+ tilingData.set_hStride(inputData.hStride);
457+ tilingData.set_wStride(inputData.wStride);
458+ tilingData.set_padH(inputData.hPad);
459+ tilingData.set_padW(inputData.wPad);
460+ tilingData.set_dilationH(inputData.hDilation);
461+ tilingData.set_dilationW(inputData.wDilation);
462+ tilingData.set_nOutputInner(splitData.nOutputInner);
463+ tilingData.set_nOutputTail(splitData.nOutputTail);
464+ tilingData.set_nOutputOuter(splitData.nOutputOuter);
465+ tilingData.set_hOutputInner(splitData.hOutputInner);
466+ tilingData.set_hOutputTail(splitData.hOutputTail);
467+ tilingData.set_hOutputOuter(splitData.hOutputOuter);
468+ tilingData.set_wOutputInner(splitData.wOutputInner);
469+ tilingData.set_wOutputTail(splitData.wOutputTail);
470+ tilingData.set_wOutputOuter(splitData.wOutputOuter);
471+ tilingData.set_cOutputInner(splitData.cOutputInner);
472+ tilingData.set_cOutputTail(splitData.cOutputTail);
473+ tilingData.set_cOutputOuter(splitData.cOutputOuter);
474+ tilingData.set_normalCoreProcessNum(splitData.normalCoreProcessNum);
475+ tilingData.set_tailCoreProcessNum(splitData.tailCoreProcessNum);
476+ tilingData.set_usedCoreNum(splitData.usedCoreNum);
477+ tilingData.set_outputBufferSize(splitData.outputBufferSize);
478+ tilingData.set_gradBufferSize(splitData.gradBufferSize);
479+ tilingData.set_argmaxBufferSize(splitData.argmaxBufferSize);
480+ tilingData.set_hProBatchSize(baseData.hProBatchSize);
481+ tilingData.set_wProBatchSize(baseData.wProBatchSize);
482+ tilingData.set_tilingKey(GetTilingKey());
483+}
484+ 
485+ge::graphStatus MaxPoolGradWithArgmaxV3NHWCTiling::DoOpTiling()
486+{
487+ DoUBTiling();
488+ DoBlockTiling();
489+ SetTilingData();
490+ PrintBaseData();
491+ PrintSplitData();
492+ return ge::GRAPH_SUCCESS;
493+}
494+ 
495+ge::graphStatus MaxPoolGradWithArgmaxV3NHWCTiling::PostTiling()
496+{
497+ context_->SetBlockDim(tilingData.get_usedCoreNum());
498+ if (tilingData.GetDataSize() > context_->GetRawTilingData()->GetCapacity()) {
499+ return ge::GRAPH_FAILED;
500+ }
501+ 
502+ tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
503+ context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
504+ return ge::GRAPH_SUCCESS;
505+}
506+ 
507+REGISTER_OPS_TILING_TEMPLATE(MaxPoolGradWithArgmaxV3, MaxPoolGradWithArgmaxV3NHWCTiling, 1);
508+ 
509+} // namespace optiling
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_nhwc_tiling.h+155-0
@@ -0,0 +1,155 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nhwc_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_GRAD_WITH_AGRMAX_V3_NHWC_TILING_H_
17+#define MAX_POOL_GRAD_WITH_AGRMAX_V3_NHWC_TILING_H_
18+ 
19+#include "max_pool_grad_with_argmax_v3_tiling_base.h"
20+ 
21+namespace optiling {
22+ 
23+BEGIN_TILING_DATA_DEF(MaxPoolGradWithArgmaxV3NHWCTilingData)
24+TILING_DATA_FIELD_DEF(int64_t, hArgmax);
25+TILING_DATA_FIELD_DEF(int64_t, wArgmax);
26+TILING_DATA_FIELD_DEF(int64_t, cOutput);
27+TILING_DATA_FIELD_DEF(int64_t, hOutput);
28+TILING_DATA_FIELD_DEF(int64_t, wOutput);
29+TILING_DATA_FIELD_DEF(int64_t, hKernel);
30+TILING_DATA_FIELD_DEF(int64_t, wKernel);
31+TILING_DATA_FIELD_DEF(int64_t, hStride);
32+TILING_DATA_FIELD_DEF(int64_t, wStride);
33+TILING_DATA_FIELD_DEF(int64_t, padH);
34+TILING_DATA_FIELD_DEF(int64_t, padW);
35+TILING_DATA_FIELD_DEF(int64_t, dilationH);
36+TILING_DATA_FIELD_DEF(int64_t, dilationW);
37+TILING_DATA_FIELD_DEF(int64_t, nOutputInner);
38+TILING_DATA_FIELD_DEF(int64_t, nOutputTail);
39+TILING_DATA_FIELD_DEF(int64_t, nOutputOuter);
40+TILING_DATA_FIELD_DEF(int64_t, hOutputInner);
41+TILING_DATA_FIELD_DEF(int64_t, hOutputTail);
42+TILING_DATA_FIELD_DEF(int64_t, hOutputOuter);
43+TILING_DATA_FIELD_DEF(int64_t, wOutputInner);
44+TILING_DATA_FIELD_DEF(int64_t, wOutputTail);
45+TILING_DATA_FIELD_DEF(int64_t, wOutputOuter);
46+TILING_DATA_FIELD_DEF(int64_t, cOutputInner);
47+TILING_DATA_FIELD_DEF(int64_t, cOutputTail);
48+TILING_DATA_FIELD_DEF(int64_t, cOutputOuter);
49+TILING_DATA_FIELD_DEF(int64_t, normalCoreProcessNum);
50+TILING_DATA_FIELD_DEF(int64_t, tailCoreProcessNum);
51+TILING_DATA_FIELD_DEF(int64_t, usedCoreNum);
52+TILING_DATA_FIELD_DEF(int64_t, outputBufferSize);
53+TILING_DATA_FIELD_DEF(int64_t, gradBufferSize);
54+TILING_DATA_FIELD_DEF(int64_t, argmaxBufferSize);
55+TILING_DATA_FIELD_DEF(int64_t, hProBatchSize);
56+TILING_DATA_FIELD_DEF(int64_t, wProBatchSize);
57+TILING_DATA_FIELD_DEF(int64_t, tilingKey);
58+END_TILING_DATA_DEF;
59+ 
60+// no check range 200 check range 201
61+// no check range int64 210 check range int64 211
62+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_200, MaxPoolGradWithArgmaxV3NHWCTilingData);
63+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_201, MaxPoolGradWithArgmaxV3NHWCTilingData);
64+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_210, MaxPoolGradWithArgmaxV3NHWCTilingData);
65+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3_211, MaxPoolGradWithArgmaxV3NHWCTilingData);
66+ 
67+struct MaxPoolGradWithArgmaxV3NHWCBaseInfo {
68+ int64_t vRegSize{0};
69+ int64_t ubBlockSize{0};
70+ int64_t inputBytes{0};
71+ int64_t indexBytes{0};
72+ int64_t availableUb{0};
73+ int64_t maxDataNumInOneBlock{0};
74+ int64_t proDataNumInOneBeatT2{0};
75+ int64_t totalCoreNum{0};
76+ int64_t coreUsedForBestPerformance{0};
77+ int64_t isPad{0};
78+ int64_t isOverlap{0};
79+ int64_t hProBatchSize{0};
80+ int64_t wProBatchSize{0};
81+ int64_t moveDataNumCacheLineT2{0};
82+};
83+ 
84+struct MaxPoolGradWithArgmaxV3NHWCSplitInfo {
85+ // DoUBTiling
86+ int64_t isCheckRange{0};
87+ 
88+ int64_t nOutputInner{0};
89+ int64_t nOutputTail{0};
90+ int64_t nOutputOuter{0};
91+ 
92+ int64_t hOutputInner{0};
93+ int64_t hOutputTail{0};
94+ int64_t hOutputOuter{0};
95+ 
96+ int64_t wOutputInner{0};
97+ int64_t wOutputTail{0};
98+ int64_t wOutputOuter{0};
99+ 
100+ int64_t cOutputInner{0};
101+ int64_t cOutputTail{0};
102+ int64_t cOutputOuter{0};
103+ 
104+ // DoBlockTiling
105+ int64_t normalCoreProcessNum{0};
106+ int64_t tailCoreProcessNum{0};
107+ int64_t usedCoreNum{0};
108+ int64_t totalBaseBlockNum{0};
109+ 
110+ // DoBufferCalculate
111+ int64_t outputBufferSize{0};
112+ int64_t gradBufferSize{0};
113+ int64_t argmaxBufferSize{0};
114+ int64_t totalBufferSize{0};
115+};
116+ 
117+class MaxPoolGradWithArgmaxV3NHWCTiling : public MaxPoolGradWithArgmaxV3BaseTiling {
118+public:
119+ explicit MaxPoolGradWithArgmaxV3NHWCTiling(gert::TilingContext* context)
120+ : MaxPoolGradWithArgmaxV3BaseTiling(context)
121+ {}
122+ 
123+ ~MaxPoolGradWithArgmaxV3NHWCTiling() override
124+ {}
125+ 
126+private:
127+ void DoUBTiling();
128+ void InitializationVars();
129+ bool TrySplitN();
130+ bool TrySplitAlignH();
131+ bool TrySplitAlignW();
132+ bool TrySplitAlignC();
133+ void SplitUnalignHWC();
134+ bool IsMeetTargetCoreNum() const;
135+ bool IsMeetUBSize();
136+ void SearchBestTiling();
137+ void DynamicAdjustmentWH();
138+ void SetTilingData();
139+ uint64_t GetTilingKey() const override;
140+ void PrintBaseData() const;
141+ void PrintSplitData() const;
142+ void DoBlockTiling();
143+ void DoBufferCalculate();
144+ bool IsCapable() override;
145+ ge::graphStatus DoOpTiling() override;
146+ ge::graphStatus PostTiling() override;
147+ 
148+ MaxPoolGradWithArgmaxV3NHWCTilingData tilingData;
149+ MaxPoolGradWithArgmaxV3NHWCBaseInfo baseData;
150+ MaxPoolGradWithArgmaxV3NHWCSplitInfo splitData;
151+};
152+ 
153+} // namespace optiling
154+ 
155+#endif
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_tiling.cpp+46-0
@@ -0,0 +1,46 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_grad_with_argmax_v3_tiling_base.h"
18+ 
19+using namespace AscendC;
20+using Ops::NN::Optiling::TilingRegistry;
21+namespace optiling {
22+ 
23+ge::graphStatus Tiling4MaxPoolGradWithArgmaxV3(gert::TilingContext* context)
24+{
25+ return TilingRegistry::GetInstance().DoTilingImpl(context);
26+}
27+ 
28+ge::graphStatus TilingPrepare4MaxPoolGradWithArgmaxV3(gert::TilingParseContext* context)
29+{
30+ fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo();
31+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr);
32+ 
33+ auto compileInfoPtr = context->GetCompiledInfo<MaxPoolGradWithArgmaxV3CompileInfo>();
34+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr);
35+ 
36+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
37+ compileInfoPtr->coreNum = ascendcPlatform.GetCoreNum();
38+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize);
39+ return ge::GRAPH_SUCCESS;
40+}
41+ 
42+IMPL_OP_OPTILING(MaxPoolGradWithArgmaxV3)
43+ .Tiling(Tiling4MaxPoolGradWithArgmaxV3)
44+ .TilingParse<MaxPoolGradWithArgmaxV3CompileInfo>(TilingPrepare4MaxPoolGradWithArgmaxV3);
45+ 
46+} // namespace optiling
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_tiling_base.cpp+377-0
@@ -0,0 +1,377 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_tiling_base.cpp
13+ * \brief
14+ */
15+ 
16+#include <cstdint>
17+#include "tiling_base/tiling_templates_registry.h"
18+#include "log/log.h"
19+#include "error_util.h"
20+#include "platform/platform_info.h"
21+#include "../../../max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_tiling.h"
22+#include "max_pool_grad_with_argmax_v3_tiling_base.h"
23+ 
24+using namespace AscendC;
25+using namespace ge;
26+ 
27+namespace optiling {
28+ 
29+static constexpr int64_t DIMS_FOUR = 4;
30+static constexpr int64_t ATTR_INDEX_KSIZE = 0;
31+static constexpr int64_t ATTR_INDEX_STRIDES = 1;
32+static constexpr int64_t ATTR_INDEX_PADS = 2;
33+static constexpr int64_t ATTR_INDEX_DTYPE = 3;
34+static constexpr int64_t ATTR_INDEX_DILATION = 4;
35+static constexpr int64_t ATTR_INDEX_CEIL_MODE = 5;
36+static constexpr int64_t ATTR_INDEX_FORMAT = 6;
37+static constexpr int64_t DIM_ZERO = 0;
38+static constexpr int64_t DIM_ONE = 1;
39+static constexpr int64_t DIM_TWO = 2;
40+static constexpr int64_t DIM_THREE = 3;
41+static constexpr int64_t DTYPE_INT32 = 3;
42+static constexpr int64_t DTYPE_INT64 = 9;
43+static constexpr int64_t INPUT_X = 0;
44+static constexpr int64_t INPUT_GRAD = 1;
45+static constexpr int64_t INPUT_ARGMAX = 2;
46+static constexpr size_t WS_SYS_SIZE = static_cast<size_t>(16 * 1024 * 1024);
47+ 
48+ge::graphStatus MaxPoolGradWithArgmaxV3BaseTiling::GetPlatformInfo()
49+{
50+ auto platformPtr = context_->GetPlatformInfo();
51+ if (platformPtr == nullptr) {
52+ auto compileInfoPtr = reinterpret_cast<const MaxPoolGradWithArgmaxV3CompileInfo*>(context_->GetCompileInfo());
53+ OP_TILING_CHECK(
54+ compileInfoPtr == nullptr, CUBE_INNER_ERR_REPORT(context_, "compile info is null"),
55+ return ge::GRAPH_FAILED);
56+ hardwareData.coreNum = compileInfoPtr->coreNum;
57+ hardwareData.ubSize = compileInfoPtr->ubSize;
58+ } else {
59+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformPtr);
60+ hardwareData.coreNum = ascendcPlatform.GetCoreNum();
61+ 
62+ uint64_t ubSizePlatform;
63+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
64+ hardwareData.ubSize = static_cast<int64_t>(ubSizePlatform);
65+ }
66+ 
67+ OP_TILING_CHECK(
68+ hardwareData.coreNum == 0, CUBE_INNER_ERR_REPORT(context_, "coreNum is 0"), return ge::GRAPH_FAILED);
69+ return ge::GRAPH_SUCCESS;
70+}
71+ 
72+static inline int64_t DivRtn(int64_t x, int64_t y)
73+{
74+ if (y == 0) {
75+ return 0;
76+ }
77+ int64_t q = x / y;
78+ int64_t r = x % y;
79+ if ((r != 0) && ((r < 0) != (y < 0))) {
80+ --q;
81+ }
82+ return q;
83+}
84+ 
85+static bool CheckGradShape(const MaxPoolGradWithArgmaxV3InputInfo& inputData)
86+{
87+ int64_t tmpH = inputData.hX + 2 * inputData.hPad - inputData.hDilation * (inputData.hKernel - 1) - 1;
88+ if (inputData.ceilMode) {
89+ tmpH += (inputData.hStride - 1);
90+ }
91+ int64_t tmpHGrad = DivRtn(tmpH, inputData.hStride) + 1;
92+ 
93+ int64_t tmpW = inputData.wX + 2 * inputData.wPad - inputData.wDilation * (inputData.wKernel - 1) - 1;
94+ if (inputData.ceilMode) {
95+ tmpW += (inputData.wStride - 1);
96+ }
97+ int64_t tmpWGrad = DivRtn(tmpW, inputData.wStride) + 1;
98+ 
99+ if (inputData.ceilMode) {
100+ if ((tmpHGrad - 1) * inputData.hStride >= inputData.hX + inputData.hPad) {
101+ tmpHGrad = tmpHGrad - 1;
102+ }
103+ if ((tmpWGrad - 1) * inputData.wStride >= inputData.wX + inputData.wPad) {
104+ tmpWGrad = tmpWGrad - 1;
105+ }
106+ }
107+ 
108+ if (tmpHGrad != inputData.hGrad || tmpWGrad != inputData.wGrad || inputData.nX != inputData.nGrad ||
109+ inputData.cX != inputData.cGrad) {
110+ OP_LOGE(
111+ "MaxPoolGradWithArgmaxV3", "grad shape expected nchw [%ld, %ld,%ld, %ld], but got [%ld, %ld,%ld, %ld]",
112+ inputData.nX, inputData.cX, tmpHGrad, tmpWGrad, inputData.nGrad, inputData.cGrad, inputData.hGrad,
113+ inputData.wGrad);
114+ return false;
115+ }
116+ return true;
117+}
118+ 
119+static inline bool IsGreaterThanInt32Max(const MaxPoolGradWithArgmaxV3InputInfo& inputData)
120+{
121+ if (inputData.indexDtype == ge::DataType::DT_INT32) {
122+ return false;
123+ }
124+ 
125+ int64_t planeSize = inputData.hX * inputData.wX;
126+ return planeSize > static_cast<int64_t>(INT32_MAX);
127+}
128+ 
129+ge::graphStatus MaxPoolGradWithArgmaxV3BaseTiling::GetShapeAttrsInfo()
130+{
131+ auto inputX = context_->GetInputShape(INPUT_X);
132+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputX);
133+ auto xShape = Ops::Base::EnsureNotScalar(inputX->GetStorageShape());
134+ 
135+ OP_TILING_CHECK(
136+ xShape.GetDimNum() != DIMS_FOUR,
137+ VECTOR_INNER_ERR_REPORT_TILIING(
138+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: input shape dim = %zu, should be equal 4",
139+ xShape.GetDimNum()),
140+ return ge::GRAPH_FAILED);
141+ 
142+ OP_TILING_CHECK(
143+ xShape.GetShapeSize() <= 0,
144+ VECTOR_INNER_ERR_REPORT_TILIING(
145+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: input shape size %ld less than zero failed",
146+ xShape.GetShapeSize()),
147+ return ge::GRAPH_FAILED);
148+ 
149+ auto inputDesc = context_->GetInputDesc(INPUT_X);
150+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputDesc);
151+ inputData.inputDtype = inputDesc->GetDataType();
152+ if (inputData.inputDtype != ge::DataType::DT_BF16 && inputData.inputDtype != ge::DataType::DT_FLOAT16 &&
153+ inputData.inputDtype != ge::DataType::DT_FLOAT) {
154+ VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: invalid dtype");
155+ return ge::GRAPH_FAILED;
156+ }
157+ 
158+ auto inputGrad = context_->GetInputShape(INPUT_GRAD);
159+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputGrad);
160+ auto gradShape = Ops::Base::EnsureNotScalar(inputGrad->GetStorageShape());
161+ 
162+ OP_TILING_CHECK(
163+ gradShape.GetShapeSize() <= 0,
164+ VECTOR_INNER_ERR_REPORT_TILIING(
165+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: grad shape size %ld less than zero failed",
166+ gradShape.GetShapeSize()),
167+ return ge::GRAPH_FAILED);
168+ 
169+ auto inputArgmax = context_->GetInputShape(INPUT_ARGMAX);
170+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputArgmax);
171+ auto argmaxShape = Ops::Base::EnsureNotScalar(inputArgmax->GetStorageShape());
172+ OP_TILING_CHECK(
173+ argmaxShape.GetShapeSize() <= 0,
174+ VECTOR_INNER_ERR_REPORT_TILIING(
175+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: argmax shape size %ld less than zero failed",
176+ argmaxShape.GetShapeSize()),
177+ return ge::GRAPH_FAILED);
178+ 
179+ auto inputArgmaxDesc = context_->GetInputDesc(INPUT_ARGMAX);
180+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputArgmaxDesc);
181+ auto argmaxDtype = inputArgmaxDesc->GetDataType();
182+ if (argmaxDtype != ge::DataType::DT_INT32 && argmaxDtype != ge::DataType::DT_INT64) {
183+ VECTOR_INNER_ERR_REPORT_TILIING(
184+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: argmax dtype only support int32, int64, but got [%s].",
185+ Ops::Base::ToString(argmaxDtype).c_str());
186+ return ge::GRAPH_FAILED;
187+ }
188+ 
189+ OP_TILING_CHECK(
190+ gradShape != argmaxShape,
191+ VECTOR_INNER_ERR_REPORT_TILIING(
192+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: argmax shape is not same as grad shape"),
193+ return ge::GRAPH_FAILED);
194+ 
195+ auto outY = context_->GetOutputShape(0);
196+ OP_CHECK_NULL_WITH_CONTEXT(context_, outY);
197+ auto yShape = Ops::Base::EnsureNotScalar(outY->GetStorageShape());
198+ OP_TILING_CHECK(
199+ yShape != xShape,
200+ VECTOR_INNER_ERR_REPORT_TILIING(
201+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: output shape is not same as input shape"),
202+ return ge::GRAPH_FAILED);
203+ 
204+ auto runtimeAttrs = context_->GetAttrs();
205+ OP_CHECK_NULL_WITH_CONTEXT(context_, runtimeAttrs);
206+ 
207+ const char* inputFormatPtr = runtimeAttrs->GetAttrPointer<char>(ATTR_INDEX_FORMAT);
208+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputFormatPtr);
209+ if (strncmp(inputFormatPtr, "NCHW", sizeof("NCHW") / sizeof(char)) == 0) {
CANN-robot
CANN-robotCANN-robot2025年12月29日

字符串比较问题: 第209行使用strncmp比较字符串，但sizeof("NCHW")/sizeof(char)计算的是5（包括空字符），而strncmp应该比较4个字符。这可能导致错误匹配。

问题类型: 字符串比较问题 文件路径: pooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_tiling_base.cpp 行号: 209 问题代码:

    if (strncmp(inputFormatPtr, "NCHW", sizeof("NCHW") / sizeof(char)) == 0) {

修改建议:

建议使用strncmp(inputFormatPtr, "NCHW", 4) == 0，或者使用strcmp如果确定字符串以空字符结尾。

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

likedislike
210+ inputData.inputFormat = ge::Format::FORMAT_NCHW;
211+ inputData.nX = xShape.GetDim(DIM_ZERO);
212+ inputData.cX = xShape.GetDim(DIM_ONE);
213+ inputData.hX = xShape.GetDim(DIM_TWO);
214+ inputData.wX = xShape.GetDim(DIM_THREE);
215+ inputData.nGrad = gradShape.GetDim(DIM_ZERO);
216+ inputData.cGrad = gradShape.GetDim(DIM_ONE);
217+ inputData.hGrad = gradShape.GetDim(DIM_TWO);
218+ inputData.wGrad = gradShape.GetDim(DIM_THREE);
219+ } else if (strncmp(inputFormatPtr, "NHWC", sizeof("NHWC") / sizeof(char)) == 0) {
220+ inputData.inputFormat = ge::Format::FORMAT_NHWC;
221+ inputData.nX = xShape.GetDim(DIM_ZERO);
222+ inputData.cX = xShape.GetDim(DIM_THREE);
223+ inputData.hX = xShape.GetDim(DIM_ONE);
224+ inputData.wX = xShape.GetDim(DIM_TWO);
225+ inputData.nGrad = gradShape.GetDim(DIM_ZERO);
226+ inputData.cGrad = gradShape.GetDim(DIM_THREE);
227+ inputData.hGrad = gradShape.GetDim(DIM_ONE);
228+ inputData.wGrad = gradShape.GetDim(DIM_TWO);
229+ } else {
230+ VECTOR_INNER_ERR_REPORT_TILIING(
231+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: input format [%s] is invalid", inputFormatPtr);
232+ return ge::GRAPH_FAILED;
233+ }
234+ 
235+ const gert::TypedContinuousVector<int64_t>* kernelSizePtr = runtimeAttrs->GetListInt(ATTR_INDEX_KSIZE);
236+ OP_CHECK_NULL_WITH_CONTEXT(context_, kernelSizePtr);
237+ inputData.hKernel = *(kernelSizePtr->GetData());
238+ inputData.wKernel = *(kernelSizePtr->GetData() + 1);
239+ OP_TILING_CHECK(
240+ inputData.hKernel <= 0 || inputData.wKernel <= 0,
241+ VECTOR_INNER_ERR_REPORT_TILIING(
242+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: kernel shape [%ld, %ld] is invalid", inputData.hKernel,
243+ inputData.wKernel),
244+ return ge::GRAPH_FAILED);
245+ 
246+ const gert::TypedContinuousVector<int64_t>* stridePtr = runtimeAttrs->GetListInt(ATTR_INDEX_STRIDES);
247+ OP_CHECK_NULL_WITH_CONTEXT(context_, stridePtr);
248+ inputData.hStride = *(stridePtr->GetData());
249+ inputData.wStride = *(stridePtr->GetData() + 1);
250+ OP_TILING_CHECK(
251+ inputData.hStride <= 0 || inputData.wStride <= 0,
252+ VECTOR_INNER_ERR_REPORT_TILIING(
253+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: stride shape [%ld, %ld] is invalid", inputData.hStride,
254+ inputData.wStride),
255+ return ge::GRAPH_FAILED);
256+ 
257+ const gert::TypedContinuousVector<int64_t>* paddingPtr = runtimeAttrs->GetListInt(ATTR_INDEX_PADS);
258+ OP_CHECK_NULL_WITH_CONTEXT(context_, paddingPtr);
259+ inputData.hPad = *(paddingPtr->GetData());
260+ inputData.wPad = *(paddingPtr->GetData() + 1);
261+ OP_TILING_CHECK(
262+ inputData.hPad > inputData.hKernel / 2 || inputData.wPad > inputData.wKernel / 2,
263+ VECTOR_INNER_ERR_REPORT_TILIING(
264+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: pad shape [%ld, %ld] is invalid", inputData.hPad,
265+ inputData.wPad),
266+ return ge::GRAPH_FAILED);
267+ 
268+ const gert::TypedContinuousVector<int64_t>* dilationPtr = runtimeAttrs->GetListInt(ATTR_INDEX_DILATION);
269+ OP_CHECK_NULL_WITH_CONTEXT(context_, dilationPtr);
270+ 
271+ inputData.hDilation = *(dilationPtr->GetData());
272+ inputData.wDilation = *(dilationPtr->GetData() + 1);
273+ OP_TILING_CHECK(
274+ inputData.hDilation <= 0 || inputData.wDilation <= 0,
275+ VECTOR_INNER_ERR_REPORT_TILIING(
276+ context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: dilation shape [%ld, %ld] is invalid",
277+ inputData.hDilation, inputData.wDilation),
278+ return ge::GRAPH_FAILED);
279+ 
280+ const bool* ceilModePtr = runtimeAttrs->GetAttrPointer<bool>(ATTR_INDEX_CEIL_MODE);
281+ OP_CHECK_NULL_WITH_CONTEXT(context_, ceilModePtr);
282+ inputData.ceilMode = *ceilModePtr;
283+ 
284+ OP_TILING_CHECK(
285+ !CheckGradShape(inputData),
286+ VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "MaxPoolGradWithArgmaxV3: grad shape is invalid"),
287+ return ge::GRAPH_FAILED);
288+ 
289+ const int* indexDtypePtr = runtimeAttrs->GetAttrPointer<int>(ATTR_INDEX_DTYPE);
290+ OP_CHECK_NULL_WITH_CONTEXT(context_, indexDtypePtr);
291+ switch (*indexDtypePtr) {
292+ case DTYPE_INT32:
293+ inputData.indexDtype = ge::DataType::DT_INT32;
294+ break;
295+ case DTYPE_INT64:
296+ inputData.indexDtype = ge::DataType::DT_INT64;
297+ break;
298+ default:
299+ inputData.indexDtype = ge::DataType::DT_INT32;
300+ break;
301+ }
302+ 
303+ if (IsGreaterThanInt32Max(inputData)) {
304+ inputData.isInt32Meet = 0;
305+ } else {
306+ inputData.isInt32Meet = 1;
307+ }
308+ 
309+ PrintInputData();
310+ return ge::GRAPH_SUCCESS;
311+}
312+ 
313+void MaxPoolGradWithArgmaxV3BaseTiling::PrintInputData() const
314+{
315+ OP_LOGD("MaxPoolGradWithArgmaxV3BaseTiling", "[MaxPoolGradWithArgmaxV3] PrintInputData start running");
316+ 
317+ std::ostringstream info;
318+ info << "inputData.hPad: " << inputData.hPad << std::endl;
319+ info << "inputData.wPad: " << inputData.wPad << std::endl;
320+ info << "inputData.hKernel: " << inputData.hKernel << std::endl;
321+ info << "inputData.wKernel: " << inputData.wKernel << std::endl;
322+ info << "inputData.hStride: " << inputData.hStride << std::endl;
323+ info << "inputData.wStride: " << inputData.wStride << std::endl;
324+ info << "inputData.hDilation: " << inputData.hDilation << std::endl;
325+ info << "inputData.wDilation: " << inputData.wDilation << std::endl;
326+ info << "inputData.ceilMode: " << inputData.ceilMode << std::endl;
327+ info << "inputData.inputDtype: " << inputData.inputDtype << std::endl;
328+ info << "inputData.indexDtype: " << inputData.indexDtype << std::endl;
329+ info << "inputData.inputFormat: " << inputData.inputFormat << std::endl;
330+ info << "inputData.nGrad: " << inputData.nGrad << std::endl;
331+ info << "inputData.cGrad: " << inputData.cGrad << std::endl;
332+ info << "inputData.hGrad: " << inputData.hGrad << std::endl;
333+ info << "inputData.wGrad: " << inputData.wGrad << std::endl;
334+ info << "inputData.nX: " << inputData.nX << std::endl;
335+ info << "inputData.cX: " << inputData.cX << std::endl;
336+ info << "inputData.hX: " << inputData.hX << std::endl;
337+ info << "inputData.wX: " << inputData.wX << std::endl;
338+ info << "inputData.isInt32Meet: " << inputData.isInt32Meet << std::endl;
339+ 
340+ OP_LOGI("MaxPoolGradWithArgmaxV3", "%s", info.str().c_str());
341+}
342+ 
343+bool MaxPoolGradWithArgmaxV3BaseTiling::IsCapable()
344+{
345+ return true;
346+}
347+ 
348+ge::graphStatus MaxPoolGradWithArgmaxV3BaseTiling::DoOpTiling()
349+{
350+ return ge::GRAPH_SUCCESS;
351+}
352+ 
353+ge::graphStatus MaxPoolGradWithArgmaxV3BaseTiling::DoLibApiTiling()
354+{
355+ return ge::GRAPH_SUCCESS;
356+}
357+ 
358+uint64_t MaxPoolGradWithArgmaxV3BaseTiling::GetTilingKey() const
359+{
360+ return 0;
361+}
362+ 
363+ge::graphStatus MaxPoolGradWithArgmaxV3BaseTiling::GetWorkspaceSize()
364+{
365+ auto sys_workspace = WS_SYS_SIZE;
366+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
367+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
368+ currentWorkspace[0] = sys_workspace;
369+ 
370+ return ge::GRAPH_SUCCESS;
371+}
372+ 
373+ge::graphStatus MaxPoolGradWithArgmaxV3BaseTiling::PostTiling()
374+{
375+ return ge::GRAPH_SUCCESS;
376+}
377+} // namespace optiling
Apooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_tiling_base.h+95-0
@@ -0,0 +1,95 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_tiling_base.h
13+ * \brief
14+ */
15+ 
16+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_V3_TILING_BASE_H_
CANN-robot
CANN-robotCANN-robot2025年12月29日

代码结构与可维护性: 头文件保护宏名称中存在拼写错误：'AGRMAX'应为'ARGMAX'。这可能导致头文件保护失效，因为其他文件中可能使用正确的拼写。

问题类型: 代码结构与可维护性 文件路径: pooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_tiling_base.h 行号: 16 问题代码:

#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_V3_TILING_BASE_H_
#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_V3_TILING_BASE_H_

修改建议:

将保护宏中的'AGRMAX'更正为'ARGMAX'，确保与文件名和功能描述一致。

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

likedislike
17+#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_GRAD_WITH_AGRMAX_V3_TILING_BASE_H_
18+ 
19+#include "register/op_def_registry.h"
20+#include "tiling/tiling_api.h"
21+#include "tiling_base/tiling_base.h"
22+#include "util/math_util.h"
23+#include "op_common/op_host/util/platform_util.h"
24+ 
25+using namespace std;
26+ 
27+namespace optiling {
28+ 
29+BEGIN_TILING_DATA_DEF(MaxPoolGradWithArgmaxV3TilingData)
30+TILING_DATA_FIELD_DEF(uint64_t, nc);
31+TILING_DATA_FIELD_DEF(uint64_t, hx);
32+TILING_DATA_FIELD_DEF(uint64_t, wx);
33+TILING_DATA_FIELD_DEF(uint64_t, kh);
34+END_TILING_DATA_DEF;
35+ 
36+REGISTER_TILING_DATA_CLASS(MaxPoolGradWithArgmaxV3, MaxPoolGradWithArgmaxV3TilingData);
37+struct MaxPoolGradWithArgmaxV3InputInfo {
38+ int64_t hPad{0};
39+ int64_t wPad{0};
40+ int64_t hStride{1};
41+ int64_t wStride{1};
42+ int64_t hKernel{1};
43+ int64_t wKernel{1};
44+ int64_t hDilation{1};
45+ int64_t wDilation{1};
46+ int64_t nX{1};
47+ int64_t cX{1};
48+ int64_t hX{1};
49+ int64_t wX{1};
50+ int64_t nGrad{1};
51+ int64_t cGrad{1};
52+ int64_t hGrad{1};
53+ int64_t wGrad{1};
54+ bool ceilMode{false};
55+ ge::DataType inputDtype{ge::DataType::DT_FLOAT};
56+ ge::DataType indexDtype{ge::DataType::DT_INT32};
57+ ge::Format inputFormat{ge::Format::FORMAT_NCHW};
58+ int64_t isInt32Meet{1};
59+};
60+ 
61+struct MaxPoolGradWithArgmaxV3HardwareInfo {
62+ int64_t coreNum{0};
63+ int64_t ubSize{0};
64+};
65+ 
66+struct MaxPoolGradWithArgmaxV3CompileInfo {
67+ uint64_t coreNum;
68+ uint64_t ubSize;
69+};
70+ 
71+class MaxPoolGradWithArgmaxV3BaseTiling : public Ops::NN::Optiling::TilingBaseClass {
72+public:
73+ explicit MaxPoolGradWithArgmaxV3BaseTiling(gert::TilingContext* context) : TilingBaseClass(context)
74+ {}
75+ 
76+ ~MaxPoolGradWithArgmaxV3BaseTiling() override
77+ {}
78+ 
79+protected:
80+ bool IsCapable() override;
81+ ge::graphStatus GetPlatformInfo() override;
82+ ge::graphStatus GetShapeAttrsInfo() override;
83+ ge::graphStatus DoOpTiling() override;
84+ ge::graphStatus DoLibApiTiling() override;
85+ uint64_t GetTilingKey() const override;
86+ ge::graphStatus GetWorkspaceSize() override;
87+ ge::graphStatus PostTiling() override;
88+ void PrintInputData() const;
89+ 
90+public:
CANN-robot
CANN-robotCANN-robot2025年12月29日

代码结构与可维护性: 类MaxPoolGradWithArgmaxV3BaseTiling将inputData和hardwareData声明为public成员变量，这破坏了封装性。外部代码可以直接修改这些内部状态，可能导致对象处于不一致的状态。

问题类型: 代码结构与可维护性 文件路径: pooling/max_pool_grad_with_argmax_v3/op_host/arch35/max_pool_grad_with_argmax_v3_tiling_base.h 行号: 90 问题代码:

public:
    MaxPoolGradWithArgmaxV3InputInfo inputData;
    MaxPoolGradWithArgmaxV3HardwareInfo hardwareData;

修改建议:

将这些成员变量改为private或protected，并提供必要的getter/setter函数来控制访问。如果确实需要公开访问，请添加注释说明理由。

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

likedislike
91+ MaxPoolGradWithArgmaxV3InputInfo inputData;
92+ MaxPoolGradWithArgmaxV3HardwareInfo hardwareData;
93+};
94+} // namespace optiling
95+#endif
Apooling/max_pool_grad_with_argmax_v3/op_host/config/ascend910_95/max_pool_grad_with_argmax_v3_binary.json+485-0
@@ -0,0 +1,485 @@
1+{
2+ "op_type": "MaxPoolGradWithArgmaxV3",
3+ "op_list": [
4+ {
5+ "bin_filename": "MaxPoolGradWithArgmaxV3_float32_int32",
6+ "inputs": [
7+ {
8+ "name": "x",
9+ "index": 0,
10+ "dtype": "float32",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [-2],
14+ "format_match_mode": "FormatAgnostic"
15+ },
16+ {
17+ "name": "grad",
18+ "index": 0,
19+ "dtype": "float32",
20+ "format": "ND",
21+ "paramType": "required",
22+ "shape": [-2],
23+ "format_match_mode": "FormatAgnostic"
24+ },
25+ {
26+ "name": "argmax",
27+ "index": 0,
28+ "dtype": "int32",
29+ "format": "ND",
30+ "paramType": "required",
31+ "shape": [-2],
32+ "format_match_mode": "FormatAgnostic"
33+ }
34+ ],
35+ "outputs": [
36+ {
37+ "name": "y",
38+ "index": 0,
39+ "dtype": "float32",
40+ "format": "ND",
41+ "paramType": "required",
42+ "shape": [-2],
43+ "format_match_mode": "FormatAgnostic"
44+ }
45+ ],
46+ "attrs": [
47+ {
48+ "name": "ksize",
49+ "dtype": "list_int",
50+ "value": []
51+ },
52+ {
53+ "name": "strides",
54+ "dtype": "list_int",
55+ "value": []
56+ },
57+ {
58+ "name": "pads",
59+ "dtype": "list_int",
60+ "value": []
61+ },
62+ {
63+ "name": "dtype",
64+ "dtype": "int",
65+ "value": 3
66+ },
67+ {
68+ "name": "dilation",
69+ "dtype": "list_int",
70+ "value": []
71+ },
72+ {
73+ "name": "ceil_mode",
74+ "dtype": "bool",
75+ "value": null
76+ },
77+ {
78+ "name": "data_format",
79+ "dtype": "string",
80+ "value": null
81+ }
82+ ]
83+ },
84+ {
85+ "bin_filename": "MaxPoolGradWithArgmaxV3_float16_int32",
86+ "inputs": [
87+ {
88+ "name": "x",
89+ "index": 0,
90+ "dtype": "float16",
91+ "format": "ND",
92+ "paramType": "required",
93+ "shape": [-2],
94+ "format_match_mode": "FormatAgnostic"
95+ },
96+ {
97+ "name": "grad",
98+ "index": 0,
99+ "dtype": "float16",
100+ "format": "ND",
101+ "paramType": "required",
102+ "shape": [-2],
103+ "format_match_mode": "FormatAgnostic"
104+ },
105+ {
106+ "name": "argmax",
107+ "index": 0,
108+ "dtype": "int32",
109+ "format": "ND",
110+ "paramType": "required",
111+ "shape": [-2],
112+ "format_match_mode": "FormatAgnostic"
113+ }
114+ ],
115+ "outputs": [
116+ {
117+ "name": "y",
118+ "index": 0,
119+ "dtype": "float16",
120+ "format": "ND",
121+ "paramType": "required",
122+ "shape": [-2],
123+ "format_match_mode": "FormatAgnostic"
124+ }
125+ ],
126+ "attrs": [
127+ {
128+ "name": "ksize",
129+ "dtype": "list_int",
130+ "value": []
131+ },
132+ {
133+ "name": "strides",
134+ "dtype": "list_int",
135+ "value": []
136+ },
137+ {
138+ "name": "pads",
139+ "dtype": "list_int",
140+ "value": []
141+ },
142+ {
143+ "name": "dtype",
144+ "dtype": "int",
145+ "value": 3
146+ },
147+ {
148+ "name": "dilation",
149+ "dtype": "list_int",
150+ "value": []
151+ },
152+ {
153+ "name": "ceil_mode",
154+ "dtype": "bool",
155+ "value": null
156+ },
157+ {
158+ "name": "data_format",
159+ "dtype": "string",
160+ "value": null
161+ }
162+ ]
163+ },
164+ {
165+ "bin_filename": "MaxPoolGradWithArgmaxV3_bfloat16_int32",
166+ "inputs": [
167+ {
168+ "name": "x",
169+ "index": 0,
170+ "dtype": "bfloat16",
171+ "format": "ND",
172+ "paramType": "required",
173+ "shape": [-2],
174+ "format_match_mode": "FormatAgnostic"
175+ },
176+ {
177+ "name": "grad",
178+ "index": 0,
179+ "dtype": "bfloat16",
180+ "format": "ND",
181+ "paramType": "required",
182+ "shape": [-2],
183+ "format_match_mode": "FormatAgnostic"
184+ },
185+ {
186+ "name": "argmax",
187+ "index": 0,
188+ "dtype": "int32",
189+ "format": "ND",
190+ "paramType": "required",
191+ "shape": [-2],
192+ "format_match_mode": "FormatAgnostic"
193+ }
194+ ],
195+ "outputs": [
196+ {
197+ "name": "y",
198+ "index": 0,
199+ "dtype": "bfloat16",
200+ "format": "ND",
201+ "paramType": "required",
202+ "shape": [-2],
203+ "format_match_mode": "FormatAgnostic"
204+ }
205+ ],
206+ "attrs": [
207+ {
208+ "name": "ksize",
209+ "dtype": "list_int",
210+ "value": []
211+ },
212+ {
213+ "name": "strides",
214+ "dtype": "list_int",
215+ "value": []
216+ },
217+ {
218+ "name": "pads",
219+ "dtype": "list_int",
220+ "value": []
221+ },
222+ {
223+ "name": "dtype",
224+ "dtype": "int",
225+ "value": 3
226+ },
227+ {
228+ "name": "dilation",
229+ "dtype": "list_int",
230+ "value": []
231+ },
232+ {
233+ "name": "ceil_mode",
234+ "dtype": "bool",
235+ "value": null
236+ },
237+ {
238+ "name": "data_format",
239+ "dtype": "string",
240+ "value": null
241+ }
242+ ]
243+ },
244+ {
245+ "bin_filename": "MaxPoolGradWithArgmaxV3_float32_int64",
246+ "inputs": [
247+ {
248+ "name": "x",
249+ "index": 0,
250+ "dtype": "float32",
251+ "format": "ND",
252+ "paramType": "required",
253+ "shape": [-2],
254+ "format_match_mode": "FormatAgnostic"
255+ },
256+ {
257+ "name": "grad",
258+ "index": 0,
259+ "dtype": "float32",
260+ "format": "ND",
261+ "paramType": "required",
262+ "shape": [-2],
263+ "format_match_mode": "FormatAgnostic"
264+ },
265+ {
266+ "name": "argmax",
267+ "index": 0,
268+ "dtype": "int64",
269+ "format": "ND",
270+ "paramType": "required",
271+ "shape": [-2],
272+ "format_match_mode": "FormatAgnostic"
273+ }
274+ ],
275+ "outputs": [
276+ {
277+ "name": "y",
278+ "index": 0,
279+ "dtype": "float32",
280+ "format": "ND",
281+ "paramType": "required",
282+ "shape": [-2],
283+ "format_match_mode": "FormatAgnostic"
284+ }
285+ ],
286+ "attrs": [
287+ {
288+ "name": "ksize",
289+ "dtype": "list_int",
290+ "value": []
291+ },
292+ {
293+ "name": "strides",
294+ "dtype": "list_int",
295+ "value": []
296+ },
297+ {
298+ "name": "pads",
299+ "dtype": "list_int",
300+ "value": []
301+ },
302+ {
303+ "name": "dtype",
304+ "dtype": "int",
305+ "value": 9
306+ },
307+ {
308+ "name": "dilation",
309+ "dtype": "list_int",
310+ "value": []
311+ },
312+ {
313+ "name": "ceil_mode",
314+ "dtype": "bool",
315+ "value": null
316+ },
317+ {
318+ "name": "data_format",
319+ "dtype": "string",
320+ "value": null
321+ }
322+ ]
323+ },
324+ {
325+ "bin_filename": "MaxPoolGradWithArgmaxV3_float16_int64",
326+ "inputs": [
327+ {
328+ "name": "x",
329+ "index": 0,
330+ "dtype": "float16",
331+ "format": "ND",
332+ "paramType": "required",
333+ "shape": [-2],
334+ "format_match_mode": "FormatAgnostic"
335+ },
336+ {
337+ "name": "grad",
338+ "index": 0,
339+ "dtype": "float16",
340+ "format": "ND",
341+ "paramType": "required",
342+ "shape": [-2],
343+ "format_match_mode": "FormatAgnostic"
344+ },
345+ {
346+ "name": "argmax",
347+ "index": 0,
348+ "dtype": "int64",
349+ "format": "ND",
350+ "paramType": "required",
351+ "shape": [-2],
352+ "format_match_mode": "FormatAgnostic"
353+ }
354+ ],
355+ "outputs": [
356+ {
357+ "name": "y",
358+ "index": 0,
359+ "dtype": "float16",
360+ "format": "ND",
361+ "paramType": "required",
362+ "shape": [-2],
363+ "format_match_mode": "FormatAgnostic"
364+ }
365+ ],
366+ "attrs": [
367+ {
368+ "name": "ksize",
369+ "dtype": "list_int",
370+ "value": []
371+ },
372+ {
373+ "name": "strides",
374+ "dtype": "list_int",
375+ "value": []
376+ },
377+ {
378+ "name": "pads",
379+ "dtype": "list_int",
380+ "value": []
381+ },
382+ {
383+ "name": "dtype",
384+ "dtype": "int",
385+ "value": 9
386+ },
387+ {
388+ "name": "dilation",
389+ "dtype": "list_int",
390+ "value": []
391+ },
392+ {
393+ "name": "ceil_mode",
394+ "dtype": "bool",
395+ "value": null
396+ },
397+ {
398+ "name": "data_format",
399+ "dtype": "string",
400+ "value": null
401+ }
402+ ]
403+ },
404+ {
405+ "bin_filename": "MaxPoolGradWithArgmaxV3_bfloat16_int64",
406+ "inputs": [
407+ {
408+ "name": "x",
409+ "index": 0,
410+ "dtype": "bfloat16",
411+ "format": "ND",
412+ "paramType": "required",
413+ "shape": [-2],
414+ "format_match_mode": "FormatAgnostic"
415+ },
416+ {
417+ "name": "grad",
418+ "index": 0,
419+ "dtype": "bfloat16",
420+ "format": "ND",
421+ "paramType": "required",
422+ "shape": [-2],
423+ "format_match_mode": "FormatAgnostic"
424+ },
425+ {
426+ "name": "argmax",
427+ "index": 0,
428+ "dtype": "int64",
429+ "format": "ND",
430+ "paramType": "required",
431+ "shape": [-2],
432+ "format_match_mode": "FormatAgnostic"
433+ }
434+ ],
435+ "outputs": [
436+ {
437+ "name": "y",
438+ "index": 0,
439+ "dtype": "bfloat16",
440+ "format": "ND",
441+ "paramType": "required",
442+ "shape": [-2],
443+ "format_match_mode": "FormatAgnostic"
444+ }
445+ ],
446+ "attrs": [
447+ {
448+ "name": "ksize",
449+ "dtype": "list_int",
450+ "value": []
451+ },
452+ {
453+ "name": "strides",
454+ "dtype": "list_int",
455+ "value": []
456+ },
457+ {
458+ "name": "pads",
459+ "dtype": "list_int",
460+ "value": []
461+ },
462+ {
463+ "name": "dtype",
464+ "dtype": "int",
465+ "value": 9
466+ },
467+ {
468+ "name": "dilation",
469+ "dtype": "list_int",
470+ "value": []
471+ },
472+ {
473+ "name": "ceil_mode",
474+ "dtype": "bool",
475+ "value": null
476+ },
477+ {
478+ "name": "data_format",
479+ "dtype": "string",
480+ "value": null
481+ }
482+ ]
483+ }
484+ ]
485+}
Apooling/max_pool_grad_with_argmax_v3/op_host/config/ascend910_95/max_pool_grad_with_argmax_v3_simplified_key.ini+13-0
@@ -0,0 +1,13 @@
1+; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2+; [某算子]
3+; default=xx
4+; ascendxx=xx
5+; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6+; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7+; 2)如果仅有default配置：各个版本按default配置
8+; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9+; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10+; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11+; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12+[MaxPoolGradWithArgmaxV3]
13+default=0
Apooling/max_pool_grad_with_argmax_v3/op_host/max_pool_grad_with_argmax_v3_def.cpp+74-0
@@ -0,0 +1,74 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3.cpp
13+ * \brief imply for max_pool_grad_with_argmax
14+ */
15+ 
16+#include <cstdint>
17+#include "register/op_def_registry.h"
18+ 
19+namespace ops {
20+constexpr int32_t DTYPE_INT32 = 3;
21+class MaxPoolGradWithArgmaxV3 : public OpDef {
22+public:
23+ explicit MaxPoolGradWithArgmaxV3(const char* name) : OpDef(name)
24+ {
25+ this->Input("x")
26+ .ParamType(REQUIRED)
27+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
28+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
29+ .UnknownShapeFormat(
30+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
31+ .AutoContiguous();
32+ this->Input("grad")
33+ .ParamType(REQUIRED)
34+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
35+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
36+ .UnknownShapeFormat(
37+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
38+ .AutoContiguous();
39+ this->Input("argmax")
40+ .ParamType(REQUIRED)
41+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
42+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
43+ .UnknownShapeFormat(
44+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
45+ .AutoContiguous();
46+ this->Output("y")
47+ .ParamType(REQUIRED)
48+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
49+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
50+ .UnknownShapeFormat(
51+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
52+ .AutoContiguous();
53+ this->Attr("ksize").AttrType(REQUIRED).ListInt();
54+ this->Attr("strides").AttrType(REQUIRED).ListInt();
55+ this->Attr("pads").AttrType(REQUIRED).ListInt();
56+ this->Attr("dtype").AttrType(OPTIONAL).Int(DTYPE_INT32);
57+ this->Attr("dilation").AttrType(OPTIONAL).ListInt({1, 1});
58+ this->Attr("ceil_mode").AttrType(OPTIONAL).Bool(false);
59+ this->Attr("data_format").AttrType(OPTIONAL).String("NCHW");
60+ 
61+ OpAICoreConfig aiCoreConfig;
62+ aiCoreConfig.DynamicCompileStaticFlag(true)
63+ .DynamicFormatFlag(false)
64+ .DynamicRankSupportFlag(true)
65+ .DynamicShapeSupportFlag(true)
66+ .NeedCheckSupportFlag(false)
67+ .PrecisionReduceFlag(true)
68+ .ExtendCfgInfo("opFile.value", "max_pool_grad_with_argmax_v3");
69+ this->AICore().AddConfig("ascend910_95", aiCoreConfig);
70+ }
71+};
72+ 
73+OP_ADD(MaxPoolGradWithArgmaxV3);
74+} // namespace ops
Apooling/max_pool_grad_with_argmax_v3/op_host/max_pool_grad_with_argmax_v3_infershape.cpp+135-0
@@ -0,0 +1,135 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_infershape.cpp
13+ * \brief
14+ */
15+#include <string>
16+#include "graph/utils/type_utils.h"
17+#include "runtime/infer_shape_context.h"
18+#include "register/op_impl_registry.h"
19+#include "log/log.h"
20+#include "util/shape_util.h"
21+#include "util/math_util.h"
22+ 
23+using namespace ge;
24+namespace ops {
25+static constexpr size_t ATTR_INDEX_KSIZE = 0;
26+static constexpr size_t ATTR_INDEX_STRIDES = 1;
27+static constexpr size_t ATTR_INDEX_PADS = 2;
28+static constexpr size_t ATTR_INDEX_DILATION = 4;
29+static constexpr size_t ATTR_INDEX_CEIL_MODE = 5;
30+static constexpr size_t ATTR_INDEX_DATA_FORMAT = 6;
31+static constexpr size_t ATTR_LIST_SHAPE_SIZE = 2;
32+static constexpr int64_t UNKNOWN_DIM_VALUE_ = -1LL;
33+ 
34+inline ge::graphStatus SetAllUnknownDim(const int64_t rank, gert::Shape* output_shape)
35+{
36+ OP_CHECK_IF(
37+ output_shape == nullptr, OP_LOGD("SetAllUnknownDim", "the output_shape is nullptr, return unsuccess"),
38+ return ge::GRAPH_FAILED);
39+ output_shape->SetDimNum(rank);
40+ for (int64_t i = 0; i < rank; ++i) {
41+ output_shape->SetDim(i, UNKNOWN_DIM_VALUE_);
42+ }
43+ OP_LOGD("SetAllUnknownDim", "set all dim = -1, output = %s", Ops::Base::ToString(*output_shape).c_str());
44+ return ge::GRAPH_SUCCESS;
45+}
46+ 
47+ge::graphStatus InferShapeForMaxPoolGradWithArgmaxV3(gert::InferShapeContext* context)
48+{
49+ if (context == nullptr) {
50+ return GRAPH_FAILED;
51+ }
52+ 
53+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPoolGradWithArgmaxV3 infershape running");
54+ auto xDesc = context->GetInputDesc(0);
55+ OP_CHECK_NULL_WITH_CONTEXT(context, xDesc);
56+ auto xOriFormat = xDesc->GetOriginFormat();
57+ OP_CHECK_IF(
58+ xOriFormat != FORMAT_ND && xOriFormat != FORMAT_NCHW && xOriFormat != FORMAT_NHWC,
59+ OP_LOGE(context->GetNodeName(), "format only supports ND, NCHW, NHWC"), return GRAPH_FAILED);
60+ 
61+ auto attrs = context->GetAttrs();
62+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
63+ 
64+ auto ksize = attrs->GetAttrPointer<gert::ContinuousVector>(ATTR_INDEX_KSIZE);
65+ OP_CHECK_NULL_WITH_CONTEXT(context, ksize);
66+ OP_CHECK_IF(
67+ ksize->GetSize() != ATTR_LIST_SHAPE_SIZE,
68+ OP_LOGE(context->GetNodeName(), "Length of ksize %lu must be 2!", ksize->GetSize()), return GRAPH_FAILED);
69+ 
70+ auto strides = attrs->GetAttrPointer<gert::ContinuousVector>(ATTR_INDEX_STRIDES);
71+ OP_CHECK_NULL_WITH_CONTEXT(context, strides);
72+ OP_CHECK_IF(
73+ strides->GetSize() != ATTR_LIST_SHAPE_SIZE,
74+ OP_LOGE(context->GetNodeName(), "Length of strides %lu must be 2!", strides->GetSize()), return GRAPH_FAILED);
75+ 
76+ auto pads = attrs->GetAttrPointer<gert::ContinuousVector>(ATTR_INDEX_PADS);
77+ OP_CHECK_NULL_WITH_CONTEXT(context, pads);
78+ OP_CHECK_IF(
79+ pads->GetSize() != ATTR_LIST_SHAPE_SIZE,
80+ OP_LOGE(context->GetNodeName(), "Length of pads %lu must be 2!", pads->GetSize()), return GRAPH_FAILED);
81+ 
82+ auto dilation = attrs->GetAttrPointer<gert::ContinuousVector>(ATTR_INDEX_DILATION);
83+ OP_CHECK_NULL_WITH_CONTEXT(context, dilation);
84+ OP_CHECK_IF(
85+ dilation->GetSize() != ATTR_LIST_SHAPE_SIZE,
86+ OP_LOGE(context->GetNodeName(), "Length of dilation %lu must be 2!", dilation->GetSize()), return GRAPH_FAILED);
87+ 
88+ auto ceil_mode = attrs->GetAttrPointer<bool>(ATTR_INDEX_CEIL_MODE);
89+ OP_CHECK_NULL_WITH_CONTEXT(context, ceil_mode);
90+ 
91+ const char* data_format = attrs->GetAttrPointer<char>(ATTR_INDEX_DATA_FORMAT); // todo 是否能匹配上
92+ OP_CHECK_NULL_WITH_CONTEXT(context, data_format);
93+ 
94+ const gert::Shape* xShape = context->GetInputShape(0);
95+ OP_CHECK_NULL_WITH_CONTEXT(context, xShape);
96+ 
97+ const gert::Shape* gradShape = context->GetInputShape(1);
98+ OP_CHECK_NULL_WITH_CONTEXT(context, gradShape);
99+ const gert::Shape* argmaxShape = context->GetInputShape(2);
100+ OP_CHECK_NULL_WITH_CONTEXT(context, argmaxShape);
101+ gert::Shape* yShape = context->GetOutputShape(0);
102+ OP_CHECK_NULL_WITH_CONTEXT(context, yShape);
103+ size_t xDimNum = xShape->GetDimNum();
104+ if (Ops::Base::IsUnknownShape(*xShape) || Ops::Base::IsUnknownShape(*gradShape) || Ops::Base::IsUnknownShape(*argmaxShape)) {
105+ SetAllUnknownDim(xDimNum, yShape);
106+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPoolGradWithArgmaxV3 infershape handle unknown rank or shape.");
107+ return ge::GRAPH_SUCCESS;
108+ }
109+ 
110+ if (Ops::Base::IsUnknownRank(*xShape)) {
111+ Ops::Base::SetUnknownRank(*yShape);
112+ return GRAPH_SUCCESS;
113+ }
114+ yShape->SetDimNum(xDimNum);
115+ *yShape = *xShape;
116+ 
117+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPoolGradWithArgmaxV3 infershape run success.");
118+ return GRAPH_SUCCESS;
119+}
120+ 
121+static ge::graphStatus InferDataTypeForMaxPoolGradWithArgmaxV3(gert::InferDataTypeContext* context)
122+{
123+ if (context == nullptr) {
124+ return GRAPH_FAILED;
125+ }
126+ 
127+ const ge::DataType xDtype = context->GetInputDataType(0);
128+ context->SetOutputDataType(0, xDtype);
129+ return GRAPH_SUCCESS;
130+}
131+ 
132+IMPL_OP_INFERSHAPE(MaxPoolGradWithArgmaxV3)
133+ .InferShape(InferShapeForMaxPoolGradWithArgmaxV3)
134+ .InferDataType(InferDataTypeForMaxPoolGradWithArgmaxV3);
135+} // namespace ops
Apooling/max_pool_grad_with_argmax_v3/op_kernel/arch35/max_pool_grad_with_argmax_v3_base.h+129-0
@@ -0,0 +1,129 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_base.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_GRAD_WITH_ARGMAX_V3_BASE_H_
17+#define MAX_POOL_GRAD_WITH_ARGMAX_V3_BASE_H_
18+ 
19+using namespace AscendC;
20+constexpr uint32_t BUFFER_NUM = 2;
21+constexpr int64_t DOUBLE = 2;
22+constexpr uint32_t HELP_BUFFER = 1024;
23+ 
24+constexpr uint32_t INDEX_TWO = 2;
25+constexpr uint32_t INDEX_THREE = 3;
26+using computeType = float;
27+ 
28+constexpr AscendC::MicroAPI::CastTrait castTraitT1ComputeType = {
29+ AscendC::MicroAPI::RegLayout::ZERO,
30+ AscendC::MicroAPI::SatMode::UNKNOWN,
31+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
32+ AscendC::RoundMode::UNKNOWN,
33+};
34+ 
35+constexpr AscendC::MicroAPI::CastTrait castTraitI64I32 = {
36+ AscendC::MicroAPI::RegLayout::ZERO,
37+ AscendC::MicroAPI::SatMode::NO_SAT,
38+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
39+ AscendC::RoundMode::CAST_ROUND,
40+};
41+ 
42+constexpr AscendC::MicroAPI::CastTrait castTraitU32U16 = {
43+ AscendC::MicroAPI::RegLayout::ZERO,
44+ AscendC::MicroAPI::SatMode::NO_SAT,
45+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
46+ AscendC::RoundMode::CAST_RINT,
47+};
48+ 
49+__aicore__ inline int64_t PStart(int64_t index, int64_t pad, int64_t kernel, int64_t dilation, int64_t stride)
50+{
51+ return (index + pad < (kernel - 1) * dilation + 1) ? 0 : (index + pad - ((kernel - 1) * dilation + 1)) / stride + 1;
52+};
53+__aicore__ inline int64_t PEnd(int64_t index, int64_t pad, int64_t stride, int64_t pooledSize)
54+{
55+ return (index + pad) / stride + 1 < pooledSize ? (index + pad) / stride + 1 : pooledSize;
56+};
57+ 
58+template <typename T2, typename T3>
59+__aicore__ inline MicroAPI::MaskReg GenT2Mask(uint32_t& maskCount)
60+{
61+ MicroAPI::MaskReg reg;
62+ if constexpr (std::is_same<T3, int32_t>::value && std::is_same<T2, int64_t>::value) {
63+ reg = AscendC::MicroAPI::UpdateMask<T2, AscendC::MicroAPI::RegTraitNumTwo>(maskCount);
64+ } else {
65+ reg = AscendC::MicroAPI::UpdateMask<T2>(maskCount);
66+ }
67+ return reg;
68+}
69+ 
70+__aicore__ inline void FilterMask(
71+ MicroAPI::MaskReg& preg, MicroAPI::RegTensor<int32_t>& hIndexReg, MicroAPI::RegTensor<int32_t>& wIndexReg,
72+ MicroAPI::RegTensor<int32_t>& zeroConstReg, MicroAPI::RegTensor<int32_t>& wMaxReg,
73+ MicroAPI::RegTensor<int32_t>& hMaxReg)
74+{
75+ AscendC::MicroAPI::MaskReg gtMask = AscendC::MicroAPI::CreateMask<int32_t, AscendC::MicroAPI::MaskPattern::ALL>();
76+ AscendC::MicroAPI::MaskReg allMask = AscendC::MicroAPI::CreateMask<int32_t, AscendC::MicroAPI::MaskPattern::ALL>();
77+ AscendC::MicroAPI::Compare<int32_t, CMPMODE::GE>(gtMask, hIndexReg, zeroConstReg, gtMask);
78+ AscendC::MicroAPI::Compare<int32_t, CMPMODE::GT>(gtMask, hMaxReg, hIndexReg, gtMask);
79+ 
80+ AscendC::MicroAPI::Compare<int32_t, CMPMODE::GE>(gtMask, wIndexReg, zeroConstReg, gtMask);
81+ AscendC::MicroAPI::Compare<int32_t, CMPMODE::GT>(gtMask, wMaxReg, wIndexReg, gtMask);
82+ AscendC::MicroAPI::MaskAnd(preg, preg, gtMask, allMask);
83+}
84+ 
85+template <typename T>
86+__aicore__ inline void GradientAcc(
87+ __local_mem__ computeType* yAddr, MicroAPI::RegTensor<computeType>& gradReg, MicroAPI::RegTensor<T>& argmaxReg,
88+ MicroAPI::MaskReg& pregArgmax)
89+{
90+ AscendC::MicroAPI::RegTensor<computeType> scatterAccResReg;
91+ AscendC::MicroAPI::DataCopyGather(
92+ scatterAccResReg, yAddr, (AscendC::MicroAPI::RegTensor<uint32_t>&)argmaxReg, pregArgmax);
93+ AscendC::MicroAPI::Add(scatterAccResReg, scatterAccResReg, gradReg, pregArgmax);
94+ AscendC::MicroAPI::DataCopyScatter(
95+ yAddr, scatterAccResReg, (AscendC::MicroAPI::RegTensor<uint32_t>&)argmaxReg, pregArgmax);
96+}
97+ 
98+template <typename T1, typename T2, typename T3>
99+__aicore__ inline void GetConCurrentInput(
100+ MicroAPI::RegTensor<T3>& argmaxReg, MicroAPI::RegTensor<computeType>& gradReg, __local_mem__ T1* gradAddr,
101+ __local_mem__ T2* argmaxAddr, MicroAPI::RegTensor<uint32_t>& parallelRegIndex, MicroAPI::MaskReg& pregT1,
102+ MicroAPI::MaskReg& pregT2)
103+{
104+ if constexpr (std::negation<std::is_same<T1, float>>::value) {
105+ AscendC::MicroAPI::RegTensor<T1> gradRegT1;
106+ AscendC::MicroAPI::RegTensor<uint16_t> parallelRegIndexU16;
107+ AscendC::MicroAPI::MaskReg allMaskU32 =
108+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
109+ AscendC::MicroAPI::Cast<uint16_t, uint32_t, castTraitU32U16>(parallelRegIndexU16, parallelRegIndex, allMaskU32);
110+ AscendC::MicroAPI::Pack(parallelRegIndexU16, (AscendC::MicroAPI::RegTensor<int32_t>&)parallelRegIndexU16);
111+ AscendC::MicroAPI::DataCopyGather(gradRegT1, gradAddr, parallelRegIndexU16, pregT1);
112+ AscendC::MicroAPI::UnPack(
113+ (AscendC::MicroAPI::RegTensor<uint32_t>&)gradRegT1, (AscendC::MicroAPI::RegTensor<uint16_t>&)gradRegT1);
114+ AscendC::MicroAPI::Cast<computeType, T1, castTraitT1ComputeType>(gradReg, gradRegT1, allMaskU32);
115+ } else {
116+ AscendC::MicroAPI::DataCopyGather(gradReg, gradAddr, parallelRegIndex, pregT1);
117+ }
118+ 
119+ if constexpr (std::is_same<T3, int32_t>::value && std::is_same<T2, int32_t>::value) {
120+ AscendC::MicroAPI::DataCopyGather(argmaxReg, argmaxAddr, parallelRegIndex, pregT2);
121+ } else if constexpr (std::is_same<T3, int32_t>::value && std::is_same<T2, int64_t>::value) {
122+ AscendC::MicroAPI::RegTensor<T2, AscendC::MicroAPI::RegTraitNumTwo> argmaxRegTwo;
123+ AscendC::MicroAPI::DataCopyGather(argmaxRegTwo, argmaxAddr, parallelRegIndex, pregT2);
124+ argmaxReg = (AscendC::MicroAPI::RegTensor<T3>&)argmaxRegTwo.reg[0];
125+ } else if constexpr (std::is_same<T3, int64_t>::value && std::is_same<T2, int64_t>::value) {
126+ AscendC::MicroAPI::DataCopyGather(argmaxReg, argmaxAddr, parallelRegIndex, pregT2);
127+ }
128+}
129+#endif // MAX_POOL_GRAD_WITH_ARGMAX_V3_BASE_H_
Apooling/max_pool_grad_with_argmax_v3/op_kernel/arch35/max_pool_grad_with_argmax_v3_nchw_kernel.h+1265-0
@@ -0,0 +1,1265 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nchw_kernel.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_GRAD_WITH_ARGMAX_V3_NCHW_KERNEL_H_
17+#define MAX_POOL_GRAD_WITH_ARGMAX_V3_NCHW_KERNEL_H_
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+#include "max_pool_grad_with_argmax_v3_base.h"
23+ 
24+namespace MaxPoolGradWithArgmaxV3NCHWNameSpace
25+{
26+ 
27+// argmaxReg输入为T3 输出为int32_t
28+template <typename T, const uint32_t IS_MUL_NC = 0>
29+__aicore__ inline void IndexConvNchw(MicroAPI::RegTensor<T>& argmaxReg, MicroAPI::RegTensor<int32_t>& hIndexReg,
30+ MicroAPI::RegTensor<int32_t>& wIndexReg, MicroAPI::RegTensor<T>& wOutputConstReg,
31+ int64_t curHIndex, int64_t curWIndex, int32_t wOutputAligned,
32+ int32_t highOutputOffset, int32_t highOutputPlaneActual, int32_t highArgmaxPlaneActual)
33+{
34+ AscendC::MicroAPI::RegTensor<T> hTmpIndexReg;
35+ AscendC::MicroAPI::RegTensor<T> wTmpIndexReg;
36+ AscendC::MicroAPI::RegTensor<T> tmpReg;
37+ AscendC::MicroAPI::MaskReg allMask = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
38+ AscendC::MicroAPI::MaskReg allMaskU32 =
39+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
40+ 
41+ AscendC::MicroAPI::Div(hTmpIndexReg, argmaxReg, wOutputConstReg, allMask);
42+ if constexpr (std::is_same<T, int64_t>::value) {
43+ AscendC::MicroAPI::Adds(tmpReg, hTmpIndexReg, T(-curHIndex), allMask);
44+ AscendC::MicroAPI::Cast<int32_t, int64_t, castTraitI64I32>(hIndexReg, tmpReg, allMask);
45+ AscendC::MicroAPI::Pack((AscendC::MicroAPI::RegTensor<uint32_t>&)hIndexReg,
46+ (AscendC::MicroAPI::RegTensor<int64_t>&)hIndexReg);
47+ } else {
48+ AscendC::MicroAPI::Adds(hIndexReg, hTmpIndexReg, T(-curHIndex), allMask);
49+ }
50+ 
51+ AscendC::MicroAPI::Mul(wTmpIndexReg, hTmpIndexReg, wOutputConstReg, allMask);
52+ AscendC::MicroAPI::Sub(wTmpIndexReg, argmaxReg, wTmpIndexReg, allMask);
53+ if constexpr (std::is_same<T, int64_t>::value) {
54+ AscendC::MicroAPI::Adds(tmpReg, wTmpIndexReg, T(-curWIndex), allMask);
55+ AscendC::MicroAPI::Cast<int32_t, int64_t, castTraitI64I32>(wIndexReg, tmpReg, allMask);
56+ AscendC::MicroAPI::Pack((AscendC::MicroAPI::RegTensor<uint32_t>&)wIndexReg,
57+ (AscendC::MicroAPI::RegTensor<int64_t>&)wIndexReg);
58+ } else {
59+ AscendC::MicroAPI::Adds(wIndexReg, wTmpIndexReg, T(-curWIndex), allMask);
60+ }
61+ 
62+ AscendC::MicroAPI::Muls((AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg, hIndexReg, T(wOutputAligned),
63+ allMaskU32);
64+ AscendC::MicroAPI::Add((AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg,
65+ (AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg, wIndexReg, allMaskU32);
66+ 
67+ AscendC::MicroAPI::Adds((AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg,
68+ (AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg, highOutputOffset, allMaskU32);
69+ 
70+ if constexpr (IS_MUL_NC == 1) {
71+ AscendC::MicroAPI::RegTensor<int32_t> highIncReg;
72+ AscendC::MicroAPI::Arange(highIncReg, 0);
73+ AscendC::MicroAPI::RegTensor<int32_t> constReg;
74+ AscendC::MicroAPI::Duplicate(constReg, highArgmaxPlaneActual);
75+ AscendC::MicroAPI::Div(highIncReg, highIncReg, constReg, allMaskU32);
76+ AscendC::MicroAPI::Muls(highIncReg, highIncReg, highOutputPlaneActual,
77+ allMaskU32);
78+ AscendC::MicroAPI::Add((AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg,
79+ (AscendC::MicroAPI::RegTensor<int32_t>&)argmaxReg, highIncReg,
80+ allMaskU32);
81+ }
82+}
83+ 
84+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
85+__aicore__ inline void DoSingleNCNchw(__local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr,
86+ __local_mem__ T2* argmaxAddr, MicroAPI::RegTensor<uint32_t>& parallelRegIndex,
87+ uint32_t argmaxMaskCount, MicroAPI::RegTensor<T3>& wOutputConstReg, int64_t curHIndex,
88+ int64_t curWIndex, int32_t wOutputAligned, int32_t highOutputOffset,
89+ MicroAPI::RegTensor<int32_t>& zeroConstReg, MicroAPI::RegTensor<int32_t>& wMaxReg,
90+ MicroAPI::RegTensor<int32_t>& hMaxReg)
91+{
92+ AscendC::MicroAPI::RegTensor<computeType> gradReg;
93+ AscendC::MicroAPI::RegTensor<T3> argmaxReg;
94+ // 相对索引
95+ AscendC::MicroAPI::RegTensor<int32_t> hIndexReg;
96+ AscendC::MicroAPI::RegTensor<int32_t> wIndexReg;
97+ 
98+ uint32_t maskT1 = argmaxMaskCount;
99+ uint32_t maskT2 = argmaxMaskCount;
100+ AscendC::MicroAPI::MaskReg pregT1 = AscendC::MicroAPI::UpdateMask<T1>(maskT1);
101+ AscendC::MicroAPI::MaskReg pregT2 = GenT2Mask<T2, T3>(maskT2);
102+ GetConCurrentInput<T1, T2, T3>(argmaxReg, gradReg, gradAddr, argmaxAddr, parallelRegIndex, pregT1, pregT2);
103+ IndexConvNchw<T3>(argmaxReg, hIndexReg, wIndexReg, wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
104+ highOutputOffset,0 ,0);
105+ uint32_t argmaxMask = argmaxMaskCount;
106+ AscendC::MicroAPI::MaskReg pregArgmax = AscendC::MicroAPI::UpdateMask<int32_t>(argmaxMask);
107+ if constexpr (IS_CHECK_RANGE == 1) {
108+ FilterMask(pregArgmax, hIndexReg, wIndexReg, zeroConstReg, wMaxReg, hMaxReg);
109+ }
110+ 
111+ GradientAcc<T3>(yAddr, gradReg, argmaxReg, pregArgmax);
112+}
113+ 
114+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
115+__aicore__ inline void DoMulNCNchw(__local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr,
116+ __local_mem__ T2* argmaxAddr, MicroAPI::RegTensor<uint32_t>& parallelRegIndex,
117+ uint32_t argmaxMaskCount, MicroAPI::RegTensor<T3>& wOutputConstReg, int64_t curHIndex,
118+ int64_t curWIndex, int32_t wOutputAligned, int32_t highOutputOffset,
119+ MicroAPI::RegTensor<int32_t>& zeroConstReg, MicroAPI::RegTensor<int32_t>& wMaxReg,
120+ MicroAPI::RegTensor<int32_t>& hMaxReg, int32_t highOutputPlaneActual, int32_t highArgmaxPlaneActual)
121+{
122+ AscendC::MicroAPI::RegTensor<computeType> gradReg;
123+ AscendC::MicroAPI::RegTensor<T3> argmaxReg;
124+ // 相对索引
125+ AscendC::MicroAPI::RegTensor<int32_t> hIndexReg;
126+ AscendC::MicroAPI::RegTensor<int32_t> wIndexReg;
127+ 
128+ uint32_t maskT1 = argmaxMaskCount;
129+ uint32_t maskT2 = argmaxMaskCount;
130+ AscendC::MicroAPI::MaskReg pregT1 = AscendC::MicroAPI::UpdateMask<T1>(maskT1);
131+ AscendC::MicroAPI::MaskReg pregT2 = GenT2Mask<T2, T3>(maskT2);
132+ GetConCurrentInput<T1, T2, T3>(argmaxReg, gradReg, gradAddr, argmaxAddr, parallelRegIndex, pregT1, pregT2);
133+ IndexConvNchw<T3, 1>(argmaxReg, hIndexReg, wIndexReg, wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
134+ highOutputOffset, highOutputPlaneActual, highArgmaxPlaneActual);
135+ uint32_t argmaxMask = argmaxMaskCount;
136+ AscendC::MicroAPI::MaskReg pregArgmax = AscendC::MicroAPI::UpdateMask<int32_t>(argmaxMask);
137+ if constexpr (IS_CHECK_RANGE == 1) {
138+ FilterMask(pregArgmax, hIndexReg, wIndexReg, zeroConstReg, wMaxReg, hMaxReg);
139+ }
140+ 
141+ GradientAcc<T3>(yAddr, gradReg, argmaxReg, pregArgmax);
142+}
143+ 
144+template <typename T>
145+__aicore__ inline void GenInitial1DIndices(MicroAPI::RegTensor<T>& indexReg, int64_t colGenRate)
146+{
147+ AscendC::MicroAPI::Arange(indexReg, 0);
148+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
149+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(colGenRate), preg);
150+}
151+ 
152+template <typename T>
153+__aicore__ inline void GenInitial2DIndices(MicroAPI::RegTensor<T>& indexReg, int64_t colGenRate, int64_t rowGenRate,
154+ int64_t colNumAligned, int64_t fullBatchColNum)
155+{
156+ AscendC::MicroAPI::Arange(indexReg, 0);
157+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg;
158+ AscendC::MicroAPI::RegTensor<T> segmentIncReg;
159+ AscendC::MicroAPI::RegTensor<T> constReg;
160+ AscendC::MicroAPI::Duplicate(constReg, T(fullBatchColNum));
161+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
162+ 
163+ AscendC::MicroAPI::Div(segmentScalarReg, indexReg, constReg, preg);
164+ 
165+ AscendC::MicroAPI::Muls(segmentIncReg, segmentScalarReg, T(fullBatchColNum), preg);
166+ AscendC::MicroAPI::Sub(segmentIncReg, indexReg, segmentIncReg, preg);
167+ 
168+ AscendC::MicroAPI::Muls(segmentIncReg, segmentIncReg, T(colGenRate), preg);
169+ AscendC::MicroAPI::Muls(segmentScalarReg, segmentScalarReg, T(rowGenRate * colNumAligned), preg);
170+ AscendC::MicroAPI::Add(indexReg, segmentScalarReg, segmentIncReg, preg);
171+}
172+ 
173+template <typename T>
174+__aicore__ inline void Gen2DIndexOne(MicroAPI::RegTensor<T>& indexReg, int64_t rowGenRate, int64_t colNumAligned)
175+{
176+ AscendC::MicroAPI::Arange(indexReg, 0);
177+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
178+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rowGenRate * colNumAligned), preg);
179+}
180+ 
181+template <typename T>
182+__aicore__ inline void GenInitial3DIndices(MicroAPI::RegTensor<T>& indexReg, int64_t colGenRate, int64_t rowGenRate,
183+ int64_t colNumAligned, int64_t fullBatchColNum, int64_t fullBatchRowNum,
184+ int64_t rowNumCount)
185+{
186+ AscendC::MicroAPI::Arange(indexReg, 0);
187+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg;
188+ AscendC::MicroAPI::RegTensor<T> segmentIncReg;
189+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg2;
190+ AscendC::MicroAPI::RegTensor<T> segmentIncReg2;
191+ AscendC::MicroAPI::RegTensor<T> constReg;
192+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
193+ 
194+ AscendC::MicroAPI::Duplicate(constReg, T(fullBatchColNum * fullBatchRowNum));
195+ AscendC::MicroAPI::Div(segmentScalarReg, indexReg, constReg, preg);
196+ AscendC::MicroAPI::Muls(segmentIncReg, segmentScalarReg, T(fullBatchColNum * fullBatchRowNum), preg);
197+ AscendC::MicroAPI::Sub(segmentIncReg, indexReg, segmentIncReg, preg);
198+ 
199+ AscendC::MicroAPI::Muls(segmentScalarReg, segmentScalarReg, T(rowNumCount * colNumAligned), preg);
200+ 
201+ AscendC::MicroAPI::Duplicate(constReg, T(fullBatchColNum));
202+ AscendC::MicroAPI::Div(segmentScalarReg2, segmentIncReg, constReg, preg);
203+ AscendC::MicroAPI::Muls(segmentIncReg2, segmentScalarReg2, T(fullBatchColNum), preg);
204+ AscendC::MicroAPI::Sub(segmentIncReg2, segmentIncReg, segmentIncReg2, preg);
205+ AscendC::MicroAPI::Muls(segmentIncReg2, segmentIncReg2, colGenRate, preg);
206+ 
207+ AscendC::MicroAPI::Muls(segmentScalarReg2, segmentScalarReg2, T(rowGenRate * colNumAligned), preg);
208+ 
209+ AscendC::MicroAPI::Add(indexReg, segmentIncReg2, segmentScalarReg2, preg);
210+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg, preg);
211+}
212+ 
213+template <typename T>
214+__aicore__ inline void Gen3DIndexOne(MicroAPI::RegTensor<T>& indexReg, int64_t rowGenRate, int64_t colNumAligned,
215+ int64_t fullBatchRowNum, int64_t rowNumCount)
216+{
217+ AscendC::MicroAPI::Arange(indexReg, 0);
218+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg;
219+ AscendC::MicroAPI::RegTensor<T> segmentIncReg;
220+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg2;
221+ AscendC::MicroAPI::RegTensor<T> segmentIncReg2;
222+ AscendC::MicroAPI::RegTensor<T> constReg;
223+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
224+ 
225+ AscendC::MicroAPI::Duplicate(constReg, T(1 * fullBatchRowNum));
226+ AscendC::MicroAPI::Div(segmentScalarReg, indexReg, constReg, preg);
227+ AscendC::MicroAPI::Muls(segmentIncReg, segmentScalarReg, T(1 * fullBatchRowNum), preg);
228+ AscendC::MicroAPI::Sub(segmentIncReg, indexReg, segmentIncReg, preg);
229+ 
230+ AscendC::MicroAPI::Muls(segmentScalarReg, segmentScalarReg, T(rowNumCount * colNumAligned), preg);
231+ 
232+ AscendC::MicroAPI::Muls(segmentIncReg, segmentIncReg, T(rowGenRate * colNumAligned), preg);
233+ 
234+ AscendC::MicroAPI::Add(indexReg, segmentIncReg, segmentScalarReg, preg);
235+}
236+ 
237+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
238+class MaxPoolGradWithArgmaxV3NCHWKernel
239+{
240+public:
241+ __aicore__ inline MaxPoolGradWithArgmaxV3NCHWKernel(void){};
242+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR grad, GM_ADDR argmax, GM_ADDR y, TPipe& pipeIn,
243+ const MaxPoolGradWithArgmaxV3NCHWTilingData& tilingData);
244+ __aicore__ inline void ParseTilingData(const MaxPoolGradWithArgmaxV3NCHWTilingData& tilingData);
245+ __aicore__ inline void Process();
246+ __aicore__ inline void ScalarCompute(int64_t loopNum);
247+ __aicore__ inline void ProcessPerLoop();
248+ __aicore__ inline void CopyIn();
249+ __aicore__ inline void Compute();
250+ __aicore__ inline void singleLineProcessVF(__local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr,
251+ __local_mem__ T2* argmaxAddr);
252+ __aicore__ inline void multipleLineProcessVF1(__local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr,
253+ __local_mem__ T2* argmaxAddr);
254+ __aicore__ inline void multipleLineProcessVF2(__local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr,
255+ __local_mem__ T2* argmaxAddr, __local_mem__ uint32_t* helpAddr);
256+ __aicore__ inline void multipleLineProcessVF2Int64(__local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr,
257+ __local_mem__ T2* argmaxAddr, __local_mem__ uint32_t* helpAddr);
258+ __aicore__ inline void ProcessNoArgmaxBlock();
259+ __aicore__ inline void CopyOut();
260+ 
261+ TPipe pipe_;
262+ TQue<QuePosition::VECIN, BUFFER_NUM> gradQue_;
263+ TQue<QuePosition::VECIN, BUFFER_NUM> argmaxQue_;
264+ TQue<QuePosition::VECOUT, BUFFER_NUM> outputQue_;
265+ TBuf<QuePosition::VECCALC> helpBuf_;
266+ 
267+ GlobalTensor<T1> gradGm_;
268+ GlobalTensor<T1> yGm_;
269+ GlobalTensor<T2> argmaxGm_;
270+ 
271+ uint32_t blockIdx_ = 0;
272+ 
273+ int64_t hArgmax_ = 1;
274+ int64_t wArgmax_ = 1;
275+ 
276+ int64_t hOutput_ = 1;
277+ int64_t wOutput_ = 1;
278+ 
279+ int64_t kernelH_ = 1;
280+ int64_t kernelW_ = 1;
281+ 
282+ int64_t strideH_ = 1;
283+ int64_t strideW_ = 1;
284+ 
285+ int64_t padH_ = 0;
286+ int64_t padW_ = 0;
287+ 
288+ int64_t dilationH_ = 1;
289+ int64_t dilationW_ = 1;
290+ 
291+ int64_t highAxisInner_ = 1;
292+ int64_t highAxisTail_ = 1;
293+ int64_t highAxisOuter_ = 1;
294+ int64_t highAxisActual_ = 1;
295+ 
296+ int64_t hOutputInner_ = 1;
297+ int64_t hOutputTail_ = 1;
298+ int64_t hOutputOuter_ = 1;
299+ int64_t hOutputActual_ = 1;
300+ 
301+ int64_t wOutputInner_ = 1;
302+ int64_t wOutputTail_ = 1;
303+ int64_t wOutputOuter_ = 1;
304+ int64_t wOutputActual_ = 1;
305+ int64_t wOutputAligned_ = 1;
306+ 
307+ int64_t normalCoreProcessNum_ = 1;
308+ int64_t tailCoreProcessNum_ = 1;
309+ int64_t curCoreProcessNum_ = 1;
310+ int64_t usedCoreNum_ = 1;
311+ 
312+ int64_t outputBufferSize_ = 1;
313+ int64_t gradBufferSize_ = 1;
314+ int64_t argmaxBufferSize_ = 1;
315+ 
316+ int64_t highAxisIndex_ = 0;
317+ int64_t hAxisIndex_ = 0;
318+ int64_t wAxisIndex_ = 0;
319+ 
320+ int64_t hArgmaxActual_ = 0;
321+ int64_t wArgmaxActual_ = 0;
322+ int64_t wArgmaxAligned_ = 0;
323+ 
324+ int64_t highAxisArgmaxOffset_ = 0;
325+ int64_t hAxisArgmaxOffset_ = 0;
326+ int64_t wAxisArgmaxOffset_ = 0;
327+ 
328+ int64_t argmaxPlaneSize_ = 1;
329+ 
330+ int64_t hProBatchSize_ = 1;
331+ int64_t wProBatchSize_ = 1;
332+ int64_t curHProBatchSize_ = 1;
333+ int64_t curWProBatchSize_ = 1;
334+ constexpr static int32_t BLOCK_SIZE = platform::GetUbBlockSize();
335+ constexpr static int32_t V_REG_SIZE = platform::GetVRegSize();
336+ 
337+ constexpr static int64_t MAX_DATA_NUM_IN_ONE_BLOCK =
338+ BLOCK_SIZE / sizeof(T1) >= BLOCK_SIZE / sizeof(T2) ? BLOCK_SIZE / sizeof(T1) : BLOCK_SIZE / sizeof(T2);
339+ constexpr static int64_t VREG_LENGTH_DATA_NUM_T2 = platform::GetVRegSize() / sizeof(T2);
340+};
341+ 
342+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
343+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::ParseTilingData(
344+ const MaxPoolGradWithArgmaxV3NCHWTilingData& tilingData)
345+{
346+ hArgmax_ = tilingData.hArgmax;
347+ wArgmax_ = tilingData.wArgmax;
348+ 
349+ hOutput_ = tilingData.hOutput;
350+ wOutput_ = tilingData.wOutput;
351+ 
352+ kernelH_ = tilingData.hKernel;
353+ kernelW_ = tilingData.wKernel;
354+ 
355+ strideH_ = tilingData.hStride;
356+ strideW_ = tilingData.wStride;
357+ 
358+ padH_ = tilingData.padH;
359+ padW_ = tilingData.padW;
360+ 
361+ dilationH_ = tilingData.dilationH;
362+ dilationW_ = tilingData.dilationW;
363+ 
364+ highAxisInner_ = tilingData.highAxisInner;
365+ highAxisTail_ = tilingData.highAxisTail;
366+ highAxisOuter_ = tilingData.highAxisOuter;
367+ 
368+ hOutputInner_ = tilingData.hOutputInner;
369+ hOutputTail_ = tilingData.hOutputTail;
370+ hOutputOuter_ = tilingData.hOutputOuter;
371+ 
372+ wOutputInner_ = tilingData.wOutputInner;
373+ wOutputTail_ = tilingData.wOutputTail;
374+ wOutputOuter_ = tilingData.wOutputOuter;
375+ 
376+ normalCoreProcessNum_ = tilingData.normalCoreProcessNum;
377+ tailCoreProcessNum_ = tilingData.tailCoreProcessNum;
378+ usedCoreNum_ = tilingData.usedCoreNum;
379+ 
380+ outputBufferSize_ = tilingData.outputBufferSize;
381+ gradBufferSize_ = tilingData.gradBufferSize;
382+ argmaxBufferSize_ = tilingData.argmaxBufferSize;
383+ 
384+ hProBatchSize_ = tilingData.hProBatchSize;
385+ wProBatchSize_ = tilingData.wProBatchSize;
386+}
387+ 
388+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
389+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::Init(
390+ GM_ADDR x, GM_ADDR grad, GM_ADDR argmax, GM_ADDR y, TPipe& pipeIn,
391+ const MaxPoolGradWithArgmaxV3NCHWTilingData& tilingData)
392+{
393+ ParseTilingData(tilingData);
394+ 
395+ blockIdx_ = GetBlockIdx();
396+ argmaxPlaneSize_ = hArgmax_ * wArgmax_;
397+ if (blockIdx_ >= usedCoreNum_) {
398+ return;
399+ }
400+ 
401+ curCoreProcessNum_ = (blockIdx_ + 1 == usedCoreNum_) ? tailCoreProcessNum_ : normalCoreProcessNum_;
402+ gradGm_.SetGlobalBuffer((__gm__ T1*)grad);
403+ argmaxGm_.SetGlobalBuffer((__gm__ T2*)argmax);
404+ yGm_.SetGlobalBuffer((__gm__ T1*)y);
405+ 
406+ pipe_ = pipeIn;
407+ pipe_.InitBuffer(outputQue_, BUFFER_NUM, outputBufferSize_);
408+ pipe_.InitBuffer(gradQue_, BUFFER_NUM, gradBufferSize_);
409+ pipe_.InitBuffer(argmaxQue_, BUFFER_NUM, argmaxBufferSize_);
410+ pipe_.InitBuffer(helpBuf_, HELP_BUFFER);
411+}
412+ 
413+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
414+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::ScalarCompute(int64_t loopNum)
415+{
416+ int64_t baseBlockIdx = blockIdx_ * normalCoreProcessNum_ + loopNum;
417+ highAxisIndex_ = baseBlockIdx / (hOutputOuter_ * wOutputOuter_);
418+ highAxisActual_ = highAxisIndex_ == (highAxisOuter_ - 1) ? highAxisTail_ : highAxisInner_;
419+ 
420+ int64_t tempTail = baseBlockIdx % (hOutputOuter_ * wOutputOuter_);
421+ hAxisIndex_ = tempTail / wOutputOuter_;
422+ hOutputActual_ = hAxisIndex_ == (hOutputOuter_ - 1) ? hOutputTail_ : hOutputInner_;
423+ 
424+ wAxisIndex_ = tempTail % wOutputOuter_;
425+ wOutputActual_ = wAxisIndex_ == (wOutputOuter_ - 1) ? wOutputTail_ : wOutputInner_;
426+ wOutputAligned_ =
427+ (wOutputActual_ + MAX_DATA_NUM_IN_ONE_BLOCK - 1) / MAX_DATA_NUM_IN_ONE_BLOCK * MAX_DATA_NUM_IN_ONE_BLOCK;
428+ 
429+ int64_t hArgmaxActualStart = PStart(hAxisIndex_ * hOutputInner_, padH_, kernelH_, dilationH_, strideH_);
430+ int64_t hArgmaxActualEnd = PEnd(hAxisIndex_ * hOutputInner_ + hOutputActual_ - 1, padH_, strideH_, hArgmax_);
431+ int64_t wArgmaxActualStart = PStart(wAxisIndex_ * wOutputInner_, padW_, kernelW_, dilationW_, strideW_);
432+ int64_t wArgmaxActualEnd = PEnd(wAxisIndex_ * wOutputInner_ + wOutputActual_ - 1, padW_, strideW_, wArgmax_);
433+ wArgmaxActual_ = wArgmaxActualEnd - wArgmaxActualStart;
434+ wArgmaxAligned_ =
435+ (wArgmaxActual_ + MAX_DATA_NUM_IN_ONE_BLOCK - 1) / MAX_DATA_NUM_IN_ONE_BLOCK * MAX_DATA_NUM_IN_ONE_BLOCK;
436+ hArgmaxActual_ = hArgmaxActualEnd - hArgmaxActualStart;
437+ 
438+ curHProBatchSize_ = hProBatchSize_ > hArgmaxActual_ ? hArgmaxActual_ : hProBatchSize_;
439+ curWProBatchSize_ = wProBatchSize_ > wArgmaxActual_ ? wArgmaxActual_ : wProBatchSize_;
440+ 
441+ highAxisArgmaxOffset_ = highAxisIndex_ * highAxisInner_ * argmaxPlaneSize_;
442+ hAxisArgmaxOffset_ = hArgmaxActualStart * wArgmax_;
443+ wAxisArgmaxOffset_ = wArgmaxActualStart;
444+}
445+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
446+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::Process()
447+{
448+ if (blockIdx_ >= usedCoreNum_) {
449+ return;
450+ }
451+ 
452+ for (int64_t loopNum = 0; loopNum < curCoreProcessNum_; loopNum++) {
453+ ScalarCompute(loopNum);
454+ ProcessPerLoop();
455+ }
456+}
457+ 
458+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
459+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::Compute()
460+{
461+ uint32_t calCount = outputBufferSize_ / sizeof(computeType);
462+ LocalTensor<computeType> yLocal = outputQue_.AllocTensor<computeType>();
463+ Duplicate(yLocal, computeType(0), calCount);
464+ 
465+ LocalTensor<T1> gradLocal = gradQue_.DeQue<T1>();
466+ LocalTensor<T2> argmaxLocal = argmaxQue_.DeQue<T2>();
467+ 
468+ __local_mem__ computeType* yAddr = (__local_mem__ computeType*)yLocal.GetPhyAddr();
469+ __local_mem__ T1* gradAddr = (__local_mem__ T1*)gradLocal.GetPhyAddr();
470+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
471+ 
472+ uint32_t wConcurrentCount = wArgmaxActual_ / curWProBatchSize_;
473+ uint32_t hConcurrentCount = hArgmaxActual_ / curHProBatchSize_;
474+ if (wConcurrentCount * DOUBLE * sizeof(T2) > V_REG_SIZE) {
475+ singleLineProcessVF(yAddr, gradAddr, argmaxAddr);
476+ } else if (wConcurrentCount * hConcurrentCount * DOUBLE * sizeof(T2) > V_REG_SIZE) {
477+ multipleLineProcessVF1(yAddr, gradAddr, argmaxAddr); // HW 并发处理
478+ } else {
479+ // NCHW 并发处理
480+ LocalTensor<uint32_t> helpTensor = helpBuf_.Get<uint32_t>();
481+ __local_mem__ uint32_t* helpAddr = (__local_mem__ uint32_t*)helpTensor.GetPhyAddr();
482+ if constexpr (std::is_same<T3, int64_t>::value) {
483+ multipleLineProcessVF2Int64(yAddr, gradAddr, argmaxAddr, helpAddr);
484+ } else {
485+ multipleLineProcessVF2(yAddr, gradAddr, argmaxAddr, helpAddr);
486+ }
487+ }
488+ 
489+ if constexpr (std::negation<std::is_same<T1, float>>::value) {
490+ Cast(yLocal.ReinterpretCast<T1>(), yLocal, RoundMode::CAST_RINT, calCount);
491+ }
492+ 
493+ outputQue_.EnQue(yLocal);
494+ gradQue_.FreeTensor(gradLocal);
495+ argmaxQue_.FreeTensor(argmaxLocal);
496+}
497+ 
498+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
499+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::ProcessNoArgmaxBlock()
500+{
501+ uint32_t calcCount = static_cast<uint32_t>(outputBufferSize_) / sizeof(T1);
502+ LocalTensor<T1> yLocal = outputQue_.AllocTensor<T1>();
503+ Duplicate(yLocal, T1(0), calcCount);
504+ outputQue_.EnQue(yLocal);
505+ CopyOut();
506+ return;
507+}
508+ 
509+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
510+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::ProcessPerLoop()
511+{
512+ if (hArgmaxActual_ <= 0 || wArgmaxActual_ <= 0) {
513+ ProcessNoArgmaxBlock(); // ceilMode为false时，最后的尾块可能是这种情况
514+ return;
515+ }
516+ 
517+ CopyIn();
518+ Compute();
519+ CopyOut();
520+}
521+ 
522+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
523+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::CopyIn()
524+{
525+ LocalTensor<T1> gradLocal = gradQue_.AllocTensor<T1>();
526+ LocalTensor<T2> argmaxLocal = argmaxQue_.AllocTensor<T2>();
527+ 
528+ int64_t argmaxGmOffset = highAxisArgmaxOffset_ + hAxisArgmaxOffset_ + wAxisArgmaxOffset_;
529+ DataCopyPadExtParams<T1> paramsT1 = {false, 0, 0, 0};
530+ LoopModeParams loopModeParamsT1;
531+ loopModeParamsT1.loop1Size = highAxisActual_;
532+ loopModeParamsT1.loop2Size = 1;
533+ loopModeParamsT1.loop1SrcStride = argmaxPlaneSize_ * sizeof(T1);
534+ loopModeParamsT1.loop2SrcStride = 0;
535+ loopModeParamsT1.loop1DstStride = hArgmaxActual_ * wArgmaxAligned_ * sizeof(T1);
536+ loopModeParamsT1.loop2DstStride = 0;
537+ 
538+ SetLoopModePara(loopModeParamsT1, DataCopyMVType::OUT_TO_UB);
539+ DataCopyExtParams copyOutParamT1 = {
540+ static_cast<uint16_t>(hArgmaxActual_),
541+ static_cast<uint32_t>(wArgmaxActual_ * sizeof(T1)),
542+ static_cast<uint32_t>((wArgmax_ - wArgmaxActual_) * sizeof(T1)),
543+ static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
544+ 
545+ DataCopyPad(gradLocal, gradGm_[argmaxGmOffset], copyOutParamT1, paramsT1);
546+
547+ DataCopyPadExtParams<T2> paramsT2 = {false, 0, 0, 0};
548+
549+ LoopModeParams loopModeParamsT2;
550+ loopModeParamsT2.loop1Size = highAxisActual_;
551+ loopModeParamsT2.loop2Size = 1;
552+ loopModeParamsT2.loop1SrcStride = argmaxPlaneSize_ * sizeof(T2);
553+ loopModeParamsT2.loop2SrcStride = 0;
554+ loopModeParamsT2.loop1DstStride = hArgmaxActual_ * wArgmaxAligned_ * sizeof(T2);
555+ loopModeParamsT2.loop2DstStride = 0;
556+ 
557+ uint32_t dstStrideT2 = (wArgmaxAligned_ - wArgmaxActual_) * sizeof(T2) / BLOCK_SIZE;
558+ SetLoopModePara(loopModeParamsT2, DataCopyMVType::OUT_TO_UB);
559+ DataCopyExtParams copyOutParamT2 = {
560+ static_cast<uint16_t>(hArgmaxActual_),
561+ static_cast<uint32_t>(wArgmaxActual_ * sizeof(T2)),
562+ static_cast<uint32_t>((wArgmax_ - wArgmaxActual_) * sizeof(T2)),
563+ static_cast<uint32_t>(dstStrideT2), static_cast<uint32_t>(0)};
564+ 
565+ DataCopyPad(argmaxLocal, argmaxGm_[argmaxGmOffset], copyOutParamT2, paramsT2);
566+ ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
567+ gradQue_.EnQue(gradLocal);
568+ argmaxQue_.EnQue(argmaxLocal);
569+}
570+ 
571+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
572+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::singleLineProcessVF(
573+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr)
574+{
575+ int64_t wOutput = wOutput_;
576+ int64_t wOutputActual = wOutputActual_;
577+ int64_t wOutputAligned = wOutputAligned_;
578+ int64_t hOutputActual = hOutputActual_;
579+ uint16_t highAxisActual = static_cast<uint16_t>(highAxisActual_);
580+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
581+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
582+ int64_t wArgmaxActual = wArgmaxActual_;
583+ int64_t wArgmaxAligned = wArgmaxAligned_;
584+ uint16_t hArgmaxActual = hArgmaxActual_;
585+ 
586+ uint16_t hProBatchSize = curHProBatchSize_;
587+ uint16_t wProBatchSize = curWProBatchSize_;
588+ 
589+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
590+ 
591+ uint16_t computeSizeT2 = V_REG_SIZE / sizeof(T2);
592+ 
593+ uint16_t repeatimes = wFullBatchCount / computeSizeT2;
594+ uint16_t wRemain = wArgmaxActual - repeatimes * wProBatchSize * computeSizeT2;
595+ 
596+ uint32_t wRemainBatchCount = wRemain / wProBatchSize;
597+ uint16_t wRemainTail = wRemain % wProBatchSize;
598+ 
599+ uint32_t one = 1;
600+ uint32_t all = computeSizeT2;
601+ 
602+ __VEC_SCOPE__
603+ {
604+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
605+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
606+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
607+ if constexpr (IS_CHECK_RANGE == 1) {
608+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
609+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
610+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
611+ }
612+ 
613+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
614+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
615+ 
616+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegIndex;
617+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
618+ 
619+ AscendC::MicroAPI::MaskReg allMaskU32 =
620+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
621+ 
622+ GenInitial1DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegIndex, wProBatchSize);
623+ 
624+ for (uint16_t highIdx = 0; highIdx < highAxisActual; ++highIdx) {
625+ uint32_t highArgmaxOffset = highIdx * hArgmaxActual * wArgmaxAligned;
626+ uint32_t highOutputOffset = highIdx * hOutputActual * wOutputAligned;
627+ for (uint16_t hIdx = 0; hIdx < hArgmaxActual; hIdx++) {
628+ for (uint16_t wRepeatIdx = 0; wRepeatIdx < repeatimes; wRepeatIdx++) {
629+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
630+ uint32_t offset = (wBatchIdx + wRepeatIdx * computeSizeT2 * wProBatchSize +
631+ hIdx * wArgmaxAligned + highArgmaxOffset);
632+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, offset, allMaskU32);
633+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
634+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, all, wOutputConstReg, curHIndex, curWIndex,
635+ wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
636+ }
637+ }
638+ // 尾段整batch 用不满mask
639+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
640+ T2 offset = (wBatchIdx + repeatimes * computeSizeT2 * wProBatchSize + hIdx * wArgmaxAligned +
641+ highArgmaxOffset);
642+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, offset, allMaskU32);
643+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
644+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, wRemainBatchCount, wOutputConstReg, curHIndex,
645+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
646+ }
647+ 
648+ // 尾段零散点
649+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
650+ T2 offset = (wBatchIdx + wRemainBatchCount * wProBatchSize +
651+ repeatimes * computeSizeT2 * wProBatchSize + hIdx * wArgmaxAligned + highArgmaxOffset);
652+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, offset, allMaskU32);
653+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, one,
654+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
655+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
656+ }
657+ }
658+ }
659+ }
660+}
661+ 
662+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
663+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::multipleLineProcessVF1(
664+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr)
665+{
666+ int64_t wOutput = wOutput_;
667+ int64_t wOutputActual = wOutputActual_;
668+ int64_t wOutputAligned = wOutputAligned_;
669+ int64_t hOutputActual = hOutputActual_;
670+ uint16_t highAxisActual = static_cast<uint16_t>(highAxisActual_);
671+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
672+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
673+ int64_t wArgmaxAligned = wArgmaxAligned_;
674+ int64_t wArgmaxActual = wArgmaxActual_;
675+ uint16_t hArgmaxActual = hArgmaxActual_;
676+ 
677+ uint16_t hProBatchSize = curHProBatchSize_;
678+ uint16_t wProBatchSize = curWProBatchSize_;
679+ 
680+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
681+ uint16_t hFullBatchCount = hArgmaxActual / hProBatchSize;
682+ uint16_t wRemainTail = wArgmaxActual % wProBatchSize;
683+ 
684+ uint16_t hConcurrentCount = V_REG_SIZE / (wFullBatchCount * sizeof(T2));
685+ 
686+ uint16_t blockConcurrentCount = hFullBatchCount / hConcurrentCount;
687+ uint16_t hRemain = hArgmaxActual - blockConcurrentCount * hConcurrentCount * hProBatchSize;
688+ 
689+ uint16_t hRemainBatchCount = hRemain / hProBatchSize;
690+ uint16_t hRemainTail = hRemain - hRemainBatchCount * hProBatchSize;
691+ 
692+ uint32_t blockOne = 1 * hConcurrentCount;
693+ uint32_t remainBatchOne = 1 * hRemainBatchCount;
694+ uint32_t remainTailOne = 1;
695+ uint32_t maskBlock = wFullBatchCount * hConcurrentCount;
696+ uint32_t maskRemainBatch = wFullBatchCount * hRemainBatchCount;
697+ uint32_t maskRemainTail = wFullBatchCount;
698+ __VEC_SCOPE__
699+ {
700+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
701+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
702+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
703+ if constexpr (IS_CHECK_RANGE == 1) {
704+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
705+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
706+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
707+ }
708+ 
709+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
710+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
711+ 
712+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegIndex;
713+ AscendC::MicroAPI::RegTensor<uint32_t> initialRegIndexOne;
714+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
715+ 
716+ AscendC::MicroAPI::MaskReg allMaskU32 =
717+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
718+ GenInitial2DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegIndex, wProBatchSize, hProBatchSize,
719+ wArgmaxAligned, wFullBatchCount);
720+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initialRegIndexOne, hProBatchSize, wArgmaxAligned);
721+ 
722+ for (uint16_t highIdx = 0; highIdx < highAxisActual; ++highIdx) {
723+ uint32_t highArgmaxOffset = highIdx * hArgmaxActual * wArgmaxAligned;
724+ uint32_t highOutputOffset = highIdx * hOutputActual * wOutputAligned;
725+ for (uint16_t hIdx = 0; hIdx < blockConcurrentCount; hIdx++) {
726+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
727+ // 整batch
728+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
729+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned +
730+ hIdx * wArgmaxAligned * hProBatchSize * hConcurrentCount + highArgmaxOffset);
731+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, offset, allMaskU32);
732+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
733+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, maskBlock, wOutputConstReg, curHIndex,
734+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
735+ }
736+ 
737+ // 尾段零散点
738+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
739+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
740+ hIdx * wArgmaxAligned * hProBatchSize * hConcurrentCount + highArgmaxOffset);
741+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndexOne, offset, allMaskU32);
742+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
743+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, blockOne, wOutputConstReg, curHIndex,
744+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
745+ }
746+ }
747+ }
748+ 
749+ // 尾行 完整hProBatch
750+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
751+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
752+ T2 offset =
753+ (wBatchIdx + hProBatchIdx * wArgmaxAligned +
754+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
755+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, offset, allMaskU32);
756+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
757+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, maskRemainBatch, wOutputConstReg, curHIndex,
758+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
759+ }
760+ 
761+ // 尾段零散点
762+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
763+ T2 offset =
764+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
765+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
766+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndexOne, offset, allMaskU32);
767+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
768+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, remainBatchOne, wOutputConstReg, curHIndex,
769+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
770+ }
771+ }
772+ // 尾行 零散hProBatch
773+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
774+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
775+ T2 offset =
776+ (wBatchIdx + hProBatchIdx * wArgmaxAligned +
777+ hRemainBatchCount * hProBatchSize * wArgmaxAligned +
778+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
779+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndex, offset, allMaskU32);
780+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
781+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, maskRemainTail, wOutputConstReg, curHIndex,
782+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
783+ }
784+ 
785+ // 尾段零散点
786+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
787+ T2 offset =
788+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
789+ hRemainBatchCount * hProBatchSize * wArgmaxAligned +
790+ blockConcurrentCount * hConcurrentCount * hProBatchSize * wArgmaxAligned + highArgmaxOffset);
791+ AscendC::MicroAPI::Adds(parallelRegIndex, initialRegIndexOne, offset, allMaskU32);
792+ DoSingleNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
793+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, remainTailOne, wOutputConstReg, curHIndex,
794+ curWIndex, wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg);
795+ }
796+ }
797+ }
798+ }
799+}
800+ 
801+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
802+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::multipleLineProcessVF2(
803+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
804+ __local_mem__ uint32_t* helpAddr)
805+{
806+ int64_t wOutput = wOutput_;
807+ int64_t wOutputActual = wOutputActual_;
808+ int64_t wOutputAligned = wOutputAligned_;
809+ int64_t hOutputActual = hOutputActual_;
810+ int32_t highOutputPlaneActual = wOutputAligned * hOutputActual;
811+ int64_t highAxisActual = highAxisActual_;
812+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
813+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
814+ int64_t wArgmaxAligned = wArgmaxAligned_;
815+ int64_t wArgmaxActual = wArgmaxActual_;
816+ uint16_t hArgmaxActual = hArgmaxActual_;
817+ 
818+ uint16_t hProBatchSize = curHProBatchSize_;
819+ uint16_t wProBatchSize = curWProBatchSize_;
820+ 
821+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
822+ uint16_t hFullBatchCount = hArgmaxActual / hProBatchSize;
823+ uint16_t wRemainTail = wArgmaxActual % wProBatchSize;
824+ uint32_t whFullBatchCount = wFullBatchCount * hFullBatchCount;
825+ 
826+ uint16_t highConcurrentCount = V_REG_SIZE / (whFullBatchCount * sizeof(T2));
827+ 
828+ uint16_t highBlockConcurrentCount = highAxisActual / highConcurrentCount;
829+ uint16_t highBlockRemainTail = highAxisActual - highBlockConcurrentCount * highConcurrentCount;
830+ 
831+ uint16_t hRemainTail = hArgmaxActual - hFullBatchCount * hProBatchSize;
832+ 
833+ uint32_t mask0 = highConcurrentCount * whFullBatchCount;
834+ uint32_t mask1 = highConcurrentCount * hFullBatchCount * 1;
835+ uint32_t mask2 = highConcurrentCount * 1 * wFullBatchCount;
836+ uint32_t mask3 = highConcurrentCount * 1 * 1;
837+ uint32_t mask4 = highBlockRemainTail * whFullBatchCount;
838+ uint32_t mask5 = highBlockRemainTail * hFullBatchCount * 1;
839+ uint32_t mask6 = highBlockRemainTail * 1 * wFullBatchCount;
840+ uint32_t mask7 = highBlockRemainTail * 1 * 1;
841+ __VEC_SCOPE__
842+ {
843+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
844+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
845+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
846+ if constexpr (IS_CHECK_RANGE == 1) {
847+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
848+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
849+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
850+ }
851+ 
852+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
853+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
854+ 
855+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
856+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
857+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
858+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
859+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
860+ 
861+ AscendC::MicroAPI::MaskReg allMaskU32 =
862+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
863+ GenInitial3DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegIndex, wProBatchSize, hProBatchSize,
864+ wArgmaxAligned, wFullBatchCount, hFullBatchCount, hArgmaxActual);
865+ Gen3DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegIndexOne, hProBatchSize, wArgmaxAligned,
866+ hFullBatchCount, hArgmaxActual);
867+ 
868+ GenInitial2DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegIndex, wProBatchSize, hArgmaxActual,
869+ wArgmaxAligned, wFullBatchCount);
870+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegIndexOne, hArgmaxActual, wArgmaxAligned);
871+ 
872+ for (uint16_t highBlockIdx = 0; highBlockIdx < highBlockConcurrentCount; ++highBlockIdx) {
873+ uint32_t highArgmaxOffset = highBlockIdx * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
874+ uint32_t highOutputOffset = highBlockIdx * highConcurrentCount * hOutputActual * wOutputAligned;
875+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
876+ // 整batch
877+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
878+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
879+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndex, offset, allMaskU32);
880+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask0,
881+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
882+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
883+ highOutputPlaneActual, whFullBatchCount);
884+ }
885+ 
886+ // 尾段零散点
887+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
888+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
889+ highArgmaxOffset);
890+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndexOne, offset, allMaskU32);
891+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask1,
892+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
893+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
894+ highOutputPlaneActual, hFullBatchCount);
895+ }
896+ }
897+ 
898+ // hRemainTail
899+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
900+ // 整batch
901+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
902+ T2 offset = (wBatchIdx + (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxAligned +
903+ highArgmaxOffset);
904+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndex, offset, allMaskU32);
905+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask2,
906+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
907+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
908+ highOutputPlaneActual, wFullBatchCount);
909+ }
910+ 
911+ // 尾段零散点
912+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
913+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount +
914+ (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
915+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndexOne, offset, allMaskU32);
916+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
917+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask3, wOutputConstReg, curHIndex, curWIndex,
918+ wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual, 1);
919+ }
920+ }
921+ }
922+ 
923+ // highBlockRemainTail
924+ uint32_t highArgmaxOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
925+ uint32_t highOutputOffset = highBlockConcurrentCount * highConcurrentCount * hOutputActual * wOutputAligned;
926+ // 整H batch
927+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
928+ // 整batch
929+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
930+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
931+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndex, offset, allMaskU32);
932+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask4,
933+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
934+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
935+ highOutputPlaneActual, whFullBatchCount);
936+ }
937+ 
938+ // 尾段零散点
939+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
940+ T2 offset =
941+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
942+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndexOne, offset, allMaskU32);
943+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask5,
944+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
945+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
946+ highOutputPlaneActual, hFullBatchCount);
947+ }
948+ }
949+ 
950+ // hRemainTail
951+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
952+ // 整batch
953+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
954+ T2 offset =
955+ (wBatchIdx + (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
956+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndex, offset, allMaskU32);
957+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask6,
958+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
959+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
960+ highOutputPlaneActual, wFullBatchCount);
961+ }
962+ 
963+ // 尾段零散点
964+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
965+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount +
966+ (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
967+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndexOne, offset, allMaskU32);
968+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
969+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask7, wOutputConstReg, curHIndex, curWIndex,
970+ wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual, 1);
971+ }
972+ }
973+ }
974+}
975+ 
976+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
977+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::multipleLineProcessVF2Int64(
978+ __local_mem__ computeType* yAddr, __local_mem__ T1* gradAddr, __local_mem__ T2* argmaxAddr,
979+ __local_mem__ uint32_t* helpAddr)
980+{
981+ int64_t wOutput = wOutput_;
982+ int64_t wOutputActual = wOutputActual_;
983+ int64_t wOutputAligned = wOutputAligned_;
984+ int64_t hOutputActual = hOutputActual_;
985+ int32_t highOutputPlaneActual = wOutputAligned * hOutputActual;
986+ int64_t highAxisActual = highAxisActual_;
987+ int64_t curHIndex = hAxisIndex_ * hOutputInner_;
988+ int64_t curWIndex = wAxisIndex_ * wOutputInner_;
989+ int64_t wArgmaxAligned = wArgmaxAligned_;
990+ int64_t wArgmaxActual = wArgmaxActual_;
991+ uint16_t hArgmaxActual = hArgmaxActual_;
992+ 
993+ uint16_t hProBatchSize = curHProBatchSize_;
994+ uint16_t wProBatchSize = curWProBatchSize_;
995+ 
996+ uint32_t wFullBatchCount = wArgmaxActual / wProBatchSize;
997+ uint16_t hFullBatchCount = hArgmaxActual / hProBatchSize;
998+ uint16_t wRemainTail = wArgmaxActual % wProBatchSize;
999+ uint32_t whFullBatchCount = wFullBatchCount * hFullBatchCount;
1000+ 
1001+ uint16_t highConcurrentCount = V_REG_SIZE / (whFullBatchCount * sizeof(T2));
1002+ 
1003+ uint16_t highBlockConcurrentCount = highAxisActual / highConcurrentCount;
1004+ uint16_t highBlockRemainTail = highAxisActual - highBlockConcurrentCount * highConcurrentCount;
1005+ 
1006+ uint16_t hRemainTail = hArgmaxActual - hFullBatchCount * hProBatchSize;
1007+ 
1008+ uint32_t mask0 = highConcurrentCount * whFullBatchCount;
1009+ uint32_t mask1 = highConcurrentCount * hFullBatchCount * 1;
1010+ uint32_t mask2 = highConcurrentCount * 1 * wFullBatchCount;
1011+ uint32_t mask3 = highConcurrentCount * 1 * 1;
1012+ uint32_t mask4 = highBlockRemainTail * whFullBatchCount;
1013+ uint32_t mask5 = highBlockRemainTail * hFullBatchCount * 1;
1014+ uint32_t mask6 = highBlockRemainTail * 1 * wFullBatchCount;
1015+ uint32_t mask7 = highBlockRemainTail * 1 * 1;
1016+ 
1017+ __VEC_SCOPE__
1018+ {
1019+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
1020+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
1021+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
1022+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
1023+ 
1024+ GenInitial3DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegIndex, wProBatchSize, hProBatchSize,
1025+ wArgmaxAligned, wFullBatchCount, hFullBatchCount, hArgmaxActual);
1026+ Gen3DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initial3DRegIndexOne, hProBatchSize, wArgmaxAligned,
1027+ hFullBatchCount, hArgmaxActual);
1028+ 
1029+ GenInitial2DIndices((AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegIndex, wProBatchSize, hArgmaxActual,
1030+ wArgmaxAligned, wFullBatchCount);
1031+ Gen2DIndexOne((AscendC::MicroAPI::RegTensor<int32_t>&)initial2DRegIndexOne, hArgmaxActual, wArgmaxAligned);
1032+ 
1033+ AscendC::MicroAPI::MaskReg allMask =
1034+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
1035+ AscendC::MicroAPI::DataCopy(helpAddr, initial3DRegIndex, allMask);
1036+ AscendC::MicroAPI::DataCopy(helpAddr + V_REG_SIZE / sizeof(uint32_t), initial3DRegIndexOne, allMask);
1037+ AscendC::MicroAPI::DataCopy(helpAddr + INDEX_TWO * V_REG_SIZE / sizeof(uint32_t), initial2DRegIndex, allMask);
1038+ AscendC::MicroAPI::DataCopy(helpAddr + INDEX_THREE * V_REG_SIZE / sizeof(uint32_t), initial2DRegIndexOne,
1039+ allMask);
1040+ }
1041+ 
1042+ __VEC_SCOPE__
1043+ {
1044+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
1045+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
1046+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
1047+ if constexpr (IS_CHECK_RANGE == 1) {
1048+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
1049+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
1050+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
1051+ }
1052+ 
1053+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
1054+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
1055+ 
1056+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
1057+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
1058+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
1059+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
1060+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
1061+ 
1062+ AscendC::MicroAPI::MaskReg allMaskU32 =
1063+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
1064+ 
1065+ AscendC::MicroAPI::DataCopy(initial3DRegIndex, helpAddr);
1066+ AscendC::MicroAPI::DataCopy(initial3DRegIndexOne, helpAddr + V_REG_SIZE / sizeof(uint32_t));
1067+ AscendC::MicroAPI::DataCopy(initial2DRegIndex, helpAddr + INDEX_TWO * V_REG_SIZE / sizeof(uint32_t));
1068+ AscendC::MicroAPI::DataCopy(initial2DRegIndexOne, helpAddr + INDEX_THREE * V_REG_SIZE / sizeof(uint32_t));
1069+ 
1070+ for (uint16_t highBlockIdx = 0; highBlockIdx < highBlockConcurrentCount; ++highBlockIdx) {
1071+ uint32_t highArgmaxOffset = highBlockIdx * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
1072+ uint32_t highOutputOffset = highBlockIdx * highConcurrentCount * hOutputActual * wOutputAligned;
1073+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
1074+ // 整batch
1075+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
1076+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
1077+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndex, offset, allMaskU32);
1078+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask0,
1079+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
1080+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
1081+ highOutputPlaneActual, whFullBatchCount);
1082+ }
1083+ 
1084+ // 尾段零散点
1085+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
1086+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned +
1087+ highArgmaxOffset);
1088+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndexOne, offset, allMaskU32);
1089+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask1,
1090+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
1091+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
1092+ highOutputPlaneActual, hFullBatchCount);
1093+ }
1094+ }
1095+ 
1096+ // hRemainTail
1097+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
1098+ // 整batch
1099+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
1100+ T2 offset = (wBatchIdx + (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxAligned +
1101+ highArgmaxOffset);
1102+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndex, offset, allMaskU32);
1103+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask2,
1104+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
1105+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
1106+ highOutputPlaneActual, wFullBatchCount);
1107+ }
1108+ 
1109+ // 尾段零散点
1110+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
1111+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount +
1112+ (hProBatchSize * hFullBatchCount + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
1113+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndexOne, offset, allMaskU32);
1114+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
1115+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask3, wOutputConstReg, curHIndex, curWIndex,
1116+ wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual, 1);
1117+ }
1118+ }
1119+ }
1120+ }
1121+ 
1122+ __VEC_SCOPE__
1123+ {
1124+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
1125+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
1126+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
1127+ if constexpr (IS_CHECK_RANGE == 1) {
1128+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
1129+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
1130+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
1131+ }
1132+ 
1133+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
1134+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
1135+ 
1136+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
1137+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
1138+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
1139+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
1140+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
1141+ 
1142+ AscendC::MicroAPI::MaskReg allMaskU32 =
1143+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
1144+ 
1145+ AscendC::MicroAPI::DataCopy(initial3DRegIndex, helpAddr);
1146+ AscendC::MicroAPI::DataCopy(initial3DRegIndexOne, helpAddr + V_REG_SIZE / sizeof(uint32_t));
1147+ AscendC::MicroAPI::DataCopy(initial2DRegIndex, helpAddr + INDEX_TWO * V_REG_SIZE / sizeof(uint32_t));
1148+ AscendC::MicroAPI::DataCopy(initial2DRegIndexOne, helpAddr + INDEX_THREE * V_REG_SIZE / sizeof(uint32_t));
1149+ 
1150+ // highBlockRemainTail
1151+ uint32_t highArgmaxOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
1152+ uint32_t highOutputOffset = highBlockConcurrentCount * highConcurrentCount * hOutputActual * wOutputAligned;
1153+ // 整H batch
1154+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hProBatchSize; hProBatchIdx++) {
1155+ // 整batch
1156+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
1157+ T2 offset = (wBatchIdx + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
1158+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndex, offset, allMaskU32);
1159+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask4,
1160+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
1161+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
1162+ highOutputPlaneActual, whFullBatchCount);
1163+ }
1164+ 
1165+ // 尾段零散点
1166+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
1167+ T2 offset =
1168+ (wBatchIdx + wProBatchSize * wFullBatchCount + hProBatchIdx * wArgmaxAligned + highArgmaxOffset);
1169+ AscendC::MicroAPI::Adds(parallelRegIndex, initial3DRegIndexOne, offset, allMaskU32);
1170+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask5,
1171+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
1172+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
1173+ highOutputPlaneActual, hFullBatchCount);
1174+ }
1175+ }
1176+ }
1177+ 
1178+ __VEC_SCOPE__
1179+ {
1180+ AscendC::MicroAPI::RegTensor<int32_t> zeroConstReg;
1181+ AscendC::MicroAPI::RegTensor<int32_t> wMaxReg;
1182+ AscendC::MicroAPI::RegTensor<int32_t> hMaxReg;
1183+ if constexpr (IS_CHECK_RANGE == 1) {
1184+ AscendC::MicroAPI::Duplicate(zeroConstReg, T2(0));
1185+ AscendC::MicroAPI::Duplicate(wMaxReg, int32_t(wOutputActual));
1186+ AscendC::MicroAPI::Duplicate(hMaxReg, int32_t(hOutputActual));
1187+ }
1188+ 
1189+ AscendC::MicroAPI::RegTensor<T3> wOutputConstReg;
1190+ AscendC::MicroAPI::Duplicate(wOutputConstReg, T3(wOutput));
1191+ 
1192+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndex;
1193+ AscendC::MicroAPI::RegTensor<uint32_t> initial3DRegIndexOne;
1194+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndex;
1195+ AscendC::MicroAPI::RegTensor<uint32_t> initial2DRegIndexOne;
1196+ AscendC::MicroAPI::RegTensor<uint32_t> parallelRegIndex;
1197+ 
1198+ AscendC::MicroAPI::MaskReg allMaskU32 =
1199+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
1200+ 
1201+ AscendC::MicroAPI::DataCopy(initial3DRegIndex, helpAddr);
1202+ AscendC::MicroAPI::DataCopy(initial3DRegIndexOne, helpAddr + V_REG_SIZE / sizeof(uint32_t));
1203+ AscendC::MicroAPI::DataCopy(initial2DRegIndex, helpAddr + INDEX_TWO * V_REG_SIZE / sizeof(uint32_t));
1204+ AscendC::MicroAPI::DataCopy(initial2DRegIndexOne, helpAddr + INDEX_THREE * V_REG_SIZE / sizeof(uint32_t));
1205+ 
1206+ // highBlockRemainTail
1207+ uint32_t highArgmaxOffset = highBlockConcurrentCount * highConcurrentCount * hArgmaxActual * wArgmaxAligned;
1208+ uint32_t highOutputOffset = highBlockConcurrentCount * highConcurrentCount * hOutputActual * wOutputAligned;
1209+ // hRemainTail
1210+ for (uint16_t hProBatchIdx = 0; hProBatchIdx < hRemainTail; hProBatchIdx++) {
1211+ // 整batch
1212+ for (uint16_t wBatchIdx = 0; wBatchIdx < wProBatchSize; wBatchIdx++) {
1213+ T2 offset =
1214+ (wBatchIdx + (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
1215+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndex, offset, allMaskU32);
1216+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask6,
1217+ wOutputConstReg, curHIndex, curWIndex, wOutputAligned,
1218+ highOutputOffset, zeroConstReg, wMaxReg, hMaxReg,
1219+ highOutputPlaneActual, wFullBatchCount);
1220+ }
1221+ 
1222+ // 尾段零散点
1223+ for (uint16_t wBatchIdx = 0; wBatchIdx < wRemainTail; wBatchIdx++) {
1224+ T2 offset = (wBatchIdx + wProBatchSize * wFullBatchCount +
1225+ (hFullBatchCount * hProBatchSize + hProBatchIdx) * wArgmaxAligned + highArgmaxOffset);
1226+ AscendC::MicroAPI::Adds(parallelRegIndex, initial2DRegIndexOne, offset, allMaskU32);
1227+ DoMulNCNchw<T1, T2, T3, IS_CHECK_RANGE>(
1228+ yAddr, gradAddr, argmaxAddr, parallelRegIndex, mask7, wOutputConstReg, curHIndex, curWIndex,
1229+ wOutputAligned, highOutputOffset, zeroConstReg, wMaxReg, hMaxReg, highOutputPlaneActual, 1);
1230+ }
1231+ }
1232+ }
1233+}
1234+ 
1235+template <typename T1, typename T2, typename T3, const uint32_t IS_CHECK_RANGE>
1236+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWKernel<T1, T2, T3, IS_CHECK_RANGE>::CopyOut()
1237+{
1238+ LocalTensor<T1> yLocal = outputQue_.DeQue<T1>();
1239+ 
1240+ int64_t outputPlaneSize = hOutput_ * wOutput_;
1241+ int64_t highOutputAxisOffset = highAxisIndex_ * highAxisInner_ * outputPlaneSize;
1242+ int64_t hOutputAxisOffset = hAxisIndex_ * hOutputInner_ * wOutput_;
1243+ int64_t wOutputAxisOffset = wAxisIndex_ * wOutputInner_;
1244+ int64_t outputGmOffset = highOutputAxisOffset + hOutputAxisOffset + wOutputAxisOffset;
1245+ 
1246+ LoopModeParams loopModeParamsT1;
1247+ loopModeParamsT1.loop1Size = highAxisActual_;
1248+ loopModeParamsT1.loop2Size = 1;
1249+ loopModeParamsT1.loop1SrcStride = hOutputActual_ * wOutputAligned_ * sizeof(T1);
1250+ loopModeParamsT1.loop2SrcStride = 0;
1251+ loopModeParamsT1.loop1DstStride = hOutput_ * wOutput_ * sizeof(T1);
1252+ loopModeParamsT1.loop2DstStride = 0;
1253+ 
1254+ SetLoopModePara(loopModeParamsT1, DataCopyMVType::UB_TO_OUT);
1255+ DataCopyExtParams copyOutParamT1 = {static_cast<uint16_t>(hOutputActual_),
1256+ static_cast<uint32_t>(wOutputActual_ * sizeof(T1)), static_cast<uint32_t>(0),
1257+ static_cast<uint32_t>((wOutput_ - wOutputActual_) * sizeof(T1)),
1258+ static_cast<uint32_t>(0)};
1259+ 
1260+ DataCopyPad(yGm_[outputGmOffset], yLocal, copyOutParamT1);
1261+ ResetLoopModePara(DataCopyMVType::UB_TO_OUT);
1262+ outputQue_.FreeTensor(yLocal);
1263+}
1264+} // namespace MaxPoolGradWithArgmaxV3NCHWNameSpace
1265+#endif // MAX_POOL_GRAD_WITH_ARGMAX_V3_NCHW_KERNEL_H_
Apooling/max_pool_grad_with_argmax_v3/op_kernel/arch35/max_pool_grad_with_argmax_v3_nchw_scalar.h+269-0
@@ -0,0 +1,269 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_grad_with_argmax_v3_nchw_scalar.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_GRAD_WITH_ARGMAX_V3_SCALAR_KERNEL_H_
17+#define MAX_POOL_GRAD_WITH_ARGMAX_V3_SCALAR_KERNEL_H_
18+#include "kernel_operator.h"
19+#include "kernel_tiling/kernel_tiling.h"
20+#include "../inc/platform.h"
21+#include "max_pool_grad_with_argmax_v3_base.h"
22+namespace MaxPoolGradWithArgmaxV3NCHWScalarNameSpace {
23+constexpr int32_t INVALID_INDEX_VALUE = -1;
24+template <typename T1, typename T2>
25+class MaxPoolGradWithArgmaxV3NCHWScalar {
26+public:
27+ __aicore__ inline MaxPoolGradWithArgmaxV3NCHWScalar(
28+ const MaxPoolGradWithArgmaxV3NCHWScalarTilingData& tilingData, TPipe& pipe)
29+ : tilingData_(tilingData), pipe_(pipe){};
30+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR grad, GM_ADDR argmax, GM_ADDR y);
31+ __aicore__ inline void Process();
32+ 
33+private:
34+ __aicore__ inline void ScalarCompute(int64_t loopNum);
35+ __aicore__ inline void ProcessPerLoop();
36+ __aicore__ inline void CopyOut();
37+ __aicore__ inline void Compute(
38+ LocalTensor<computeType>& yLocal, int64_t argmaxNcActual, int64_t argmaxHActual, int64_t argmaxWActual,
39+ int64_t argmaxLoopIndex);
40+ __aicore__ inline uint32_t ConvertIndexToUBIndex(int64_t indexValue);
41+ __aicore__ inline void ComputeActualOffset(
42+ int64_t loopIndex, int64_t& argmaxNcActual, int64_t& argmaxHActual, int64_t& argmaxWActual,
43+ int64_t& argmaxGmOffset, int64_t& argmaxNcIndex);
44+ __aicore__ inline void CopyInArgmaxGrad(
45+ int64_t& argmaxNcActual, int64_t& argmaxHActual, int64_t& argmaxWActual, int64_t& argmaxGmOffset);
46+ __aicore__ inline uint32_t ConvertIndexToUBIndex(int64_t indexValue, int64_t innerNcIndex, int64_t argmaxLoopIndex);
47+ 
48+private:
49+ const MaxPoolGradWithArgmaxV3NCHWScalarTilingData& tilingData_;
50+ TPipe& pipe_;
51+ TQue<QuePosition::VECIN, BUFFER_NUM> gradQue_;
52+ TQue<QuePosition::VECIN, BUFFER_NUM> argmaxQue_;
53+ TQue<QuePosition::VECOUT, BUFFER_NUM> outputQue_;
54+ GlobalTensor<T1> gradGm_;
55+ GlobalTensor<T1> yGm_;
56+ GlobalTensor<T2> argmaxGm_;
57+ int64_t blockIdx_ = 0;
58+ int64_t curCoreProcessNum_ = 1;
59+ int64_t argmaxPlaneSize_ = 1;
60+ int64_t highAxisIndex_ = 0;
61+ int64_t highAxisActual_ = 0;
62+ int64_t hAxisIndex_ = 0;
63+ int64_t hOutputActual_ = 0;
64+ int64_t wAxisIndex_ = 0;
65+ int64_t wOutputActual_ = 0;
66+ int64_t highAxisArgmaxOffset_ = 0;
67+ int64_t hAxisArgmaxOffset_ = 0;
68+ int64_t wAxisArgmaxOffset_ = 0;
69+};
70+template <typename T1, typename T2>
71+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::Init(
72+ GM_ADDR x, GM_ADDR grad, GM_ADDR argmax, GM_ADDR y)
73+{
74+ blockIdx_ = GetBlockIdx();
75+ argmaxPlaneSize_ = tilingData_.hArgmax * tilingData_.wArgmax;
76+ curCoreProcessNum_ =
77+ (blockIdx_ + 1 == tilingData_.usedCoreNum) ? tilingData_.tailCoreProcessNum : tilingData_.normalCoreProcessNum;
78+ gradGm_.SetGlobalBuffer((__gm__ T1*)grad);
79+ argmaxGm_.SetGlobalBuffer((__gm__ T2*)argmax);
80+ yGm_.SetGlobalBuffer((__gm__ T1*)y);
81+ 
82+ pipe_.InitBuffer(outputQue_, BUFFER_NUM, tilingData_.outputBufferSize);
83+ pipe_.InitBuffer(gradQue_, BUFFER_NUM, tilingData_.gradBufferSize);
84+ pipe_.InitBuffer(argmaxQue_, BUFFER_NUM, tilingData_.argmaxBufferSize);
85+ return;
86+}
87+ 
88+template <typename T1, typename T2>
89+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::Process()
90+{
91+ if (blockIdx_ >= tilingData_.usedCoreNum) {
92+ return;
93+ }
94+ for (int64_t loopNum = 0; loopNum < curCoreProcessNum_; loopNum++) {
95+ ScalarCompute(loopNum);
96+ ProcessPerLoop();
97+ }
98+ return;
99+}
100+template <typename T1, typename T2>
101+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::ScalarCompute(int64_t loopNum)
102+{
103+ int64_t baseBlockIdx = blockIdx_ * tilingData_.normalCoreProcessNum + loopNum;
104+ highAxisIndex_ = baseBlockIdx / (tilingData_.hOutputOuter * tilingData_.wOutputOuter);
105+ highAxisActual_ =
106+ (highAxisIndex_ == (tilingData_.highAxisOuter - 1) ? tilingData_.highAxisTail : tilingData_.highAxisInner);
107+ 
108+ int64_t tempTail = baseBlockIdx - highAxisIndex_ * tilingData_.hOutputOuter * tilingData_.wOutputOuter;
109+ hAxisIndex_ = tempTail / tilingData_.wOutputOuter;
110+ hOutputActual_ =
111+ (hAxisIndex_ == (tilingData_.hOutputOuter - 1) ? tilingData_.hOutputTail : tilingData_.hOutputInner);
112+ 
113+ wAxisIndex_ = tempTail - hAxisIndex_ * tilingData_.wOutputOuter;
114+ wOutputActual_ =
115+ (wAxisIndex_ == (tilingData_.wOutputOuter - 1) ? tilingData_.wOutputTail : tilingData_.wOutputInner);
116+ 
117+ int64_t hArgmaxActualStart = PStart(
118+ hAxisIndex_ * tilingData_.hOutputInner, tilingData_.padH, tilingData_.hKernel, tilingData_.dilationH,
119+ tilingData_.hStride);
120+ int64_t wArgmaxActualStart = PStart(
121+ wAxisIndex_ * tilingData_.wOutputInner, tilingData_.padW, tilingData_.wKernel, tilingData_.dilationW,
122+ tilingData_.wStride);
123+ 
124+ highAxisArgmaxOffset_ = highAxisIndex_ * tilingData_.highAxisInner * argmaxPlaneSize_;
125+ hAxisArgmaxOffset_ = hArgmaxActualStart * tilingData_.wArgmax;
126+ wAxisArgmaxOffset_ = wArgmaxActualStart;
127+ return;
128+}
129+template <typename T1, typename T2>
130+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::ComputeActualOffset(
131+ int64_t loopIndex, int64_t& argmaxNcActual, int64_t& argmaxHActual, int64_t& argmaxWActual, int64_t& argmaxGmOffset,
132+ int64_t& argmaxNcIndex)
133+{
134+ argmaxNcIndex = loopIndex / (tilingData_.argmaxHOuter * tilingData_.argmaxWOuter);
135+ argmaxNcActual =
136+ (argmaxNcIndex == (tilingData_.argmaxNcOuter - 1) ? tilingData_.argmaxNcTail : tilingData_.argmaxNcInner);
137+ int64_t remain = loopIndex - argmaxNcIndex * tilingData_.argmaxHOuter * tilingData_.argmaxWOuter;
138+ int64_t argmaxHIndex = remain / tilingData_.argmaxWOuter;
139+ argmaxHActual =
140+ (argmaxHIndex == (tilingData_.argmaxHOuter - 1) ? tilingData_.argmaxHTail : tilingData_.argmaxHInner);
141+ int64_t argmaxWIndex = remain - argmaxHIndex * tilingData_.argmaxWOuter;
142+ argmaxWActual =
143+ (argmaxWIndex == (tilingData_.argmaxWOuter - 1) ? tilingData_.argmaxWTail : tilingData_.argmaxWInner);
144+ argmaxGmOffset = highAxisArgmaxOffset_ + hAxisArgmaxOffset_ + wAxisArgmaxOffset_ +
145+ argmaxNcIndex * tilingData_.argmaxNcInner * tilingData_.hArgmax * tilingData_.wArgmax +
146+ argmaxHIndex * tilingData_.argmaxHInner * tilingData_.wArgmax +
147+ argmaxWIndex * tilingData_.argmaxWInner;
148+ return;
149+}
150+template <typename T1, typename T2>
151+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::ProcessPerLoop()
152+{
153+ uint32_t calCount = static_cast<uint32_t>(tilingData_.outputBufferSize) / sizeof(computeType);
154+ LocalTensor<computeType> yLocal = outputQue_.AllocTensor<computeType>();
155+ Duplicate(yLocal, computeType(0), calCount);
156+ int64_t argmaxNcActual = 0;
157+ int64_t argmaxHActual = 0;
158+ int64_t argmaxWActual = 0;
159+ int64_t argmaxGmOffset = 0;
160+ int64_t argmaxNcIndex = 0;
161+ for (int64_t i = 0; i < tilingData_.argmaxInnerLoop; i++) {
162+ ComputeActualOffset(i, argmaxNcActual, argmaxHActual, argmaxWActual, argmaxGmOffset, argmaxNcIndex);
163+ CopyInArgmaxGrad(argmaxNcActual, argmaxHActual, argmaxWActual, argmaxGmOffset);
164+ Compute(yLocal, argmaxNcActual, argmaxHActual, argmaxWActual, argmaxNcIndex);
165+ }
166+ if constexpr (std::negation<std::is_same<T1, float>>::value) {
167+ Cast(yLocal.ReinterpretCast<T1>(), yLocal, RoundMode::CAST_RINT, calCount);
168+ }
169+ outputQue_.EnQue(yLocal);
170+ CopyOut();
171+ return;
172+}
173+template <typename T1, typename T2>
174+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::CopyInArgmaxGrad(
175+ int64_t& argmaxNcActual, int64_t& argmaxHActual, int64_t& argmaxWActual, int64_t& argmaxGmOffset)
176+{
177+ LocalTensor<T1> gradLocal = gradQue_.AllocTensor<T1>();
178+ LocalTensor<T2> argmaxLocal = argmaxQue_.AllocTensor<T2>();
179+ DataCopyPadExtParams<T1> paramsT1 = {false, 0, 0, 0};
180+ DataCopyExtParams copyInParamT1 = {
181+ static_cast<uint16_t>(1), static_cast<uint32_t>(argmaxNcActual * argmaxHActual * argmaxWActual * sizeof(T1)),
182+ static_cast<uint32_t>(0), static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
183+ DataCopyPad(gradLocal, gradGm_[argmaxGmOffset], copyInParamT1, paramsT1);
184+ 
185+ DataCopyPadExtParams<T2> paramsT2 = {false, 0, 0, 0};
186+ DataCopyExtParams copyInParamT2 = {
187+ static_cast<uint16_t>(1), static_cast<uint32_t>(argmaxNcActual * argmaxHActual * argmaxWActual * sizeof(T2)),
188+ static_cast<uint32_t>(0), static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
189+ DataCopyPad(argmaxLocal, argmaxGm_[argmaxGmOffset], copyInParamT2, paramsT2);
190+ gradQue_.EnQue(gradLocal);
191+ argmaxQue_.EnQue(argmaxLocal);
192+ return;
193+}
194+template <typename T1, typename T2>
195+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::CopyOut()
196+{
197+ LocalTensor<T1> yLocal = outputQue_.DeQue<T1>();
198+ event_t eventIDSToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
199+ SetFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
200+ WaitFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
201+ int64_t outputPlaneSize = tilingData_.hOutput * tilingData_.wOutput;
202+ int64_t highOutputAxisOffset = highAxisIndex_ * tilingData_.highAxisInner * outputPlaneSize;
203+ int64_t hOutputAxisOffset = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.wOutput;
204+ int64_t wOutputAxisOffset = wAxisIndex_ * tilingData_.wOutputInner;
205+ int64_t outputGmOffset = highOutputAxisOffset + hOutputAxisOffset + wOutputAxisOffset;
206+ 
207+ DataCopyExtParams copyOutParamT1 = {
208+ static_cast<uint16_t>(1), static_cast<uint32_t>(highAxisActual_ * hOutputActual_ * wOutputActual_ * sizeof(T1)),
209+ static_cast<uint32_t>(0), static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
210+ DataCopyPad(yGm_[outputGmOffset], yLocal, copyOutParamT1);
211+ outputQue_.FreeTensor(yLocal);
212+}
213+template <typename T1, typename T2>
214+__aicore__ inline void MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::Compute(
215+ LocalTensor<computeType>& yLocal, int64_t argmaxNcActual, int64_t argmaxHActual, int64_t argmaxWActual,
216+ int64_t argmaxLoopIndex)
217+{
218+ LocalTensor<T1> gradLocal = gradQue_.DeQue<T1>();
219+ LocalTensor<T2> argmaxLocal = argmaxQue_.DeQue<T2>();
220+ event_t eventIDMTE2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
221+ SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);
222+ WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);
223+ computeType gradValue;
224+ int64_t argmaxCountInner = argmaxHActual * argmaxWActual;
225+ for (int64_t i = 0; i < argmaxNcActual; i++) {
226+ int64_t argmaxOffsetHigh = i * argmaxCountInner;
227+ for (int64_t j = 0; j < argmaxCountInner; j++) {
228+ int64_t indexValue = argmaxLocal.GetValue(argmaxOffsetHigh + j);
229+ int32_t outputIndexInUB = ConvertIndexToUBIndex(indexValue, i, argmaxLoopIndex);
230+ if (outputIndexInUB == INVALID_INDEX_VALUE) {
231+ continue;
232+ }
233+ if constexpr (std::is_same<T1, bfloat16_t>::value) {
234+ gradValue = ToFloat(gradLocal.GetValue(argmaxOffsetHigh + j));
235+ } else {
236+ gradValue = static_cast<computeType>(gradLocal.GetValue(argmaxOffsetHigh + j));
237+ }
238+ computeType ubValue = yLocal.GetValue(outputIndexInUB);
239+ yLocal.SetValue(outputIndexInUB, gradValue + ubValue);
240+ }
241+ }
242+ gradQue_.FreeTensor(gradLocal);
243+ argmaxQue_.FreeTensor(argmaxLocal);
244+ return;
245+}
246+template <typename T1, typename T2>
247+__aicore__ inline uint32_t MaxPoolGradWithArgmaxV3NCHWScalar<T1, T2>::ConvertIndexToUBIndex(
248+ int64_t indexValue, int64_t innerNcIndex, int64_t argmaxLoopIndex)
249+{
250+ int64_t curHStartIndex = hAxisIndex_ * tilingData_.hOutputInner;
251+ int64_t curWStartIndex = wAxisIndex_ * tilingData_.wOutputInner;
252+ 
253+ int64_t curHEndIndex = curHStartIndex + hOutputActual_;
254+ int64_t curWEndIndex = curWStartIndex + wOutputActual_;
255+ 
256+ int64_t relativeHIndex = indexValue / tilingData_.wOutput;
257+ int64_t relativeWIndex = indexValue - relativeHIndex * tilingData_.wOutput;
258+ 
259+ if (relativeHIndex < curHStartIndex || relativeHIndex > curHEndIndex || relativeWIndex < curWStartIndex ||
260+ relativeWIndex > curWEndIndex) {
261+ return INVALID_INDEX_VALUE;
262+ }
263+ int32_t ubIndex = (relativeHIndex - curHStartIndex) * wOutputActual_ + (relativeWIndex - curWStartIndex) +
264+ (argmaxLoopIndex * tilingData_.argmaxNcInner + innerNcIndex) * tilingData_.hOutputInner *
265+ tilingData_.wOutputInner;
266+ return ubIndex;
267+}
268+} // namespace MaxPoolGradWithArgmaxV3NCHWScalarNameSpace
269+#endif
Apooling/max_pool_grad_with_argmax_v3/op_kernel/arch35/max_pool_grad_with_argmax_v3_nhwc_kernel.h+1155-0文件内容审核中，请稍后刷新重试
Apooling/max_pool_grad_with_argmax_v3/op_kernel/max_pool_grad_with_argmax_v3.cpp+87-0
@@ -0,0 +1,87 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file max_pool_grad_with_argmax_v3.cpp
13+ * \brief max_pool_grad_with_argmax_v3 implied
14+ */
15+ 
16+#include <cstdint>
17+#include "kernel_operator.h"
18+#include "kernel_tiling/kernel_tiling.h"
19+ 
20+#include "./arch35/max_pool_grad_with_argmax_v3_nchw_kernel.h"
21+#include "./arch35/max_pool_grad_with_argmax_v3_nhwc_kernel.h"
22+#include "./arch35/max_pool_grad_with_argmax_v3_nchw_scalar.h"
23+ 
24+#define NO_CHECK_RANGE_TILING_KEY_NCHW 100
25+#define CHECK_RANGE_TILING_KEY_NCHW 101
26+#define NO_CHECK_RANGE_TILING_KEY_NHWC 200
27+#define CHECK_RANGE_TILING_KEY_NHWC 201
28+#define CHECK_RANGE_TILING_KEY_NCHW_SCALAR 301
29+ 
30+#define NO_CHECK_RANGE_TILING_KEY_NCHW_INT64 110
31+#define CHECK_RANGE_TILING_KEY_NCHW_INT64 111
32+#define NO_CHECK_RANGE_TILING_KEY_NHWC_INT64 210
33+#define CHECK_RANGE_TILING_KEY_NHWC_INT64 211
34+ 
35+extern "C" __global__ __aicore__ void max_pool_grad_with_argmax_v3(
36+ GM_ADDR x, GM_ADDR grad, GM_ADDR argmax, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
37+{
38+ AscendC::TPipe pipeBase;
39+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
40+ if (TILING_KEY_IS(NO_CHECK_RANGE_TILING_KEY_NCHW)) {
41+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NCHWTilingData, tilingDataIn, tiling);
42+ MaxPoolGradWithArgmaxV3NCHWNameSpace::MaxPoolGradWithArgmaxV3NCHWKernel<DTYPE_X, DTYPE_ARGMAX, int32_t, 0> op;
43+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
44+ op.Process();
45+ } else if (TILING_KEY_IS(CHECK_RANGE_TILING_KEY_NCHW)) {
46+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NCHWTilingData, tilingDataIn, tiling);
47+ MaxPoolGradWithArgmaxV3NCHWNameSpace::MaxPoolGradWithArgmaxV3NCHWKernel<DTYPE_X, DTYPE_ARGMAX, int32_t, 1> op;
48+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
49+ op.Process();
50+ } else if (TILING_KEY_IS(NO_CHECK_RANGE_TILING_KEY_NHWC)) {
51+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NHWCTilingData, tilingDataIn, tiling);
52+ MaxPoolGradWithArgmaxV3NHWCNameSpace::MaxPoolGradWithArgmaxV3KernelNHWC<DTYPE_X, DTYPE_ARGMAX, int32_t, 0> op;
53+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
54+ op.Process();
55+ } else if (TILING_KEY_IS(CHECK_RANGE_TILING_KEY_NHWC)) {
56+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NHWCTilingData, tilingDataIn, tiling);
57+ MaxPoolGradWithArgmaxV3NHWCNameSpace::MaxPoolGradWithArgmaxV3KernelNHWC<DTYPE_X, DTYPE_ARGMAX, int32_t, 1> op;
58+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
59+ op.Process();
60+ } else if (TILING_KEY_IS(NO_CHECK_RANGE_TILING_KEY_NCHW_INT64)) {
61+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NCHWTilingData, tilingDataIn, tiling);
62+ MaxPoolGradWithArgmaxV3NCHWNameSpace::MaxPoolGradWithArgmaxV3NCHWKernel<DTYPE_X, DTYPE_ARGMAX, int64_t, 0> op;
63+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
64+ op.Process();
65+ } else if (TILING_KEY_IS(CHECK_RANGE_TILING_KEY_NCHW_INT64)) {
66+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NCHWTilingData, tilingDataIn, tiling);
67+ MaxPoolGradWithArgmaxV3NCHWNameSpace::MaxPoolGradWithArgmaxV3NCHWKernel<DTYPE_X, DTYPE_ARGMAX, int64_t, 1> op;
68+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
69+ op.Process();
70+ } else if (TILING_KEY_IS(NO_CHECK_RANGE_TILING_KEY_NHWC_INT64)) {
71+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NHWCTilingData, tilingDataIn, tiling);
72+ MaxPoolGradWithArgmaxV3NHWCNameSpace::MaxPoolGradWithArgmaxV3KernelNHWC<DTYPE_X, DTYPE_ARGMAX, int64_t, 0> op;
73+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
74+ op.Process();
75+ } else if (TILING_KEY_IS(CHECK_RANGE_TILING_KEY_NHWC_INT64)) {
76+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NHWCTilingData, tilingDataIn, tiling);
77+ MaxPoolGradWithArgmaxV3NHWCNameSpace::MaxPoolGradWithArgmaxV3KernelNHWC<DTYPE_X, DTYPE_ARGMAX, int64_t, 1> op;
78+ op.Init(x, grad, argmax, y, pipeBase, tilingDataIn);
79+ op.Process();
80+ } else if (TILING_KEY_IS(CHECK_RANGE_TILING_KEY_NCHW_SCALAR)) {
81+ GET_TILING_DATA_WITH_STRUCT(MaxPoolGradWithArgmaxV3NCHWScalarTilingData, tilingDataIn, tiling);
82+ MaxPoolGradWithArgmaxV3NCHWScalarNameSpace::MaxPoolGradWithArgmaxV3NCHWScalar<DTYPE_X, DTYPE_ARGMAX> op(
83+ tilingDataIn, pipeBase);
84+ op.Init(x, grad, argmax, y);
85+ op.Process();
86+ }
87+}
Rloss/mse_loss/op_host/CMakeLists.txt→pooling/max_pool_grad_with_argmax_v3/tests/CMakeLists.txt+16-11
@@ -1,11 +1,16 @@
1-#1+# This program is free software, you can redistribute it and/or modify.
2-# Copyright (c) 2025 Huawei Technologies Co., Ltd.2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This file is a part of the CANN Open Software.
4-# CANN Open Software License Agreement Version 2.0 (the "License").4+# Licensed under 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, 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.7+# 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.8+#/
9-#/9+ 
10- 10+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
11-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE mse_loss ACLNNTYPE aclnn_exclude DEPENDENCIES mse_loss_v2)11+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
12+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
13+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Rpooling/max_pool_grad_with_argmax_v3/op_host/CMakeLists.txt→pooling/max_pool_grad_with_argmax_v3/tests/ut/CMakeLists.txt+16-11
@@ -1,11 +1,16 @@
1-# -----------------------------------------------------------------------------------------------------------1+# This program is free software, you can redistribute it and/or modify.
2-# Copyright (c) 2025 Huawei Technologies Co., Ltd.2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+# This file is a part of the CANN Open Software.
4-# CANN Open Software License Agreement Version 2.0 (the "License").4+# Licensed under 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, 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.7+# 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.8+#/
9-# -----------------------------------------------------------------------------------------------------------9+ 
10- 10+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
11-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE max_pool_grad_with_argmax_v3 ACLNNTYPE aclnn_exclude)11+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
12+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
13+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Apooling/max_pool_grad_with_argmax_v3/tests/ut/op_host/CMakeLists.txt+14-0
@@ -0,0 +1,14 @@
1+# This program is free software, you can redistribute it and/or modify.
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+# This file is a part of the CANN Open Software.
4+# Licensed under 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, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+# See LICENSE in the root of the software repository for the full text of the License.
8+#/
9+ 
10+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
11+if(UT_TEST_ALL OR OP_HOST_UT)
12+ add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13+ add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14+endif()
Apooling/max_pool_grad_with_argmax_v3/tests/ut/op_host/test_max_pool_grad_with_argmax_v3_infershape.cpp+179-0
@@ -0,0 +1,179 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include "exe_graph/runtime/storage_format.h"
13+#include "exe_graph/runtime/storage_shape.h"
14+#include <gtest/gtest.h>
15+#include "kernel_run_context_facker.h"
16+#include "register/op_impl_registry.h"
17+#include "log/log.h"
18+#include "platform/platform_info.h"
19+#include "../../../op_graph/max_pool_grad_with_argmax_v3_proto.h"
20+ 
21+namespace {
22+template <typename T>
23+std::string Shape2String(const T& shape)
24+{
25+ std::ostringstream oss;
26+ oss << "[";
27+ if (shape.GetDimNum() > 0) {
28+ for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) {
29+ oss << shape.GetDim(i) << ", ";
30+ }
31+ oss << shape.GetDim(shape.GetDimNum() - 1);
32+ }
33+ oss << "]";
34+ return oss.str();
35+}
36+ 
37+class MaxPoolGradWithArgmaxV3Infer : public testing::Test {
38+protected:
39+ static void SetUpTestCase()
40+ {
41+ std::cout << "MaxPoolGradWithArgmaxV3InferTest SetUp" << std::endl;
42+ }
43+ 
44+ static void TearDownTestCase()
45+ {
46+ std::cout << "max_pool_grad_with_argmax_v3_infer_test TearDown" << std::endl;
47+ }
48+};
49+ 
50+TEST_F(MaxPoolGradWithArgmaxV3Infer, max_pool_grad_with_argmax_v3_infershape_test_01)
51+{
52+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolGradWithArgmaxV3")->infer_shape;
53+ 
54+ gert::StorageShape xShape = {{4, 512, 16, 16}, {}};
55+ gert::StorageShape gradShape = {{}, {}};
56+ gert::StorageShape yShape = {{}, {}};
57+ gert::StorageShape indicesShape = {{}, {}};
58+ auto holder = gert::InferShapeContextFaker()
59+ .NodeIoNum(3, 1)
60+ .IrInstanceNum({1, 1, 1})
61+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
62+ .NodeInputTd(1, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
63+ .NodeInputTd(2, ge::DT_INT32, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
64+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
65+ .NodeAttrs(
66+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
67+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 2})},
68+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
69+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
70+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
71+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(true)},
72+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
73+ .InputShapes({&xShape, &gradShape, &indicesShape})
74+ .OutputShapes({&yShape})
75+ .Build();
76+ 
77+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
78+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
79+ ASSERT_EQ(Shape2String(*output), "[4, 512, 16, 16]");
80+}
81+ 
82+TEST_F(MaxPoolGradWithArgmaxV3Infer, max_pool_grad_with_argmax_v3_inferdtype_test_01)
83+{
84+ auto data_type_func = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolGradWithArgmaxV3")->infer_datatype;
85+ 
86+ if (data_type_func != nullptr) {
87+ ge::DataType input_ref = ge::DT_FLOAT;
88+ ge::DataType input_ref1 = ge::DT_INT32;
89+ ge::DataType output_ref = ge::DT_FLOAT;
90+ auto context_holder = gert::InferDataTypeContextFaker()
91+ .NodeIoNum(3, 1)
92+ .IrInstanceNum({1, 1, 1})
93+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
94+ .NodeInputTd(1, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
95+ .NodeInputTd(2, ge::DT_INT32, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
96+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
97+ .InputDataTypes({&input_ref, &input_ref, &input_ref1})
98+ .OutputDataTypes({&output_ref})
99+ .NodeAttrs(
100+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 3})},
101+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 3})},
102+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
103+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
104+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({9, 5})},
105+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
106+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
107+ .Build();
108+ auto context = context_holder.GetContext<gert::InferDataTypeContext>();
109+ EXPECT_EQ(data_type_func(context), ge::GRAPH_SUCCESS);
110+ ASSERT_NE(context, nullptr);
111+ EXPECT_EQ(context->GetOutputDataType(0), output_ref);
112+ }
113+}
114+ 
115+TEST_F(MaxPoolGradWithArgmaxV3Infer, max_pool_grad_with_argmax_v3_infershape_test_02)
116+{
117+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolGradWithArgmaxV3")->infer_shape;
118+ 
119+ gert::StorageShape xShape = {{4, 512, 16, 16}, {}};
120+ gert::StorageShape gradShape = {{4, 512, 16, 16}, {}};
121+ gert::StorageShape yShape = {{}, {}};
122+ gert::StorageShape indicesShape = {{4, 512, -1, 16}, {}};
123+ auto holder = gert::InferShapeContextFaker()
124+ .NodeIoNum(3, 1)
125+ .IrInstanceNum({1, 1, 1})
126+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
127+ .NodeInputTd(1, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
128+ .NodeInputTd(2, ge::DT_INT32, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
129+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
130+ .NodeAttrs(
131+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
132+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 2})},
133+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
134+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
135+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
136+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(true)},
137+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
138+ .InputShapes({&xShape, &gradShape, &indicesShape})
139+ .OutputShapes({&yShape})
140+ .Build();
141+ 
142+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
143+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
144+ ASSERT_EQ(Shape2String(*output), "[-1, -1, -1, -1]");
145+}
146+ 
147+TEST_F(MaxPoolGradWithArgmaxV3Infer, max_pool_grad_with_argmax_v3_infershape_test_03)
148+{
149+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolGradWithArgmaxV3")->infer_shape;
150+ 
151+ gert::StorageShape xShape = {{4, 512, 16, 16}, {}};
152+ gert::StorageShape gradShape = {{4, -1, 16, 16}, {}};
153+ gert::StorageShape yShape = {{}, {}};
154+ gert::StorageShape indicesShape = {{4, 512, 16, 16}, {}};
155+ auto holder = gert::InferShapeContextFaker()
156+ .NodeIoNum(3, 1)
157+ .IrInstanceNum({1, 1, 1})
158+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
159+ .NodeInputTd(1, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
160+ .NodeInputTd(2, ge::DT_INT32, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
161+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
162+ .NodeAttrs(
163+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
164+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 2})},
165+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
166+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
167+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
168+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(true)},
169+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
170+ .InputShapes({&xShape, &gradShape, &indicesShape})
171+ .OutputShapes({&yShape})
172+ .Build();
173+ 
174+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
175+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
176+ ASSERT_EQ(Shape2String(*output), "[-1, -1, -1, -1]");
177+}
178+ 
179+} // namespace
Apooling/max_pool_grad_with_argmax_v3/tests/ut/op_host/test_max_pool_grad_with_argmax_v3_tiling.cpp+498-0
@@ -0,0 +1,498 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include <fstream>
13+#include <vector>
14+#include <gtest/gtest.h>
15+#include "log/log.h"
16+#include "kernel_run_context_facker.h"
17+#include "test_cube_util.h"
18+#include "exe_graph/runtime/storage_format.h"
19+#include "exe_graph/runtime/storage_shape.h"
20+#include "platform/platform_infos_def.h"
21+#include "ut_op_util.h"
22+#include "../../../op_host/arch35/max_pool_grad_with_argmax_v3_tiling_base.h"
23+ 
24+using namespace ut_util;
25+using namespace std;
26+using namespace ge;
27+ 
28+class MaxPoolGradWithArgmaxV3Tiling : public testing::Test {
29+protected:
30+ static void SetUpTestCase()
31+ {
32+ std::cout << "MaxPoolGradWithArgmaxV3Tiling SetUp" << std::endl;
33+ }
34+ 
35+ static void TearDownTestCase()
36+ {
37+ std::cout << "MaxPoolGradWithArgmaxV3Tiling TearDown" << std::endl;
38+ }
39+};
40+ 
41+static void ExecuteTestCase(
42+ gert::StorageShape xShape, gert::StorageShape yShape, gert::StorageShape gradShape, gert::StorageShape argmaxShape,
43+ std::vector<int64_t> ksize, std::vector<int64_t> strides, std::vector<int64_t> pads, std::vector<int64_t> dilation,
44+ ge::DataType dtype, int64_t index_dtype, ge::DataType index_dtype_enum, bool ceil_mode, std::string data_format,
45+ uint64_t except_tilingkey, std::string expect)
46+{
47+ dlog_setlevel(0, 0, 0);
48+ 
49+ string compile_info_string = R"({
50+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
51+ "Intrinsic_fix_pipe_l0c2out": false,
52+ "Intrinsic_data_move_l12ub": true,
53+ "Intrinsic_data_move_l0c2ub": true,
54+ "Intrinsic_data_move_out2l1_nd2nz": false,
55+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
56+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
57+ "CORE_NUM": 64}
58+ })";
59+ map<string, string> soc_infos;
60+ map<string, string> aicore_spec;
61+ map<string, string> intrinsics;
62+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
63+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend910_95"}};
64+ // platform info
65+ fe::PlatFormInfos platform_info;
66+ platform_info.Init();
67+ // compile info
68+ optiling::MaxPoolGradWithArgmaxV3CompileInfo compile_info;
69+ 
70+ std::string op_type("MaxPoolGradWithArgmaxV3");
71+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
72+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
73+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
74+ 
75+ // tilingParseFunc simulate
76+ auto kernel_holder =
77+ gert::KernelRunContextFaker()
78+ .KernelIONum(2, 1)
79+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
80+ .Outputs({&compile_info})
81+ .Build();
82+ 
83+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
84+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
85+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
86+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
87+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
88+ "AICoreintrinsicDtypeMap", intrinsics);
89+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
90+ "version", soc_version_infos);
91+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
92+ 
93+ // tilingFunc simulate
94+ auto param = gert::TilingData::CreateCap(4096);
95+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
96+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
97+ ASSERT_NE(param, nullptr);
98+ auto holder = gert::TilingContextFaker()
99+ .SetOpType(op_type)
100+ .NodeIoNum(3, 1)
101+ .IrInstanceNum({1, 1, 1})
102+ .InputShapes({&xShape, &gradShape, &argmaxShape})
103+ .OutputShapes({&yShape})
104+ .CompileInfo(&compile_info)
105+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
106+ .NodeInputTd(0, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
107+ .NodeInputTd(1, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
108+ .NodeInputTd(2, index_dtype_enum, ge::FORMAT_ND, ge::FORMAT_ND)
109+ .NodeOutputTd(0, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
110+ .NodeAttrs(
111+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(ksize)},
112+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(strides)},
113+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(pads)},
114+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(index_dtype)},
115+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(dilation)},
116+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(ceil_mode)},
117+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>(data_format)}})
118+ .TilingData(param.get())
119+ .Workspace(ws_size)
120+ .Build();
121+ 
122+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
123+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
124+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
125+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
126+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
127+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
128+ 
129+ // workspaces nullptr return failed
130+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
131+ auto tiling_key = tiling_context->GetTilingKey();
132+ ASSERT_EQ(tiling_key, except_tilingkey);
133+ // auto tilingData = tiling_context->GetRawTilingData();
134+ // ASSERT_NE(tilingData, nullptr);
135+ // dlog_setlevel(0, 3, 0);
136+ // auto tiling_data_result = to_string<int64_t>(tilingData->GetData(), tilingData->GetDataSize());
137+ // std::cout<<tiling_data_result<<std::endl;
138+ // EXPECT_EQ(tiling_data_result, expect);
139+}
140+ 
141+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test1)
142+{
143+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
144+ gert::StorageShape yShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
145+ ;
146+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
147+ gert::StorageShape gradShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
148+ std::vector<int64_t> ksize = {64, 64};
149+ std::vector<int64_t> strides = {64, 64};
150+ std::vector<int64_t> pads = {0, 0};
151+ std::vector<int64_t> dilation = {1, 1};
152+ ge::DataType dtype = ge::DT_FLOAT;
153+ ge::DataType dtype_index = ge::DT_INT32;
154+ int64_t index_dtype = 3;
155+ bool ceil_mode = false;
156+ std::string data_format = "NCHW";
157+ uint64_t except_tilingkey = 301;
158+ std::string expect =
159+ "1 1 64 64 64 64 64 64 0 0 1 1 1 1 6 64 64 1 64 64 1 1 1 6 61440 30720 30720 1 1 1 1 1 1 1 1 1 1 ";
160+ ExecuteTestCase(
161+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
162+ ceil_mode, data_format, except_tilingkey, expect);
163+}
164+ 
165+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test2)
166+{
167+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
168+ gert::StorageShape yShape = xShape;
169+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
170+ gert::StorageShape gradShape = argmaxShape;
171+ std::vector<int64_t> ksize = {64, 64};
172+ std::vector<int64_t> strides = {64, 64};
173+ std::vector<int64_t> pads = {0, 0};
174+ std::vector<int64_t> dilation = {1, 1};
175+ ge::DataType dtype = ge::DT_FLOAT16;
176+ ge::DataType dtype_index = ge::DT_INT32;
177+ int64_t index_dtype = 3;
178+ bool ceil_mode = false;
179+ std::string data_format = "NCHW";
180+ uint64_t except_tilingkey = 301;
181+ std::string expect =
182+ "1 1 64 64 64 64 64 64 0 0 1 1 1 1 6 64 64 1 64 64 1 1 1 6 61440 20480 40960 1 1 1 1 1 1 1 1 1 1 ";
183+ ExecuteTestCase(
184+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
185+ ceil_mode, data_format, except_tilingkey, expect);
186+}
187+ 
188+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test3)
189+{
190+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
191+ gert::StorageShape yShape = xShape;
192+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
193+ gert::StorageShape gradShape = argmaxShape;
194+ std::vector<int64_t> ksize = {64, 64};
195+ std::vector<int64_t> strides = {64, 64};
196+ std::vector<int64_t> pads = {0, 0};
197+ std::vector<int64_t> dilation = {1, 1};
198+ ge::DataType dtype = ge::DT_BF16;
199+ ge::DataType dtype_index = ge::DT_INT32;
200+ int64_t index_dtype = 3;
201+ bool ceil_mode = false;
202+ std::string data_format = "NCHW";
203+ uint64_t except_tilingkey = 301;
204+ std::string expect =
205+ "1 1 64 64 64 64 64 64 0 0 1 1 1 1 6 64 64 1 64 64 1 1 1 6 61440 20480 40960 1 1 1 1 1 1 1 1 1 1 ";
206+ ExecuteTestCase(
207+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
208+ ceil_mode, data_format, except_tilingkey, expect);
209+}
210+ 
211+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test4)
212+{
213+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
214+ gert::StorageShape yShape = xShape;
215+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
216+ gert::StorageShape gradShape = argmaxShape;
217+ std::vector<int64_t> ksize = {64, 64};
218+ std::vector<int64_t> strides = {64, 64};
219+ std::vector<int64_t> pads = {0, 0};
220+ std::vector<int64_t> dilation = {1, 1};
221+ ge::DataType dtype = ge::DT_FLOAT;
222+ ge::DataType dtype_index = ge::DT_INT64;
223+ int64_t index_dtype = 9;
224+ bool ceil_mode = false;
225+ std::string data_format = "NCHW";
226+ uint64_t except_tilingkey = 301;
227+ std::string expect =
228+ "1 1 64 64 64 64 64 64 0 0 1 1 1 1 6 64 64 1 64 64 1 1 1 6 61440 20480 40960 1 1 1 1 1 1 1 1 1 1 ";
229+ ExecuteTestCase(
230+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
231+ ceil_mode, data_format, except_tilingkey, expect);
232+}
233+ 
234+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test5)
235+{
236+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
237+ gert::StorageShape yShape = xShape;
238+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
239+ gert::StorageShape gradShape = argmaxShape;
240+ std::vector<int64_t> ksize = {64, 64};
241+ std::vector<int64_t> strides = {64, 64};
242+ std::vector<int64_t> pads = {0, 0};
243+ std::vector<int64_t> dilation = {1, 1};
244+ ge::DataType dtype = ge::DT_FLOAT16;
245+ ge::DataType dtype_index = ge::DT_INT64;
246+ int64_t index_dtype = 9;
247+ bool ceil_mode = false;
248+ std::string data_format = "NCHW";
249+ uint64_t except_tilingkey = 301;
250+ std::string expect =
251+ "1 1 64 64 64 64 64 64 0 0 1 1 1 1 6 64 64 1 64 64 1 1 1 6 61440 12288 49152 1 1 1 1 1 1 1 1 1 1 ";
252+ ExecuteTestCase(
253+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
254+ ceil_mode, data_format, except_tilingkey, expect);
255+}
256+ 
257+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test6)
258+{
259+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
260+ gert::StorageShape yShape = xShape;
261+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
262+ gert::StorageShape gradShape = argmaxShape;
263+ std::vector<int64_t> ksize = {64, 64};
264+ std::vector<int64_t> strides = {64, 64};
265+ std::vector<int64_t> pads = {0, 0};
266+ std::vector<int64_t> dilation = {1, 1};
267+ ge::DataType dtype = ge::DT_BF16;
268+ ge::DataType dtype_index = ge::DT_INT64;
269+ int64_t index_dtype = 9;
270+ bool ceil_mode = false;
271+ std::string data_format = "NCHW";
272+ uint64_t except_tilingkey = 301;
273+ std::string expect =
274+ "1 1 64 64 64 64 64 64 0 0 1 1 1 1 6 64 64 1 64 64 1 1 1 6 61440 12288 49152 1 1 1 1 1 1 1 1 1 1 ";
275+ ExecuteTestCase(
276+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
277+ ceil_mode, data_format, except_tilingkey, expect);
278+}
279+ 
280+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test7)
281+{
282+ gert::StorageShape xShape = {{33, 12, 2, 22}, {33, 12, 2, 22}};
283+ gert::StorageShape yShape = xShape;
284+ gert::StorageShape argmaxShape = {{33, 12, 1, 4}, {33, 12, 1, 4}};
285+ gert::StorageShape gradShape = argmaxShape;
286+ std::vector<int64_t> ksize = {2, 6};
287+ std::vector<int64_t> strides = {2, 6};
288+ std::vector<int64_t> pads = {0, 0};
289+ std::vector<int64_t> dilation = {1, 1};
290+ ge::DataType dtype = ge::DT_FLOAT;
291+ ge::DataType dtype_index = ge::DT_INT32;
292+ int64_t index_dtype = 3;
293+ bool ceil_mode = true;
294+ std::string data_format = "NCHW";
295+ uint64_t except_tilingkey = 100;
296+ std::string expect = "1 4 2 22 2 6 2 6 0 0 1 1 6 6 66 2 2 1 22 22 1 2 2 33 1152 384 384 1 1 100 ";
297+ ExecuteTestCase(
298+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
299+ ceil_mode, data_format, except_tilingkey, expect);
300+}
301+ 
302+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test8)
303+{
304+ gert::StorageShape xShape = {{33, 12, 2, 22}, {33, 12, 2, 22}};
305+ gert::StorageShape yShape = xShape;
306+ gert::StorageShape argmaxShape = {{33, 12, 1, 4}, {33, 12, 1, 4}};
307+ gert::StorageShape gradShape = argmaxShape;
308+ std::vector<int64_t> ksize = {2, 6};
309+ std::vector<int64_t> strides = {2, 6};
310+ std::vector<int64_t> pads = {0, 0};
311+ std::vector<int64_t> dilation = {1, 1};
312+ ge::DataType dtype = ge::DT_FLOAT16;
313+ ge::DataType dtype_index = ge::DT_INT32;
314+ int64_t index_dtype = 3;
315+ bool ceil_mode = true;
316+ std::string data_format = "NCHW";
317+ uint64_t except_tilingkey = 100;
318+ std::string expect = "1 4 2 22 2 6 2 6 0 0 1 1 6 6 66 2 2 1 22 22 1 2 2 33 1536 384 768 1 1 100 ";
319+ ExecuteTestCase(
320+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
321+ ceil_mode, data_format, except_tilingkey, expect);
322+}
323+ 
324+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test9)
325+{
326+ gert::StorageShape xShape = {{33, 12, 2, 22}, {33, 12, 2, 22}};
327+ gert::StorageShape yShape = xShape;
328+ gert::StorageShape argmaxShape = {{33, 12, 1, 4}, {33, 12, 1, 4}};
329+ gert::StorageShape gradShape = argmaxShape;
330+ std::vector<int64_t> ksize = {2, 6};
331+ std::vector<int64_t> strides = {2, 6};
332+ std::vector<int64_t> pads = {0, 0};
333+ std::vector<int64_t> dilation = {1, 1};
334+ ge::DataType dtype = ge::DT_BF16;
335+ ge::DataType dtype_index = ge::DT_INT32;
336+ int64_t index_dtype = 3;
337+ bool ceil_mode = true;
338+ std::string data_format = "NCHW";
339+ uint64_t except_tilingkey = 100;
340+ std::string expect = "1 4 2 22 2 6 2 6 0 0 1 1 6 6 66 2 2 1 22 22 1 2 2 33 1536 384 768 1 1 100 ";
341+ ExecuteTestCase(
342+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
343+ ceil_mode, data_format, except_tilingkey, expect);
344+}
345+ 
346+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test10)
347+{
348+ gert::StorageShape xShape = {{33, 12, 2, 22}, {33, 12, 2, 22}};
349+ gert::StorageShape yShape = xShape;
350+ gert::StorageShape argmaxShape = {{33, 12, 1, 4}, {33, 12, 1, 4}};
351+ gert::StorageShape gradShape = argmaxShape;
352+ std::vector<int64_t> ksize = {2, 6};
353+ std::vector<int64_t> strides = {2, 6};
354+ std::vector<int64_t> pads = {0, 0};
355+ std::vector<int64_t> dilation = {1, 1};
356+ ge::DataType dtype = ge::DT_FLOAT;
357+ ge::DataType dtype_index = ge::DT_INT64;
358+ int64_t index_dtype = 9;
359+ bool ceil_mode = true;
360+ std::string data_format = "NCHW";
361+ uint64_t except_tilingkey = 100;
362+ std::string expect = "1 4 2 22 2 6 2 6 0 0 1 1 6 6 66 2 2 1 22 22 1 2 2 33 1152 384 768 1 1 100 ";
363+ ExecuteTestCase(
364+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
365+ ceil_mode, data_format, except_tilingkey, expect);
366+}
367+ 
368+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test11)
369+{
370+ gert::StorageShape xShape = {{33, 12, 2, 22}, {33, 12, 2, 22}};
371+ gert::StorageShape yShape = xShape;
372+ gert::StorageShape argmaxShape = {{33, 12, 1, 4}, {33, 12, 1, 4}};
373+ gert::StorageShape gradShape = argmaxShape;
374+ std::vector<int64_t> ksize = {2, 6};
375+ std::vector<int64_t> strides = {2, 6};
376+ std::vector<int64_t> pads = {0, 0};
377+ std::vector<int64_t> dilation = {1, 1};
378+ ge::DataType dtype = ge::DT_FLOAT16;
379+ ge::DataType dtype_index = ge::DT_INT64;
380+ int64_t index_dtype = 9;
381+ bool ceil_mode = true;
382+ std::string data_format = "NCHW";
383+ uint64_t except_tilingkey = 100;
384+ std::string expect = "1 4 2 22 2 6 2 6 0 0 1 1 6 6 66 2 2 1 22 22 1 2 2 33 1536 384 1536 1 1 100 ";
385+ ExecuteTestCase(
386+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
387+ ceil_mode, data_format, except_tilingkey, expect);
388+}
389+ 
390+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NCHW_Test12)
391+{
392+ gert::StorageShape xShape = {{33, 12, 2, 22}, {33, 12, 2, 22}};
393+ gert::StorageShape yShape = xShape;
394+ gert::StorageShape argmaxShape = {{33, 12, 1, 4}, {33, 12, 1, 4}};
395+ gert::StorageShape gradShape = argmaxShape;
396+ std::vector<int64_t> ksize = {2, 6};
397+ std::vector<int64_t> strides = {2, 6};
398+ std::vector<int64_t> pads = {0, 0};
399+ std::vector<int64_t> dilation = {1, 1};
400+ ge::DataType dtype = ge::DT_BF16;
401+ ge::DataType dtype_index = ge::DT_INT64;
402+ int64_t index_dtype = 9;
403+ bool ceil_mode = true;
404+ std::string data_format = "NCHW";
405+ uint64_t except_tilingkey = 100;
406+ std::string expect = "1 4 2 22 2 6 2 6 0 0 1 1 6 6 66 2 2 1 22 22 1 2 2 33 1536 384 1536 1 1 100 ";
407+ ExecuteTestCase(
408+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
409+ ceil_mode, data_format, except_tilingkey, expect);
410+}
411+ 
412+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NHWC_Test1)
413+{
414+ gert::StorageShape xShape = {{16, 18, 30, 2}, {16, 18, 30, 2}};
415+ gert::StorageShape yShape = xShape;
416+ gert::StorageShape argmaxShape = {{16, 10, 16, 2}, {16, 10, 16, 2}};
417+ gert::StorageShape gradShape = argmaxShape;
418+ std::vector<int64_t> ksize = {2, 2};
419+ std::vector<int64_t> strides = {2, 2};
420+ std::vector<int64_t> pads = {1, 1};
421+ std::vector<int64_t> dilation = {1, 1};
422+ ge::DataType dtype = ge::DT_FLOAT;
423+ ge::DataType dtype_index = ge::DT_INT32;
424+ int64_t index_dtype = 3;
425+ bool ceil_mode = false;
426+ std::string data_format = "NHWC";
427+ uint64_t except_tilingkey = 201;
428+ std::string expect = "10 16 2 18 30 2 2 2 2 1 1 1 1 1 1 16 5 3 4 30 30 1 2 2 1 1 1 64 4800 1792 1792 1 1 201 ";
429+ ExecuteTestCase(
430+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
431+ ceil_mode, data_format, except_tilingkey, expect);
432+}
433+ 
434+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NHWC_Test2)
435+{
436+ gert::StorageShape xShape = {{16, 18, 30, 2}, {16, 18, 30, 2}};
437+ gert::StorageShape yShape = xShape;
438+ gert::StorageShape argmaxShape = {{16, 10, 16, 2}, {16, 10, 16, 2}};
439+ gert::StorageShape gradShape = argmaxShape;
440+ std::vector<int64_t> ksize = {2, 2};
441+ std::vector<int64_t> strides = {2, 2};
442+ std::vector<int64_t> pads = {1, 1};
443+ std::vector<int64_t> dilation = {1, 1};
444+ ge::DataType dtype = ge::DT_FLOAT16;
445+ ge::DataType dtype_index = ge::DT_INT32;
446+ int64_t index_dtype = 3;
447+ bool ceil_mode = false;
448+ std::string data_format = "NHWC";
449+ uint64_t except_tilingkey = 201;
450+ std::string expect = "10 16 2 18 30 2 2 2 2 1 1 1 1 1 1 16 5 3 4 30 30 1 2 2 1 1 1 64 9600 1792 3328 1 1 201 ";
451+ ExecuteTestCase(
452+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
453+ ceil_mode, data_format, except_tilingkey, expect);
454+}
455+ 
456+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NHWC_Test3)
457+{
458+ gert::StorageShape xShape = {{16, 18, 30, 2}, {16, 18, 30, 2}};
459+ gert::StorageShape yShape = xShape;
460+ gert::StorageShape argmaxShape = {{16, 10, 16, 2}, {16, 10, 16, 2}};
461+ gert::StorageShape gradShape = argmaxShape;
462+ std::vector<int64_t> ksize = {2, 2};
463+ std::vector<int64_t> strides = {2, 2};
464+ std::vector<int64_t> pads = {1, 1};
465+ std::vector<int64_t> dilation = {1, 1};
466+ ge::DataType dtype = ge::DT_BF16;
467+ ge::DataType dtype_index = ge::DT_INT32;
468+ int64_t index_dtype = 3;
469+ bool ceil_mode = false;
470+ std::string data_format = "NHWC";
471+ uint64_t except_tilingkey = 201;
472+ std::string expect = "10 16 2 18 30 2 2 2 2 1 1 1 1 1 1 16 5 3 4 30 30 1 2 2 1 1 1 64 9600 1792 3328 1 1 201 ";
473+ ExecuteTestCase(
474+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
475+ ceil_mode, data_format, except_tilingkey, expect);
476+}
477+ 
478+TEST_F(MaxPoolGradWithArgmaxV3Tiling, MaxPoolGradWithArgmaxV3Tiling_NHWC_Test4)
479+{
480+ gert::StorageShape xShape = {{2, 64, 64, 3}, {2, 64, 64, 3}};
481+ gert::StorageShape yShape = xShape;
482+ gert::StorageShape argmaxShape = {{2, 1, 1, 3}, {2, 1, 1, 3}};
483+ gert::StorageShape gradShape = argmaxShape;
484+ std::vector<int64_t> ksize = {64, 64};
485+ std::vector<int64_t> strides = {64, 64};
486+ std::vector<int64_t> pads = {0, 0};
487+ std::vector<int64_t> dilation = {1, 1};
488+ ge::DataType dtype = ge::DT_FLOAT;
489+ ge::DataType dtype_index = ge::DT_INT32;
490+ int64_t index_dtype = 3;
491+ bool ceil_mode = false;
492+ std::string data_format = "NHWC";
493+ uint64_t except_tilingkey = 201;
494+ std::string expect = "1 1 3 64 64 64 64 64 64 0 0 1 1 1 1 2 2 2 32 64 64 1 3 3 1 1 1 64 4096 384 384 1 1 201 ";
495+ ExecuteTestCase(
496+ xShape, yShape, gradShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, dtype_index,
497+ ceil_mode, data_format, except_tilingkey, expect);
498+}
Mpooling/max_pool_with_argmax_v3/CMakeLists.txt+6-9
@@ -8,12 +8,9 @@
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-file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11+# 设置算子定义时支持的芯片类型
12-if(NOT ENABLE_TEST AND NOT BENCHMARK)12+set(SUPPORT_COMPUTE_UNIT "ascend910_95")
13- list(REMOVE_ITEM CURRENT_DIRS tests)13+# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14-endif()14+set(SUPPORT_TILING_DIR "arch35")
15-foreach(SUB_DIR ${CURRENT_DIRS})15+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE max_pool_with_argmax_v3
16- if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")16+ ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
17- add_subdirectory(${SUB_DIR})
18- endif()
19-endforeach()
Apooling/max_pool_with_argmax_v3/examples/arch35/test_aclnn_max_pool_with_argmax_v3.cpp+171-0
@@ -0,0 +1,171 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+#include <iostream>
11+#include <vector>
12+#include "acl/acl.h"
13+#include "aclnnop/aclnn_max_pool2d_with_indices.h"
14+ 
15+#define CHECK_RET(cond, return_expr) \
16+ do { \
17+ if (!(cond)) { \
18+ return_expr; \
19+ } \
20+ } while (0)
21+ 
22+#define LOG_PRINT(message, ...) \
23+ do { \
24+ printf(message, ##__VA_ARGS__); \
25+ } while (0)
26+ 
27+int64_t GetShapeSize(const std::vector<int64_t>& shape) {
28+ int64_t shapeSize = 1;
29+ for (auto i : shape) {
30+ shapeSize *= i;
31+ }
32+ return shapeSize;
33+}
34+ 
35+int Init(int32_t deviceId, aclrtStream* stream) {
36+ // 固定写法，资源初始化
37+ auto ret = aclInit(nullptr);
38+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
39+ ret = aclrtSetDevice(deviceId);
40+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
41+ ret = aclrtCreateStream(stream);
42+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
43+ return 0;
44+}
45+ 
46+template <typename T>
47+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
48+ aclDataType dataType, aclTensor** tensor) {
49+ auto size = GetShapeSize(shape) * sizeof(T);
50+ // 调用aclrtMalloc申请Device侧内存
51+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
52+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
53+ 
54+ // 调用aclrtMemcpy将Host侧数据拷贝到Device侧内存上
55+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
56+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
57+ 
58+ // 计算连续tensor的strides
59+ std::vector<int64_t> strides(shape.size(), 1);
60+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
61+ strides[i] = shape[i + 1] * strides[i + 1];
62+ }
63+ 
64+ // 调用aclCreateTensor接口创建aclTensor
65+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_NCHW,
66+ shape.data(), shape.size(), *deviceAddr);
67+ return 0;
68+}
69+ 
70+int main() {
71+ // 1. （固定写法）device/stream初始化，参考acl API手册
72+ // 根据自己的实际device填写deviceId
73+ int32_t deviceId = 0;
74+ aclrtStream stream;
75+ auto ret = Init(deviceId, &stream);
76+ // check根据自己的需要处理
77+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
78+ 
79+ // 2. 构造输入与输出，需要根据API的接口自定义构造
80+ std::vector<int64_t> selfShape = {1, 1, 4, 3};
81+ std::vector<int64_t> outShape = {1, 1, 2, 1};
82+ std::vector<int64_t> indicesShape = {1, 1, 2, 1};
83+ std::vector<int64_t> kernelSizeData = {2, 2};
84+ std::vector<int64_t> strideData = {2, 2};
85+ std::vector<int64_t> paddingData = {0, 0};
86+ std::vector<int64_t> dilationData = {1, 1};
87+ void* selfDeviceAddr = nullptr;
88+ void* outDeviceAddr = nullptr;
89+ void* indicesDeviceAddr = nullptr;
90+ aclTensor* self = nullptr;
91+ aclTensor* out = nullptr;
92+ aclTensor* indices = nullptr;
93+ std::vector<float> selfHostData = {0.0850, -0.5147, -0.0212, -0.5654, -0.3222, 0.5847, 1.7510, 0.9954, 0.1842, 0.8392, 0.4835, 0.9213};
94+ std::vector<float> outHostData = {0, 0};
95+ std::vector<int32_t> indicesHostData = {0, 0};
96+ 
97+ // 创建self aclTensor
98+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
99+ CHECK_RET(ret == ACL_SUCCESS, return ret);
100+ // 创建out aclTensor
101+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
102+ CHECK_RET(ret == ACL_SUCCESS, return ret);
103+ // 创建indices aclTensor
104+ ret = CreateAclTensor(indicesHostData, indicesShape, &indicesDeviceAddr, aclDataType::ACL_INT32, &indices);
105+ CHECK_RET(ret == ACL_SUCCESS, return ret);
106+ 
107+ // 创建输入数组
108+ aclIntArray* kernelSize = aclCreateIntArray(kernelSizeData.data(), 2);
109+ aclIntArray* stride = aclCreateIntArray(strideData.data(), 2);
110+ aclIntArray* padding = aclCreateIntArray(paddingData.data(), 2);
111+ aclIntArray* dilation = aclCreateIntArray(dilationData.data(), 2);
112+ const bool ceilMode = false;
113+ 
114+ uint64_t workspaceSize = 0;
115+ aclOpExecutor* executor;
116+ 
117+ // aclnnMaxPool2dWithIndices接口调用示例
118+ // 3. 调用CANN算子库API，需要修改为具体的API名称
119+ // 调用aclnnMaxPool2dWithIndices第一段接口
120+ ret = aclnnMaxPool2dWithIndicesGetWorkspaceSize(self, kernelSize, stride, padding, dilation, ceilMode, out, indices, &workspaceSize, &executor);
121+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaxPool2dWithIndicesGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
122+ // 根据第一段接口计算出的workspaceSize申请device内存
123+ void* workspaceAddr = nullptr;
124+ if (workspaceSize > 0) {
125+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
126+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
127+ }
128+ // 调用aclnnMaxPool2dWithIndices第二段接口
129+ ret = aclnnMaxPool2dWithIndices(workspaceAddr, workspaceSize, executor, stream);
130+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaxPool2dWithIndices failed. ERROR: %d\n", ret); return ret);
131+ 
132+ // 4. （固定写法）同步等待任务执行结束
133+ ret = aclrtSynchronizeStream(stream);
134+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
135+ 
136+ // 5. 获取输出的值，将Device侧内存上的结果拷贝至Host侧，需要根据具体API的接口定义修改
137+ auto size = GetShapeSize(outShape);
138+ std::vector<float> resultData(size, 0);
139+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
140+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
141+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy out result from device to host failed. ERROR: %d\n", ret); return ret);
142+ for (int64_t i = 0; i < size; i++) {
143+ LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
144+ }
145+ 
146+ size = GetShapeSize(indicesShape);
147+ std::vector<int> indicesResultData(size, 0);
148+ ret = aclrtMemcpy(indicesResultData.data(), indicesResultData.size() * sizeof(indicesResultData[0]), indicesDeviceAddr,
149+ size * sizeof(indicesResultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
150+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy indices result from device to host failed. ERROR: %d\n", ret); return ret);
151+ for (int64_t i = 0; i < size; i++) {
152+ LOG_PRINT("result[%ld] is: %d\n", i, indicesResultData[i]);
153+ }
154+ 
155+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
156+ aclDestroyTensor(self);
157+ aclDestroyTensor(out);
158+ aclDestroyTensor(indices);
159+ 
160+ // 7. 释放device资源，需要根据具体API的接口定义修改
161+ aclrtFree(selfDeviceAddr);
162+ aclrtFree(outDeviceAddr);
163+ aclrtFree(indicesDeviceAddr);
164+ if (workspaceSize > 0) {
165+ aclrtFree(workspaceAddr);
166+ }
167+ aclrtDestroyStream(stream);
168+ aclrtResetDevice(deviceId);
169+ aclFinalize();
170+ return 0;
171+}
Rpooling/max_pool_with_argmax_v3/op_host/op_api/max_pool_with_argmax_v3.cpp→pooling/max_pool_with_argmax_v3/op_api/max_pool_with_argmax_v3.cpp+1-1
@@ -6,7 +6,7 @@
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#include "max_pool_with_argmax_v3.h"11#include "max_pool_with_argmax_v3.h"
12#include "opdev/data_type_utils.h"12#include "opdev/data_type_utils.h"
Rpooling/max_pool_with_argmax_v3/op_host/op_api/max_pool_with_argmax_v3.h→pooling/max_pool_with_argmax_v3/op_api/max_pool_with_argmax_v3.h+1-1
@@ -6,7 +6,7 @@
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#ifndef OP_API_INC_LEVEL0_MAX_POOL_WITH_ARGMAX_V3_H_11#ifndef OP_API_INC_LEVEL0_MAX_POOL_WITH_ARGMAX_V3_H_
12#define OP_API_INC_LEVEL0_MAX_POOL_WITH_ARGMAX_V3_H_12#define OP_API_INC_LEVEL0_MAX_POOL_WITH_ARGMAX_V3_H_
Apooling/max_pool_with_argmax_v3/op_graph/max_pool_with_argmax_v3_proto.h+70-0
@@ -0,0 +1,70 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_proto.h
13+ * \brief
14+ */
15+#ifndef OPS_BUILT_IN_OP_PROTO_INC_NN_POOLING_OPS_H_
16+#define OPS_BUILT_IN_OP_PROTO_INC_NN_POOLING_OPS_H_
17+ 
18+#include "graph/operator_reg.h"
19+#include "graph/operator.h"
20+ 
21+namespace ge {
22+ 
23+/**
24+* @brief Performs max pooling on the input and outputs both max values and indices.
25+ 
26+* @par Inputs:
27+* One input:
28+* x: A tensor of type bfloat16, float16, float32, the shape is `[batch, channels, height_in, width_in]` or
29+ `[batch, height_in, width_in, channels]`.
30+ 
31+* @par Attributes:
32+* @li ksize: A required list of int64 values,
33+* specifying the size of the window for each dimension of the input tensor.
34+* A list that has length 2.
35+* @li strides: A required list of int64 values,
36+* specifying the stride of the sliding window for each dimension of the input tensor.
37+* A list that has length 2.
38+* @li pads: A required list of int64 values,
39+* specifying the pad of the input feature map. No default value.
40+* A list that has length 2:
41+* 0 <= pads[0] <= (ksize[0]//2), 0 <= pads[1] <= (ksize[1]//2).
42+* @li dilation: A list that has length 2, default value is {1,1}.
43+* @li dtype: An optional int. default value is 3. (3 is int32, 9 is int64)
44+* @li ceil_mode: When true, will use ceil instead of floor to compute the output shape, defaults to false.
45+* @li data_format: The value can be "NCHW" or "NHWC", defaults to "NCHW".
46+ 
47+* @par Outputs:
48+* @li y: A tensor has the same type and format as input "x", the shape is `[batch, channels, height_out, width_out]` or
49+ `[batch, height_out, width_out, channels]`.
50+* @li argmax: A tensor of type is int64 or int32, the shape is `[batch, channels, height_out, width_out]` or
51+ `[batch, height_out, width_out, channels]`.
52+ 
53+* @par Third-party framework compatibility
54+* Compatible with the PyTorch operator max_pool2d_with_indices.
55+*/
56+REG_OP(MaxPoolWithArgmaxV3)
57+ .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
58+ .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT32, DT_BF16}))
59+ .OUTPUT(argmax, TensorType({DT_INT32, DT_INT64}))
60+ .REQUIRED_ATTR(ksize, ListInt)
61+ .REQUIRED_ATTR(strides, ListInt)
62+ .REQUIRED_ATTR(pads, ListInt)
63+ .ATTR(dtype, Int, 3)
64+ .ATTR(dilation, ListInt, {1, 1})
65+ .ATTR(ceil_mode, Bool, false)
66+ .ATTR(data_format, String, "NCHW")
67+ .OP_END_FACTORY_REG(MaxPoolWithArgmaxV3)
68+ 
69+} // namespace ge
70+#endif // OPS_BUILT_IN_OP_PROTO_INC_NN_POOLING_OPS_H
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_big_kernel_mul_core_tiling.cpp+192-0
@@ -0,0 +1,192 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_big_kernel_mul_core_tiling.cpp
13+ * \brief big kernel imply for max_pool_with_argmax
14+ */
15+ 
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_with_argmax_v3_big_kernel_mul_core_tiling.h"
18+ 
19+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT32 = 400001;
20+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT64 = 400002;
21+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT32 = 400003;
22+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT64 = 400004;
23+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT32 = 400005;
24+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT64 = 400006;
25+static constexpr uint32_t SPLIT_CORE_THRESHOLD = 16;
26+static constexpr uint64_t MIN_SIZE_THRESHOLD = 1024;
27+static constexpr uint64_t MIN_KERNEL_WIDTH_THRESHOLD = 128;
28+static constexpr uint32_t VALUE_WORKSPACE_SIZE = 64 * 4;
29+static constexpr uint32_t INDEX_WORKSPACE_SIZE = 64 * 8;
30+static constexpr int64_t MAX_VALUE_BUFFER_LENGTH = 256;
31+static constexpr int64_t MAX_INDEX_BUFFER_LENGTH = 512;
32+static constexpr int64_t MASK_RATIO = 8;
33+static constexpr int64_t DOUBLE = 2;
34+static constexpr int64_t TRIPPLE = 3;
35+static constexpr int64_t UB_CONST = 65;
36+static constexpr int64_t ALIGN_VALUE = 64;
37+static constexpr uint32_t WS_SYS_SIZE = 16 * 1024 * 1024;
38+using namespace AscendC;
39+ 
40+namespace optiling
41+{
42+ 
43+bool MaxPoolWithArgmaxV3BigKernelMulCoreTiling::IsCapable()
44+{
45+ if (inputData.inputFormat != ge::Format::FORMAT_NCHW) {
46+ return false;
47+ }
48+ if (inputData.dilation[H_DIM] != 1 || inputData.dilation[W_DIM] != 1) {
49+ return false;
50+ }
51+ if (inputData.pad[H_DIM] != 0 || inputData.pad[W_DIM] != 0) {
52+ return false;
53+ }
54+ totalIdx = inputData.batches * inputData.outShape[H_DIM] * inputData.outShape[W_DIM];
55+ if (coreNum == 0) {
56+ return false;
57+ }
58+ uint32_t factor = totalIdx / coreNum;
59+ if ((factor > 0) || (factor == 0 && totalIdx >= SPLIT_CORE_THRESHOLD)) {
60+ return false;
61+ }
62+ uint64_t kernelSize = inputData.kernelSize[H_DIM] * inputData.kernelSize[W_DIM];
63+ if (kernelSize < MIN_SIZE_THRESHOLD) {
64+ return false;
65+ }
66+ return true;
67+}
68+ 
69+uint64_t MaxPoolWithArgmaxV3BigKernelMulCoreTiling::GetTilingKey() const
70+{
71+ if (inputData.indexDtype == ge::DataType::DT_INT32) {
72+ if (dtype == ge::DataType::DT_BF16) {
73+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT32;
74+ } else if (dtype == ge::DataType::DT_FLOAT16) {
75+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT32;
76+ } else {
77+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT32;
78+ }
79+ } else {
80+ if (dtype == ge::DataType::DT_BF16) {
81+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT64;
82+ } else if (dtype == ge::DataType::DT_FLOAT16) {
83+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT64;
84+ } else {
85+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT64;
86+ }
87+ }
88+}
89+ 
90+void MaxPoolWithArgmaxV3BigKernelMulCoreTiling::DoUBTiling()
91+{
92+ maxCountLength =
93+ (ubSize - TRIPPLE * MAX_VALUE_BUFFER_LENGTH - DOUBLE * MAX_INDEX_BUFFER_LENGTH) * MASK_RATIO / UB_CONST;
94+ maxCountLength = maxCountLength / ALIGN_VALUE * ALIGN_VALUE;
95+ coreNums = totalIdx;
96+ uint64_t coreNumRatio = coreNum / coreNums;
97+ if (inputData.kernelSize[H_DIM] > coreNumRatio / DOUBLE ||
98+ inputData.kernelSize[W_DIM] < MIN_KERNEL_WIDTH_THRESHOLD) {
99+ kernelBlockFactorH = Ops::Base::CeilDiv(inputData.kernelSize[H_DIM], coreNumRatio);
100+ multiCoreNum = Ops::Base::CeilDiv(inputData.kernelSize[H_DIM], kernelBlockFactorH);
101+ tailKernelBlockFactorH = inputData.kernelSize[H_DIM] - (multiCoreNum - 1) * kernelBlockFactorH;
102+ splitW = 0;
103+ } else {
104+ splitSlice = coreNumRatio / inputData.kernelSize[H_DIM];
105+ wSplitSize = max(Ops::Base::CeilDiv(inputData.kernelSize[W_DIM], splitSlice), MIN_KERNEL_WIDTH_THRESHOLD);
106+ splitSlice = Ops::Base::CeilDiv(inputData.kernelSize[W_DIM], wSplitSize);
107+ splitW = 1;
108+ tailWSplitSize = inputData.kernelSize[W_DIM] - (splitSlice - 1) * wSplitSize;
109+ multiCoreNum = inputData.kernelSize[H_DIM] * splitSlice;
110+ }
111+}
112+ 
113+void MaxPoolWithArgmaxV3BigKernelMulCoreTiling::SetTilingData()
114+{
115+ tiling.set_hInDim(inputData.inputShape[H_DIM]);
116+ tiling.set_wInDim(inputData.inputShape[W_DIM]);
117+ tiling.set_hOutDim(inputData.outShape[H_DIM]);
118+ tiling.set_wOutDim(inputData.outShape[W_DIM]);
119+ tiling.set_kH(inputData.kernelSize[H_DIM]);
120+ tiling.set_kW(inputData.kernelSize[W_DIM]);
121+ tiling.set_sH(inputData.stride[H_DIM]);
122+ tiling.set_sW(inputData.stride[W_DIM]);
123+ tiling.set_pH(inputData.pad[H_DIM]);
124+ tiling.set_pW(inputData.pad[W_DIM]);
125+ tiling.set_dH(inputData.dilation[H_DIM]);
126+ tiling.set_dW(inputData.dilation[W_DIM]);
127+ tiling.set_coreNums(coreNums);
128+ tiling.set_multiCoreNum(multiCoreNum);
129+ tiling.set_kernelBlockFactorH(kernelBlockFactorH);
130+ tiling.set_tailKernelBlockFactorH(tailKernelBlockFactorH);
131+ tiling.set_splitW(splitW);
132+ tiling.set_wSplitSize(wSplitSize);
133+ tiling.set_tailWSplitSize(tailWSplitSize);
134+ tiling.set_splitSlice(splitSlice);
135+ tiling.set_maxCountLength(maxCountLength);
136+ tiling.set_valueBufferLength(MAX_VALUE_BUFFER_LENGTH);
137+ tiling.set_indexBufferLength(MAX_INDEX_BUFFER_LENGTH);
138+}
139+ 
140+ge::graphStatus MaxPoolWithArgmaxV3BigKernelMulCoreTiling::DoOpTiling()
141+{
142+ DoUBTiling();
143+ SetTilingData();
144+ return ge::GRAPH_SUCCESS;
145+}
146+ge::graphStatus MaxPoolWithArgmaxV3BigKernelMulCoreTiling::GetWorkspaceSize()
147+{
148+ auto sysWorkspace = WS_SYS_SIZE + VALUE_WORKSPACE_SIZE + INDEX_WORKSPACE_SIZE;
149+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
150+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
151+ currentWorkspace[0] = sysWorkspace;
152+ return ge::GRAPH_SUCCESS;
153+}
154+ge::graphStatus MaxPoolWithArgmaxV3BigKernelMulCoreTiling::PostTiling()
155+{
156+ context_->SetBlockDim(coreNums * multiCoreNum);
157+ tiling.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
158+ context_->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
159+ return ge::GRAPH_SUCCESS;
160+}
161+void MaxPoolWithArgmaxV3BigKernelMulCoreTiling::DumpTilingInfo()
162+{
163+ std::string str;
164+ str += " hInDim:" + std::to_string(tiling.get_hInDim());
165+ str += " wInDim:" + std::to_string(tiling.get_wInDim());
166+ str += " hOutDim:" + std::to_string(tiling.get_hOutDim());
167+ str += " wOutDim:" + std::to_string(tiling.get_wOutDim());
168+ str += " kH:" + std::to_string(tiling.get_kH());
169+ str += " kW:" + std::to_string(tiling.get_kW());
170+ str += " sH:" + std::to_string(tiling.get_sH());
171+ str += " sW:" + std::to_string(tiling.get_sW());
172+ str += " pH:" + std::to_string(tiling.get_pH());
173+ str += " pW:" + std::to_string(tiling.get_pW());
174+ str += " dH:" + std::to_string(tiling.get_dH());
175+ str += " dW:" + std::to_string(tiling.get_dW());
176+ str += " coreNums:" + std::to_string(tiling.get_coreNums());
177+ str += " multiCoreNum:" + std::to_string(tiling.get_multiCoreNum());
178+ str += " kernelBlockFactorH:" + std::to_string(tiling.get_kernelBlockFactorH());
179+ str += " tailKernelBlockFactorH:" + std::to_string(tiling.get_tailKernelBlockFactorH());
180+ str += " splitW:" + std::to_string(tiling.get_splitW());
181+ str += " wSplitSize:" + std::to_string(tiling.get_wSplitSize());
182+ str += " tailWSplitSize:" + std::to_string(tiling.get_tailWSplitSize());
183+ str += " splitSlice:" + std::to_string(tiling.get_splitSlice());
184+ str += " maxCountLength:" + std::to_string(tiling.get_maxCountLength());
185+ str += " valueBufferLength:" + std::to_string(tiling.get_valueBufferLength());
186+ str += " indexBufferLength:" + std::to_string(tiling.get_indexBufferLength());
187+ OP_LOGI(context_, "%s", str.c_str());
188+}
189+ 
190+REGISTER_OPS_TILING_TEMPLATE(MaxPoolWithArgmaxV3, MaxPoolWithArgmaxV3BigKernelMulCoreTiling, 4);
191+ 
192+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_big_kernel_mul_core_tiling.h+87-0
@@ -0,0 +1,87 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_big_kernel_mul_core_tiling.h
13+ * \brief big kernel imply for max_pool_with_argmax
14+ */
15+ 
16+#ifndef CANN_MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_MUL_CORE_TILING_H
17+#define CANN_MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_MUL_CORE_TILING_H
18+ 
19+#include "max_pool_with_argmax_v3_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+BEGIN_TILING_DATA_DEF(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData)
24+TILING_DATA_FIELD_DEF(int64_t, hInDim);
25+TILING_DATA_FIELD_DEF(int64_t, wInDim);
26+TILING_DATA_FIELD_DEF(int64_t, hOutDim);
27+TILING_DATA_FIELD_DEF(int64_t, wOutDim);
28+TILING_DATA_FIELD_DEF(int64_t, kW);
29+TILING_DATA_FIELD_DEF(int64_t, kH);
30+TILING_DATA_FIELD_DEF(int64_t, sW);
31+TILING_DATA_FIELD_DEF(int64_t, sH);
32+TILING_DATA_FIELD_DEF(int64_t, pW);
33+TILING_DATA_FIELD_DEF(int64_t, pH);
34+TILING_DATA_FIELD_DEF(int64_t, dW);
35+TILING_DATA_FIELD_DEF(int64_t, dH);
36+TILING_DATA_FIELD_DEF(int64_t, coreNums);
37+TILING_DATA_FIELD_DEF(int64_t, multiCoreNum);
38+TILING_DATA_FIELD_DEF(int64_t, kernelBlockFactorH);
39+TILING_DATA_FIELD_DEF(int64_t, tailKernelBlockFactorH);
40+TILING_DATA_FIELD_DEF(int64_t, splitW);
41+TILING_DATA_FIELD_DEF(int64_t, wSplitSize);
42+TILING_DATA_FIELD_DEF(int64_t, tailWSplitSize);
43+TILING_DATA_FIELD_DEF(int64_t, splitSlice);
44+TILING_DATA_FIELD_DEF(int64_t, maxCountLength);
45+TILING_DATA_FIELD_DEF(int64_t, valueBufferLength);
46+TILING_DATA_FIELD_DEF(int64_t, indexBufferLength);
47+END_TILING_DATA_DEF;
48+ 
49+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_400001, MaxPoolWithArgmaxV3BigKernelMulCoreTilingData);
50+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_400002, MaxPoolWithArgmaxV3BigKernelMulCoreTilingData);
51+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_400003, MaxPoolWithArgmaxV3BigKernelMulCoreTilingData);
52+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_400004, MaxPoolWithArgmaxV3BigKernelMulCoreTilingData);
53+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_400005, MaxPoolWithArgmaxV3BigKernelMulCoreTilingData);
54+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_400006, MaxPoolWithArgmaxV3BigKernelMulCoreTilingData);
55+ 
56+class MaxPoolWithArgmaxV3BigKernelMulCoreTiling : public MaxPoolWithArgmaxV3BaseTiling {
57+public:
58+ explicit MaxPoolWithArgmaxV3BigKernelMulCoreTiling(gert::TilingContext* context)
59+ : MaxPoolWithArgmaxV3BaseTiling(context)
60+ {}
61+ ~MaxPoolWithArgmaxV3BigKernelMulCoreTiling() override
62+ {}
63+ 
64+private:
65+ void DoUBTiling();
66+ void SetTilingData();
67+ uint64_t GetTilingKey() const;
68+ bool IsCapable() override;
69+ ge::graphStatus DoOpTiling() override;
70+ ge::graphStatus PostTiling() override;
71+ ge::graphStatus GetWorkspaceSize() override;
72+ void DumpTilingInfo() override;
73+ MaxPoolWithArgmaxV3BigKernelMulCoreTilingData tiling;
74+ uint32_t totalIdx{0};
75+ uint32_t coreNums{0};
76+ int64_t multiCoreNum{0};
77+ uint64_t kernelBlockFactorH{0};
78+ int64_t tailKernelBlockFactorH{0};
79+ int64_t splitW{0};
80+ uint64_t wSplitSize{0};
81+ int64_t tailWSplitSize{0};
82+ uint64_t splitSlice{0};
83+ int64_t maxCountLength{0};
84+};
85+ 
86+} // namespace optiling
87+#endif // CANN_MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_MUL_CORE_TILING_H
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_big_kernel_tiling.cpp+130-0
@@ -0,0 +1,130 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_big_kernel_tiling.cpp
13+ * \brief big kernel imply for max_pool_with_argmax
14+ */
15+ 
16+#include "op_util.h"
17+#include "platform/platform_info.h"
18+#include "tiling_base/tiling_templates_registry.h"
19+#include "max_pool_with_argmax_v3_big_kernel_tiling.h"
20+ 
21+static constexpr uint64_t MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_NCHW = 311110;
22+ 
23+static constexpr int64_t OUT_BUFFER_LEN = 1024;
24+static constexpr int64_t BUFFER_NUM = 2;
25+static constexpr int64_t MIN_COUNT = 1024;
26+static constexpr int64_t BYTES_FOUR = 4;
27+static constexpr int64_t BYTES_EIGHT = 8;
28+static constexpr int64_t KW_THRESHOLD = 128;
29+using namespace AscendC;
30+ 
31+namespace optiling {
32+ 
33+bool MaxPoolWithArgmaxV3BigKernelTiling::IsCapable()
34+{
35+ int64_t ubAvailable = ubSize - (BYTES_FOUR + BYTES_EIGHT) * OUT_BUFFER_LEN;
36+ maxCount_ = ubAvailable / BUFFER_NUM;
37+ int64_t vRegSize = Ops::Base::GetVRegSize(context_);
38+ maxCount_ = Ops::Base::FloorAlign(maxCount_, vRegSize);
39+ int64_t dtypeSize = ge::GetSizeByDataType(dtype);
40+ OP_CHECK_IF(
41+ dtypeSize <= 0, OP_LOGE(context_, "dtypeSize must be greater than 0, dtypeSize: %ld", dtypeSize), return false);
42+ maxCount_ = maxCount_ / dtypeSize;
43+ if (inputData.dilation[H_DIM] == 1 && inputData.dilation[W_DIM] == 1 && maxCount_ > MIN_COUNT &&
44+ inputData.inputFormat == ge::Format::FORMAT_NCHW && inputData.kernelSize[W_DIM] * dtypeSize > KW_THRESHOLD) {
45+ return true;
46+ }
47+ return false;
48+}
49+ 
50+uint64_t MaxPoolWithArgmaxV3BigKernelTiling::GetTilingKey() const
51+{
52+ return MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_NCHW;
53+}
54+ 
55+void MaxPoolWithArgmaxV3BigKernelTiling::DoUBTiling()
56+{
57+ totalIdx_ = inputData.batches * inputData.outShape[H_DIM] * inputData.outShape[W_DIM];
58+ // coreNum已在tiling_base中校验过非0
59+ blockFactor_ = totalIdx_ / coreNum;
60+ blockTail_ = totalIdx_ % coreNum;
61+ if (blockFactor_ == 0) {
62+ coreNums_ = totalIdx_;
63+ } else {
64+ coreNums_ = coreNum;
65+ }
66+ isSigOut_ = (inputData.outShape[H_DIM] == 1 && inputData.outShape[W_DIM] == 1) ? 1 : 0;
67+}
68+ 
69+void MaxPoolWithArgmaxV3BigKernelTiling::SetTilingData()
70+{
71+ tiling.set_hInDim(inputData.inputShape[H_DIM]);
72+ tiling.set_wInDim(inputData.inputShape[W_DIM]);
73+ tiling.set_hOutDim(inputData.outShape[H_DIM]);
74+ tiling.set_wOutDim(inputData.outShape[W_DIM]);
75+ tiling.set_kH(inputData.kernelSize[H_DIM]);
76+ tiling.set_kW(inputData.kernelSize[W_DIM]);
77+ tiling.set_sH(inputData.stride[H_DIM]);
78+ tiling.set_sW(inputData.stride[W_DIM]);
79+ tiling.set_pH(inputData.pad[H_DIM]);
80+ tiling.set_pW(inputData.pad[W_DIM]);
81+ tiling.set_dH(inputData.dilation[H_DIM]);
82+ tiling.set_dW(inputData.dilation[W_DIM]);
83+ tiling.set_blockFactor(blockFactor_);
84+ tiling.set_blockTail(blockTail_);
85+ tiling.set_totalIdx(totalIdx_);
86+ tiling.set_coreNums(coreNums_);
87+ tiling.set_maxCount(maxCount_);
88+ tiling.set_isSigOut(isSigOut_);
89+}
90+ 
91+ge::graphStatus MaxPoolWithArgmaxV3BigKernelTiling::DoOpTiling()
92+{
93+ DoUBTiling();
94+ SetTilingData();
95+ return ge::GRAPH_SUCCESS;
96+}
97+ 
98+ge::graphStatus MaxPoolWithArgmaxV3BigKernelTiling::PostTiling()
99+{
100+ context_->SetBlockDim(coreNums_);
101+ tiling.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
102+ context_->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
103+ return ge::GRAPH_SUCCESS;
104+}
105+ 
106+void MaxPoolWithArgmaxV3BigKernelTiling::DumpTilingInfo()
107+{
108+ std::string str;
109+ str += " hInDim:" + std::to_string(tiling.get_hInDim());
110+ str += " wInDim:" + std::to_string(tiling.get_wInDim());
111+ str += " hOutDim:" + std::to_string(tiling.get_hOutDim());
112+ str += " wOutDim:" + std::to_string(tiling.get_wOutDim());
113+ str += " kH:" + std::to_string(tiling.get_kH());
114+ str += " kW:" + std::to_string(tiling.get_kW());
115+ str += " sH:" + std::to_string(tiling.get_sH());
116+ str += " sW:" + std::to_string(tiling.get_sW());
117+ str += " pH:" + std::to_string(tiling.get_pH());
118+ str += " pW:" + std::to_string(tiling.get_pW());
119+ str += " blockFactor:" + std::to_string(tiling.get_blockFactor());
120+ str += " blockTail:" + std::to_string(tiling.get_blockTail());
121+ str += " totalIdx:" + std::to_string(tiling.get_totalIdx());
122+ str += " coreNums:" + std::to_string(tiling.get_coreNums());
123+ str += " maxCount:" + std::to_string(tiling.get_maxCount());
124+ str += " isSigOut:" + std::to_string(tiling.get_isSigOut());
125+ OP_LOGI(context_, "%s", str.c_str());
126+}
127+ 
128+REGISTER_OPS_TILING_TEMPLATE(MaxPoolWithArgmaxV3, MaxPoolWithArgmaxV3BigKernelTiling, 6);
129+ 
130+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_big_kernel_tiling.h+71-0
@@ -0,0 +1,71 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_big_kernel_tiling.h
13+ * \brief big kernel imply for max_pool_with_argmax
14+ */
15+ 
16+#ifndef CANN_MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_TILING_H
17+#define CANN_MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_TILING_H
18+ 
19+#include "max_pool_with_argmax_v3_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+BEGIN_TILING_DATA_DEF(MaxPoolWithArgmaxV3BigKernelTilingData)
24+TILING_DATA_FIELD_DEF(int64_t, hInDim);
25+TILING_DATA_FIELD_DEF(int64_t, wInDim);
26+TILING_DATA_FIELD_DEF(int64_t, hOutDim);
27+TILING_DATA_FIELD_DEF(int64_t, wOutDim);
28+TILING_DATA_FIELD_DEF(int64_t, kW);
29+TILING_DATA_FIELD_DEF(int64_t, kH);
30+TILING_DATA_FIELD_DEF(int64_t, sW);
31+TILING_DATA_FIELD_DEF(int64_t, sH);
32+TILING_DATA_FIELD_DEF(int64_t, pW);
33+TILING_DATA_FIELD_DEF(int64_t, pH);
34+TILING_DATA_FIELD_DEF(int64_t, dW);
35+TILING_DATA_FIELD_DEF(int64_t, dH);
36+TILING_DATA_FIELD_DEF(int64_t, blockFactor);
37+TILING_DATA_FIELD_DEF(int64_t, blockTail);
38+TILING_DATA_FIELD_DEF(int64_t, totalIdx);
39+TILING_DATA_FIELD_DEF(int64_t, coreNums);
40+TILING_DATA_FIELD_DEF(int64_t, maxCount);
41+TILING_DATA_FIELD_DEF(int64_t, isSigOut);
42+END_TILING_DATA_DEF;
43+ 
44+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_311110, MaxPoolWithArgmaxV3BigKernelTilingData);
45+ 
46+class MaxPoolWithArgmaxV3BigKernelTiling : public MaxPoolWithArgmaxV3BaseTiling {
47+public:
48+ explicit MaxPoolWithArgmaxV3BigKernelTiling(gert::TilingContext* context) : MaxPoolWithArgmaxV3BaseTiling(context)
49+ {}
50+ ~MaxPoolWithArgmaxV3BigKernelTiling() override
51+ {}
52+ 
53+private:
54+ void DoUBTiling();
55+ void SetTilingData();
56+ uint64_t GetTilingKey() const;
57+ bool IsCapable() override;
58+ ge::graphStatus DoOpTiling() override;
59+ ge::graphStatus PostTiling() override;
60+ void DumpTilingInfo() override;
61+ MaxPoolWithArgmaxV3BigKernelTilingData tiling;
62+ int64_t totalIdx_{0};
63+ int64_t blockFactor_{0};
64+ int64_t blockTail_{0};
65+ int64_t maxCount_{0};
66+ int64_t isSigOut_{0};
67+ int64_t coreNums_{0};
68+};
69+ 
70+} // namespace optiling
71+#endif // CANN_MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_TILING_H
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_gather_tiling.cpp+317-0
@@ -0,0 +1,317 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_gather_tiling.cpp
13+ * \brief
14+ */
15+#include "platform/platform_info.h"
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_with_argmax_v3_gather_tiling.h"
18+ 
19+namespace 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 = 0;
25+static constexpr int64_t HELPER_BUFFER_SIZE = 1024;
26+static constexpr int64_t NO_PADDING_TILING_KEY = 300001;
27+static constexpr int64_t PADDING_TILING_KEY = 300002;
28+static constexpr int64_t MAX_BANDWIDTH_COEFFICIENTS = 2;
29+static constexpr int64_t DOUBLE = 2;
30+static constexpr int64_t CACHE_LINE_SIZE = 128;
31+static constexpr int64_t MIN_OUTPUT_THRESHOLD = 32;
32+static constexpr int64_t DILATION_THRESHOLD = 1;
33+ 
34+void MaxPoolWithArgmaxV3GatherTiling::InitializationVars()
35+{
36+ baseData_.inputBytes = dtype == ge::DT_FLOAT ? FLOAT32_SIZE : FLOAT16_SIZE;
37+ baseData_.indexBytes = inputData.indexDtype == ge::DT_INT32 ? INT32_SIZE : INT64_SIZE;
38+ baseData_.availableUb = ubSize - UB_RESVERVED_SIZE;
39+ baseData_.totalCoreNum = coreNum;
40+ baseData_.coreUsedForBestPerformance = baseData_.inputBytes == FLOAT32_SIZE ?
41+ baseData_.totalCoreNum / MAX_BANDWIDTH_COEFFICIENTS :
42+ baseData_.totalCoreNum;
43+ 
44+ baseData_.padTop = inputData.pad[H_DIM];
45+ baseData_.padLeft = inputData.pad[W_DIM];
46+ baseData_.hInput = inputData.inputShape[H_DIM];
47+ baseData_.wInput = inputData.inputShape[W_DIM];
48+ baseData_.hOutput = inputData.outShape[H_DIM];
49+ baseData_.wOutput = inputData.outShape[W_DIM];
50+ baseData_.hStride = inputData.stride[H_DIM];
51+ baseData_.wStride = inputData.stride[W_DIM];
52+ baseData_.hKernel = inputData.kernelSize[H_DIM];
53+ baseData_.wKernel = inputData.kernelSize[W_DIM];
54+ baseData_.highAxisTotal = inputData.batches;
55+ baseData_.hDilation = inputData.dilation[H_DIM];
56+ baseData_.wDilation = inputData.dilation[W_DIM];
57+ baseData_.isPad = 0;
58+ if (baseData_.padTop != 0 || baseData_.padLeft != 0) {
59+ baseData_.isPad = 1;
60+ }
61+ 
62+ if (inputData.ceilMode && baseData_.isPad == 0) {
63+ if (((baseData_.wOutput - 1) * baseData_.wStride + baseData_.wKernel) != baseData_.wInput ||
64+ ((baseData_.hOutput - 1) * baseData_.hStride + baseData_.hKernel) != baseData_.hInput) {
65+ baseData_.isPad = 1;
66+ }
67+ }
68+ 
69+ baseData_.oneBlockNumT1 = Ops::Base::GetUbBlockSize(context_) / baseData_.inputBytes;
70+ baseData_.oneBlockNumT2 = Ops::Base::GetUbBlockSize(context_) / baseData_.indexBytes;
71+}
72+ 
73+bool MaxPoolWithArgmaxV3GatherTiling::IsCapable()
74+{
75+ if (inputData.dilation[H_DIM] > DILATION_THRESHOLD || inputData.dilation[W_DIM] > DILATION_THRESHOLD ||
76+ inputData.inputFormat != ge::Format::FORMAT_NCHW) {
77+ return false;
78+ }
79+ 
80+ InitializationVars();
81+ if (baseData_.wKernel * baseData_.inputBytes >= CACHE_LINE_SIZE) {
82+ return false;
83+ }
84+ 
85+ splitData_.hOutputInner = 1;
86+ splitData_.wOutputInner = 1;
87+ splitData_.highAxisInner = 1;
88+ DoBufferCalculate();
89+ return splitData_.totalBufferSize <= baseData_.availableUb / MIN_OUTPUT_THRESHOLD;
90+}
91+ 
92+uint64_t MaxPoolWithArgmaxV3GatherTiling::GetTilingKey() const
93+{
94+ uint64_t tilingKey = NO_PADDING_TILING_KEY;
95+ if (baseData_.isPad == 1) {
96+ tilingKey = PADDING_TILING_KEY;
97+ }
98+ return tilingKey;
99+}
100+ 
101+void MaxPoolWithArgmaxV3GatherTiling::DoBufferCalculate()
102+{
103+ splitData_.hInputInner =
104+ (splitData_.hOutputInner - 1) * baseData_.hStride + (baseData_.hKernel - 1) * baseData_.hDilation + 1;
105+ splitData_.wInputInner =
106+ (splitData_.wOutputInner - 1) * baseData_.wStride + (baseData_.wKernel - 1) * baseData_.wDilation + 1;
107+ int64_t maxDataNumInOneBlock = std::max(baseData_.oneBlockNumT1, baseData_.oneBlockNumT2);
108+ int64_t wInputInnerAligned = Ops::Base::CeilAlign(splitData_.wInputInner, baseData_.oneBlockNumT1);
109+ int64_t wOutputInnerAligned = Ops::Base::CeilAlign(splitData_.wOutputInner, maxDataNumInOneBlock);
110+ 
111+ int64_t inputBufferSize =
112+ splitData_.highAxisInner * splitData_.hInputInner * wInputInnerAligned * baseData_.inputBytes;
113+ splitData_.inputBufferSize = inputBufferSize;
114+ // pad情况下COPY IN的UB地址不一定32字节对齐
115+ if (baseData_.isPad == 1) {
116+ inputBufferSize *= DOUBLE;
117+ }
118+ int64_t outputDataSize = splitData_.highAxisInner * splitData_.hOutputInner * wOutputInnerAligned;
119+ splitData_.maxValueBufferSize = outputDataSize * baseData_.inputBytes;
120+ splitData_.argmaxBufferSize = outputDataSize * baseData_.indexBytes;
121+ 
122+ int64_t tmpTotalBufferSize =
123+ inputBufferSize + splitData_.maxValueBufferSize + splitData_.argmaxBufferSize + HELPER_BUFFER_SIZE;
124+ 
125+ splitData_.totalBufferSize = tmpTotalBufferSize * DOUBLE;
126+ if (baseData_.isPad == 1) {
127+ splitData_.totalBufferSize -= splitData_.inputBufferSize;
128+ }
129+}
130+ 
131+bool MaxPoolWithArgmaxV3GatherTiling::IsMeetTargetCoreNum() const
132+{
133+ int64_t tmpWOutputOuter = Ops::Base::CeilDiv(baseData_.wOutput, splitData_.wOutputInner);
134+ int64_t tmpHOutputOuter = Ops::Base::CeilDiv(baseData_.hOutput, splitData_.hOutputInner);
135+ int64_t tmpNOutputOuter = Ops::Base::CeilDiv(baseData_.highAxisTotal, splitData_.highAxisInner);
136+ return tmpWOutputOuter * tmpHOutputOuter * tmpNOutputOuter >= baseData_.coreUsedForBestPerformance;
137+}
138+ 
139+bool MaxPoolWithArgmaxV3GatherTiling::IsMeetUBSize()
140+{
141+ DoBufferCalculate();
142+ return splitData_.totalBufferSize <= baseData_.availableUb;
143+}
144+ 
145+void MaxPoolWithArgmaxV3GatherTiling::BinarySearch(int64_t start, int64_t end, int64_t* value)
146+{
147+ int64_t left = start;
148+ int64_t right = end;
149+ int64_t bestSplit = 1;
150+ 
151+ while (left <= right) {
152+ int64_t mid = left + (right - left) / DOUBLE;
153+ *value = mid;
154+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
155+ bestSplit = mid;
156+ left = mid + 1;
157+ } else {
158+ right = mid - 1;
159+ }
160+ }
161+ 
162+ *value = bestSplit;
163+}
164+ 
165+bool MaxPoolWithArgmaxV3GatherTiling::TrySplitNC()
166+{
167+ splitData_.hOutputInner = baseData_.hOutput;
168+ splitData_.wOutputInner = baseData_.wOutput;
169+ 
170+ splitData_.highAxisInner = Ops::Base::CeilDiv(baseData_.highAxisTotal, baseData_.coreUsedForBestPerformance);
171+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
172+ return true;
173+ }
174+ 
175+ splitData_.highAxisInner = 1;
176+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
177+ BinarySearch(1, baseData_.highAxisTotal, &splitData_.highAxisInner);
178+ return true;
179+ } else {
180+ return false;
181+ }
182+}
183+ 
184+bool MaxPoolWithArgmaxV3GatherTiling::TrySplitH()
185+{
186+ splitData_.highAxisInner = 1;
187+ splitData_.wOutputInner = baseData_.wOutput;
188+ splitData_.hOutputInner = 1;
189+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
190+ BinarySearch(1, baseData_.hOutput, &splitData_.hOutputInner);
191+ return true;
192+ } else {
193+ return false;
194+ }
195+}
196+ 
197+bool MaxPoolWithArgmaxV3GatherTiling::TrySplitW()
198+{
199+ splitData_.highAxisInner = 1;
200+ splitData_.hOutputInner = 1;
201+ 
202+ splitData_.wOutputInner = 1;
203+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
204+ BinarySearch(1, baseData_.wOutput, &splitData_.wOutputInner);
205+ return true;
206+ } else {
207+ return false;
208+ }
209+}
210+ 
211+void MaxPoolWithArgmaxV3GatherTiling::SearchBestTiling()
212+{
213+ if (TrySplitNC()) {
214+ return;
215+ }
216+ if (TrySplitH()) {
217+ return;
218+ }
219+ if (TrySplitW()) {
220+ return;
221+ }
222+}
223+ 
224+void MaxPoolWithArgmaxV3GatherTiling::DoUBTiling()
225+{
226+ SearchBestTiling();
227+ DoBufferCalculate();
228+ 
229+ splitData_.wOutputOuter = Ops::Base::CeilDiv(baseData_.wOutput, splitData_.wOutputInner);
230+ int64_t tempWOutputTail = baseData_.wOutput % splitData_.wOutputInner;
231+ splitData_.wOutputTail = tempWOutputTail == 0 ? splitData_.wOutputInner : tempWOutputTail;
232+ 
233+ splitData_.hOutputOuter = Ops::Base::CeilDiv(baseData_.hOutput, splitData_.hOutputInner);
234+ int64_t tempHOutputTail = baseData_.hOutput % splitData_.hOutputInner;
235+ splitData_.hOutputTail = tempHOutputTail == 0 ? splitData_.hOutputInner : tempHOutputTail;
236+ 
237+ splitData_.highAxisOuter = Ops::Base::CeilDiv(baseData_.highAxisTotal, splitData_.highAxisInner);
238+ int64_t tempNOutputTail = baseData_.highAxisTotal % splitData_.highAxisInner;
239+ splitData_.highAxisTail = tempNOutputTail == 0 ? splitData_.highAxisInner : tempNOutputTail;
240+}
241+ 
242+void MaxPoolWithArgmaxV3GatherTiling::DoBlockTiling()
243+{
244+ splitData_.totalBaseBlockNum = splitData_.highAxisOuter * splitData_.hOutputOuter * splitData_.wOutputOuter;
245+ splitData_.normalCoreProcessNum = Ops::Base::CeilDiv(splitData_.totalBaseBlockNum, baseData_.totalCoreNum);
246+ splitData_.usedCoreNum = Ops::Base::CeilDiv(splitData_.totalBaseBlockNum, splitData_.normalCoreProcessNum);
247+ splitData_.tailCoreProcessNum =
248+ splitData_.totalBaseBlockNum - splitData_.normalCoreProcessNum * (splitData_.usedCoreNum - 1);
249+}
250+ 
251+void MaxPoolWithArgmaxV3GatherTiling::PrintBaseData() const
252+{
253+ OP_LOGI("PrintBaseData", "%s", baseData_.ToString().c_str());
254+}
255+ 
256+void MaxPoolWithArgmaxV3GatherTiling::PrintSplitData() const
257+{
258+ OP_LOGI("PrintSplitData", "%s", splitData_.ToString().c_str());
259+}
260+ 
261+void MaxPoolWithArgmaxV3GatherTiling::SetTilingData()
262+{
263+ tilingData_.set_hInput(baseData_.hInput);
264+ tilingData_.set_wInput(baseData_.wInput);
265+ tilingData_.set_hOutput(baseData_.hOutput);
266+ tilingData_.set_wOutput(baseData_.wOutput);
267+ tilingData_.set_hKernel(baseData_.hKernel);
268+ tilingData_.set_wKernel(baseData_.wKernel);
269+ tilingData_.set_hStride(baseData_.hStride);
270+ tilingData_.set_wStride(baseData_.wStride);
271+ tilingData_.set_padTop(baseData_.padTop);
272+ tilingData_.set_padLeft(baseData_.padLeft);
273+ tilingData_.set_highAxisInner(splitData_.highAxisInner);
274+ tilingData_.set_highAxisTail(splitData_.highAxisTail);
275+ tilingData_.set_highAxisOuter(splitData_.highAxisOuter);
276+ tilingData_.set_hOutputInner(splitData_.hOutputInner);
277+ tilingData_.set_hOutputTail(splitData_.hOutputTail);
278+ tilingData_.set_hOutputOuter(splitData_.hOutputOuter);
279+ tilingData_.set_wOutputInner(splitData_.wOutputInner);
280+ tilingData_.set_wOutputTail(splitData_.wOutputTail);
281+ tilingData_.set_wOutputOuter(splitData_.wOutputOuter);
282+ tilingData_.set_normalCoreProcessNum(splitData_.normalCoreProcessNum);
283+ tilingData_.set_tailCoreProcessNum(splitData_.tailCoreProcessNum);
284+ tilingData_.set_usedCoreNum(splitData_.usedCoreNum);
285+ tilingData_.set_inputBufferSize(splitData_.inputBufferSize);
286+ tilingData_.set_maxValueBufferSize(splitData_.maxValueBufferSize);
287+ tilingData_.set_argmaxBufferSize(splitData_.argmaxBufferSize);
288+ tilingData_.set_isPad(baseData_.isPad);
289+ tilingData_.set_hDilation(baseData_.hDilation);
290+ tilingData_.set_wDilation(baseData_.wDilation);
291+}
292+ 
293+ge::graphStatus MaxPoolWithArgmaxV3GatherTiling::DoOpTiling()
294+{
295+ DoUBTiling();
296+ DoBlockTiling();
297+ SetTilingData();
298+ PrintBaseData();
299+ PrintSplitData();
300+ return ge::GRAPH_SUCCESS;
301+}
302+ 
303+ge::graphStatus MaxPoolWithArgmaxV3GatherTiling::PostTiling()
304+{
305+ context_->SetBlockDim(tilingData_.get_usedCoreNum());
306+ if (tilingData_.GetDataSize() > context_->GetRawTilingData()->GetCapacity()) {
307+ return ge::GRAPH_FAILED;
308+ }
309+ 
310+ tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
311+ context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
312+ return ge::GRAPH_SUCCESS;
313+}
314+ 
315+REGISTER_OPS_TILING_TEMPLATE(MaxPoolWithArgmaxV3, MaxPoolWithArgmaxV3GatherTiling, 0);
316+ 
317+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_gather_tiling.h+178-0
@@ -0,0 +1,178 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_gather_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_WITH_AGRMAX_V3_GATHER_TILING_H_
17+#define MAX_POOL_WITH_AGRMAX_V3_GATHER_TILING_H_
18+ 
19+#include "max_pool_with_argmax_v3_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+BEGIN_TILING_DATA_DEF(MaxPoolWithArgmaxV3GatherTilingData)
24+TILING_DATA_FIELD_DEF(int64_t, hInput);
25+TILING_DATA_FIELD_DEF(int64_t, wInput);
26+TILING_DATA_FIELD_DEF(int64_t, hOutput);
27+TILING_DATA_FIELD_DEF(int64_t, wOutput);
28+TILING_DATA_FIELD_DEF(int64_t, hKernel);
29+TILING_DATA_FIELD_DEF(int64_t, wKernel);
30+TILING_DATA_FIELD_DEF(int64_t, hStride);
31+TILING_DATA_FIELD_DEF(int64_t, wStride);
32+TILING_DATA_FIELD_DEF(int64_t, padLeft);
33+TILING_DATA_FIELD_DEF(int64_t, padTop);
34+TILING_DATA_FIELD_DEF(int64_t, highAxisInner);
35+TILING_DATA_FIELD_DEF(int64_t, highAxisTail);
36+TILING_DATA_FIELD_DEF(int64_t, highAxisOuter);
37+TILING_DATA_FIELD_DEF(int64_t, hOutputInner);
38+TILING_DATA_FIELD_DEF(int64_t, hOutputTail);
39+TILING_DATA_FIELD_DEF(int64_t, hOutputOuter);
40+TILING_DATA_FIELD_DEF(int64_t, wOutputInner);
41+TILING_DATA_FIELD_DEF(int64_t, wOutputTail);
42+TILING_DATA_FIELD_DEF(int64_t, wOutputOuter);
43+TILING_DATA_FIELD_DEF(int64_t, normalCoreProcessNum);
44+TILING_DATA_FIELD_DEF(int64_t, tailCoreProcessNum);
45+TILING_DATA_FIELD_DEF(int64_t, usedCoreNum);
46+TILING_DATA_FIELD_DEF(int64_t, inputBufferSize);
47+TILING_DATA_FIELD_DEF(int64_t, maxValueBufferSize);
48+TILING_DATA_FIELD_DEF(int64_t, argmaxBufferSize);
49+TILING_DATA_FIELD_DEF(int64_t, isPad);
50+TILING_DATA_FIELD_DEF(int64_t, hDilation);
51+TILING_DATA_FIELD_DEF(int64_t, wDilation);
52+END_TILING_DATA_DEF;
53+ 
54+// 300001 - no padding, 300002 - padding
55+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_300001, MaxPoolWithArgmaxV3GatherTilingData);
56+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_300002, MaxPoolWithArgmaxV3GatherTilingData);
57+ 
58+struct MaxPoolWithArgmaxV3GatherBaseInfo {
59+ int64_t inputBytes = 0;
60+ int64_t indexBytes = 0;
61+ int64_t availableUb = 0;
62+ int64_t totalCoreNum = 0;
63+ int64_t oneBlockNumT1 = 0;
64+ int64_t oneBlockNumT2 = 0;
65+ int64_t coreUsedForBestPerformance = 0;
66+ 
67+ int64_t padTop = 0;
68+ int64_t padLeft = 0;
69+ int64_t hStride = 0;
70+ int64_t wStride = 0;
71+ int64_t hKernel = 0;
72+ int64_t wKernel = 0;
73+ int64_t hInput = 0;
74+ int64_t wInput = 0;
75+ int64_t hOutput = 0;
76+ int64_t wOutput = 0;
77+ int64_t highAxisTotal = 0;
78+ int64_t isPad = 0;
79+ int64_t hDilation = 0;
80+ int64_t wDilation = 0;
81+ std::string ToString() const
82+ {
83+ std::stringstream info;
84+ info << "MaxPoolWithArgmaxV3GatherBaseInfo {";
85+ info << "inputBytes:" << inputBytes << ",indexBytes:" << indexBytes << ",availableUb:" << availableUb
86+ << ",totalCoreNum:" << totalCoreNum << ",coreUsedForBestPerformance:" << coreUsedForBestPerformance
87+ << ",padTop:" << padTop << ",padLeft:" << padLeft << ",hStride:" << hStride << ",wStride:" << wStride
88+ << ",hKernel:" << hKernel << ",wKernel:" << wKernel << ",nInput:" << highAxisTotal << ",hInput:" << hInput
89+ << ",wInput:" << wInput << ",hOutput:" << hOutput << ",wOutput:" << wOutput << ",isPad:" << isPad
90+ << ",hDilation:" << hDilation << ",wDilation:" << wDilation;
91+ info << " }";
92+ return info.str();
93+ }
94+};
95+ 
96+struct MaxPoolWithArgmaxV3GatherSplitInfo {
97+ // InitializationVars
98+ int64_t highAxisInner = 0;
99+ int64_t highAxisTail = 0;
100+ int64_t highAxisOuter = 0;
101+ int64_t highAxisAligned = 0;
102+ 
103+ // DoUBTiling
104+ int64_t hOutputInner = 0;
105+ int64_t hOutputTail = 0;
106+ int64_t hOutputOuter = 0;
107+ int64_t wOutputInner = 0;
108+ int64_t wOutputTail = 0;
109+ int64_t wOutputOuter = 0;
110+ 
111+ // DoBlockTiling
112+ int64_t normalCoreProcessNum = 0;
113+ int64_t tailCoreProcessNum = 0;
114+ int64_t usedCoreNum = 0;
115+ int64_t totalBaseBlockNum = 0;
116+ 
117+ // DoBufferCalculate
118+ int64_t hInputInner = 0;
119+ int64_t wInputInner = 0;
120+ int64_t baseBlockPlaneSizeAligned = 0;
121+ int64_t inputBufferSize = 0;
122+ int64_t maxValueBufferSize = 0;
123+ int64_t argmaxBufferSize = 0;
124+ int64_t totalBufferSize = 0;
125+ std::string ToString() const
126+ {
127+ std::stringstream info;
128+ info << "MaxPoolWithArgmaxV3NhwcSplitInfo {";
129+ info << "highAxisInner:" << highAxisInner << ",highAxisTail:" << highAxisTail
130+ << ",highAxisOuter:" << highAxisOuter << ",hOutputInner:" << hOutputInner << ",hOutputTail:" << hOutputTail
131+ << ",hOutputOuter:" << hOutputOuter << ",wOutputInner:" << wOutputInner << ",wOutputTail:" << wOutputTail
132+ << ",wOutputOuter:" << wOutputOuter << ",normalCoreProcessNum:" << normalCoreProcessNum
133+ << ",tailCoreProcessNum:" << tailCoreProcessNum << ",usedCoreNum:" << usedCoreNum
134+ << ",totalBaseBlockNum:" << totalBaseBlockNum << ",hInputInner:" << hInputInner
135+ << ",wInputInner:" << wInputInner << ",inputBufferSize:" << inputBufferSize
136+ << ",maxValueBufferSize:" << maxValueBufferSize << ",argmaxBufferSize:" << argmaxBufferSize
137+ << ",totalBufferSize:" << totalBufferSize;
138+ info << " }";
139+ return info.str();
140+ }
141+};
142+ 
143+class MaxPoolWithArgmaxV3GatherTiling : public MaxPoolWithArgmaxV3BaseTiling {
144+public:
145+ explicit MaxPoolWithArgmaxV3GatherTiling(gert::TilingContext* context) : MaxPoolWithArgmaxV3BaseTiling(context)
146+ {}
147+ 
148+ ~MaxPoolWithArgmaxV3GatherTiling() override
149+ {}
150+ 
151+private:
152+ void DoUBTiling();
153+ void InitializationVars();
154+ bool IsMeetTargetCoreNum() const;
155+ void SearchBestTiling();
156+ bool IsMeetUBSize();
157+ void SetTilingData();
158+ void BinarySearch(int64_t start, int64_t end, int64_t* value);
159+ bool TrySplitNC();
160+ bool TrySplitH();
161+ bool TrySplitW();
162+ uint64_t GetTilingKey() const override;
163+ void PrintBaseData() const;
164+ void PrintSplitData() const;
165+ void DoBlockTiling();
166+ void DoBufferCalculate();
167+ bool IsCapable() override;
168+ ge::graphStatus DoOpTiling() override;
169+ ge::graphStatus PostTiling() override;
170+ 
171+ MaxPoolWithArgmaxV3GatherTilingData tilingData_;
172+ MaxPoolWithArgmaxV3GatherBaseInfo baseData_;
173+ MaxPoolWithArgmaxV3GatherSplitInfo splitData_;
174+};
175+ 
176+} // namespace optiling
177+ 
178+#endif
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_nhwc_tiling.cpp+414-0
@@ -0,0 +1,414 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_nhwc_tiling.cpp
13+ * \brief
14+ */
15+#include "platform/platform_info.h"
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_with_argmax_v3_nhwc_tiling.h"
18+#include "op_common/op_host/util/platform_util.h"
19+ 
20+namespace optiling {
21+static constexpr int64_t FLOAT16_SIZE = 2;
22+static constexpr int64_t FLOAT32_SIZE = 4;
23+static constexpr int64_t INT32_SIZE = 4;
24+static constexpr int64_t INT64_SIZE = 8;
25+static constexpr int64_t UB_RESVERVED_SIZE = 0;
26+static constexpr int64_t HELPER_BUFFER_SIZE = 1024;
27+ 
28+static constexpr uint64_t SMALL_C_NO_PADDING_TILING_KEY = 700001;
29+static constexpr uint64_t SMALL_C_PADDING_TILING_KEY = 700002;
30+static constexpr uint64_t LARGE_C_NO_PADDING_TILING_KEY = 800001;
31+static constexpr uint64_t LARGE_C_PADDING_TILING_KEY = 800002;
32+ 
33+static constexpr int64_t DOUBLE = 2;
34+ 
35+static constexpr int64_t TEMPLATE_MODE_SMALL_C = 1;
36+static constexpr int64_t TEMPLATE_MODE_LARGE_C = 2;
37+ 
38+void MaxPoolWithArgmaxV3NhwcTiling::InitializationVars()
39+{
40+ baseData_.inputBytes = dtype == ge::DT_FLOAT ? FLOAT32_SIZE : FLOAT16_SIZE;
41+ baseData_.indexBytes = inputData.indexDtype == ge::DT_INT32 ? INT32_SIZE : INT64_SIZE;
42+ baseData_.availableUb = static_cast<int64_t>(ubSize) - UB_RESVERVED_SIZE;
43+ baseData_.totalCoreNum = static_cast<int64_t>(coreNum);
44+ baseData_.coreUsedForBestPerformance = baseData_.totalCoreNum;
45+ 
46+ baseData_.padTop = inputData.pad[H_DIM];
47+ baseData_.padLeft = inputData.pad[W_DIM];
48+ baseData_.hDilation = inputData.dilation[H_DIM];
49+ baseData_.wDilation = inputData.dilation[W_DIM];
50+ baseData_.hStride = inputData.stride[H_DIM];
51+ baseData_.wStride = inputData.stride[W_DIM];
52+ baseData_.hKernel = inputData.kernelSize[H_DIM];
53+ baseData_.wKernel = inputData.kernelSize[W_DIM];
54+ 
55+ baseData_.nInput = inputData.nInput;
56+ baseData_.hInput = inputData.inputShape[H_DIM];
57+ baseData_.wInput = inputData.inputShape[W_DIM];
58+ baseData_.cInput = inputData.cInput;
59+ 
60+ baseData_.hOutput = inputData.outShape[H_DIM];
61+ baseData_.wOutput = inputData.outShape[W_DIM];
62+ 
63+ baseData_.isPad = 0;
64+ if (baseData_.padTop != 0 || baseData_.padLeft != 0) {
65+ baseData_.isPad = 1;
66+ }
67+ 
68+ if (inputData.ceilMode && baseData_.isPad == 0) {
69+ if (((baseData_.wOutput - 1) * baseData_.wStride + baseData_.wKernel) != baseData_.wInput ||
70+ ((baseData_.hOutput - 1) * baseData_.hStride + baseData_.hKernel) != baseData_.hInput) {
71+ baseData_.isPad = 1;
72+ }
73+ }
74+ 
75+ // 并发度按照索引类型计算 索引类型位宽 >= 输入类型位宽
76+ baseData_.concurrentCount = Ops::Base::GetVRegSize(context_) / baseData_.indexBytes;
77+ baseData_.templateMode =
78+ DOUBLE * baseData_.cInput > baseData_.concurrentCount ? TEMPLATE_MODE_LARGE_C : TEMPLATE_MODE_SMALL_C;
79+}
80+ 
81+bool MaxPoolWithArgmaxV3NhwcTiling::IsCapable()
82+{
83+ if (inputData.dilation[H_DIM] != 1 || inputData.dilation[W_DIM] != 1 ||
84+ inputData.inputFormat != ge::Format::FORMAT_NHWC) {
85+ return false;
86+ }
87+ 
88+ InitializationVars();
89+ return true;
90+}
91+ 
92+uint64_t MaxPoolWithArgmaxV3NhwcTiling::GetTilingKey() const
93+{
94+ uint64_t tilingKey = 0;
95+ switch (baseData_.templateMode) {
96+ case TEMPLATE_MODE_SMALL_C:
97+ if (baseData_.isPad == 1) {
98+ tilingKey = SMALL_C_PADDING_TILING_KEY;
99+ } else {
100+ tilingKey = SMALL_C_NO_PADDING_TILING_KEY;
101+ }
102+ break;
103+ case TEMPLATE_MODE_LARGE_C:
104+ if (baseData_.isPad == 1) {
105+ tilingKey = LARGE_C_PADDING_TILING_KEY;
106+ } else {
107+ tilingKey = LARGE_C_NO_PADDING_TILING_KEY;
108+ }
109+ break;
110+ default:
111+ break;
112+ }
113+ 
114+ return tilingKey;
115+}
116+ 
117+void MaxPoolWithArgmaxV3NhwcTiling::DoBufferCalculate()
118+{
119+ if (splitData_.hKernelInner == 0 && splitData_.wKernelInner == 0) {
120+ splitData_.hInputInner = (splitData_.hOutputInner - 1) * baseData_.hStride + baseData_.hKernel;
121+ splitData_.wInputInner = (splitData_.wOutputInner - 1) * baseData_.wStride + baseData_.wKernel;
122+ } else {
123+ splitData_.hInputInner = splitData_.hKernelInner;
124+ splitData_.wInputInner = splitData_.wKernelInner;
125+ }
126+ 
127+ int64_t oneBlockNumT1 = Ops::Base::GetUbBlockSize(context_) / baseData_.inputBytes;
128+ int64_t oneBlockNumT2 = Ops::Base::GetUbBlockSize(context_) / baseData_.indexBytes;
129+ int64_t maxDataNumInOneBlock = std::max(oneBlockNumT1, oneBlockNumT2);
130+ int64_t cOutputInnerAligned = Ops::Base::CeilAlign(splitData_.cOutputInner, maxDataNumInOneBlock);
131+ 
132+ splitData_.inputBufferSize = splitData_.nOutputInner * splitData_.hInputInner * splitData_.wInputInner *
133+ cOutputInnerAligned * baseData_.inputBytes;
134+ splitData_.maxValueBufferSize = splitData_.nOutputInner * splitData_.hOutputInner * splitData_.wOutputInner *
135+ cOutputInnerAligned * baseData_.inputBytes;
136+ splitData_.argmaxBufferSize = splitData_.nOutputInner * splitData_.hOutputInner * splitData_.wOutputInner *
137+ cOutputInnerAligned * baseData_.indexBytes;
138+ 
139+ int64_t tmpTotalBufferSize =
140+ splitData_.inputBufferSize + splitData_.maxValueBufferSize + splitData_.argmaxBufferSize + HELPER_BUFFER_SIZE;
141+ splitData_.totalBufferSize = tmpTotalBufferSize * DOUBLE;
142+}
143+ 
144+bool MaxPoolWithArgmaxV3NhwcTiling::IsMeetTargetCoreNum() const
145+{
146+ int64_t tmpWOutputOuter = Ops::Base::CeilDiv(baseData_.wOutput, splitData_.wOutputInner);
147+ int64_t tmpHOutputOuter = Ops::Base::CeilDiv(baseData_.hOutput, splitData_.hOutputInner);
148+ int64_t tmpNOutputOuter = Ops::Base::CeilDiv(baseData_.nInput, splitData_.nOutputInner);
149+ int64_t tmpCOutputOuter = Ops::Base::CeilDiv(baseData_.cInput, splitData_.cOutputInner);
150+ return tmpWOutputOuter * tmpHOutputOuter * tmpNOutputOuter * tmpCOutputOuter >=
151+ baseData_.coreUsedForBestPerformance;
152+}
153+ 
154+bool MaxPoolWithArgmaxV3NhwcTiling::IsMeetUBSize()
155+{
156+ DoBufferCalculate();
157+ return splitData_.totalBufferSize <= baseData_.availableUb;
158+}
159+ 
160+void MaxPoolWithArgmaxV3NhwcTiling::BinarySearch(int64_t start, int64_t end, int64_t* value, int64_t rate)
161+{
162+ int64_t left = start;
163+ int64_t right = end;
164+ int64_t bestSplit = 1;
165+ 
166+ while (left <= right) {
167+ int64_t mid = left + (right - left) / DOUBLE;
168+ *value = mid * rate;
169+ 
170+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
171+ bestSplit = mid;
172+ left = mid + 1;
173+ } else {
174+ right = mid - 1;
175+ }
176+ }
177+ 
178+ *value = bestSplit * rate;
179+}
180+ 
181+bool MaxPoolWithArgmaxV3NhwcTiling::TrySplitN()
182+{
183+ splitData_.hOutputInner = baseData_.hOutput;
184+ splitData_.wOutputInner = baseData_.wOutput;
185+ splitData_.cOutputInner = baseData_.cInput;
186+ 
187+ splitData_.nOutputInner = Ops::Base::CeilDiv(baseData_.nInput, baseData_.coreUsedForBestPerformance);
188+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
189+ return true;
190+ }
191+ 
192+ splitData_.nOutputInner = 1;
193+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
194+ BinarySearch(1, baseData_.nInput, &splitData_.nOutputInner);
195+ return true;
196+ }
197+ 
198+ return false;
199+}
200+ 
201+bool MaxPoolWithArgmaxV3NhwcTiling::TrySplitH()
202+{
203+ splitData_.nOutputInner = 1;
204+ splitData_.wOutputInner = baseData_.wOutput;
205+ splitData_.cOutputInner = baseData_.cInput;
206+ 
207+ splitData_.hOutputInner = 1;
208+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
209+ BinarySearch(1, baseData_.hOutput, &splitData_.hOutputInner);
210+ return true;
211+ }
212+ 
213+ return false;
214+}
215+ 
216+bool MaxPoolWithArgmaxV3NhwcTiling::TrySplitW()
217+{
218+ splitData_.nOutputInner = 1;
219+ splitData_.hOutputInner = 1;
220+ splitData_.cOutputInner = baseData_.cInput;
221+ 
222+ splitData_.wOutputInner = 1;
223+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
224+ BinarySearch(1, baseData_.wOutput, &splitData_.wOutputInner);
225+ return true;
226+ }
227+ 
228+ return false;
229+}
230+ 
231+void MaxPoolWithArgmaxV3NhwcTiling::SplitC()
232+{
233+ // NHW全切为1，此时还需要切C可能场景：
234+ // （1） C 超大或者kernel超大 (UB 不满足) （2） N*H*W超小场景（核数不满足） （3） 都不满足
235+ splitData_.nOutputInner = 1;
236+ splitData_.hOutputInner = 1;
237+ splitData_.wOutputInner = 1;
238+ 
239+ int64_t tmpC = baseData_.cInput < baseData_.concurrentCount ? baseData_.cInput : baseData_.concurrentCount;
240+ splitData_.cOutputInner = tmpC;
241+ if (IsMeetUBSize() && IsMeetTargetCoreNum()) {
242+ BinarySearch(
243+ 1, Ops::Base::CeilDiv(baseData_.cInput / DOUBLE, baseData_.concurrentCount), &splitData_.cOutputInner,
244+ baseData_.concurrentCount);
245+ }
246+}
247+ 
248+void MaxPoolWithArgmaxV3NhwcTiling::SplitKernel()
249+{
250+ splitData_.wKernelInner = baseData_.wKernel;
251+ splitData_.hKernelInner = baseData_.hKernel;
252+ splitData_.wKernelOuter = 1;
253+ splitData_.hKernelOuter = 1;
254+ while (splitData_.hKernelInner != 1 || splitData_.wKernelInner != 1) {
255+ if (!IsMeetUBSize()) {
256+ DynamicAdjustmentKernelWH();
257+ } else {
258+ break;
259+ }
260+ }
261+ 
262+ splitData_.hKernelOuter = Ops::Base::CeilDiv(baseData_.hKernel, splitData_.hKernelInner);
263+ int64_t tempHKernelTail = baseData_.hKernel % splitData_.hKernelInner;
264+ splitData_.hKernelTail = tempHKernelTail == 0 ? splitData_.hKernelInner : tempHKernelTail;
265+ 
266+ splitData_.wKernelOuter = Ops::Base::CeilDiv(baseData_.wKernel, splitData_.wKernelInner);
267+ int64_t tempWKernelTail = baseData_.wKernel % splitData_.wKernelInner;
268+ splitData_.wKernelTail = tempWKernelTail == 0 ? splitData_.wKernelInner : tempWKernelTail;
269+}
270+ 
271+void MaxPoolWithArgmaxV3NhwcTiling::DynamicAdjustmentKernelWH()
272+{
273+ if (splitData_.hKernelInner == 1) {
274+ splitData_.wKernelOuter++;
275+ splitData_.wKernelInner = Ops::Base::CeilDiv(baseData_.wKernel, splitData_.wKernelOuter);
276+ } else {
277+ splitData_.hKernelOuter++;
278+ splitData_.hKernelInner = Ops::Base::CeilDiv(baseData_.hKernel, splitData_.hKernelOuter);
279+ }
280+}
281+ 
282+void MaxPoolWithArgmaxV3NhwcTiling::SearchBestTiling()
283+{
284+ if (TrySplitN()) {
285+ return;
286+ }
287+ 
288+ if (TrySplitH()) {
289+ return;
290+ }
291+ 
292+ if (TrySplitW()) {
293+ return;
294+ }
295+ 
296+ SplitC();
297+ if (!IsMeetUBSize()) {
298+ // 超大kernel场景 C * H * W = 32768 左右
299+ splitData_.isSplitKernel = 1;
300+ SplitKernel();
301+ }
302+}
303+ 
304+void MaxPoolWithArgmaxV3NhwcTiling::DoUBTiling()
305+{
306+ // 切输出,反算输入
307+ SearchBestTiling();
308+ DoBufferCalculate();
309+ splitData_.wOutputOuter = Ops::Base::CeilDiv(baseData_.wOutput, splitData_.wOutputInner);
310+ int64_t tempWOutputTail = baseData_.wOutput % splitData_.wOutputInner;
311+ splitData_.wOutputTail = tempWOutputTail == 0 ? splitData_.wOutputInner : tempWOutputTail;
312+ 
313+ splitData_.hOutputOuter = Ops::Base::CeilDiv(baseData_.hOutput, splitData_.hOutputInner);
314+ int64_t tempHOutputTail = baseData_.hOutput % splitData_.hOutputInner;
315+ splitData_.hOutputTail = tempHOutputTail == 0 ? splitData_.hOutputInner : tempHOutputTail;
316+ 
317+ splitData_.nOutputOuter = Ops::Base::CeilDiv(baseData_.nInput, splitData_.nOutputInner);
318+ int64_t tempNOutputTail = baseData_.nInput % splitData_.nOutputInner;
319+ splitData_.nOutputTail = tempNOutputTail == 0 ? splitData_.nOutputInner : tempNOutputTail;
320+ 
321+ splitData_.cOutputOuter = Ops::Base::CeilDiv(baseData_.cInput, splitData_.cOutputInner);
322+ int64_t tempCOutputTail = baseData_.cInput % splitData_.cOutputInner;
323+ splitData_.cOutputTail = tempCOutputTail == 0 ? splitData_.cOutputInner : tempCOutputTail;
324+}
325+ 
326+void MaxPoolWithArgmaxV3NhwcTiling::DoBlockTiling()
327+{
328+ splitData_.totalBaseBlockNum =
329+ splitData_.nOutputOuter * splitData_.cOutputOuter * splitData_.hOutputOuter * splitData_.wOutputOuter;
330+ splitData_.normalCoreProcessNum = Ops::Base::CeilDiv(splitData_.totalBaseBlockNum, baseData_.totalCoreNum);
331+ splitData_.usedCoreNum = Ops::Base::CeilDiv(splitData_.totalBaseBlockNum, splitData_.normalCoreProcessNum);
332+ splitData_.tailCoreProcessNum =
333+ splitData_.totalBaseBlockNum - splitData_.normalCoreProcessNum * (splitData_.usedCoreNum - 1);
334+}
335+ 
336+void MaxPoolWithArgmaxV3NhwcTiling::RerouteTemplateBySplit()
337+{
338+ if (splitData_.isSplitKernel == 1 || (splitData_.wOutputInner == 1 && splitData_.hOutputInner == 1)) {
339+ OP_LOGD("MaxPoolWithArgmaxV3Nhwc", "[GetTilingKey] split kernel or single kernel scenario to large c template");
340+ baseData_.templateMode = TEMPLATE_MODE_LARGE_C;
341+ }
342+}
343+ 
344+void MaxPoolWithArgmaxV3NhwcTiling::SetTilingData()
345+{
346+ tilingData_.set_cInput(baseData_.cInput);
347+ tilingData_.set_hInput(baseData_.hInput);
348+ tilingData_.set_wInput(baseData_.wInput);
349+ tilingData_.set_hOutput(baseData_.hOutput);
350+ tilingData_.set_wOutput(baseData_.wOutput);
351+ tilingData_.set_hKernel(baseData_.hKernel);
352+ tilingData_.set_wKernel(baseData_.wKernel);
353+ tilingData_.set_hStride(baseData_.hStride);
354+ tilingData_.set_wStride(baseData_.wStride);
355+ tilingData_.set_padTop(baseData_.padTop);
356+ tilingData_.set_padLeft(baseData_.padLeft);
357+ tilingData_.set_hDilation(baseData_.hDilation);
358+ tilingData_.set_wDilation(baseData_.wDilation);
359+ tilingData_.set_nOutputInner(splitData_.nOutputInner);
360+ tilingData_.set_nOutputTail(splitData_.nOutputTail);
361+ tilingData_.set_nOutputOuter(splitData_.nOutputOuter);
362+ tilingData_.set_hOutputInner(splitData_.hOutputInner);
363+ tilingData_.set_hOutputTail(splitData_.hOutputTail);
364+ tilingData_.set_hOutputOuter(splitData_.hOutputOuter);
365+ tilingData_.set_wOutputInner(splitData_.wOutputInner);
366+ tilingData_.set_wOutputTail(splitData_.wOutputTail);
367+ tilingData_.set_wOutputOuter(splitData_.wOutputOuter);
368+ tilingData_.set_cOutputInner(splitData_.cOutputInner);
369+ tilingData_.set_cOutputTail(splitData_.cOutputTail);
370+ tilingData_.set_cOutputOuter(splitData_.cOutputOuter);
371+ tilingData_.set_normalCoreProcessNum(splitData_.normalCoreProcessNum);
372+ tilingData_.set_tailCoreProcessNum(splitData_.tailCoreProcessNum);
373+ tilingData_.set_usedCoreNum(splitData_.usedCoreNum);
374+ tilingData_.set_inputBufferSize(splitData_.inputBufferSize);
375+ tilingData_.set_maxValueBufferSize(splitData_.maxValueBufferSize);
376+ tilingData_.set_argmaxBufferSize(splitData_.argmaxBufferSize);
377+ tilingData_.set_isPad(baseData_.isPad);
378+ tilingData_.set_isSplitKernel(splitData_.isSplitKernel);
379+ tilingData_.set_hKernelInner(splitData_.hKernelInner);
380+ tilingData_.set_hKernelTail(splitData_.hKernelTail);
381+ tilingData_.set_hKernelOuter(splitData_.hKernelOuter);
382+ tilingData_.set_wKernelInner(splitData_.wKernelInner);
383+ tilingData_.set_wKernelTail(splitData_.wKernelTail);
384+ tilingData_.set_wKernelOuter(splitData_.wKernelOuter);
385+ tilingData_.set_wKernelOuter(splitData_.wKernelOuter);
386+ tilingData_.set_tilingKey(GetTilingKey());
387+}
388+ 
389+ge::graphStatus MaxPoolWithArgmaxV3NhwcTiling::DoOpTiling()
390+{
391+ DoUBTiling();
392+ RerouteTemplateBySplit();
393+ DoBlockTiling();
394+ SetTilingData();
395+ OP_LOGI("PrintBaseData", "%s", baseData_.ToString().c_str());
396+ OP_LOGI("PrintSplitData", "%s", splitData_.ToString().c_str());
397+ return ge::GRAPH_SUCCESS;
398+}
399+ 
400+ge::graphStatus MaxPoolWithArgmaxV3NhwcTiling::PostTiling()
401+{
402+ context_->SetBlockDim(tilingData_.get_usedCoreNum());
403+ if (tilingData_.GetDataSize() > context_->GetRawTilingData()->GetCapacity()) {
404+ return ge::GRAPH_FAILED;
405+ }
406+ 
407+ tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
408+ context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
409+ return ge::GRAPH_SUCCESS;
410+}
411+ 
412+REGISTER_OPS_TILING_TEMPLATE(MaxPoolWithArgmaxV3, MaxPoolWithArgmaxV3NhwcTiling, 20);
413+ 
414+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_nhwc_tiling.h+213-0
@@ -0,0 +1,213 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_nhwc_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_WITH_AGRMAX_V3_NHWC_TILING_H_
17+#define MAX_POOL_WITH_AGRMAX_V3_NHWC_TILING_H_
18+ 
19+#include "max_pool_with_argmax_v3_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+BEGIN_TILING_DATA_DEF(MaxPoolWithArgmaxV3NhwcTilingData)
24+TILING_DATA_FIELD_DEF(int64_t, cInput);
25+TILING_DATA_FIELD_DEF(int64_t, hInput);
26+TILING_DATA_FIELD_DEF(int64_t, wInput);
27+TILING_DATA_FIELD_DEF(int64_t, hOutput);
28+TILING_DATA_FIELD_DEF(int64_t, wOutput);
29+TILING_DATA_FIELD_DEF(int64_t, hKernel);
30+TILING_DATA_FIELD_DEF(int64_t, wKernel);
31+TILING_DATA_FIELD_DEF(int64_t, hStride);
32+TILING_DATA_FIELD_DEF(int64_t, wStride);
33+TILING_DATA_FIELD_DEF(int64_t, padLeft);
34+TILING_DATA_FIELD_DEF(int64_t, padTop);
35+TILING_DATA_FIELD_DEF(int64_t, hDilation);
36+TILING_DATA_FIELD_DEF(int64_t, wDilation);
37+TILING_DATA_FIELD_DEF(int64_t, nOutputInner);
38+TILING_DATA_FIELD_DEF(int64_t, nOutputTail);
39+TILING_DATA_FIELD_DEF(int64_t, nOutputOuter);
40+TILING_DATA_FIELD_DEF(int64_t, hOutputInner);
41+TILING_DATA_FIELD_DEF(int64_t, hOutputTail);
42+TILING_DATA_FIELD_DEF(int64_t, hOutputOuter);
43+TILING_DATA_FIELD_DEF(int64_t, wOutputInner);
44+TILING_DATA_FIELD_DEF(int64_t, wOutputTail);
45+TILING_DATA_FIELD_DEF(int64_t, wOutputOuter);
46+TILING_DATA_FIELD_DEF(int64_t, cOutputInner);
47+TILING_DATA_FIELD_DEF(int64_t, cOutputTail);
48+TILING_DATA_FIELD_DEF(int64_t, cOutputOuter);
49+TILING_DATA_FIELD_DEF(int64_t, normalCoreProcessNum);
50+TILING_DATA_FIELD_DEF(int64_t, tailCoreProcessNum);
51+TILING_DATA_FIELD_DEF(int64_t, usedCoreNum);
52+TILING_DATA_FIELD_DEF(int64_t, inputBufferSize);
53+TILING_DATA_FIELD_DEF(int64_t, maxValueBufferSize);
54+TILING_DATA_FIELD_DEF(int64_t, argmaxBufferSize);
55+TILING_DATA_FIELD_DEF(int64_t, isPad);
56+TILING_DATA_FIELD_DEF(int64_t, isSplitKernel);
57+TILING_DATA_FIELD_DEF(int64_t, hKernelInner);
58+TILING_DATA_FIELD_DEF(int64_t, hKernelTail);
59+TILING_DATA_FIELD_DEF(int64_t, hKernelOuter);
60+TILING_DATA_FIELD_DEF(int64_t, wKernelInner);
61+TILING_DATA_FIELD_DEF(int64_t, wKernelTail);
62+TILING_DATA_FIELD_DEF(int64_t, wKernelOuter);
63+TILING_DATA_FIELD_DEF(int64_t, tilingKey);
64+END_TILING_DATA_DEF;
65+ 
66+// small c 700001 - no padding, 700002 - padding
67+// large c 800001 - no padding, 800002 - padding
68+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_700001, MaxPoolWithArgmaxV3NhwcTilingData);
69+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_700002, MaxPoolWithArgmaxV3NhwcTilingData);
70+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_800001, MaxPoolWithArgmaxV3NhwcTilingData);
71+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_800002, MaxPoolWithArgmaxV3NhwcTilingData);
72+ 
73+struct MaxPoolWithArgmaxV3NhwcBaseInfo {
74+ int64_t inputBytes{0};
75+ int64_t indexBytes{0};
76+ int64_t availableUb{0};
77+ int64_t totalCoreNum{0};
78+ int64_t coreUsedForBestPerformance{0};
79+ int64_t concurrentCount{0};
80+ int64_t templateMode{0};
81+ 
82+ int64_t padTop{0};
83+ int64_t padLeft{0};
84+ int64_t hDilation{0};
85+ int64_t wDilation{0};
86+ int64_t hStride{0};
87+ int64_t wStride{0};
88+ int64_t hKernel{0};
89+ int64_t wKernel{0};
90+ int64_t nInput{0};
91+ int64_t hInput{0};
92+ int64_t wInput{0};
93+ int64_t cInput{0};
94+ int64_t hOutput{0};
95+ int64_t wOutput{0};
96+ 
97+ int64_t isPad{0};
98+ 
99+ std::string ToString() const
100+ {
101+ std::stringstream info;
102+ info << "MaxPoolWithArgmaxV3NhwcBaseInfo {";
103+ info << "inputBytes:" << inputBytes << ",indexBytes:" << indexBytes << ",availableUb:" << availableUb
104+ << ",totalCoreNum:" << totalCoreNum << ",coreUsedForBestPerformance:" << coreUsedForBestPerformance
105+ << ",concurrentCount:" << concurrentCount << ",templateMode:" << templateMode << ",padTop:" << padTop
106+ << ",padLeft:" << padLeft << ",hDilation:" << hDilation << ",wDilation:" << wDilation
107+ << ",hStride:" << hStride << ",wStride:" << wStride << ",hKernel:" << hKernel << ",wKernel:" << wKernel
108+ << ",nInput:" << nInput << ",hInput:" << hInput << ",wInput:" << wInput << ",cInput:" << cInput
109+ << ",hOutput:" << hOutput << ",wOutput:" << wOutput << ",isPad:" << isPad;
110+ info << " }";
111+ return info.str();
112+ }
113+};
114+ 
115+struct MaxPoolWithArgmaxV3NhwcSplitInfo {
116+ // DoUBTiling
117+ int64_t nOutputInner{0};
118+ int64_t nOutputTail{0};
119+ int64_t nOutputOuter{0};
120+ 
121+ int64_t hOutputInner{0};
122+ int64_t hOutputTail{0};
123+ int64_t hOutputOuter{0};
124+ 
125+ int64_t wOutputInner{0};
126+ int64_t wOutputTail{0};
127+ int64_t wOutputOuter{0};
128+ 
129+ int64_t cOutputInner{0};
130+ int64_t cOutputTail{0};
131+ int64_t cOutputOuter{0};
132+ 
133+ int64_t isSplitKernel{0};
134+ int64_t hKernelInner{0};
135+ int64_t hKernelTail{0};
136+ int64_t hKernelOuter{0};
137+ int64_t wKernelInner{0};
138+ int64_t wKernelTail{0};
139+ int64_t wKernelOuter{0};
140+ 
141+ // DoBlockTiling
142+ int64_t normalCoreProcessNum{0};
143+ int64_t tailCoreProcessNum{0};
144+ int64_t usedCoreNum{0};
145+ int64_t totalBaseBlockNum{0};
146+ 
147+ // DoBufferCalculate
148+ int64_t hInputInner{0};
149+ int64_t wInputInner{0};
150+ int64_t inputBufferSize{0};
151+ int64_t maxValueBufferSize{0};
152+ int64_t argmaxBufferSize{0};
153+ int64_t totalBufferSize{0};
154+ 
155+ std::string ToString() const
156+ {
157+ std::stringstream info;
158+ info << "MaxPoolWithArgmaxV3NhwcSplitInfo {";
159+ info << "nOutputInner:" << nOutputInner << ",nOutputTail:" << nOutputTail << ",nOutputOuter:" << nOutputOuter
160+ << ",hOutputInner:" << hOutputInner << ",hOutputTail:" << hOutputTail << ",hOutputOuter:" << hOutputOuter
161+ << ",wOutputInner:" << wOutputInner << ",wOutputTail:" << wOutputTail << ",wOutputOuter:" << wOutputOuter
162+ << ",cOutputInner:" << cOutputInner << ",cOutputTail:" << cOutputTail << ",cOutputOuter:" << cOutputOuter
163+ << ",isSplitKernel:" << isSplitKernel << ",hKernelInner:" << hKernelInner << ",hKernelTail:" << hKernelTail
164+ << ",hKernelOuter:" << hKernelOuter << ",wKernelInner:" << wKernelInner << ",wKernelTail:" << wKernelTail
165+ << ",wKernelOuter:" << wKernelOuter << ",normalCoreProcessNum:" << normalCoreProcessNum
166+ << ",tailCoreProcessNum:" << tailCoreProcessNum << ",usedCoreNum:" << usedCoreNum
167+ << ",totalBaseBlockNum:" << totalBaseBlockNum << ",hInputInner:" << hInputInner
168+ << ",wInputInner:" << wInputInner << ",inputBufferSize:" << inputBufferSize
169+ << ",maxValueBufferSize:" << maxValueBufferSize << ",argmaxBufferSize:" << argmaxBufferSize
170+ << ",totalBufferSize:" << totalBufferSize;
171+ info << " }";
172+ return info.str();
173+ }
174+};
175+ 
176+class MaxPoolWithArgmaxV3NhwcTiling : public MaxPoolWithArgmaxV3BaseTiling {
177+public:
178+ explicit MaxPoolWithArgmaxV3NhwcTiling(gert::TilingContext* context) : MaxPoolWithArgmaxV3BaseTiling(context)
179+ {}
180+ 
181+ ~MaxPoolWithArgmaxV3NhwcTiling() override
182+ {}
183+ 
184+private:
185+ void DoUBTiling();
186+ void InitializationVars();
187+ bool IsMeetTargetCoreNum() const;
188+ bool IsMeetUBSize();
189+ void SearchBestTiling();
190+ void BinarySearch(int64_t start, int64_t end, int64_t* value, int64_t rate = 1);
191+ bool TrySplitN();
192+ bool TrySplitH();
193+ bool TrySplitW();
194+ void SplitC();
195+ void SplitKernel();
196+ void DynamicAdjustmentKernelWH();
197+ void SetTilingData();
198+ uint64_t GetTilingKey() const override;
199+ void DoBlockTiling();
200+ void RerouteTemplateBySplit();
201+ void DoBufferCalculate();
202+ bool IsCapable() override;
203+ ge::graphStatus DoOpTiling() override;
204+ ge::graphStatus PostTiling() override;
205+ 
206+ MaxPoolWithArgmaxV3NhwcTilingData tilingData_;
207+ MaxPoolWithArgmaxV3NhwcBaseInfo baseData_;
208+ MaxPoolWithArgmaxV3NhwcSplitInfo splitData_;
209+};
210+ 
211+} // namespace optiling
212+ 
213+#endif
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_simt_tiling.cpp+175-0
@@ -0,0 +1,175 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_simt_tiling.cpp
13+ * \brief
14+ */
15+#include <cctype>
16+#include <algorithm>
17+#include "tiling_base/tiling_templates_registry.h"
18+#include "max_pool_with_argmax_v3_simt_tiling.h"
19+ 
20+using namespace AscendC;
21+using namespace ge;
22+ 
23+namespace optiling {
24+ 
25+ge::graphStatus MaxPoolWithArgmaxV3TilingSIMT::GetShapeAttrsInfo()
26+{
27+ auto runtimeAttrs = context_->GetAttrs();
28+ const char* data_format = runtimeAttrs->GetAttrPointer<char>(FORMAT_POS);
29+ OP_CHECK_NULL_WITH_CONTEXT(context_, data_format);
30+ inputData.data_format = data_format;
31+ std::transform(
32+ inputData.data_format.begin(), inputData.data_format.end(), inputData.data_format.begin(),
33+ [](unsigned char c) { return std::tolower(c); });
34+ OP_CHECK_IF(
35+ !(inputData.data_format == "nchw" || inputData.data_format == "nhwc"),
36+ OP_LOGE(context_, "ATTR data_format is %s ,expect [NCHW] or [NHWC].", data_format), return ge::GRAPH_FAILED);
37+ auto inputX = context_->GetInputShape(0);
38+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputX);
39+ auto inputShape = Ops::Base::EnsureNotScalar(inputX->GetStorageShape());
40+ 
41+ auto outX = context_->GetOutputShape(0);
42+ OP_CHECK_NULL_WITH_CONTEXT(context_, outX);
43+ auto outShape = Ops::Base::EnsureNotScalar(outX->GetStorageShape());
44+ 
45+ auto indicesX = context_->GetOutputShape(1);
46+ OP_CHECK_NULL_WITH_CONTEXT(context_, indicesX);
47+ auto indicesShape = Ops::Base::EnsureNotScalar(indicesX->GetStorageShape());
48+ 
49+ if (inputShape.GetDimNum() != NCHW_DIMS) {
50+ OP_LOGE(
51+ context_->GetNodeName(), "MaxPoolWithArgmaxV3: input shape dim = %zu, should be equal 4",
52+ inputShape.GetDimNum());
53+ return ge::GRAPH_FAILED;
54+ }
55+ if (inputData.data_format == "nhwc") {
56+ nDimPos = 0;
57+ cDimPos = 3;
58+ hDimPos = 1;
59+ wDimPos = 2;
60+ }
61+ inputData.inputShape = array<uint64_t, NCHW_DIMS>{
62+ uint64_t(inputShape.GetDim(nDimPos)), uint64_t(inputShape.GetDim(cDimPos)),
63+ uint64_t(inputShape.GetDim(hDimPos)), uint64_t(inputShape.GetDim(wDimPos))};
64+ inputData.outShape = array<uint64_t, NCHW_DIMS>{
65+ uint64_t(inputShape.GetDim(nDimPos)), uint64_t(inputShape.GetDim(cDimPos)), uint64_t(outShape.GetDim(hDimPos)),
66+ uint64_t(outShape.GetDim(wDimPos))};
67+ auto inputDesc = context_->GetInputDesc(0);
68+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputDesc);
69+ dtype = inputDesc->GetDataType();
70+ if (dtype != ge::DataType::DT_BF16 && dtype != ge::DataType::DT_FLOAT16 && dtype != ge::DataType::DT_FLOAT) {
71+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: invalid dtype");
72+ return ge::GRAPH_FAILED;
73+ }
74+ if (indicesShape != outShape) {
75+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: indices shape and values shape is different");
76+ return ge::GRAPH_FAILED;
77+ }
78+ OP_CHECK_NULL_WITH_CONTEXT(context_, runtimeAttrs);
79+ const gert::TypedContinuousVector<int64_t>* kernelSize = runtimeAttrs->GetListInt(KERNEL_POS);
80+ OP_CHECK_NULL_WITH_CONTEXT(context_, kernelSize);
81+ inputData.kernelSize =
82+ array<uint64_t, HW_DIMS>{uint64_t(*(kernelSize->GetData())), uint64_t(*(kernelSize->GetData() + 1))};
83+ const gert::TypedContinuousVector<int64_t>* stride = runtimeAttrs->GetListInt(STRIDE_POS);
84+ OP_CHECK_NULL_WITH_CONTEXT(context_, stride);
85+ inputData.stride = array<uint64_t, HW_DIMS>{uint64_t(*(stride->GetData())), uint64_t(*(stride->GetData() + 1))};
86+ const gert::TypedContinuousVector<int64_t>* padding = runtimeAttrs->GetListInt(PADDING_POS);
87+ OP_CHECK_NULL_WITH_CONTEXT(context_, padding);
88+ inputData.pad = array<uint64_t, HW_DIMS>{uint64_t(*(padding->GetData())), uint64_t(*(padding->GetData() + 1))};
89+ const gert::TypedContinuousVector<int64_t>* dilation = runtimeAttrs->GetListInt(DILATION_POS);
90+ OP_CHECK_NULL_WITH_CONTEXT(context_, dilation);
91+ inputData.dilation =
92+ array<uint64_t, HW_DIMS>{uint64_t(*(dilation->GetData())), uint64_t(*(dilation->GetData() + 1))};
93+ inputData.ceilMode = *runtimeAttrs->GetAttrPointer<bool>(CEIL_POS);
94+ return ge::GRAPH_SUCCESS;
95+}
96+ 
97+ge::graphStatus MaxPoolWithArgmaxV3TilingSIMT::DoOpTiling()
98+{
99+ tiling.set_nDim(inputData.inputShape[N_DIM_]);
100+ tiling.set_cDim(inputData.inputShape[C_DIM_]);
101+ tiling.set_hInDim(inputData.inputShape[H_DIM_]);
102+ tiling.set_wInDim(inputData.inputShape[W_DIM_]);
103+ tiling.set_hOutDim(inputData.outShape[H_DIM_]);
104+ tiling.set_wOutDim(inputData.outShape[W_DIM_]);
105+ tiling.set_kSizeH(inputData.kernelSize[H_IDX_]);
106+ tiling.set_kSizeW(inputData.kernelSize[W_IDX_]);
107+ tiling.set_stridesH(inputData.stride[H_IDX_]);
108+ tiling.set_stridesW(inputData.stride[W_IDX_]);
109+ tiling.set_padH(inputData.pad[H_IDX_]);
110+ tiling.set_padW(inputData.pad[W_IDX_]);
111+ tiling.set_dilationH(inputData.dilation[H_IDX_]);
112+ tiling.set_dilationW(inputData.dilation[W_IDX_]);
113+ tiling.set_ceilMode(inputData.ceilMode);
114+ outputDataCount = tiling.get_nDim() * tiling.get_cDim() * tiling.get_hOutDim() * tiling.get_wOutDim();
115+ int64_t threads = std::min(outputDataCount, MAX_THREAD_NUM);
116+ int64_t blockNum = Ops::Base::CeilDiv(outputDataCount, threads);
117+ blockNum = std::min(blockNum, static_cast<int64_t>(coreNum));
118+ context_->SetBlockDim(blockNum);
119+ tiling.set_threadNums(threads);
120+ tiling.set_blockNums(blockNum);
121+ OP_LOGI(context_->GetNodeName(), "%s", ToString(tiling).c_str());
122+ return ge::GRAPH_SUCCESS;
123+}
124+ 
125+uint64_t MaxPoolWithArgmaxV3TilingSIMT::GetTilingKey() const
126+{
127+ if (inputData.data_format == "nchw" && outputDataCount <= MAX_INT32) {
128+ return SIMT_NCHW_TILING_KEY_INT32;
129+ } else if (inputData.data_format == "nhwc" && outputDataCount <= MAX_INT32) {
130+ return SIMT_NHWC_TILING_KEY_INT32;
131+ } else if (inputData.data_format == "nchw" && outputDataCount > MAX_INT32) {
132+ return SIMT_NCHW_TILING_KEY_INT64;
133+ } else if (inputData.data_format == "nhwc" && outputDataCount > MAX_INT32) {
134+ return SIMT_NHWC_TILING_KEY_INT64;
135+ }
136+ return SIMT_NCHW_TILING_KEY_INT32;
137+}
138+ 
139+ge::graphStatus MaxPoolWithArgmaxV3TilingSIMT::PostTiling()
140+{
141+ OP_CHECK_IF(
142+ context_->GetRawTilingData()->GetCapacity() < tiling.GetDataSize(),
143+ OP_LOGE(
144+ context_, "tiling data's[%zu] is larger than capacity[%zu].", tiling.GetDataSize(),
145+ context_->GetRawTilingData()->GetCapacity()),
146+ return ge::GRAPH_FAILED);
147+ tiling.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
148+ context_->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
149+ return ge::GRAPH_SUCCESS;
150+}
151+ 
152+std::string MaxPoolWithArgmaxV3TilingSIMT::ToString(MaxPoolWithArgmaxV3SimtTilingData& tiling)
153+{
154+ std::string str;
155+ str += " threadNums:" + std::to_string(tiling.get_threadNums());
156+ str += " blockNums:" + std::to_string(tiling.get_blockNums());
157+ str += " nDim:" + std::to_string(tiling.get_nDim());
158+ str += " cDim:" + std::to_string(tiling.get_cDim());
159+ str += " hInDim:" + std::to_string(tiling.get_hInDim());
160+ str += " wInDim:" + std::to_string(tiling.get_wInDim());
161+ str += " hOutDim:" + std::to_string(tiling.get_hOutDim());
162+ str += " wOutDim:" + std::to_string(tiling.get_wOutDim());
163+ str += " kSizeH:" + std::to_string(tiling.get_kSizeH());
164+ str += " kSizeW:" + std::to_string(tiling.get_kSizeW());
165+ str += " stridesH:" + std::to_string(tiling.get_stridesH());
166+ str += " stridesW:" + std::to_string(tiling.get_stridesW());
167+ str += " padH:" + std::to_string(tiling.get_padH());
168+ str += " padW:" + std::to_string(tiling.get_padW());
169+ str += " dilationH:" + std::to_string(tiling.get_dilationH());
170+ str += " dilationW:" + std::to_string(tiling.get_dilationW());
171+ str += " ceilMode:" + std::to_string(tiling.get_ceilMode());
172+ return str;
173+}
174+REGISTER_OPS_TILING_TEMPLATE(MaxPoolWithArgmaxV3, MaxPoolWithArgmaxV3TilingSIMT, 100);
175+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_simt_tiling.h+110-0
@@ -0,0 +1,110 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_simt_tiling.h
13+ * \brief simt imply for max_pool_with_argmax
14+ */
15+ 
16+#ifndef CANN_MAX_POOL_WITH_ARGMAX_V3_SIMT_TILING_H
17+#define CANN_MAX_POOL_WITH_ARGMAX_V3_SIMT_TILING_H
18+ 
19+#include "max_pool_with_argmax_v3_tiling.h"
20+ 
21+namespace optiling {
22+const int CHW_DIMS = 3;
23+const int NCHW_DIMS = 4;
24+const int KERNEL_POS = 0;
25+const int STRIDE_POS = 1;
26+const int PADDING_POS = 2;
27+const int DILATION_POS = 4;
28+const int CEIL_POS = 5;
29+const int FORMAT_POS = 6;
30+const int N_DIM_ = 0;
31+const int C_DIM_ = 1;
32+const int H_DIM_ = 2;
33+const int W_DIM_ = 3;
34+const int H_IDX_ = 0;
35+const int W_IDX_ = 1;
36+constexpr int64_t MAX_INT32 = 2147483647;
37+constexpr uint64_t SIMT_NCHW_TILING_KEY_INT32 = 500001;
38+constexpr uint64_t SIMT_NHWC_TILING_KEY_INT32 = 500002;
39+constexpr uint64_t SIMT_NCHW_TILING_KEY_INT64 = 500011;
40+constexpr uint64_t SIMT_NHWC_TILING_KEY_INT64 = 500012;
41+constexpr int64_t MAX_THREAD_NUM = 256;
42+constexpr size_t SYS_WORKSPACE_SIZE = 16 * 1024 * 1024;
43+ 
44+BEGIN_TILING_DATA_DEF(MaxPoolWithArgmaxV3SimtTilingData)
45+TILING_DATA_FIELD_DEF(int64_t, threadNums);
46+TILING_DATA_FIELD_DEF(int64_t, blockNums);
47+TILING_DATA_FIELD_DEF(int64_t, nDim);
48+TILING_DATA_FIELD_DEF(int64_t, cDim);
49+TILING_DATA_FIELD_DEF(int64_t, hInDim);
50+TILING_DATA_FIELD_DEF(int64_t, wInDim);
51+TILING_DATA_FIELD_DEF(int64_t, hOutDim);
52+TILING_DATA_FIELD_DEF(int64_t, wOutDim);
53+TILING_DATA_FIELD_DEF(int64_t, kSizeH);
54+TILING_DATA_FIELD_DEF(int64_t, kSizeW);
55+TILING_DATA_FIELD_DEF(int64_t, stridesH);
56+TILING_DATA_FIELD_DEF(int64_t, stridesW);
57+TILING_DATA_FIELD_DEF(int64_t, padH);
58+TILING_DATA_FIELD_DEF(int64_t, padW);
59+TILING_DATA_FIELD_DEF(int64_t, dilationH);
60+TILING_DATA_FIELD_DEF(int64_t, dilationW);
61+TILING_DATA_FIELD_DEF(int64_t, ceilMode);
62+END_TILING_DATA_DEF;
63+// 500001 for NCHW 500002
64+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_500001, MaxPoolWithArgmaxV3SimtTilingData);
65+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3_500002, MaxPoolWithArgmaxV3SimtTilingData);
66+ 
67+struct InputSIMTInfo {
68+ array<uint64_t, NCHW_DIMS> inputShape;
69+ array<uint64_t, NCHW_DIMS> outShape;
70+ array<uint64_t, HW_DIMS> kernelSize;
71+ array<uint64_t, HW_DIMS> stride;
72+ array<uint64_t, HW_DIMS> pad;
73+ array<uint64_t, HW_DIMS> dilation;
74+ bool ceilMode;
75+ std::string data_format;
76+};
77+ 
78+class MaxPoolWithArgmaxV3TilingSIMT : public MaxPoolWithArgmaxV3BaseTiling {
79+public:
80+ explicit MaxPoolWithArgmaxV3TilingSIMT(gert::TilingContext* context) : MaxPoolWithArgmaxV3BaseTiling(context)
81+ {}
82+ 
83+ ~MaxPoolWithArgmaxV3TilingSIMT() override
84+ {}
85+ 
86+protected:
87+ // 获取INPUT/OUTPUT/ATTR信息
88+ ge::graphStatus GetShapeAttrsInfo() override;
89+ // 计算数据切分TilingData
90+ ge::graphStatus DoOpTiling() override;
91+ // 计算TilingKey
92+ uint64_t GetTilingKey() const;
93+ // 保存Tiling数据
94+ ge::graphStatus PostTiling() override;
95+ // tiling信息打屏
96+ std::string ToString(MaxPoolWithArgmaxV3SimtTilingData& tilingData);
97+ 
98+private:
99+ uint64_t GenerateTilingKey(uint64_t innerKey);
100+ MaxPoolWithArgmaxV3SimtTilingData tiling;
101+ InputSIMTInfo inputData;
102+ int nDimPos = 0;
103+ int cDimPos = 1;
104+ int hDimPos = 2;
105+ int wDimPos = 3;
106+ int64_t outputDataCount = 0;
107+};
108+ 
109+} // namespace optiling
110+#endif // MAX_POOL_WITH_ARGMAX_V3_SIMT_TILING_H
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_tiling.cpp+46-0
@@ -0,0 +1,46 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "max_pool_with_argmax_v3_tiling.h"
18+ 
19+using namespace AscendC;
20+using Ops::NN::Optiling::TilingRegistry;
21+namespace optiling {
22+ 
23+ge::graphStatus Tiling4MaxPoolWithArgmaxV3(gert::TilingContext* context)
24+{
25+ return TilingRegistry::GetInstance().DoTilingImpl(context);
26+}
27+ 
28+ge::graphStatus TilingPrepare4MaxPoolWithArgmaxV3(gert::TilingParseContext* context)
29+{
30+ fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo();
31+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr);
32+ 
33+ auto compileInfoPtr = context->GetCompiledInfo<MaxPoolWithArgmaxV3CompileInfo>();
34+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr);
35+ 
36+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
37+ compileInfoPtr->coreNum = ascendcPlatform.GetCoreNum();
38+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize);
39+ return ge::GRAPH_SUCCESS;
40+}
41+ 
42+IMPL_OP_OPTILING(MaxPoolWithArgmaxV3)
43+ .Tiling(Tiling4MaxPoolWithArgmaxV3)
44+ .TilingParse<MaxPoolWithArgmaxV3CompileInfo>(TilingPrepare4MaxPoolWithArgmaxV3);
45+ 
46+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_tiling.h+102-0
@@ -0,0 +1,102 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_WITH_AGRMAX_V3_TILING_BASE_H_
17+#define AIR_CXX_RUNTIME_V2_OP_IMPL_MAX_POOL_WITH_AGRMAX_V3_TILING_BASE_H_
18+ 
19+#include <array>
20+ 
21+#include <cstdint>
22+#include <vector>
23+#include <string>
24+#include "register/op_def_registry.h"
25+#include "register/tilingdata_base.h"
26+#include "tiling_base/tiling_base.h"
27+#include "tiling/tiling_api.h"
28+#include "util/math_util.h"
29+#include "atvoss/broadcast/broadcast_tiling.h"
30+#include "op_common/op_host/util/platform_util.h"
31+ 
32+using namespace std;
33+ 
34+namespace optiling {
35+const int HW_DIMS = 2;
36+const int HW_PAD_DIMS = 5;
37+const int MAX_CORE_NUM = 64;
38+const uint32_t H_DIM = 0;
39+const uint32_t W_DIM = 1;
40+const uint32_t MAX_DIV = 2;
41+const uint32_t NCHW_CONV_ADDR_LIST_SIZE = 16;
42+const uint32_t MIN_TRANSPOSE_ROWS = 16;
43+const uint32_t INT64_FP32 = 2;
44+const uint32_t BINARY_SEARCH_COEFF = 2;
45+const uint32_t BLOCK_LEN_FP32 = 8;
46+const uint32_t BLOCK_LEN_FP16 = 16;
47+ 
48+BEGIN_TILING_DATA_DEF(MaxPoolWithArgmaxV3TilingData)
49+TILING_DATA_FIELD_DEF(uint64_t, nc);
50+TILING_DATA_FIELD_DEF(uint64_t, hx);
51+TILING_DATA_FIELD_DEF(uint64_t, wx);
52+TILING_DATA_FIELD_DEF(uint64_t, kh);
53+END_TILING_DATA_DEF;
54+ 
55+REGISTER_TILING_DATA_CLASS(MaxPoolWithArgmaxV3, MaxPoolWithArgmaxV3TilingData);
56+ 
57+struct InputInfo {
58+ uint64_t batches;
59+ array<uint64_t, HW_DIMS> inputShape;
60+ array<uint64_t, HW_DIMS> outShape;
61+ array<uint64_t, HW_DIMS> kernelSize;
62+ array<uint64_t, HW_DIMS> stride;
63+ array<uint64_t, HW_DIMS> pad;
64+ array<uint64_t, HW_DIMS> dilation;
65+ bool ceilMode;
66+ ge::DataType indexDtype;
67+ ge::Format inputFormat;
68+ uint64_t nInput;
69+ uint64_t cInput;
70+};
71+ 
72+struct MaxPoolWithArgmaxV3CompileInfo {
73+ uint64_t coreNum;
74+ uint64_t ubSize;
75+};
76+ 
77+class MaxPoolWithArgmaxV3BaseTiling : public Ops::NN::Optiling::TilingBaseClass {
78+public:
79+ explicit MaxPoolWithArgmaxV3BaseTiling(gert::TilingContext* context) : TilingBaseClass(context)
80+ {}
81+ 
82+ ~MaxPoolWithArgmaxV3BaseTiling() override
83+ {}
84+ 
85+protected:
86+ bool IsCapable() override;
87+ ge::graphStatus GetPlatformInfo() override;
88+ ge::graphStatus GetShapeAttrsInfo() override;
89+ ge::graphStatus DoOpTiling() override;
90+ ge::graphStatus DoLibApiTiling() override;
91+ uint64_t GetTilingKey() const override;
92+ ge::graphStatus GetWorkspaceSize() override;
93+ ge::graphStatus PostTiling() override;
94+ 
95+public:
96+ InputInfo inputData;
97+ ge::DataType dtype = ge::DataType::DT_FLOAT;
98+ uint32_t coreNum = 1;
99+ uint32_t ubSize = 0;
100+};
101+} // namespace optiling
102+#endif
Apooling/max_pool_with_argmax_v3/op_host/arch35/max_pool_with_argmax_v3_tiling_base.cpp+251-0
@@ -0,0 +1,251 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_tiling_base.cpp
13+ * \brief
14+ */
15+ 
16+#include "tiling_base/tiling_templates_registry.h"
17+#include "log/log.h"
18+#include "error_util.h"
19+#include "platform/platform_info.h"
20+#include "max_pool_with_argmax_v3_tiling.h"
21+ 
22+using namespace AscendC;
23+using namespace ge;
24+ 
25+namespace optiling {
26+ 
27+const int INPUT_IDX_X = 0;
28+const int NCHW_DIMS = 4;
29+const int KERNEL_POS = 0;
30+const int STRIDE_POS = 1;
31+const int PADDING_POS = 2;
32+const int DTYPE_POS = 3;
33+const int DILATION_POS = 4;
34+const int CEIL_POS = 5;
35+const int FORMAT_POS = 6;
36+const int WS_SYS_SIZE = 16 * 1024 * 1024;
37+static const int MP_MAX_2D_DIM_ZERO = 0;
38+static const int MP_MAX_2D_DIM_ONE = 1;
39+static const int MP_MAX_2D_DIM_TWO = 2;
40+static const int MP_MAX_2D_DIM_THREE = 3;
41+static const int64_t MP_MAX_2D_TYPE_INT32 = 3;
42+static const int64_t MP_MAX_2D_TYPE_INT64 = 9;
43+ 
44+ge::graphStatus MaxPoolWithArgmaxV3BaseTiling::GetPlatformInfo()
45+{
46+ auto platformPtr = context_->GetPlatformInfo();
47+ if (platformPtr == nullptr) {
48+ auto compileInfoPtr = reinterpret_cast<const MaxPoolWithArgmaxV3CompileInfo*>(context_->GetCompileInfo());
49+ OP_CHECK_IF(
50+ compileInfoPtr == nullptr, CUBE_INNER_ERR_REPORT(context_, "compile info is null"),
51+ return ge::GRAPH_FAILED);
52+ coreNum = compileInfoPtr->coreNum;
53+ 
54+ ubSize = compileInfoPtr->ubSize;
55+ } else {
56+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformPtr);
57+ coreNum = ascendcPlatform.GetCoreNum();
58+ 
59+ uint64_t ubSizePlatform;
60+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
61+ ubSize = static_cast<int64_t>(ubSizePlatform);
62+ }
63+ OP_CHECK_IF(coreNum == 0, CUBE_INNER_ERR_REPORT(context_, "coreNum is 0"), return ge::GRAPH_FAILED);
64+ 
65+ return ge::GRAPH_SUCCESS;
66+}
67+ 
68+ge::graphStatus MaxPoolWithArgmaxV3BaseTiling::GetShapeAttrsInfo()
69+{
70+ auto inputX = context_->GetInputShape(0);
71+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputX);
72+ auto inputShape = Ops::Base::EnsureNotScalar(inputX->GetStorageShape());
73+ 
74+ OP_CHECK_IF(
75+ inputShape.GetDimNum() != NCHW_DIMS,
76+ OP_LOGE(
77+ context_->GetNodeName(), "MaxPoolWithArgmaxV3: input shape dim = %zu, should be equal 4",
78+ inputShape.GetDimNum()),
79+ return ge::GRAPH_FAILED);
80+ OP_CHECK_IF(
81+ inputShape.GetShapeSize() <= 0,
82+ OP_LOGE(
83+ context_->GetNodeName(), "MaxPoolWithArgmaxV3: input shape size %ld less than zero failed",
84+ inputShape.GetShapeSize()),
85+ return ge::GRAPH_FAILED);
86+ auto inputDesc = context_->GetInputDesc(0);
87+ OP_CHECK_NULL_WITH_CONTEXT(context_, inputDesc);
88+ dtype = inputDesc->GetDataType();
89+ if (dtype != ge::DataType::DT_BF16 && dtype != ge::DataType::DT_FLOAT16 && dtype != ge::DataType::DT_FLOAT) {
90+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: invalid dtype");
91+ return ge::GRAPH_FAILED;
92+ }
93+ 
94+ auto outX = context_->GetOutputShape(0);
95+ OP_CHECK_NULL_WITH_CONTEXT(context_, outX);
96+ auto outShape = Ops::Base::EnsureNotScalar(outX->GetStorageShape());
97+ auto indicesX = context_->GetOutputShape(1);
98+ OP_CHECK_NULL_WITH_CONTEXT(context_, indicesX);
99+ auto indicesShape = Ops::Base::EnsureNotScalar(indicesX->GetStorageShape());
100+ if (indicesShape != outShape) {
101+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: indices shape and values shape is different");
102+ return ge::GRAPH_FAILED;
103+ }
104+ auto runtimeAttrs = context_->GetAttrs();
105+ OP_CHECK_NULL_WITH_CONTEXT(context_, runtimeAttrs);
106+ 
107+ std::string inputFormatStr("NCHW");
108+ const char* inputFormat = runtimeAttrs->GetAttrPointer<char>(FORMAT_POS);
109+ if (inputFormat != nullptr) {
110+ inputFormatStr = inputFormat;
111+ }
112+ int h_dim = MP_MAX_2D_DIM_TWO, w_dim = MP_MAX_2D_DIM_THREE;
113+ if (inputFormatStr == "NCHW") {
114+ inputData.inputFormat = ge::Format::FORMAT_NCHW;
115+ inputData.batches = inputShape.GetDim(MP_MAX_2D_DIM_ZERO) * inputShape.GetDim(MP_MAX_2D_DIM_ONE);
116+ inputData.nInput = inputShape.GetDim(MP_MAX_2D_DIM_ZERO);
117+ inputData.cInput = inputShape.GetDim(MP_MAX_2D_DIM_ONE);
118+ } else if (inputFormatStr == "NHWC") {
119+ inputData.inputFormat = ge::Format::FORMAT_NHWC;
120+ h_dim = MP_MAX_2D_DIM_ONE;
121+ w_dim = MP_MAX_2D_DIM_TWO;
122+ inputData.batches = inputShape.GetDim(MP_MAX_2D_DIM_ZERO) * inputShape.GetDim(MP_MAX_2D_DIM_THREE);
123+ inputData.nInput = inputShape.GetDim(MP_MAX_2D_DIM_ZERO);
124+ inputData.cInput = inputShape.GetDim(MP_MAX_2D_DIM_THREE);
125+ } else {
126+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: not support format %s", inputFormatStr.c_str());
127+ return ge::GRAPH_FAILED;
128+ }
129+ 
130+ OP_CHECK_IF(
131+ outShape.GetDim(h_dim) < 1 || outShape.GetDim(w_dim) < 1,
132+ OP_LOGE(
133+ context_->GetNodeName(), "MaxPoolWithArgmaxV3: output shape [%ld, %ld] not support", outShape.GetDim(h_dim),
134+ outShape.GetDim(w_dim)),
135+ return ge::GRAPH_FAILED);
136+ 
137+ inputData.inputShape =
138+ array<uint64_t, HW_DIMS>{uint64_t(inputShape.GetDim(h_dim)), uint64_t(inputShape.GetDim(w_dim))};
139+ inputData.outShape = array<uint64_t, HW_DIMS>{uint64_t(outShape.GetDim(h_dim)), uint64_t(outShape.GetDim(w_dim))};
140+ 
141+ int32_t hValue = 0;
142+ int32_t wValue = 0;
143+ const gert::TypedContinuousVector<int64_t>* kernelSize = runtimeAttrs->GetListInt(KERNEL_POS);
144+ OP_CHECK_NULL_WITH_CONTEXT(context_, kernelSize);
145+ hValue = *(kernelSize->GetData());
146+ wValue = *(kernelSize->GetData() + 1);
147+ inputData.kernelSize = array<uint64_t, HW_DIMS>{uint64_t(hValue), uint64_t(wValue)};
148+ OP_CHECK_IF(
149+ hValue <= 0 || wValue <= 0,
150+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: not support kernel shape [%d, %d]", hValue, wValue),
151+ return ge::GRAPH_FAILED);
152+ 
153+ int32_t khValue = hValue;
154+ int32_t kwValue = wValue;
155+ const gert::TypedContinuousVector<int64_t>* stride = runtimeAttrs->GetListInt(STRIDE_POS);
156+ OP_CHECK_NULL_WITH_CONTEXT(context_, stride);
157+ hValue = *(stride->GetData());
158+ wValue = *(stride->GetData() + 1);
159+ inputData.stride = array<uint64_t, HW_DIMS>{uint64_t(hValue), uint64_t(wValue)};
160+ OP_CHECK_IF(
161+ hValue <= 0 || wValue <= 0,
162+ OP_LOGE(context_->GetNodeName(), "MaxPoolWithArgmaxV3: not support stride shape [%d, %d]", hValue, wValue),
163+ return ge::GRAPH_FAILED);
164+ 
165+ const gert::TypedContinuousVector<int64_t>* padding = runtimeAttrs->GetListInt(PADDING_POS);
166+ OP_CHECK_NULL_WITH_CONTEXT(context_, padding);
167+ hValue = *(padding->GetData());
168+ wValue = *(padding->GetData() + 1);
169+ inputData.pad = array<uint64_t, HW_DIMS>{uint64_t(hValue), uint64_t(wValue)};
170+ OP_CHECK_IF(
171+ hValue > khValue / 2 || wValue > kwValue / 2,
172+ OP_LOGE(
173+ context_->GetNodeName(), "MaxPoolWithArgmaxV3: not support pad shape [%d, %d] kernel shape [%d, %d]",
174+ hValue, wValue, khValue, kwValue),
175+ return ge::GRAPH_FAILED);
176+ 
177+ inputData.dilation = array<uint64_t, HW_DIMS>{1, 1};
178+ hValue = 1;
179+ wValue = 1;
180+ const gert::TypedContinuousVector<int64_t>* dilation = runtimeAttrs->GetListInt(DILATION_POS);
181+ if (dilation != nullptr) {
182+ hValue = *(dilation->GetData());
183+ wValue = *(dilation->GetData() + 1);
184+ inputData.dilation = array<uint64_t, HW_DIMS>{uint64_t(hValue), uint64_t(wValue)};
185+ OP_CHECK_IF(
186+ hValue <= 0 || wValue <= 0,
187+ OP_LOGE(
188+ context_->GetNodeName(), "MaxPoolWithArgmaxV3: not support dilation shape [%d, %d]", hValue, wValue),
189+ return ge::GRAPH_FAILED);
190+ }
191+ 
192+ inputData.ceilMode = false;
193+ const bool* ceilModePtr = runtimeAttrs->GetAttrPointer<bool>(CEIL_POS);
194+ if (ceilModePtr != nullptr) {
195+ inputData.ceilMode = *ceilModePtr;
196+ }
197+ 
198+ int indexDtype = 3;
199+ const int* indexDtypePtr = runtimeAttrs->GetAttrPointer<int>(DTYPE_POS);
200+ if (indexDtypePtr != nullptr) {
201+ indexDtype = *indexDtypePtr;
202+ }
203+ switch (indexDtype) {
204+ case MP_MAX_2D_TYPE_INT32:
205+ inputData.indexDtype = ge::DataType::DT_INT32;
206+ break;
207+ case MP_MAX_2D_TYPE_INT64:
208+ inputData.indexDtype = ge::DataType::DT_INT64;
209+ break;
210+ default:
211+ inputData.indexDtype = ge::DataType::DT_INT32;
212+ break;
213+ }
214+ return ge::GRAPH_SUCCESS;
215+}
216+ 
217+bool MaxPoolWithArgmaxV3BaseTiling::IsCapable()
218+{
219+ return true;
220+}
221+ 
222+ge::graphStatus MaxPoolWithArgmaxV3BaseTiling::DoOpTiling()
223+{
224+ return ge::GRAPH_SUCCESS;
225+}
226+ 
227+ge::graphStatus MaxPoolWithArgmaxV3BaseTiling::DoLibApiTiling()
228+{
229+ return ge::GRAPH_SUCCESS;
230+}
231+ 
232+uint64_t MaxPoolWithArgmaxV3BaseTiling::GetTilingKey() const
233+{
234+ return 0;
235+}
236+ 
237+ge::graphStatus MaxPoolWithArgmaxV3BaseTiling::GetWorkspaceSize()
238+{
239+ auto sys_workspace = WS_SYS_SIZE;
240+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
241+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
242+ currentWorkspace[0] = sys_workspace;
243+ 
244+ return ge::GRAPH_SUCCESS;
245+}
246+ 
247+ge::graphStatus MaxPoolWithArgmaxV3BaseTiling::PostTiling()
248+{
249+ return ge::GRAPH_SUCCESS;
250+}
251+} // namespace optiling
Apooling/max_pool_with_argmax_v3/op_host/config/ascend910_95/max_pool_with_argmax_v3_binary.json+431-0
@@ -0,0 +1,431 @@
1+{
2+ "op_type": "MaxPoolWithArgmaxV3",
3+ "op_list": [
4+ {
5+ "bin_filename": "MaxPoolWithArgmaxV3_float32_int32",
6+ "inputs": [
7+ {
8+ "name": "x",
9+ "index": 0,
10+ "dtype": "float32",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [-2],
14+ "format_match_mode": "FormatAgnostic"
15+ }
16+ ],
17+ "outputs": [
18+ {
19+ "name": "y",
20+ "index": 0,
21+ "dtype": "float32",
22+ "format": "ND",
23+ "paramType": "required",
24+ "shape": [-2],
25+ "format_match_mode": "FormatAgnostic"
26+ },
27+ {
28+ "name": "argmax",
29+ "index": 0,
30+ "dtype": "int32",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [-2],
34+ "format_match_mode": "FormatAgnostic"
35+ }
36+ ],
37+ "attrs": [
38+ {
39+ "name": "ksize",
40+ "dtype": "list_int",
41+ "value": []
42+ },
43+ {
44+ "name": "strides",
45+ "dtype": "list_int",
46+ "value": []
47+ },
48+ {
49+ "name": "pads",
50+ "dtype": "list_int",
51+ "value": []
52+ },
53+ {
54+ "name": "dtype",
55+ "dtype": "int",
56+ "value": 3
57+ },
58+ {
59+ "name": "dilation",
60+ "dtype": "list_int",
61+ "value": []
62+ },
63+ {
64+ "name": "ceil_mode",
65+ "dtype": "bool",
66+ "value": null
67+ },
68+ {
69+ "name": "data_format",
70+ "dtype": "string",
71+ "value": null
72+ }
73+ ]
74+ },
75+ {
76+ "bin_filename": "MaxPoolWithArgmaxV3_float16_int32",
77+ "inputs": [
78+ {
79+ "name": "x",
80+ "index": 0,
81+ "dtype": "float16",
82+ "format": "ND",
83+ "paramType": "required",
84+ "shape": [-2],
85+ "format_match_mode": "FormatAgnostic"
86+ }
87+ ],
88+ "outputs": [
89+ {
90+ "name": "y",
91+ "index": 0,
92+ "dtype": "float16",
93+ "format": "ND",
94+ "paramType": "required",
95+ "shape": [-2],
96+ "format_match_mode": "FormatAgnostic"
97+ },
98+ {
99+ "name": "argmax",
100+ "index": 0,
101+ "dtype": "int32",
102+ "format": "ND",
103+ "paramType": "required",
104+ "shape": [-2],
105+ "format_match_mode": "FormatAgnostic"
106+ }
107+ ],
108+ "attrs": [
109+ {
110+ "name": "ksize",
111+ "dtype": "list_int",
112+ "value": []
113+ },
114+ {
115+ "name": "strides",
116+ "dtype": "list_int",
117+ "value": []
118+ },
119+ {
120+ "name": "pads",
121+ "dtype": "list_int",
122+ "value": []
123+ },
124+ {
125+ "name": "dtype",
126+ "dtype": "int",
127+ "value": 3
128+ },
129+ {
130+ "name": "dilation",
131+ "dtype": "list_int",
132+ "value": []
133+ },
134+ {
135+ "name": "ceil_mode",
136+ "dtype": "bool",
137+ "value": null
138+ },
139+ {
140+ "name": "data_format",
141+ "dtype": "string",
142+ "value": null
143+ }
144+ ]
145+ },
146+ {
147+ "bin_filename": "MaxPoolWithArgmaxV3_bfloat16_int32",
148+ "inputs": [
149+ {
150+ "name": "x",
151+ "index": 0,
152+ "dtype": "bfloat16",
153+ "format": "ND",
154+ "paramType": "required",
155+ "shape": [-2],
156+ "format_match_mode": "FormatAgnostic"
157+ }
158+ ],
159+ "outputs": [
160+ {
161+ "name": "y",
162+ "index": 0,
163+ "dtype": "bfloat16",
164+ "format": "ND",
165+ "paramType": "required",
166+ "shape": [-2],
167+ "format_match_mode": "FormatAgnostic"
168+ },
169+ {
170+ "name": "argmax",
171+ "index": 0,
172+ "dtype": "int32",
173+ "format": "ND",
174+ "paramType": "required",
175+ "shape": [-2],
176+ "format_match_mode": "FormatAgnostic"
177+ }
178+ ],
179+ "attrs": [
180+ {
181+ "name": "ksize",
182+ "dtype": "list_int",
183+ "value": []
184+ },
185+ {
186+ "name": "strides",
187+ "dtype": "list_int",
188+ "value": []
189+ },
190+ {
191+ "name": "pads",
192+ "dtype": "list_int",
193+ "value": []
194+ },
195+ {
196+ "name": "dtype",
197+ "dtype": "int",
198+ "value": 3
199+ },
200+ {
201+ "name": "dilation",
202+ "dtype": "list_int",
203+ "value": []
204+ },
205+ {
206+ "name": "ceil_mode",
207+ "dtype": "bool",
208+ "value": null
209+ },
210+ {
211+ "name": "data_format",
212+ "dtype": "string",
213+ "value": null
214+ }
215+ ]
216+ },
217+ {
218+ "bin_filename": "MaxPoolWithArgmaxV3_float32_int64",
219+ "inputs": [
220+ {
221+ "name": "x",
222+ "index": 0,
223+ "dtype": "float32",
224+ "format": "ND",
225+ "paramType": "required",
226+ "shape": [-2],
227+ "format_match_mode": "FormatAgnostic"
228+ }
229+ ],
230+ "outputs": [
231+ {
232+ "name": "y",
233+ "index": 0,
234+ "dtype": "float32",
235+ "format": "ND",
236+ "paramType": "required",
237+ "shape": [-2],
238+ "format_match_mode": "FormatAgnostic"
239+ },
240+ {
241+ "name": "argmax",
242+ "index": 0,
243+ "dtype": "int64",
244+ "format": "ND",
245+ "paramType": "required",
246+ "shape": [-2],
247+ "format_match_mode": "FormatAgnostic"
248+ }
249+ ],
250+ "attrs": [
251+ {
252+ "name": "ksize",
253+ "dtype": "list_int",
254+ "value": []
255+ },
256+ {
257+ "name": "strides",
258+ "dtype": "list_int",
259+ "value": []
260+ },
261+ {
262+ "name": "pads",
263+ "dtype": "list_int",
264+ "value": []
265+ },
266+ {
267+ "name": "dtype",
268+ "dtype": "int",
269+ "value": 9
270+ },
271+ {
272+ "name": "dilation",
273+ "dtype": "list_int",
274+ "value": []
275+ },
276+ {
277+ "name": "ceil_mode",
278+ "dtype": "bool",
279+ "value": null
280+ },
281+ {
282+ "name": "data_format",
283+ "dtype": "string",
284+ "value": null
285+ }
286+ ]
287+ },
288+ {
289+ "bin_filename": "MaxPoolWithArgmaxV3_float16_int64",
290+ "inputs": [
291+ {
292+ "name": "x",
293+ "index": 0,
294+ "dtype": "float16",
295+ "format": "ND",
296+ "paramType": "required",
297+ "shape": [-2],
298+ "format_match_mode": "FormatAgnostic"
299+ }
300+ ],
301+ "outputs": [
302+ {
303+ "name": "y",
304+ "index": 0,
305+ "dtype": "float16",
306+ "format": "ND",
307+ "paramType": "required",
308+ "shape": [-2],
309+ "format_match_mode": "FormatAgnostic"
310+ },
311+ {
312+ "name": "argmax",
313+ "index": 0,
314+ "dtype": "int64",
315+ "format": "ND",
316+ "paramType": "required",
317+ "shape": [-2],
318+ "format_match_mode": "FormatAgnostic"
319+ }
320+ ],
321+ "attrs": [
322+ {
323+ "name": "ksize",
324+ "dtype": "list_int",
325+ "value": []
326+ },
327+ {
328+ "name": "strides",
329+ "dtype": "list_int",
330+ "value": []
331+ },
332+ {
333+ "name": "pads",
334+ "dtype": "list_int",
335+ "value": []
336+ },
337+ {
338+ "name": "dtype",
339+ "dtype": "int",
340+ "value": 9
341+ },
342+ {
343+ "name": "dilation",
344+ "dtype": "list_int",
345+ "value": []
346+ },
347+ {
348+ "name": "ceil_mode",
349+ "dtype": "bool",
350+ "value": null
351+ },
352+ {
353+ "name": "data_format",
354+ "dtype": "string",
355+ "value": null
356+ }
357+ ]
358+ },
359+ {
360+ "bin_filename": "MaxPoolWithArgmaxV3_bfloat16_int64",
361+ "inputs": [
362+ {
363+ "name": "x",
364+ "index": 0,
365+ "dtype": "bfloat16",
366+ "format": "ND",
367+ "paramType": "required",
368+ "shape": [-2],
369+ "format_match_mode": "FormatAgnostic"
370+ }
371+ ],
372+ "outputs": [
373+ {
374+ "name": "y",
375+ "index": 0,
376+ "dtype": "bfloat16",
377+ "format": "ND",
378+ "paramType": "required",
379+ "shape": [-2],
380+ "format_match_mode": "FormatAgnostic"
381+ },
382+ {
383+ "name": "argmax",
384+ "index": 0,
385+ "dtype": "int64",
386+ "format": "ND",
387+ "paramType": "required",
388+ "shape": [-2],
389+ "format_match_mode": "FormatAgnostic"
390+ }
391+ ],
392+ "attrs": [
393+ {
394+ "name": "ksize",
395+ "dtype": "list_int",
396+ "value": []
397+ },
398+ {
399+ "name": "strides",
400+ "dtype": "list_int",
401+ "value": []
402+ },
403+ {
404+ "name": "pads",
405+ "dtype": "list_int",
406+ "value": []
407+ },
408+ {
409+ "name": "dtype",
410+ "dtype": "int",
411+ "value": 9
412+ },
413+ {
414+ "name": "dilation",
415+ "dtype": "list_int",
416+ "value": []
417+ },
418+ {
419+ "name": "ceil_mode",
420+ "dtype": "bool",
421+ "value": null
422+ },
423+ {
424+ "name": "data_format",
425+ "dtype": "string",
426+ "value": null
427+ }
428+ ]
429+ }
430+ ]
431+}
Apooling/max_pool_with_argmax_v3/op_host/config/ascend910_95/max_pool_with_argmax_v3_simplified_key.ini+13-0
@@ -0,0 +1,13 @@
1+; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2+; [某算子]
3+; default=xx
4+; ascendxx=xx
5+; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6+; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7+; 2)如果仅有default配置：各个版本按default配置
8+; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9+; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10+; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11+; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12+[MaxPoolWithArgmaxV3]
13+default=0
Apooling/max_pool_with_argmax_v3/op_host/max_pool_with_argmax_v3_def.cpp+66-0
@@ -0,0 +1,66 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_def.cpp
13+ * \brief imply for max_pool_with_argmax
14+ */
15+ 
16+#include "register/op_def_registry.h"
17+ 
18+namespace ops {
19+constexpr int DTYPE_INT32 = 3;
20+class MaxPoolWithArgmaxV3 : public OpDef {
21+public:
22+ explicit MaxPoolWithArgmaxV3(const char* name) : OpDef(name)
23+ {
24+ this->Input("x")
25+ .ParamType(REQUIRED)
26+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
27+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
28+ .UnknownShapeFormat(
29+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
30+ .AutoContiguous();
31+ this->Output("y")
32+ .ParamType(REQUIRED)
33+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
34+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
35+ .UnknownShapeFormat(
36+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
37+ .AutoContiguous();
38+ this->Output("argmax")
39+ .ParamType(REQUIRED)
40+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
41+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
42+ .UnknownShapeFormat(
43+ {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
44+ .AutoContiguous();
45+ this->Attr("ksize").AttrType(REQUIRED).ListInt();
46+ this->Attr("strides").AttrType(REQUIRED).ListInt();
47+ this->Attr("pads").AttrType(REQUIRED).ListInt();
48+ this->Attr("dtype").AttrType(OPTIONAL).Int(DTYPE_INT32);
49+ this->Attr("dilation").AttrType(OPTIONAL).ListInt({1, 1});
50+ this->Attr("ceil_mode").AttrType(OPTIONAL).Bool(false);
51+ this->Attr("data_format").AttrType(OPTIONAL).String("NCHW");
52+ 
53+ OpAICoreConfig aiCoreConfig;
54+ aiCoreConfig.DynamicCompileStaticFlag(true)
55+ .DynamicFormatFlag(false)
56+ .DynamicRankSupportFlag(true)
57+ .DynamicShapeSupportFlag(true)
58+ .NeedCheckSupportFlag(false)
59+ .PrecisionReduceFlag(true)
60+ .ExtendCfgInfo("opFile.value", "max_pool_with_argmax_v3");
61+ this->AICore().AddConfig("ascend910_95", aiCoreConfig);
62+ }
63+};
64+ 
65+OP_ADD(MaxPoolWithArgmaxV3);
66+} // namespace ops
Apooling/max_pool_with_argmax_v3/op_host/max_pool_with_argmax_v3_infershape.cpp+204-0
@@ -0,0 +1,204 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_infershape.cpp
13+ * \brief
14+ */
15+#include <string>
16+#include "graph/utils/type_utils.h"
17+#include "runtime/infer_shape_context.h"
18+#include "register/op_impl_registry.h"
19+#include "log/log.h"
20+#include "util/shape_util.h"
21+#include "util/math_util.h"
22+ 
23+using namespace ge;
24+namespace ops {
25+static constexpr size_t INDEX_KSIZE = 0;
26+static constexpr size_t INDEX_STRIDES = 1;
27+static constexpr size_t INDEX_PADS = 2;
28+static constexpr size_t INDEX_DTYPE = 3;
29+static constexpr size_t INDEX_DILATION = 4;
30+static constexpr size_t INDEX_CEIL_MODE = 5;
31+static constexpr size_t INDEX_DATA_FORMAT = 6;
32+static constexpr size_t ATTR_LIST_SHAPE_SIZE = 2;
33+static constexpr size_t INDEX_OUT_MAX = 0;
34+static constexpr size_t INDEX_OUT_INDICES = 1;
35+static constexpr size_t C0_SIZE = 16;
36+static constexpr size_t PARAM_NUM = 4;
37+static constexpr size_t PARAM_H_DIM = 0;
38+static constexpr size_t PARAM_W_DIM = 1;
39+static constexpr size_t SHAPE_H_DIM = 2;
40+static constexpr size_t SHAPE_W_DIM = 3;
41+static constexpr size_t CHW_DIM = 3;
42+static constexpr size_t INT32_DTYPE = 3;
43+static constexpr size_t INT64_DTYPE = 9;
44+ 
45+static int64_t DivRtn(int64_t x, int64_t y)
46+{
47+ if (y == 0) {
48+ OP_LOGE("MaxPoolWithArgmaxV3", "y value cannot be zero.");
49+ return GRAPH_FAILED;
50+ }
51+ int64_t q = x / y;
52+ int64_t r = x % y;
53+ if ((r != 0) && ((r < 0) != (y < 0))) {
54+ --q;
55+ }
56+ return q;
57+}
58+ 
59+static void UpdateMaxShape(
60+ const int64_t (&param)[PARAM_NUM], bool ceil_mode, const int64_t& dim_size, int64_t& out_max_shape)
61+{
62+ int64_t ksize = param[INDEX_KSIZE];
63+ int64_t strides = param[INDEX_STRIDES];
64+ int64_t pad = param[INDEX_PADS];
65+ int64_t dilation = param[PARAM_NUM - 1];
66+ int64_t exact_size = dim_size + 2 * pad - dilation * (ksize - 1) - 1 + (ceil_mode ? (strides - 1) : 0);
67+ out_max_shape = DivRtn(exact_size, strides) + 1;
68+ if (ceil_mode) {
69+ if ((out_max_shape - 1) * strides >= dim_size + pad) {
70+ out_max_shape = out_max_shape - 1;
71+ }
72+ }
73+}
74+ 
75+ge::graphStatus InferShapeForMaxPoolWithArgmaxV3(gert::InferShapeContext* context)
76+{
77+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPoolWithArgmaxV3 infershape running");
78+ auto src_td = context->GetInputDesc(0);
79+ OP_CHECK_NULL_WITH_CONTEXT(context, src_td);
80+ auto input_format = src_td->GetOriginFormat();
81+ auto indices_td = context->GetOutputDesc(INDEX_OUT_INDICES);
82+ OP_CHECK_NULL_WITH_CONTEXT(context, indices_td);
83+ auto indices_dtype = indices_td->GetDataType();
84+ OP_LOGD(context->GetNodeName(), "indices_dtype = %d", indices_dtype);
85+ 
86+ OP_CHECK_IF(
87+ input_format != FORMAT_ND && input_format != FORMAT_NCHW && input_format != FORMAT_NHWC,
88+ OP_LOGE(context->GetNodeName(), "format only supports ND, NCHW, NHWC"), return GRAPH_FAILED);
89+ 
90+ size_t param_h_dim = PARAM_H_DIM;
91+ size_t param_w_dim = PARAM_W_DIM;
92+ size_t input_h_dim = SHAPE_H_DIM;
93+ size_t input_w_dim = SHAPE_W_DIM;
94+ 
95+ auto attrs = context->GetAttrs();
96+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
97+ 
98+ auto ksize = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_KSIZE);
99+ OP_CHECK_NULL_WITH_CONTEXT(context, ksize);
100+ OP_CHECK_IF(
101+ ksize->GetSize() != ATTR_LIST_SHAPE_SIZE,
102+ OP_LOGE(context->GetNodeName(), "Length of ksize %lu must be 2!", ksize->GetSize()), return GRAPH_FAILED);
103+ auto ksize_data = reinterpret_cast<const int64_t*>(ksize->GetData());
104+ 
105+ auto strides = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_STRIDES);
106+ OP_CHECK_NULL_WITH_CONTEXT(context, strides);
107+ OP_CHECK_IF(
108+ strides->GetSize() != ATTR_LIST_SHAPE_SIZE,
109+ OP_LOGE(context->GetNodeName(), "Length of strides %lu must be 2!", strides->GetSize()), return GRAPH_FAILED);
110+ auto strides_data = reinterpret_cast<const int64_t*>(strides->GetData());
111+ 
112+ auto pads = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_PADS);
113+ OP_CHECK_NULL_WITH_CONTEXT(context, pads);
114+ OP_CHECK_IF(
115+ pads->GetSize() != ATTR_LIST_SHAPE_SIZE,
116+ OP_LOGE(context->GetNodeName(), "Length of pads %lu must be 2!", pads->GetSize()), return GRAPH_FAILED);
117+ auto pads_data = reinterpret_cast<const int64_t*>(pads->GetData());
118+ 
119+ auto dilation = attrs->GetAttrPointer<gert::ContinuousVector>(INDEX_DILATION);
120+ OP_CHECK_NULL_WITH_CONTEXT(context, dilation);
121+ OP_CHECK_IF(
122+ dilation->GetSize() != ATTR_LIST_SHAPE_SIZE,
123+ OP_LOGE(context->GetNodeName(), "Length of dilation %lu must be 2!", dilation->GetSize()), return GRAPH_FAILED);
124+ auto dilation_data = reinterpret_cast<const int64_t*>(dilation->GetData());
125+ 
126+ auto ceil_mode = attrs->GetAttrPointer<bool>(INDEX_CEIL_MODE);
127+ OP_CHECK_NULL_WITH_CONTEXT(context, ceil_mode);
128+ 
129+ const char* data_format = attrs->GetAttrPointer<char>(INDEX_DATA_FORMAT);
130+ OP_CHECK_NULL_WITH_CONTEXT(context, data_format);
131+ 
132+ std::string data_format_str = data_format;
133+ if (data_format_str == "NHWC") {
134+ input_h_dim = input_h_dim - 1UL;
135+ input_w_dim = input_w_dim - 1UL;
136+ }
137+ 
138+ const gert::Shape* in_shape = context->GetInputShape(0);
139+ OP_CHECK_NULL_WITH_CONTEXT(context, in_shape);
140+ gert::Shape* out_max_shape = context->GetOutputShape(INDEX_OUT_MAX);
141+ OP_CHECK_NULL_WITH_CONTEXT(context, out_max_shape);
142+ *out_max_shape = *in_shape;
143+ gert::Shape* out_indices_shape = context->GetOutputShape(INDEX_OUT_INDICES);
144+ OP_CHECK_NULL_WITH_CONTEXT(context, out_indices_shape);
145+ *out_indices_shape = *in_shape;
146+ 
147+ if (Ops::Base::IsUnknownRank(*in_shape) || Ops::Base::IsUnknownShape(*in_shape)) {
148+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPoolWithArgmaxV3 infershape handle unknown rank or shape.");
149+ return ge::GRAPH_SUCCESS;
150+ }
151+ 
152+ size_t dim_num = in_shape->GetDimNum();
153+ int64_t max_dim = 0;
154+ if (dim_num == CHW_DIM) {
155+ input_h_dim = input_h_dim - 1UL;
156+ input_w_dim = input_w_dim - 1UL;
157+ }
158+ for (size_t i = 0; i < dim_num; i++) {
159+ int64_t input_dim = in_shape->GetDim(i);
160+ if (i == input_h_dim) {
161+ int64_t param[PARAM_NUM] = {
162+ ksize_data[param_h_dim], strides_data[param_h_dim], pads_data[param_h_dim], dilation_data[param_h_dim]};
163+ UpdateMaxShape(param, *ceil_mode, input_dim, max_dim);
164+ out_max_shape->SetDim(i, max_dim);
165+ out_indices_shape->SetDim(i, max_dim);
166+ } else if (i == input_w_dim) {
167+ int64_t param[PARAM_NUM] = {
168+ ksize_data[param_w_dim], strides_data[param_w_dim], pads_data[param_w_dim], dilation_data[param_w_dim]};
169+ UpdateMaxShape(param, *ceil_mode, input_dim, max_dim);
170+ out_max_shape->SetDim(i, max_dim);
171+ out_indices_shape->SetDim(i, max_dim);
172+ } else {
173+ out_max_shape->SetDim(i, input_dim);
174+ out_indices_shape->SetDim(i, input_dim);
175+ }
176+ }
177+ OP_LOGD(context->GetNodeName(), "runtime2.0 MaxPoolWithArgmaxV3 infershape run success.");
178+ return GRAPH_SUCCESS;
179+}
180+ 
181+static ge::graphStatus InferDataTypeForMaxPoolWithArgmaxV3(gert::InferDataTypeContext* context)
182+{
183+ if (context == nullptr) {
184+ return GRAPH_FAILED;
185+ }
186+ 
187+ const ge::DataType x = context->GetInputDataType(0);
188+ context->SetOutputDataType(INDEX_OUT_MAX, x);
189+ 
190+ auto attrsPtr = context->GetAttrs();
191+ OP_CHECK_NULL_WITH_CONTEXT(context, attrsPtr);
192+ const int64_t* dstDtype = attrsPtr->GetAttrPointer<int64_t>(INDEX_DTYPE);
193+ OP_CHECK_NULL_WITH_CONTEXT(context, dstDtype);
194+ ge::DataType indicesDtype = *dstDtype == INT32_DTYPE ? ge::DT_INT32 : ge::DT_INT64;
195+ 
196+ context->SetOutputDataType(INDEX_OUT_INDICES, indicesDtype);
197+ 
198+ return GRAPH_SUCCESS;
199+}
200+ 
201+IMPL_OP_INFERSHAPE(MaxPoolWithArgmaxV3)
202+ .InferShape(InferShapeForMaxPoolWithArgmaxV3)
203+ .InferDataType(InferDataTypeForMaxPoolWithArgmaxV3);
204+} // namespace ops
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_base.h+124-0
@@ -0,0 +1,124 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_base.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_WITH_ARGMAX_V3_BASE_H_
17+#define MAX_POOL_WITH_ARGMAX_V3_BASE_H_
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+ 
23+using namespace AscendC;
24+ 
25+// 默认 rate1D = 1 生成 0 1 2 3 ... rate1D = 0 生成 0 0 0 0 ...
26+template <typename T>
27+__aicore__ inline void GenGatterIndex2D(MicroAPI::RegTensor<T>& indexReg, T rate2D, T num1D, T rate1D = 1)
28+{
29+ AscendC::MicroAPI::Arange(indexReg, 0);
30+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg;
31+ AscendC::MicroAPI::RegTensor<T> tmpReg;
32+ AscendC::MicroAPI::RegTensor<T> constReg;
33+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
34+ AscendC::MicroAPI::Duplicate(constReg, T(num1D));
35+ AscendC::MicroAPI::Div(segmentScalarReg, indexReg, constReg, preg);
36+ AscendC::MicroAPI::Muls(tmpReg, segmentScalarReg, T(num1D), preg);
37+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
38+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rate1D), preg);
39+ AscendC::MicroAPI::Muls(segmentScalarReg, segmentScalarReg, T(rate2D), preg);
40+ 
41+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg, preg);
42+}
43+ 
44+template <typename T>
45+__aicore__ inline void GenGatterIndex3D(
46+ MicroAPI::RegTensor<T>& indexReg, T rate3D, T num2D, T rate2D, T num1D, T rate1D = 1)
47+{
48+ AscendC::MicroAPI::Arange(indexReg, 0);
49+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg;
50+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg2;
51+ AscendC::MicroAPI::RegTensor<T> tmpReg;
52+ AscendC::MicroAPI::RegTensor<T> constReg;
53+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
54+ AscendC::MicroAPI::Duplicate(constReg, T(num2D));
55+ AscendC::MicroAPI::Div(segmentScalarReg2, indexReg, constReg, preg);
56+ AscendC::MicroAPI::Muls(tmpReg, segmentScalarReg2, T(num2D), preg);
57+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
58+ AscendC::MicroAPI::Muls(segmentScalarReg2, segmentScalarReg2, T(rate3D), preg);
59+ 
60+ AscendC::MicroAPI::Duplicate(constReg, T(num1D));
61+ AscendC::MicroAPI::Div(segmentScalarReg, indexReg, constReg, preg);
62+ AscendC::MicroAPI::Muls(tmpReg, segmentScalarReg, T(num1D), preg);
63+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
64+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rate1D), preg);
65+ AscendC::MicroAPI::Muls(segmentScalarReg, segmentScalarReg, T(rate2D), preg);
66+ 
67+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg, preg);
68+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg2, preg);
69+}
70+ 
71+template <typename T>
72+__aicore__ inline void GenGatterIndex4D(
73+ MicroAPI::RegTensor<T>& indexReg, T rate4D, T num3D, T rate3D, T num2D, T rate2D, T num1D, T rate1D = 1)
74+{
75+ AscendC::MicroAPI::Arange(indexReg, 0);
76+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg;
77+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg2;
78+ AscendC::MicroAPI::RegTensor<T> segmentScalarReg3;
79+ AscendC::MicroAPI::RegTensor<T> tmpReg;
80+ AscendC::MicroAPI::RegTensor<T> constReg;
81+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
82+ AscendC::MicroAPI::Duplicate(constReg, T(num3D));
83+ AscendC::MicroAPI::Div(segmentScalarReg3, indexReg, constReg, preg);
84+ AscendC::MicroAPI::Muls(tmpReg, segmentScalarReg3, T(num3D), preg);
85+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
86+ AscendC::MicroAPI::Muls(segmentScalarReg3, segmentScalarReg3, T(rate4D), preg);
87+ 
88+ AscendC::MicroAPI::Duplicate(constReg, T(num2D));
89+ AscendC::MicroAPI::Div(segmentScalarReg2, indexReg, constReg, preg);
90+ AscendC::MicroAPI::Muls(tmpReg, segmentScalarReg2, T(num2D), preg);
91+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
92+ AscendC::MicroAPI::Muls(segmentScalarReg2, segmentScalarReg2, T(rate3D), preg);
93+ 
94+ AscendC::MicroAPI::Duplicate(constReg, T(num1D));
95+ AscendC::MicroAPI::Div(segmentScalarReg, indexReg, constReg, preg);
96+ AscendC::MicroAPI::Muls(tmpReg, segmentScalarReg, T(num1D), preg);
97+ AscendC::MicroAPI::Sub(indexReg, indexReg, tmpReg, preg);
98+ AscendC::MicroAPI::Muls(indexReg, indexReg, T(rate1D), preg);
99+ AscendC::MicroAPI::Muls(segmentScalarReg, segmentScalarReg, T(rate2D), preg);
100+ 
101+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg, preg);
102+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg2, preg);
103+ AscendC::MicroAPI::Add(indexReg, indexReg, segmentScalarReg3, preg);
104+}
105+ 
106+template <typename T>
107+__aicore__ inline void DuplicateNegInfReg(MicroAPI::RegTensor<T>& negInfReg)
108+{
109+ // -inf
110+ constexpr uint32_t FLOAT32_NEG_INF = 0xFF800000;
111+ constexpr uint16_t FLOAT16_NEG_INF = 0xFC00;
112+ constexpr uint16_t BFLOAT16_NEG_INF = 0xFF80;
113+ using computeType = std::conditional_t<std::is_same<T, float>::value, uint32_t, uint16_t>;
114+ 
115+ if constexpr (std::is_same<T, float>::value) {
116+ AscendC::MicroAPI::Duplicate((AscendC::MicroAPI::RegTensor<computeType>&)negInfReg, (FLOAT32_NEG_INF));
117+ } else if constexpr (std::is_same<T, half>::value) {
118+ AscendC::MicroAPI::Duplicate((AscendC::MicroAPI::RegTensor<computeType>&)negInfReg, (FLOAT16_NEG_INF));
119+ } else {
120+ AscendC::MicroAPI::Duplicate((AscendC::MicroAPI::RegTensor<computeType>&)negInfReg, (BFLOAT16_NEG_INF));
121+ }
122+}
123+ 
124+#endif // MAX_POOL_WITH_ARGMAX_V3_BASE_H_
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_big_kernel.h+562-0
@@ -0,0 +1,562 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_big_kernel.h
13+ * \brief
14+ */
15+#ifndef MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_H_
16+#define MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_H_
17+ 
18+#include "max_pool_with_argmax_v3_base.h"
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+#include "../inc/kernel_utils.h"
23+ 
24+namespace MaxPoolWithArgMaxV3BigKernel {
25+using namespace AscendC;
26+ 
27+constexpr int32_t BUFFER_NUM = 2;
28+constexpr uint32_t FLOAT32_NEG_INF = 0xFF800000; // -inf 0xFF800000
29+constexpr int32_t OUT_BUFFER_LEN = 1024;
30+constexpr int32_t EIGHT = 8;
31+constexpr int32_t FOUR = 4;
32+constexpr int32_t TWO = 2;
33+ 
34+constexpr MicroAPI::CastTrait castTraitB322B16 = {
35+ MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
36+ 
37+constexpr MicroAPI::CastTrait castTraitB162B32 = {
38+ MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
39+ 
40+constexpr MicroAPI::CastTrait castTraitB322B64 = {
41+ MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
42+ 
43+template <typename T, typename U>
44+__aicore__ inline void StoreOneElement(
45+ const __local_mem__ void* output, MicroAPI::RegTensor<U>& src, MicroAPI::MaskReg& preg, uint32_t offset)
46+{
47+ if constexpr (IsSameType<T, half>::value) {
48+ MicroAPI::RegTensor<half> xFp16;
49+ MicroAPI::Cast<half, float, castTraitB322B16>(xFp16, src, preg);
50+ MicroAPI::DataCopy<half, MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B16>(
51+ (__local_mem__ half*)(output) + offset, xFp16, preg);
52+ } else if constexpr (IsSameType<T, bfloat16_t>::value) {
53+ MicroAPI::RegTensor<bfloat16_t> xBf16;
54+ MicroAPI::Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg);
55+ MicroAPI::DataCopy<bfloat16_t, MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B16>(
56+ (__local_mem__ bfloat16_t*)(output) + offset, xBf16, preg);
57+ } else if constexpr (sizeof(T) == FOUR) {
58+ MicroAPI::DataCopy<float, MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(
59+ ((__local_mem__ float*)output) + offset, (MicroAPI::RegTensor<float>&)src, preg);
60+ } else {
61+ MicroAPI::UnalignReg u0;
62+ auto dstAddr = (__local_mem__ T*)(output) + offset;
63+ MicroAPI::DataCopyUnAlign(dstAddr, src, u0, 1);
64+ MicroAPI::DataCopyUnAlignPost(dstAddr, u0, 0);
65+ }
66+}
67+ 
68+template <typename T, typename U>
69+__aicore__ inline void LoadOneElement(
70+ const __local_mem__ void* input, MicroAPI::RegTensor<U>& dst, MicroAPI::MaskReg& preg, uint32_t offset)
71+{
72+ if constexpr (IsSameType<T, half>::value) {
73+ MicroAPI::RegTensor<half> xFp16;
74+ MicroAPI::DataCopy<half, MicroAPI::LoadDist::DIST_BRC_B16>(xFp16, (__local_mem__ half*)(input) + offset);
75+ MicroAPI::Cast<float, half, castTraitB162B32>(dst, xFp16, preg);
76+ } else if constexpr (IsSameType<T, bfloat16_t>::value) {
77+ MicroAPI::RegTensor<bfloat16_t> xBf16;
78+ MicroAPI::DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_BRC_B16>(
79+ xBf16, (__local_mem__ bfloat16_t*)(input) + offset);
80+ MicroAPI::Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg);
81+ } else if constexpr (sizeof(T) == FOUR) {
82+ MicroAPI::DataCopy<T, MicroAPI::LoadDist::DIST_BRC_B32>(dst, ((__local_mem__ T*)(input)) + offset);
83+ } else {
84+ MicroAPI::UnalignReg u0;
85+ auto srcAddr = (__local_mem__ T*)(input) + offset;
86+ MicroAPI::DataCopyUnAlignPre(u0, srcAddr);
87+ MicroAPI::DataCopyUnAlign(dst, u0, srcAddr, 1);
88+ }
89+}
90+ 
91+template <typename T>
92+__aicore__ inline void LoadOneTensor(
93+ const __local_mem__ void* input, MicroAPI::RegTensor<float>& dst, MicroAPI::MaskReg& preg,
94+ MicroAPI::AddrReg& offset)
95+{
96+ if constexpr (IsSameType<T, half>::value) {
97+ MicroAPI::RegTensor<half> xFp16;
98+ DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, (__local_mem__ half*)(input), offset);
99+ Cast<float, half, castTraitB162B32>(dst, xFp16, preg);
100+ } else if constexpr (IsSameType<T, bfloat16_t>::value) {
101+ MicroAPI::RegTensor<bfloat16_t> xBf16;
102+ MicroAPI::DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(
103+ xBf16, (__local_mem__ bfloat16_t*)(input), offset);
104+ MicroAPI::Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg);
105+ } else {
106+ MicroAPI::DataCopy(dst, (__local_mem__ float*)(input), offset);
107+ }
108+}
109+ 
110+template <typename T, bool SPLITKW>
111+__aicore__ inline void CalcRealIndex(
112+ MicroAPI::RegTensor<T>& resIndex, MicroAPI::RegTensor<int32_t>& index, int64_t curKw, int64_t inputW,
113+ int64_t offset)
114+{
115+ MicroAPI::MaskReg pregOneIndex = MicroAPI::CreateMask<int32_t, MicroAPI::MaskPattern::VL1>();
116+ 
117+ MicroAPI::RegTensor<T> indexCast;
118+ if constexpr (IsSameType<T, int64_t>::value) {
119+ MicroAPI::Cast<int64_t, int32_t, castTraitB322B64>(indexCast, index, pregOneIndex);
120+ } else {
121+ MicroAPI::Copy(indexCast, index, pregOneIndex);
122+ }
123+ if constexpr (SPLITKW) {
124+ MicroAPI::Adds(resIndex, indexCast, (T)offset, pregOneIndex);
125+ } else {
126+ MicroAPI::RegTensor<T> wLen;
127+ MicroAPI::RegTensor<T> v0;
128+ MicroAPI::RegTensor<T> v1;
129+ MicroAPI::Duplicate(wLen, (T)curKw, pregOneIndex);
130+ MicroAPI::Div(v0, indexCast, wLen, pregOneIndex);
131+ MicroAPI::Muls(resIndex, v0, inputW, pregOneIndex);
132+ MicroAPI::Adds(resIndex, resIndex, (T)offset, pregOneIndex);
133+ MicroAPI::Mul(wLen, wLen, v0, pregOneIndex);
134+ MicroAPI::Sub(v0, indexCast, wLen, pregOneIndex);
135+ MicroAPI::Add(resIndex, resIndex, v0, pregOneIndex);
136+ }
137+}
138+ 
139+template <typename T>
140+__aicore__ inline void DuplicateNegInf(const __local_mem__ void* dstAddr, uint32_t calNum, uint32_t offset)
141+{
142+ MicroAPI::RegTensor<T> v0;
143+ MicroAPI::UnalignReg u0;
144+ DuplicateNegInfReg<T>(v0);
145+ __local_mem__ T* addr = (__local_mem__ T*)dstAddr + offset;
146+ MicroAPI::DataCopyUnAlign(addr, v0, u0, calNum);
147+ MicroAPI::DataCopyUnAlignPost(addr, u0, 0);
148+ MicroAPI::LocalMemBar<MicroAPI::MemType::VEC_STORE, MicroAPI::MemType::VEC_LOAD>();
149+}
150+ 
151+template <typename T>
152+__aicore__ inline void ReduceMaxWithIndex(
153+ MicroAPI::RegTensor<T>& dst, MicroAPI::RegTensor<int32_t>& dstIndex, MicroAPI::RegTensor<T>& src,
154+ MicroAPI::RegTensor<int32_t>& srcIndex, int32_t indexPadValue)
155+{
156+ // select first max value or last nan from one reg
157+ MicroAPI::MaskReg maskAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
158+ MicroAPI::MaskReg notNanMaskReg;
159+ MicroAPI::MaskReg nanMaskReg;
160+ MicroAPI::RegTensor<T> vd1;
161+ MicroAPI::RegTensor<T> vd2;
162+ MicroAPI::RegTensor<int32_t> nanIndex;
163+ MicroAPI::Duplicate(nanIndex, indexPadValue);
164+ MicroAPI::Compare<T, CMPMODE::NE>(nanMaskReg, src, src, maskAll); // nan mask
165+ MicroAPI::MaskNot(notNanMaskReg, nanMaskReg, maskAll); // not nan mask
166+ MicroAPI::Select(nanIndex, srcIndex, nanIndex, nanMaskReg); // nan index
167+ MicroAPI::ReduceMax(nanIndex, nanIndex, maskAll); // max nan index
168+ MicroAPI::ReduceMax(vd1, src, notNanMaskReg); // max value
169+ MicroAPI::Duplicate(vd2, vd1, maskAll); // max value
170+ MicroAPI::Compare<T, CMPMODE::EQ>(notNanMaskReg, src, vd2, maskAll); // nan mask
171+ MicroAPI::ReduceMin(dstIndex, srcIndex, notNanMaskReg); // not nan max index
172+ MicroAPI::CompareScalar<int32_t, CMPMODE::NE>(nanMaskReg, nanIndex, indexPadValue, maskAll); // nan
173+ MicroAPI::Select(dstIndex, nanIndex, dstIndex, nanMaskReg);
174+ MicroAPI::Duplicate(dstIndex, dstIndex, maskAll);
175+ MicroAPI::Compare<int32_t, CMPMODE::EQ>(notNanMaskReg, dstIndex, srcIndex, maskAll);
176+ MicroAPI::ReduceMax(dst, src, notNanMaskReg); // max value
177+ // all value in the kernel is -inf
178+ MicroAPI::CompareScalar<int32_t, CMPMODE::EQ>(notNanMaskReg, dstIndex, indexPadValue, maskAll);
179+ MicroAPI::Duplicate(nanIndex, static_cast<int32_t>(0));
180+ MicroAPI::Select(dstIndex, nanIndex, dstIndex, notNanMaskReg);
181+}
182+ 
183+template <typename T, typename U, typename TINDEX>
184+__aicore__ inline void MergeMaxRes(
185+ MicroAPI::RegTensor<U>& res, MicroAPI::RegTensor<TINDEX>& realResIndex, const __local_mem__ T* dstLocalAddr,
186+ const __local_mem__ TINDEX* indexLocalAddr, int32_t offset)
187+{
188+ // merge cur result with pre result
189+ MicroAPI::MaskReg maskAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
190+ MicroAPI::MaskReg notNanMaskReg;
191+ MicroAPI::MaskReg nanMaskReg;
192+ MicroAPI::MaskReg pregOne = MicroAPI::CreateMask<U, MicroAPI::MaskPattern::VL1>();
193+ MicroAPI::RegTensor<U> lastRes;
194+ MicroAPI::RegTensor<TINDEX> lastResIndex;
195+ LoadOneElement<T, U>(dstLocalAddr, lastRes, pregOne, offset);
196+ MicroAPI::Compare<U, CMPMODE::NE>(nanMaskReg, res, res, maskAll); // cur nan
197+ MicroAPI::Compare<U, CMPMODE::GT>(notNanMaskReg, res, lastRes, maskAll); // cur large > last
198+ MicroAPI::MaskXor(notNanMaskReg, notNanMaskReg, nanMaskReg, maskAll); // gt & nan
199+ MicroAPI::Select(res, res, lastRes, notNanMaskReg); // nan index
200+ LoadOneElement<TINDEX, TINDEX>(indexLocalAddr, lastResIndex, pregOne, offset);
201+ MicroAPI::Select(realResIndex, realResIndex, lastResIndex, notNanMaskReg); // nan index
202+ MicroAPI::LocalMemBar<MicroAPI::MemType::VEC_LOAD, MicroAPI::MemType::VEC_STORE>();
203+}
204+ 
205+template <typename T1, typename T2, typename TINDEX>
206+class MaxPoolWithArgmaxV3BigKernel {
207+public:
208+ __aicore__ inline MaxPoolWithArgmaxV3BigKernel(
209+ TPipe* pipe, const MaxPoolWithArgmaxV3BigKernelTilingData* __restrict tiling)
210+ : pipe_(pipe), tilingData_(tiling){};
211+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR indices);
212+ __aicore__ inline void Process();
213+ 
214+private:
215+ __aicore__ inline void CalcKernelSize(int64_t curIdx, int64_t& curkH, int64_t& curkW, int64_t& curInOffset);
216+ template <bool SPLIT_KERNEL>
217+ __aicore__ inline void BaseCompute(int64_t beginIdx, int64_t endIdx, int64_t maxCount);
218+ __aicore__ inline void CopyInSingleRow(int64_t offset, int64_t blockLen);
219+ __aicore__ inline void CopyInMultiRows(int64_t offset, int64_t blockLen, int64_t blockCount);
220+ __aicore__ inline void CopyMaxOut(int64_t curIdx);
221+ __aicore__ inline void NoSplitKernelProcess(
222+ int32_t localCurIdx, int64_t curkH, int64_t curkW, int64_t curInOffset, int64_t maxCount);
223+ __aicore__ inline void SplitKernelProcess(
224+ int32_t localCurIdx, int64_t curkH, int64_t curkW, int64_t curInOffset, int64_t maxCount);
225+ template <bool MERGE, bool SPLITKW>
226+ __aicore__ inline void ComputeSingle(int32_t localCurIdx, int64_t dataCount, int64_t offset, int64_t curKw);
227+ template <bool CLEAR>
228+ __aicore__ inline void InitOutLocal(int32_t localCurIdx);
229+ __aicore__ inline int64_t min(int64_t a, int64_t b)
230+ {
231+ return (a > b) ? b : a;
232+ }
233+ 
234+ TPipe* pipe_;
235+ // 输入队列
236+ TQue<QuePosition::VECIN, BUFFER_NUM> inputQue_;
237+ // 输出ub
238+ TBuf<> maxUBOutput_;
239+ TBuf<> indexUBOutput_;
240+ 
241+ GlobalTensor<T1> xGm_;
242+ GlobalTensor<T1> maxGm_;
243+ GlobalTensor<TINDEX> indicesGm_;
244+ 
245+ const MaxPoolWithArgmaxV3BigKernelTilingData* tilingData_;
246+ 
247+ int64_t inHW_ = 1;
248+ int64_t curOriginH_ = 0;
249+ int64_t curOriginW_ = 0;
250+ int64_t curOriginIndex_ = 0;
251+ int64_t beginIdx_ = 0;
252+ int64_t endIdx_ = 0;
253+};
254+ 
255+template <typename T1, typename T2, typename TINDEX>
256+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR indices)
257+{
258+ inHW_ = tilingData_->hInDim * tilingData_->wInDim;
259+ if (GetBlockIdx() < tilingData_->blockTail) {
260+ beginIdx_ = GetBlockIdx() * (tilingData_->blockFactor + 1);
261+ endIdx_ = beginIdx_ + tilingData_->blockFactor + 1;
262+ } else {
263+ beginIdx_ = GetBlockIdx() * tilingData_->blockFactor + tilingData_->blockTail;
264+ endIdx_ = beginIdx_ + tilingData_->blockFactor;
265+ }
266+ // GM
267+ xGm_.SetGlobalBuffer((__gm__ T1*)x);
268+ maxGm_.SetGlobalBuffer((__gm__ T1*)y);
269+ indicesGm_.SetGlobalBuffer((__gm__ TINDEX*)indices);
270+ 
271+ pipe_->InitBuffer(inputQue_, BUFFER_NUM, tilingData_->maxCount * sizeof(T1));
272+ pipe_->InitBuffer(maxUBOutput_, OUT_BUFFER_LEN * sizeof(T1));
273+ pipe_->InitBuffer(indexUBOutput_, OUT_BUFFER_LEN * sizeof(TINDEX));
274+}
275+ 
276+template <typename T1, typename T2, typename TINDEX>
277+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::Process()
278+{
279+ if (tilingData_->kH * tilingData_->kW <= tilingData_->maxCount) {
280+ BaseCompute<false>(beginIdx_, endIdx_, tilingData_->maxCount);
281+ } else {
282+ BaseCompute<true>(beginIdx_, endIdx_, tilingData_->maxCount);
283+ }
284+}
285+ 
286+template <typename T1, typename T2, typename TINDEX>
287+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::CalcKernelSize(
288+ int64_t curIdx, int64_t& curkH, int64_t& curkW, int64_t& curInOffset)
289+{
290+ if (tilingData_->isSigOut) {
291+ curInOffset = curIdx * inHW_;
292+ curOriginIndex_ = 0;
293+ curkH = min(tilingData_->kH - tilingData_->pH, tilingData_->hInDim);
294+ curkW = min(tilingData_->kW - tilingData_->pW, tilingData_->wInDim);
295+ return;
296+ }
297+ int64_t outHW = tilingData_->hOutDim * tilingData_->wOutDim;
298+ int64_t cur2D = curIdx % outHW;
299+ int64_t curNc = curIdx / outHW;
300+ int64_t curHo = cur2D / tilingData_->wOutDim;
301+ int64_t curWo = cur2D % tilingData_->wOutDim;
302+ 
303+ curOriginH_ = tilingData_->sH * curHo - tilingData_->pH;
304+ if (curOriginH_ < 0) {
305+ curkH = min(tilingData_->kH + curOriginH_, tilingData_->hInDim);
306+ curOriginH_ = 0;
307+ } else {
308+ curkH = min(tilingData_->hInDim - curOriginH_, tilingData_->kH);
309+ }
310+ 
311+ curOriginW_ = tilingData_->sW * curWo - tilingData_->pW;
312+ if (curOriginW_ < 0) {
313+ curkW = min(tilingData_->kW + curOriginW_, tilingData_->wInDim);
314+ curOriginW_ = 0;
315+ } else {
316+ curkW = min(tilingData_->wInDim - curOriginW_, tilingData_->kW);
317+ }
318+ 
319+ curOriginIndex_ = curOriginH_ * tilingData_->wInDim + curOriginW_;
320+ curInOffset = curNc * inHW_ + curOriginIndex_;
321+}
322+ 
323+template <typename T1, typename T2, typename TINDEX>
324+template <bool SPLIT_KERNEL>
325+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::BaseCompute(
326+ int64_t beginIdx, int64_t endIdx, int64_t maxCount)
327+{
328+ int64_t curkH = 1;
329+ int64_t curkW = 1;
330+ int64_t curInOffset = 0;
331+ // current blockdim range
332+ for (int64_t idx = beginIdx; idx < endIdx; idx++) {
333+ CalcKernelSize(idx, curkH, curkW, curInOffset);
334+ constexpr int32_t maxLocalLen = OUT_BUFFER_LEN;
335+ int32_t localCurIdx = (idx - beginIdx) % maxLocalLen;
336+ if constexpr (SPLIT_KERNEL) {
337+ InitOutLocal<true>(localCurIdx);
338+ SplitKernelProcess(localCurIdx, curkH, curkW, curInOffset, maxCount);
339+ } else {
340+ InitOutLocal<false>(localCurIdx);
341+ NoSplitKernelProcess(localCurIdx, curkH, curkW, curInOffset, maxCount);
342+ }
343+ CopyMaxOut(idx);
344+ }
345+}
346+ 
347+template <typename T1, typename T2, typename TINDEX>
348+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::CopyInSingleRow(int64_t offset, int64_t blockLen)
349+{
350+ LocalTensor<T1> xLocal = inputQue_.AllocTensor<T1>();
351+ 
352+ DataCopyPadExtParams<T1> padExtParams;
353+ padExtParams.isPad = false;
354+ padExtParams.leftPadding = 0;
355+ padExtParams.rightPadding = 0;
356+ padExtParams.paddingValue = 0;
357+ 
358+ DataCopyExtParams extParams;
359+ extParams.blockCount = 1;
360+ extParams.blockLen = blockLen * sizeof(T1);
361+ extParams.srcStride = 0;
362+ extParams.dstStride = 0;
363+ DataCopyPad(xLocal, xGm_[offset], extParams, padExtParams);
364+ inputQue_.EnQue(xLocal);
365+}
366+ 
367+template <typename T1, typename T2, typename TINDEX>
368+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::CopyInMultiRows(
369+ int64_t offset, int64_t blockLen, int64_t blockCount)
370+{
371+ LocalTensor<T1> xLocal = inputQue_.AllocTensor<T1>();
372+ 
373+ DataCopyPadExtParams<T1> padExtParams;
374+ padExtParams.isPad = false;
375+ padExtParams.leftPadding = 0;
376+ padExtParams.rightPadding = 0;
377+ padExtParams.paddingValue = 0;
378+ 
379+ DataCopyExtParams extParams;
380+ extParams.blockCount = blockCount;
381+ extParams.blockLen = blockLen * sizeof(T1);
382+ extParams.srcStride = (tilingData_->wInDim - blockLen) * sizeof(T1);
383+ extParams.dstStride = 0;
384+ DataCopyPad<T1, PaddingMode::Compact>(xLocal, xGm_[offset], extParams, padExtParams);
385+ inputQue_.EnQue(xLocal);
386+}
387+ 
388+template <typename T1, typename T2, typename TINDEX>
389+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::CopyMaxOut(int64_t curIdx)
390+{
391+ constexpr int32_t maxLocalLen = OUT_BUFFER_LEN;
392+ int32_t localCurIdx = (curIdx - beginIdx_) % maxLocalLen;
393+ 
394+ if (localCurIdx == maxLocalLen - 1 || curIdx == endIdx_ - 1) {
395+ LocalTensor<T1> maxOutLocal = maxUBOutput_.Get<T1>();
396+ LocalTensor<TINDEX> indexLocal = indexUBOutput_.Get<TINDEX>();
397+ 
398+ DataCopyExtParams extParams;
399+ extParams.blockCount = 1;
400+ extParams.blockLen = (localCurIdx + 1) * sizeof(T1);
401+ extParams.srcStride = 0;
402+ extParams.dstStride = 0;
403+ event_t eventIdVtoMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
404+ SetFlag<HardEvent::V_MTE3>(eventIdVtoMTE3);
405+ WaitFlag<HardEvent::V_MTE3>(eventIdVtoMTE3);
406+ DataCopyPad(maxGm_[curIdx - localCurIdx], maxOutLocal, extParams);
407+ extParams.blockLen = (localCurIdx + 1) * sizeof(TINDEX);
408+ DataCopyPad(indicesGm_[curIdx - localCurIdx], indexLocal, extParams);
409+ }
410+}
411+ 
412+template <typename T1, typename T2, typename TINDEX>
413+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::NoSplitKernelProcess(
414+ int32_t localCurIdx, int64_t curkH, int64_t curkW, int64_t curInOffset, int64_t maxCount)
415+{
416+ CopyInMultiRows(curInOffset, curkW, curkH);
417+ ComputeSingle<false, false>(localCurIdx, curkW * curkH, curOriginIndex_, curkW);
418+}
419+ 
420+template <typename T1, typename T2, typename TINDEX>
421+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::SplitKernelProcess(
422+ int32_t localCurIdx, int64_t curkH, int64_t curkW, int64_t curInOffset, int64_t maxCount)
423+{
424+ int64_t realIndex = 0;
425+ int64_t inputOffset = curInOffset;
426+ int64_t kernelOffset = curOriginIndex_;
427+ int64_t maxIndex = 0;
428+ 
429+ if (curkW <= maxCount) {
430+ // 整行搬入
431+ int64_t hFactor = maxCount / curkW;
432+ int64_t hLoops = (curkH + hFactor - 1) / hFactor;
433+ int64_t hTail = curkH - (hLoops - 1) * hFactor;
434+ 
435+ for (int64_t hLoop = 0; hLoop < hLoops; hLoop++) {
436+ int32_t curhFactor = hLoop == hLoops - 1 ? hTail : hFactor;
437+ CopyInMultiRows(inputOffset, curkW, curhFactor);
438+ ComputeSingle<true, false>(localCurIdx, curkW * curhFactor, kernelOffset, curkW);
439+ inputOffset += curhFactor * tilingData_->wInDim;
440+ kernelOffset += curhFactor * tilingData_->wInDim;
441+ }
442+ } else {
443+ // 单行很大，单行循环搬
444+ int64_t hLoops = curkH;
445+ int64_t wFactor = maxCount;
446+ int64_t wLoops = (curkW + wFactor - 1) / wFactor;
447+ int64_t wTail = curkW - (wLoops - 1) * wFactor;
448+ 
449+ for (int64_t hLoop = 0; hLoop < hLoops; hLoop++) {
450+ inputOffset = curInOffset + hLoop * tilingData_->wInDim;
451+ kernelOffset = curOriginIndex_ + hLoop * tilingData_->wInDim;
452+ for (int64_t wLoop = 0; wLoop < wLoops; wLoop++) {
453+ int32_t curFactor = wLoop == wLoops - 1 ? wTail : wFactor;
454+ CopyInSingleRow(inputOffset, curFactor);
455+ ComputeSingle<true, true>(localCurIdx, curFactor, kernelOffset, curkW);
456+ inputOffset += curFactor;
457+ kernelOffset += curFactor;
458+ }
459+ }
460+ }
461+}
462+ 
463+template <typename T1, typename T2, typename TINDEX>
464+template <bool CLEAR>
465+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::InitOutLocal(int32_t localCurIdx)
466+{
467+ if (localCurIdx != 0) {
468+ return;
469+ }
470+ 
471+ constexpr int32_t maxLocalLen = OUT_BUFFER_LEN;
472+ event_t eventIdMTE3toV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
473+ SetFlag<HardEvent::MTE3_V>(eventIdMTE3toV);
474+ WaitFlag<HardEvent::MTE3_V>(eventIdMTE3toV);
475+ 
476+ if constexpr (!CLEAR) { // kerel 全载场景无需merge，因此无需初始化output
477+ return;
478+ }
479+ LocalTensor<T1> maxOutLocal = maxUBOutput_.Get<T1>();
480+ __local_mem__ T1* dstAddr = (__local_mem__ T1*)maxOutLocal.GetPhyAddr();
481+ constexpr uint32_t repeatElm = platform::GetVRegSize() / sizeof(T1);
482+ uint16_t repeatTimes = CeilDivision(maxLocalLen, repeatElm);
483+ uint32_t num = maxLocalLen;
484+ __VEC_SCOPE__
485+ {
486+ MicroAPI::RegTensor<T1> v0;
487+ DuplicateNegInfReg<T1>(v0);
488+ for (uint16_t i = 0; i < repeatTimes; i++) {
489+ MicroAPI::MaskReg p0 = MicroAPI::UpdateMask<T1>(num);
490+ MicroAPI::AddrReg offsetReg = MicroAPI::CreateAddrReg<T1>(i, repeatElm);
491+ MicroAPI::DataCopy(dstAddr, v0, offsetReg, p0);
492+ }
493+ }
494+}
495+ 
496+template <typename T1, typename T2, typename TINDEX>
497+template <bool MERGE, bool SPLITKW>
498+__aicore__ inline void MaxPoolWithArgmaxV3BigKernel<T1, T2, TINDEX>::ComputeSingle(
499+ int32_t localCurIdx, int64_t dataCount, int64_t offset, int64_t curKw)
500+{
501+ LocalTensor<T1> maxOutLocal = maxUBOutput_.Get<T1>();
502+ LocalTensor<TINDEX> indexLocal = indexUBOutput_.Get<TINDEX>();
503+ LocalTensor<T1> xLocal = inputQue_.DeQue<T1>();
504+ union {
505+ T2 f;
506+ unsigned int i;
507+ } minValue;
508+ minValue.i = FLOAT32_NEG_INF;
509+ __local_mem__ T1* xLocalAddr = (__local_mem__ T1*)xLocal.GetPhyAddr();
510+ __local_mem__ T1* dstLocalAddr = (__local_mem__ T1*)maxOutLocal.GetPhyAddr();
511+ __local_mem__ TINDEX* indexLocalAddr = (__local_mem__ TINDEX*)indexLocal.GetPhyAddr();
512+ constexpr int32_t padIndex = -1;
513+ 
514+ constexpr uint32_t repeatElm = platform::GetVRegSize() / sizeof(T2);
515+ uint16_t repeatTimes = CeilDivision(dataCount, repeatElm);
516+ uint32_t num = repeatTimes * repeatElm; // 需要vreg_len对齐
517+ uint32_t padNum = num - dataCount;
518+ TINDEX inputW = tilingData_->wInDim;
519+ __VEC_SCOPE__
520+ {
521+ DuplicateNegInf<T1>(xLocalAddr, padNum, dataCount);
522+ MicroAPI::RegTensor<T2> vd0;
523+ MicroAPI::RegTensor<T2> vd1;
524+ MicroAPI::RegTensor<T2> vd2;
525+ MicroAPI::RegTensor<T2> vd3;
526+ MicroAPI::RegTensor<T2> res;
527+ MicroAPI::RegTensor<int32_t> resIndex;
528+ MicroAPI::RegTensor<int32_t> index;
529+ MicroAPI::MaskReg cmpMaskNanReg;
530+ MicroAPI::MaskReg cmpMaskReg;
531+ MicroAPI::MaskReg maskAll = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
532+ 
533+ MicroAPI::Duplicate(resIndex, padIndex);
534+ MicroAPI::Duplicate(res, minValue.f);
535+ MicroAPI::Arange(index, 0);
536+ for (uint16_t i = 0; i < repeatTimes; i++) {
537+ MicroAPI::MaskReg p0 = MicroAPI::UpdateMask<T2>(num);
538+ MicroAPI::AddrReg offset = MicroAPI::CreateAddrReg<T1>(i, repeatElm);
539+ LoadOneTensor<T1>(xLocalAddr, vd0, p0, offset);
540+ MicroAPI::Compare<T2, CMPMODE::NE>(cmpMaskNanReg, vd0, vd0, maskAll); // cur nan
541+ MicroAPI::Compare<T2, CMPMODE::GT>(cmpMaskReg, vd0, res, maskAll); // cur large > last
542+ MicroAPI::MaskXor(cmpMaskReg, cmpMaskReg, cmpMaskNanReg, maskAll); // gt & nan
543+ MicroAPI::Select(res, vd0, res, cmpMaskReg);
544+ MicroAPI::Select(resIndex, index, resIndex, cmpMaskReg);
545+ MicroAPI::Adds(index, index, repeatElm, maskAll);
546+ }
547+ ReduceMaxWithIndex<T2>(res, index, res, resIndex, padIndex);
548+ MicroAPI::MaskReg pregOne = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::VL1>();
549+ MicroAPI::RegTensor<TINDEX> realResIndex;
550+ CalcRealIndex<TINDEX, SPLITKW>(realResIndex, index, curKw, inputW, offset);
551+ if constexpr (MERGE) {
552+ // merge cur result with last result
553+ MergeMaxRes<T1, T2, TINDEX>(res, realResIndex, dstLocalAddr, indexLocalAddr, localCurIdx);
554+ }
555+ StoreOneElement<TINDEX, TINDEX>(indexLocalAddr, realResIndex, pregOne, localCurIdx);
556+ StoreOneElement<T1, T2>(dstLocalAddr, res, pregOne, localCurIdx);
557+ }
558+ inputQue_.FreeTensor<T1>(xLocal);
559+}
560+ 
561+} // namespace MaxPoolWithArgMaxV3BigKernel
562+#endif // MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_H_
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_big_kernel_mul_core.h+728-0
@@ -0,0 +1,728 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_big_kernel_mul_core.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_MUL_CORE_H_
17+#define MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_MUL_CORE_H_
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+ 
23+namespace MaxPoolWithArgmaxV3BigKernelMulCore {
24+using namespace AscendC;
25+constexpr int32_t BUFFER_NUM = 1;
26+constexpr int64_t REPEAT_DATA = 256;
27+constexpr uint16_t FLOAT16_NEG_INF = 64512; // -inf 0xFC00
28+constexpr uint16_t FLOAT16_INF = 31744; // inf 0x7C00
29+constexpr uint16_t FLOAT16_NAN_END = 32768; // 0x8000
30+constexpr int32_t FLOAT32_NEG_INF = -2139095040; // -inf 0xFF800000
31+constexpr int32_t FLOAT32_INF = 2139095040; // inf 0x7F800000
32+constexpr int32_t FLOAT32_NEG_ZERO = -2147483648; // -0
33+constexpr uint32_t VALUE_WORKSPACE_SIZE = 64 * 4;
34+constexpr uint32_t INDEX_WORKSPACE_SIZE = 64 * 8;
35+constexpr int64_t MASK_RATIO = 8;
36+ 
37+template <typename T>
38+class InnerComputer {
39+public:
40+ __aicore__ inline void Compute(
41+ const LocalTensor<T>& xLocal, LocalTensor<float>& castToFP32, TBuf<>& maxUB, TBuf<>& workLocalUB,
42+ uint32_t dataCount)
43+ {
44+ LocalTensor<T> maxOutLocal = maxUB.Get<T>();
45+ ReduceMax<T>(maxOutLocal, xLocal, xLocal, dataCount, true);
46+ // pipev
47+ }
48+ 
49+ __aicore__ inline void GetMask(
50+ const LocalTensor<T>& xLocal, LocalTensor<float>& castToFP32, LocalTensor<uint8_t>& mask, uint32_t dataCount)
51+ {
52+ uint32_t dataCountAlign = (dataCount + REPEAT_DATA - 1) / REPEAT_DATA * REPEAT_DATA;
53+ if (dataCountAlign > dataCount) {
54+ Duplicate(xLocal[dataCount], T(0), dataCountAlign - dataCount);
55+ // pipev
56+ }
57+ Compare(mask, xLocal, xLocal, CMPMODE::EQ, dataCountAlign);
58+ // pipev
59+ Not(mask, mask, dataCountAlign / MASK_RATIO);
60+ // pipev
61+ }
62+};
63+ 
64+template <>
65+class InnerComputer<bfloat16_t> {
66+public:
67+ __aicore__ inline void Compute(
68+ const LocalTensor<bfloat16_t>& xLocal, LocalTensor<float>& castToFP32, TBuf<>& maxUB, TBuf<>& workLocalUB,
69+ uint32_t dataCount)
70+ {
71+ LocalTensor<float> maxOutLocal = maxUB.Get<float>();
72+ Cast(castToFP32, xLocal, RoundMode::CAST_NONE, dataCount);
73+ // pipev
74+ ReduceMax<float>(maxOutLocal, castToFP32, castToFP32, dataCount, true);
75+ // pipev
76+ }
77+ 
78+ __aicore__ inline void GetMask(
79+ const LocalTensor<bfloat16_t>& xLocal, LocalTensor<float>& castToFP32, LocalTensor<uint8_t>& mask,
80+ uint32_t dataCount)
81+ {
82+ uint32_t dataCountAlign = (dataCount + REPEAT_DATA - 1) / REPEAT_DATA * REPEAT_DATA;
83+ if (dataCountAlign > dataCount) {
84+ Duplicate(castToFP32[dataCount], float(0), dataCountAlign - dataCount);
85+ // pipev
86+ }
87+ Compare(mask, castToFP32, castToFP32, CMPMODE::EQ, dataCountAlign);
88+ // pipev
89+ Not(mask, mask, dataCountAlign / MASK_RATIO);
90+ // pipev
91+ }
92+};
93+ 
94+template <typename T1, typename T2, typename TINDEX>
95+class MaxPoolWithArgmaxV3BigKernelMulCore {
96+public:
97+ __aicore__ inline MaxPoolWithArgmaxV3BigKernelMulCore(){};
98+ __aicore__ inline void Init(
99+ GM_ADDR x, GM_ADDR y, GM_ADDR indices, GM_ADDR workspace, TPipe* pipe_in,
100+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tiling);
101+ __aicore__ inline void Process();
102+ 
103+private:
104+ __aicore__ inline void Prepare(int64_t curIdx, int64_t innerBlockIdx);
105+ __aicore__ inline void BaseCompute(int64_t curIdx);
106+ __aicore__ inline int64_t HwCopyInput(
107+ int64_t offset, int64_t blockCount, int64_t blockLen, int64_t blockLenAlign, int64_t srcStride);
108+ __aicore__ inline int32_t Compute(int64_t dataCount);
109+ __aicore__ inline int64_t RestoreIndex(int32_t index, int64_t hLen, int64_t wLen);
110+ __aicore__ inline void CopyMaxOut(int64_t curIdx);
111+ __aicore__ inline void CopyIndicesOut(int64_t maxIndex, int64_t curIdx);
112+ __aicore__ inline void NaNIndicesInit(LocalTensor<float> indicesLocal);
113+ __aicore__ inline void GetIndexWithLastNan(
114+ LocalTensor<float> indicesMaxLocal, LocalTensor<uint8_t> maskNanLocal, int64_t dataCount, int32_t& index);
115+ __aicore__ inline int64_t AllWInKernelProcess();
116+ __aicore__ inline void UpdateMax(int64_t curMaxIndex, T2& maxValue, int64_t& maxIndice, bool first);
117+ __aicore__ inline int32_t KernelRealIndex(int32_t index, int64_t blockLen, int64_t blockLenAlign);
118+ __aicore__ inline void CopyOut(int64_t idx, int32_t index);
119+ __aicore__ inline void ComputeMulCore(int32_t& index);
120+ __aicore__ inline void CopyInMulCore(int64_t startIdx);
121+ __aicore__ inline void SplitW(
122+ int64_t blockLen, int64_t alignBlockLen, int64_t strStride, int64_t& maxValueIndex, T2& value);
123+ __aicore__ inline int64_t CeilValue(int64_t inputValue, int64_t upperValue)
124+ {
125+ if (upperValue == 0) {
126+ return inputValue;
127+ }
128+ return (inputValue + upperValue - 1) / upperValue * upperValue;
129+ }
130+ 
131+ __aicore__ inline int64_t Min(int64_t a, int64_t b)
132+ {
133+ return (a > b) ? b : a;
134+ }
135+ 
136+ __aicore__ inline bool IsNan(T2 value)
137+ {
138+ if (std::is_same<T2, half>::value) {
139+ uint16_t nan = *reinterpret_cast<uint16_t*>(&value);
140+ if ((nan > FLOAT16_INF && nan < FLOAT16_NAN_END) || nan > FLOAT16_NEG_INF) {
141+ return true;
142+ }
143+ } else {
144+ int32_t nan = *reinterpret_cast<int32_t*>(&value);
145+ if ((nan != FLOAT32_NEG_ZERO) && (nan > FLOAT32_INF || nan < FLOAT32_NEG_INF)) {
146+ return true;
147+ }
148+ }
149+ return false;
150+ }
151+ 
152+ TPipe* pipe;
153+ // 输入队列
154+ TQue<QuePosition::VECIN, BUFFER_NUM> inputQue;
155+ // 最大值ub
156+ TBuf<> maxUB;
157+ TBuf<> maxUBOutput;
158+ // indices初始下标
159+ TBuf<> indicesInitUB;
160+ // Compare结果mask
161+ TBuf<> maskNanUB;
162+ // nan场景最大值和下标
163+ TBuf<> nanMaxIndexUB;
164+ TBuf<> nanMaxIndexUBOutput;
165+ TBuf<> castBuff;
166+ 
167+ GlobalTensor<T1> xGm, maxGm;
168+ GlobalTensor<TINDEX> indicesGm;
169+ GlobalTensor<T2> maxValueWorkspaceGm;
170+ GlobalTensor<TINDEX> maxValueIndexWorkspaceGm;
171+ 
172+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* tilingData;
173+ 
174+ uint32_t cBlockIdx = 0;
175+ 
176+ int64_t inHW = 1;
177+ int64_t outHW = 1;
178+ int64_t curNc = 0;
179+ int64_t curOriginH = 0;
180+ int64_t curOriginW = 0;
181+ int64_t curOriginIndex = 0;
182+ int64_t curkH = 1;
183+ int64_t curkW = 1;
184+ int64_t curInOffset = 0;
185+ T1 minT1 = 0;
186+ T2 minT2 = 0;
187+ int32_t inputXQueOffset = 0;
188+ int64_t curKernelBlockFactorH = 0;
189+ int64_t curWSplitSize = 0;
190+ 
191+ constexpr static int64_t BYTE_T1 = sizeof(T1);
192+ constexpr static int64_t BLOCK_DATA = platform::GetUbBlockSize();
193+ constexpr static int64_t BLOCK_NUM_T1 = BLOCK_DATA / sizeof(T1);
194+ constexpr static int64_t REPEAT_NUM_T1 = REPEAT_DATA / sizeof(T1);
195+ constexpr static int64_t REPEAT_NUM_T2 = REPEAT_DATA / sizeof(T2);
196+ constexpr static int64_t BLOCK_NUM_T2 = BLOCK_DATA / sizeof(T2);
197+};
198+ 
199+template <typename T1, typename T2, typename TINDEX>
200+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::Init(
201+ GM_ADDR x, GM_ADDR y, GM_ADDR indices, GM_ADDR workspace, TPipe* pipe_in,
202+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tiling)
203+{
204+ pipe = pipe_in;
205+ tilingData = tiling;
206+ // base info
207+ cBlockIdx = GetBlockIdx();
208+ inHW = tilingData->hInDim * tilingData->wInDim;
209+ outHW = tilingData->hOutDim * tilingData->wOutDim;
210+ 
211+ uint32_t MIN_FLOAT32 = 0xFF800000; // -inf
212+ uint16_t MIN_FLOAT16 = 0xFC00; // -inf
213+ uint16_t MIN_BFLOAT16 = 0xFF80; // -inf
214+ if (std::is_same<T1, float>::value) {
215+ minT1 = *reinterpret_cast<T1*>(&MIN_FLOAT32);
216+ minT2 = *reinterpret_cast<T2*>(&MIN_FLOAT32);
217+ } else if (std::is_same<T1, half>::value) {
218+ minT1 = *reinterpret_cast<T1*>(&MIN_FLOAT16);
219+ minT2 = *reinterpret_cast<T2*>(&MIN_FLOAT16);
220+ } else if (std::is_same<T1, bfloat16_t>::value) {
221+ minT1 = *reinterpret_cast<T1*>(&MIN_BFLOAT16);
222+ minT2 = *reinterpret_cast<T2*>(&MIN_FLOAT32);
223+ }
224+ 
225+ // GM
226+ xGm.SetGlobalBuffer((__gm__ T1*)x);
227+ maxGm.SetGlobalBuffer((__gm__ T1*)y);
228+ indicesGm.SetGlobalBuffer((__gm__ TINDEX*)indices);
229+ maxValueWorkspaceGm.SetGlobalBuffer((__gm__ T2*)workspace);
230+ maxValueIndexWorkspaceGm.SetGlobalBuffer(reinterpret_cast<__gm__ TINDEX*>(workspace[VALUE_WORKSPACE_SIZE]));
231+ 
232+ pipe->InitBuffer(
233+ inputQue, BUFFER_NUM,
234+ tilingData->maxCountLength * sizeof(float)); // 原地cast 并复用为 nan index 的列表
235+ pipe->InitBuffer(maxUB, tilingData->valueBufferLength); // next do 256 参数化
236+ pipe->InitBuffer(maxUBOutput, tilingData->valueBufferLength); // next do 256 参数化
237+ pipe->InitBuffer(indicesInitUB, tilingData->maxCountLength * sizeof(float));
238+ pipe->InitBuffer(maskNanUB, tilingData->maxCountLength / MASK_RATIO); // 复用为reducemax时候的临时空间 worklocal
239+ pipe->InitBuffer(nanMaxIndexUB, tilingData->valueBufferLength);
240+ pipe->InitBuffer(nanMaxIndexUBOutput, tilingData->indexBufferLength);
241+ pipe->InitBuffer(castBuff, tilingData->indexBufferLength);
242+ 
243+ if (std::is_same<T1, bfloat16_t>::value) {
244+ inputXQueOffset = tilingData->maxCountLength; // inputQue的偏移 默认是0，bf16要做原地cast，输入放到后半部分
245+ }
246+}
247+ 
248+template <typename T1, typename T2, typename TINDEX>
249+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::Process()
250+{
251+ if (cBlockIdx >= tilingData->multiCoreNum * tilingData->coreNums) {
252+ return;
253+ }
254+ // init indices
255+ LocalTensor<float> indicesLocal = indicesInitUB.Get<float>();
256+ NaNIndicesInit(indicesLocal);
257+ 
258+ int64_t idx = cBlockIdx / tilingData->multiCoreNum;
259+ int64_t innerBlockIdx = cBlockIdx % tilingData->multiCoreNum;
260+ int64_t startIdx = idx * tilingData->multiCoreNum;
261+ 
262+ if (tilingData->splitW == 0) {
263+ if ((innerBlockIdx + 1) == tilingData->multiCoreNum) {
264+ curKernelBlockFactorH = tilingData->tailKernelBlockFactorH;
265+ } else {
266+ curKernelBlockFactorH = tilingData->kernelBlockFactorH;
267+ }
268+ } else {
269+ if ((innerBlockIdx + 1) % tilingData->splitSlice == 0) {
270+ curWSplitSize = tilingData->tailWSplitSize;
271+ } else {
272+ curWSplitSize = tilingData->wSplitSize;
273+ }
274+ }
275+ int32_t index = 0;
276+ Prepare(idx, innerBlockIdx);
277+ BaseCompute(idx);
278+ SyncAll();
279+ CopyInMulCore(startIdx);
280+ ComputeMulCore(index);
281+ CopyOut(idx, index);
282+}
283+template <typename T1, typename T2, typename TINDEX>
284+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::CopyOut(int64_t idx, int32_t index)
285+{
286+ if (cBlockIdx % tilingData->multiCoreNum != 0) {
287+ return;
288+ }
289+ LocalTensor<TINDEX> indicesResult = nanMaxIndexUBOutput.Get<TINDEX>();
290+ DataCopyExtParams extParams;
291+ extParams.blockCount = 1;
292+ extParams.blockLen = 1 * sizeof(T1);
293+ extParams.srcStride = 0;
294+ extParams.dstStride = 0;
295+ if (std::is_same<T1, bfloat16_t>::value) {
296+ LocalTensor<float> maxOutLocal = maxUBOutput.Get<float>();
297+ LocalTensor<T1> castBuffLocal = castBuff.Get<T1>();
298+ Cast(castBuffLocal, maxOutLocal, RoundMode::CAST_RINT, MASK_RATIO);
299+ event_t eventIdVtoMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
300+ SetFlag<HardEvent::V_MTE3>(eventIdVtoMTE3);
301+ WaitFlag<HardEvent::V_MTE3>(eventIdVtoMTE3);
302+ DataCopyPad(maxGm[idx], castBuffLocal[0], extParams);
303+ } else {
304+ LocalTensor<T1> maxOutLocal = maxUBOutput.Get<T1>();
305+ DataCopyPad(maxGm[idx], maxOutLocal[0], extParams);
306+ }
307+ extParams.blockLen = 1 * sizeof(TINDEX);
308+ indicesResult.SetValue(0, indicesResult.GetValue(index));
309+ 
310+ event_t eventIdStoMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
311+ SetFlag<HardEvent::S_MTE3>(eventIdStoMTE3);
312+ WaitFlag<HardEvent::S_MTE3>(eventIdStoMTE3);
313+ 
314+ DataCopyPad(indicesGm[idx], indicesResult[0], extParams);
315+ return;
316+}
317+template <typename T1, typename T2, typename TINDEX>
318+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::ComputeMulCore(int32_t& index)
319+{
320+ if (cBlockIdx % tilingData->multiCoreNum != 0) {
321+ return;
322+ }
323+ event_t eventIdMTE2toV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
324+ SetFlag<HardEvent::MTE2_V>(eventIdMTE2toV);
325+ WaitFlag<HardEvent::MTE2_V>(eventIdMTE2toV);
326+ LocalTensor<T2> maxUBLocal = maxUB.Get<T2>();
327+ LocalTensor<uint8_t> maskNanLocal = maskNanUB.Get<uint8_t>();
328+ LocalTensor<float> castBuffLocal = castBuff.Get<float>();
329+ LocalTensor<T2> maxOutLocal = maxUBOutput.Get<T2>();
330+ ReduceMax<T2>(maxOutLocal, maxUBLocal, maxUBLocal, tilingData->multiCoreNum, true);
331+ 
332+ event_t eventIdVtoS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
333+ SetFlag<HardEvent::V_S>(eventIdVtoS);
334+ WaitFlag<HardEvent::V_S>(eventIdVtoS);
335+ 
336+ T2 maxIndex = maxOutLocal.GetValue(1);
337+ T2 maxValue = maxOutLocal.GetValue(0);
338+ if (std::is_same<T2, half>::value) {
339+ int16_t indexInt16 = *reinterpret_cast<int16_t*>(&maxIndex);
340+ index = static_cast<int32_t>(indexInt16);
341+ } else {
342+ index = *reinterpret_cast<int32_t*>(&maxIndex);
343+ }
344+ 
345+ if (IsNan(maxValue)) {
346+ uint32_t dataCountAlign = CeilValue(tilingData->multiCoreNum, REPEAT_NUM_T2);
347+ uint32_t alignLen = CeilValue(tilingData->multiCoreNum, BLOCK_NUM_T2);
348+ if (dataCountAlign > alignLen) {
349+ Duplicate(maxUBLocal[alignLen], T2(0), dataCountAlign - alignLen);
350+ }
351+ Compare(maskNanLocal, maxUBLocal, maxUBLocal, CMPMODE::EQ, dataCountAlign);
352+ Not(maskNanLocal, maskNanLocal, dataCountAlign / MASK_RATIO);
353+ 
354+ LocalTensor<float> indicesLocal = indicesInitUB.Get<float>();
355+ LocalTensor<float> nanMaxIndex = nanMaxIndexUB.Get<float>();
356+ Select(
357+ castBuffLocal, maskNanLocal, indicesLocal, float(-1), SELMODE::VSEL_TENSOR_SCALAR_MODE,
358+ tilingData->multiCoreNum);
359+ ReduceMax<float>(nanMaxIndex, castBuffLocal, castBuffLocal, tilingData->multiCoreNum, false);
360+ 
361+ event_t eventIdVtoS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
362+ SetFlag<HardEvent::V_S>(eventIdVtoS);
363+ WaitFlag<HardEvent::V_S>(eventIdVtoS);
364+ index = ScalarCast<float, int32_t, RoundMode::CAST_ROUND>(nanMaxIndex.GetValue(0));
365+ }
366+ return;
367+}
368+template <typename T1, typename T2, typename TINDEX>
369+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::CopyInMulCore(int64_t startIdx)
370+{
371+ if (cBlockIdx % tilingData->multiCoreNum != 0) {
372+ return;
373+ }
374+ int32_t alignValue = CeilValue(tilingData->multiCoreNum, BLOCK_NUM_T2);
375+ LocalTensor<T2> maxValue = maxUB.Get<T2>();
376+ DataCopyPadExtParams<T2> padExtParams;
377+ padExtParams.isPad = true;
378+ padExtParams.leftPadding = 0;
379+ padExtParams.rightPadding = alignValue - tilingData->multiCoreNum;
380+ padExtParams.paddingValue = 0;
381+ 
382+ DataCopyExtParams extParams;
383+ extParams.blockCount = 1;
384+ extParams.blockLen = tilingData->multiCoreNum * sizeof(T2);
385+ extParams.srcStride = 0;
386+ extParams.dstStride = 0;
387+ extParams.rsv = 0;
388+ DataCopyPad(maxValue, maxValueWorkspaceGm[startIdx], extParams, padExtParams);
389+ 
390+ LocalTensor<TINDEX> maxValueIndices = nanMaxIndexUBOutput.Get<TINDEX>();
391+ DataCopyExtParams indicesExtParams;
392+ indicesExtParams.blockCount = 1;
393+ indicesExtParams.blockLen = tilingData->multiCoreNum * sizeof(TINDEX);
394+ indicesExtParams.srcStride = 0;
395+ indicesExtParams.dstStride = 0;
396+ indicesExtParams.rsv = 0;
397+ 
398+ DataCopyPadExtParams<TINDEX> indicesPadExtParams;
399+ indicesPadExtParams.isPad = false;
400+ indicesPadExtParams.leftPadding = 0;
401+ indicesPadExtParams.rightPadding = 0;
402+ indicesPadExtParams.paddingValue = 0;
403+ DataCopyPad(maxValueIndices, maxValueIndexWorkspaceGm[startIdx], indicesExtParams, indicesPadExtParams);
404+ return;
405+}
406+template <typename T1, typename T2, typename TINDEX>
407+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::Prepare(
408+ int64_t curIdx, int64_t innerBlockIdx)
409+{
410+ int64_t cur2D = curIdx % outHW;
411+ int64_t curNc = curIdx / outHW;
412+ int64_t curHo = cur2D / tilingData->wOutDim;
413+ int64_t curWo = cur2D % tilingData->wOutDim;
414+ 
415+ curOriginH = tilingData->sH * curHo - tilingData->pH;
416+ if (curOriginH < 0) {
417+ curkH = Min(tilingData->kH + curOriginH, tilingData->hInDim);
418+ curOriginH = 0;
419+ } else {
420+ curkH = Min(tilingData->hInDim - curOriginH, tilingData->kH);
421+ }
422+ 
423+ curOriginW = tilingData->sW * curWo - tilingData->pW;
424+ if (curOriginW < 0) {
425+ curkW = Min(tilingData->kW + curOriginW, tilingData->wInDim);
426+ curOriginW = 0;
427+ } else {
428+ curkW = Min(tilingData->wInDim - curOriginW, tilingData->kW);
429+ }
430+ 
431+ if (tilingData->splitW == 0) {
432+ curOriginIndex =
433+ (curOriginH + innerBlockIdx * tilingData->kernelBlockFactorH) * tilingData->wInDim + curOriginW;
434+ } else {
435+ curOriginIndex = (curOriginH + innerBlockIdx / tilingData->splitSlice) * tilingData->wInDim + curOriginW +
436+ innerBlockIdx % tilingData->splitSlice * tilingData->wSplitSize;
437+ }
438+ curInOffset = curNc * inHW + curOriginIndex;
439+ return;
440+}
441+ 
442+template <typename T1, typename T2, typename TINDEX>
443+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::BaseCompute(int64_t curIdx)
444+{
445+ int64_t realIndex = AllWInKernelProcess();
446+ CopyMaxOut(curIdx);
447+ CopyIndicesOut(realIndex, curIdx);
448+}
449+ 
450+template <typename T1, typename T2, typename TINDEX>
451+__aicore__ inline int64_t MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::HwCopyInput(
452+ int64_t offset, int64_t blockCount, int64_t blockLen, int64_t blockLenAlign, int64_t srcStride)
453+{
454+ LocalTensor<T1> xLocal = inputQue.AllocTensor<T1>();
455+ int64_t alignNum = blockLenAlign - blockLen;
456+ DataCopyPadExtParams<T1> padExtParams;
457+ padExtParams.isPad = alignNum != 0;
458+ padExtParams.leftPadding = 0;
459+ padExtParams.rightPadding = padExtParams.isPad ? alignNum : 0;
460+ padExtParams.paddingValue = minT1;
461+ 
462+ DataCopyExtParams extParams;
463+ extParams.blockCount = blockCount;
464+ extParams.blockLen = blockLen * sizeof(T1);
465+ extParams.srcStride = srcStride * sizeof(T1);
466+ extParams.dstStride = 0;
467+ 
468+ DataCopyPad(xLocal[inputXQueOffset], xGm[offset], extParams, padExtParams);
469+ inputQue.EnQue(xLocal);
470+ return blockCount * blockLenAlign;
471+}
472+ 
473+template <typename T1, typename T2, typename TINDEX>
474+__aicore__ inline int32_t MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::Compute(int64_t dataCount)
475+{
476+ LocalTensor<T1> xLocal = inputQue.DeQue<T1>();
477+ LocalTensor<float> castToFP32 = xLocal.template ReinterpretCast<float>();
478+ InnerComputer<T1> computer;
479+ 
480+ computer.Compute(xLocal[inputXQueOffset], castToFP32, maxUB, maskNanUB, dataCount);
481+ 
482+ event_t eventIdVtoS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
483+ SetFlag<HardEvent::V_S>(eventIdVtoS);
484+ WaitFlag<HardEvent::V_S>(eventIdVtoS);
485+ 
486+ int32_t index = 0;
487+ LocalTensor<uint8_t> maskNanLocal = maskNanUB.Get<uint8_t>();
488+ // 输入为fp16时
489+ if (std::is_same<T1, half>::value) {
490+ LocalTensor<half> maxOutLocal = maxUB.Get<half>();
491+ half maxIndex = maxOutLocal.GetValue(1);
492+ int16_t indexInt16 = *reinterpret_cast<int16_t*>(&maxIndex);
493+ index = static_cast<int32_t>(indexInt16);
494+ 
495+ half maxValue = maxOutLocal.GetValue(0);
496+ if (IsNan(maxValue)) {
497+ computer.GetMask(xLocal[inputXQueOffset], castToFP32, maskNanLocal, dataCount);
498+ GetIndexWithLastNan(castToFP32, maskNanLocal, dataCount, index);
499+ }
500+ } else {
501+ LocalTensor<float> maxOutLocal = maxUB.Get<float>();
502+ float maxIndex = maxOutLocal.GetValue(1);
503+ index = *reinterpret_cast<int32_t*>(&maxIndex);
504+ float maxValue = maxOutLocal.GetValue(0);
505+ if (IsNan(maxValue)) {
506+ computer.GetMask(xLocal[inputXQueOffset], castToFP32, maskNanLocal, dataCount);
507+ GetIndexWithLastNan(castToFP32, maskNanLocal, dataCount, index);
508+ }
509+ }
510+ inputQue.FreeTensor<T1>(xLocal);
511+ 
512+ return index;
513+}
514+ 
515+template <typename T1, typename T2, typename TINDEX>
516+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::GetIndexWithLastNan(
517+ LocalTensor<float> indicesMaxLocal, LocalTensor<uint8_t> maskNanLocal, int64_t dataCount, int32_t& index)
518+{
519+ LocalTensor<float> indicesLocal = indicesInitUB.Get<float>();
520+ Select(indicesMaxLocal, maskNanLocal, indicesLocal, float(-1), SELMODE::VSEL_TENSOR_SCALAR_MODE, dataCount);
521+ 
522+ LocalTensor<float> nanMaxIndex = nanMaxIndexUB.Get<float>();
523+ 
524+ ReduceMax<float>(nanMaxIndex, indicesMaxLocal, indicesMaxLocal, dataCount, false);
525+ 
526+ event_t eventIdVtoS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
527+ SetFlag<HardEvent::V_S>(eventIdVtoS);
528+ WaitFlag<HardEvent::V_S>(eventIdVtoS);
529+ 
530+ index = ScalarCast<float, int32_t, RoundMode::CAST_ROUND>(nanMaxIndex.GetValue(0));
531+}
532+ 
533+template <typename T1, typename T2, typename TINDEX>
534+__aicore__ inline int64_t MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::RestoreIndex(
535+ int32_t index, int64_t hLen, int64_t wLen)
536+{
537+ int64_t realIndex = 0;
538+ if (tilingData->splitW == 0) {
539+ int64_t alignBlockLen = CeilValue(curkW, BLOCK_NUM_T1);
540+ realIndex = curOriginIndex + index / alignBlockLen * tilingData->wInDim + index % alignBlockLen;
541+ } else {
542+ realIndex = curOriginIndex + index;
543+ }
544+ return realIndex;
545+}
546+ 
547+template <typename T1, typename T2, typename TINDEX>
548+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::CopyMaxOut(int64_t curIdx)
549+{
550+ DataCopyExtParams extParams;
551+ extParams.blockCount = 1;
552+ extParams.blockLen = 1 * sizeof(T2);
553+ extParams.srcStride = 0;
554+ extParams.dstStride = 0;
555+ 
556+ LocalTensor<T2> maxValueResult = maxUB.Get<T2>();
557+ 
558+ event_t eventIdVtoMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
559+ SetFlag<HardEvent::V_MTE3>(eventIdVtoMTE3);
560+ WaitFlag<HardEvent::V_MTE3>(eventIdVtoMTE3);
561+ 
562+ DataCopyPad(maxValueWorkspaceGm[cBlockIdx], maxValueResult[0], extParams);
563+}
564+ 
565+template <typename T1, typename T2, typename TINDEX>
566+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::CopyIndicesOut(
567+ int64_t maxIndex, int64_t curIdx)
568+{
569+ DataCopyExtParams extParams;
570+ extParams.blockCount = 1;
571+ extParams.blockLen = sizeof(int32_t);
572+ extParams.srcStride = 0;
573+ extParams.dstStride = 0;
574+ 
575+ LocalTensor<TINDEX> indexTensor = nanMaxIndexUBOutput.Get<TINDEX>();
576+ 
577+ if (std::is_same<TINDEX, int32_t>::value) {
578+ indexTensor.SetValue(0, maxIndex);
579+ } else {
580+ extParams.blockLen = sizeof(int64_t);
581+ indexTensor.SetValue(0, maxIndex);
582+ indexTensor.SetValue(1, maxIndex >> 32);
583+ }
584+ 
585+ event_t eventIdStoMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
586+ SetFlag<HardEvent::S_MTE3>(eventIdStoMTE3);
587+ WaitFlag<HardEvent::S_MTE3>(eventIdStoMTE3);
588+ DataCopyPad(maxValueIndexWorkspaceGm[cBlockIdx], indexTensor, extParams);
589+}
590+ 
591+template <typename T1, typename T2, typename TINDEX>
592+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::NaNIndicesInit(
593+ LocalTensor<float> indicesLocal)
594+{
595+ int64_t hAlignkW = CeilValue(tilingData->kW, BLOCK_NUM_T1) * tilingData->kH;
596+ int32_t InitIndicesNum = (hAlignkW > tilingData->maxCountLength) ? tilingData->maxCountLength : hAlignkW;
597+ CreateVecIndex(indicesLocal, 0.0f, InitIndicesNum);
598+}
599+template <typename T1, typename T2, typename TINDEX>
600+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::SplitW(
601+ int64_t blockLen, int64_t alignBlockLen, int64_t strStride, int64_t& maxValueIndex, T2& value)
602+{
603+ int64_t kernelOffset = 0;
604+ int64_t inputOffset = curInOffset;
605+ int64_t maxIndex = 0;
606+ T2 maxValue = 0;
607+ if (alignBlockLen <= tilingData->maxCountLength) {
608+ int64_t eachLoopLine = tilingData->maxCountLength / alignBlockLen;
609+ int64_t loop = (curKernelBlockFactorH + eachLoopLine - 1) / eachLoopLine;
610+ int64_t tailLoopLine = curKernelBlockFactorH - (loop - 1) * eachLoopLine;
611+ for (int64_t hwLoop = 0; hwLoop < loop; hwLoop++) {
612+ int64_t blockCount = (hwLoop == loop - 1 ? tailLoopLine : eachLoopLine);
613+ int32_t dataCount = HwCopyInput(inputOffset, blockCount, blockLen, alignBlockLen, strStride);
614+ int32_t index = Compute(dataCount);
615+ index = KernelRealIndex(index, blockLen, alignBlockLen);
616+ bool first = (hwLoop == 0);
617+ UpdateMax(kernelOffset + index, maxValue, maxIndex, first);
618+ inputOffset += blockCount * tilingData->wInDim;
619+ kernelOffset += blockCount * blockLen;
620+ }
621+ } else {
622+ int64_t loopWCount = (blockLen + tilingData->maxCountLength - 1) / tilingData->maxCountLength;
623+ int64_t tailLoopWSize = blockLen - (loopWCount - 1) * tilingData->maxCountLength;
624+ for (int64_t w = 0; w < curKernelBlockFactorH; w++) {
625+ for (int64_t eachLoopW = 0; eachLoopW < loopWCount; eachLoopW++) {
626+ int64_t wBlockLen = (eachLoopW == (loopWCount - 1) ? tailLoopWSize : tilingData->maxCountLength);
627+ int64_t alignWBlockLen =
628+ (eachLoopW == (loopWCount - 1) ? CeilValue(tailLoopWSize, BLOCK_NUM_T1) :
629+ tilingData->maxCountLength);
630+ int32_t dataCount = HwCopyInput(inputOffset, 1, wBlockLen, alignWBlockLen, 0);
631+ int32_t index = Compute(dataCount);
632+ bool first = (eachLoopW == 0 && w == 0);
633+ UpdateMax(kernelOffset + index, maxValue, maxIndex, first);
634+ inputOffset += wBlockLen;
635+ kernelOffset += wBlockLen;
636+ }
637+ inputOffset += tilingData->wInDim - blockLen;
638+ }
639+ }
640+ maxValueIndex = maxIndex;
641+ value = maxValue;
642+}
643+template <typename T1, typename T2, typename TINDEX>
644+__aicore__ inline int64_t MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::AllWInKernelProcess()
645+{
646+ int64_t realIndex = 0;
647+ int64_t blockCount = 0;
648+ int64_t alignBlockLen = 0;
649+ int64_t blockLen = 0;
650+ int64_t strStride = 0;
651+ if (tilingData->splitW == 0) {
652+ blockCount = curKernelBlockFactorH;
653+ blockLen = curkW;
654+ alignBlockLen = CeilValue(curkW, BLOCK_NUM_T1);
655+ strStride = (tilingData->wInDim - blockLen);
656+ } else {
657+ blockCount = 1;
658+ blockLen = curWSplitSize;
659+ alignBlockLen = CeilValue(curWSplitSize, BLOCK_NUM_T1);
660+ strStride = 0;
661+ }
662+ int64_t inputOffset = curInOffset;
663+ int64_t kernelOffset = 0;
664+ T2 maxValue = 0;
665+ int64_t maxIndex = 0;
666+ if (blockCount * alignBlockLen <= tilingData->maxCountLength) {
667+ int32_t dataCount = HwCopyInput(curInOffset, blockCount, blockLen, alignBlockLen, strStride);
668+ int32_t index = Compute(dataCount);
669+ index = KernelRealIndex(index, blockLen, alignBlockLen);
670+ realIndex = RestoreIndex(index, curkH, curkW);
671+ } else {
672+ if (tilingData->splitW == 0) {
673+ SplitW(blockLen, alignBlockLen, strStride, maxIndex, maxValue);
674+ } else {
675+ int64_t loop = (curWSplitSize + tilingData->maxCountLength - 1) / tilingData->maxCountLength;
676+ int64_t eachBlockLen = tilingData->maxCountLength;
677+ int64_t tailBlockLen = curWSplitSize - (loop - 1) * eachBlockLen;
678+ for (int64_t hwLoop = 0; hwLoop < loop; hwLoop++) {
679+ blockLen = (hwLoop == loop - 1 ? tailBlockLen : eachBlockLen);
680+ alignBlockLen = (hwLoop == loop - 1 ? CeilValue(tailBlockLen, BLOCK_NUM_T1) : eachBlockLen);
681+ int32_t dataCount = HwCopyInput(inputOffset, blockCount, blockLen, alignBlockLen, strStride);
682+ int32_t index = Compute(dataCount);
683+ bool first = (hwLoop == 0);
684+ UpdateMax(kernelOffset + index, maxValue, maxIndex, first);
685+ inputOffset += blockLen;
686+ kernelOffset += blockLen;
687+ }
688+ }
689+ realIndex = RestoreIndex(maxIndex, curkH, curkW);
690+ LocalTensor<T2> maxOutLocal = maxUB.Get<T2>();
691+ maxOutLocal.SetValue(0, maxValue);
692+ }
693+ return realIndex;
694+}
695+template <typename T1, typename T2, typename TINDEX>
696+__aicore__ inline void MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::UpdateMax(
697+ int64_t curMaxIndex, T2& maxValue, int64_t& maxIndex, bool first)
698+{
699+ LocalTensor<T2> maxOutLocal = maxUB.Get<T2>();
700+ T2 curMaxValue = maxOutLocal.GetValue(0);
701+ if (first) {
702+ maxIndex = curMaxIndex;
703+ maxValue = curMaxValue;
704+ return;
705+ }
706+ if (IsNan(curMaxValue)) {
707+ maxIndex = curMaxIndex;
708+ maxValue = curMaxValue;
709+ } else if (curMaxValue > maxValue) {
710+ maxIndex = curMaxIndex;
711+ maxValue = curMaxValue;
712+ }
713+}
714+ 
715+template <typename T1, typename T2, typename TINDEX>
716+__aicore__ inline int32_t MaxPoolWithArgmaxV3BigKernelMulCore<T1, T2, TINDEX>::KernelRealIndex(
717+ int32_t index, int64_t blockLen, int64_t blockLenAlign)
718+{
719+ int64_t alignNum = blockLenAlign - blockLen;
720+ if (alignNum != 0) {
721+ return index - alignNum * (index / blockLenAlign);
722+ } else {
723+ return index;
724+ }
725+}
726+ 
727+} // namespace MaxPoolWithArgmaxV3BigKernelMulCore
728+#endif // MAX_POOL_WITH_ARGMAX_V3_BIG_KERNEL_MUL_CORE_H_
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_gather_kernel.h+659-0
@@ -0,0 +1,659 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_gather_kernel.h
13+ * \brief
14+ */
15+#ifndef MAX_POOL_WITH_ARGMAX_V3_GATHER_KERNEL_H_
16+#define MAX_POOL_WITH_ARGMAX_V3_GATHER_KERNEL_H_
17+ 
18+#include "max_pool_with_argmax_v3_base.h"
19+ 
20+namespace MaxPoolWithArgmaxV3GatherNameSpace {
21+using namespace AscendC;
22+constexpr uint32_t BUFFER_NUM = 2;
23+constexpr int64_t HELPER_BUFFER_SIZE = 1024;
24+constexpr int64_t THREE_DIM = 3;
25+constexpr int64_t RATIO = 2;
26+constexpr uint16_t B32 = 4;
27+ 
28+template <typename T1, typename T2, const uint32_t IS_PAD = 0>
29+class MaxPoolWithArgmaxV3GatherKernel {
30+public:
31+ __aicore__ inline MaxPoolWithArgmaxV3GatherKernel(
32+ TPipe& pipeIn, const MaxPoolWithArgmaxV3GatherTilingData& tilingData)
33+ : pipe_(pipeIn), tilingData_(tilingData){};
34+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR argmax);
35+ __aicore__ inline void Process();
36+ 
37+private:
38+ __aicore__ inline void ScalarCompute(int64_t loopNum);
39+ __aicore__ inline void ProcessPerLoop();
40+ __aicore__ inline void CopyIn();
41+ __aicore__ inline void Compute();
42+ __aicore__ inline void CopyOut();
43+ __aicore__ inline void DupBufferNegInf(__local_mem__ T1* dstAddr, uint32_t repeatElm, uint16_t loop, uint32_t tail);
44+ __aicore__ inline void CopyToCalcBuffer(
45+ __local_mem__ T1* dstAddr, __local_mem__ T1* srcAddr, uint16_t batch, uint16_t rows, uint16_t loopCols,
46+ uint16_t tailCols, uint32_t repeatElm, uint32_t srcBatchStride, uint32_t srcRowStride, uint32_t dstBatchStride,
47+ uint32_t dstRowStride, uint32_t dstRowOffset, uint32_t dstColOffset);
48+ __aicore__ inline void DupAndCopyToCalcBuffer(__local_mem__ T1* dstAddr, __local_mem__ T1* srcAddr);
49+ __aicore__ inline void ConvertIndexWithoutPadAlign(
50+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, T2 left, T2 wInputActualNoPad, T2 hIndexBase,
51+ MicroAPI::RegTensor<T2>& dstReg, int32_t ncInputOffset);
52+ __aicore__ inline void ConvertIndexWithoutPadAlignNc(
53+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, T2 left, T2 wInputActualNoPad, T2 hIndexBase,
54+ MicroAPI::RegTensor<T2>& dstReg, int32_t ncInputOffset, int32_t ncOutputCount, int32_t inputNcSize);
55+ __aicore__ inline void ProcessW(
56+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, int32_t hOffset, uint16_t wStrideOffset,
57+ MicroAPI::RegTensor<int32_t>& indexReg, uint16_t hKernel, uint16_t wKernel, uint16_t repeatElem,
58+ int32_t outputOffset, MicroAPI::RegTensor<int32_t>& maxIndexReg, uint32_t hDilation, uint32_t wDilation);
59+ __aicore__ inline void SingleRowGather(
60+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr);
61+ __aicore__ inline void MultiRowGather(
62+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr);
63+ __aicore__ inline void MultiNcGather(
64+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr);
65+ 
66+private:
67+ TPipe& pipe_;
68+ const MaxPoolWithArgmaxV3GatherTilingData& tilingData_;
69+ TQue<QuePosition::VECIN, BUFFER_NUM> inputQue_;
70+ TQue<QuePosition::VECOUT, BUFFER_NUM> maxValueQue_;
71+ TQue<QuePosition::VECOUT, BUFFER_NUM> argmaxQue_;
72+ TBuf<TPosition::VECCALC> inputCalcBuff_;
73+ 
74+ GlobalTensor<T1> xGm_;
75+ GlobalTensor<T1> yGm_;
76+ GlobalTensor<T2> argmaxGm_;
77+ 
78+ uint32_t blockIdx_ = 0;
79+ int64_t highAxisActual_ = 0;
80+ int64_t hOutputActual_ = 0;
81+ int64_t wOutputActual_ = 0;
82+ int64_t curCoreProcessNum_ = 0;
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+ 
91+ int64_t highAxisIndex_ = 0;
92+ int64_t hAxisIndex_ = 0;
93+ int64_t wAxisIndex_ = 0;
94+ 
95+ int64_t highInputAxisOffset_ = 0;
96+ int64_t hInputAxisOffset_ = 0;
97+ int64_t wInputAxisOffset_ = 0;
98+ 
99+ int64_t hInputActualNoPad_ = 0;
100+ int64_t wInputActualNoPad_ = 0;
101+ int64_t wOutputActualAligned_ = 0;
102+ 
103+ constexpr static int32_t blockSize_ = platform::GetUbBlockSize();
104+ constexpr static int64_t maxDataNumOneBlock_ =
105+ blockSize_ / sizeof(T1) >= blockSize_ / sizeof(T2) ? blockSize_ / sizeof(T1) : blockSize_ / sizeof(T2);
106+ constexpr static uint16_t vlT2_ = platform::GetVRegSize() / sizeof(T2);
107+ constexpr static uint16_t vlT1_ = platform::GetVRegSize() / sizeof(T1);
108+};
109+ 
110+template <typename T1, typename T2, const uint32_t IS_PAD>
111+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR argmax)
112+{
113+ blockIdx_ = GetBlockIdx();
114+ if (blockIdx_ >= tilingData_.usedCoreNum) {
115+ return;
116+ }
117+ curCoreProcessNum_ =
118+ (blockIdx_ + 1 == tilingData_.usedCoreNum) ? tilingData_.tailCoreProcessNum : tilingData_.normalCoreProcessNum;
119+ xGm_.SetGlobalBuffer((__gm__ T1*)x);
120+ yGm_.SetGlobalBuffer((__gm__ T1*)y);
121+ argmaxGm_.SetGlobalBuffer((__gm__ T2*)argmax);
122+ 
123+ pipe_.InitBuffer(inputQue_, BUFFER_NUM, tilingData_.inputBufferSize);
124+ if constexpr (IS_PAD == 1) {
125+ pipe_.InitBuffer(inputCalcBuff_, tilingData_.inputBufferSize);
126+ }
127+ pipe_.InitBuffer(maxValueQue_, BUFFER_NUM, tilingData_.maxValueBufferSize);
128+ pipe_.InitBuffer(argmaxQue_, BUFFER_NUM, tilingData_.argmaxBufferSize);
129+ return;
130+}
131+ 
132+template <typename T1, typename T2, const uint32_t IS_PAD>
133+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::Process()
134+{
135+ if (blockIdx_ >= tilingData_.usedCoreNum) {
136+ return;
137+ }
138+ for (int64_t loopNum = 0; loopNum < curCoreProcessNum_; loopNum++) {
139+ ScalarCompute(loopNum);
140+ ProcessPerLoop();
141+ }
142+}
143+template <typename T1, typename T2, const uint32_t IS_PAD>
144+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::ScalarCompute(int64_t loopNum)
145+{
146+ int64_t baseBlockIdx = blockIdx_ * tilingData_.normalCoreProcessNum + loopNum;
147+ highAxisIndex_ = baseBlockIdx / (tilingData_.hOutputOuter * tilingData_.wOutputOuter);
148+ highAxisActual_ =
149+ highAxisIndex_ == (tilingData_.highAxisOuter - 1) ? tilingData_.highAxisTail : tilingData_.highAxisInner;
150+ int64_t tempTail = baseBlockIdx % (tilingData_.hOutputOuter * tilingData_.wOutputOuter);
151+ 
152+ hAxisIndex_ = tempTail / tilingData_.wOutputOuter;
153+ hOutputActual_ = hAxisIndex_ == (tilingData_.hOutputOuter - 1) ? tilingData_.hOutputTail : tilingData_.hOutputInner;
154+ 
155+ wAxisIndex_ = tempTail % tilingData_.wOutputOuter;
156+ wOutputActual_ = wAxisIndex_ == (tilingData_.wOutputOuter - 1) ? tilingData_.wOutputTail : tilingData_.wOutputInner;
157+ wOutputActualAligned_ = CeilDivision(wOutputActual_, maxDataNumOneBlock_) * maxDataNumOneBlock_;
158+ 
159+ hInputActualPad_ =
160+ (hOutputActual_ - 1) * tilingData_.hStride + (tilingData_.hKernel - 1) * tilingData_.hDilation + 1;
161+ wInputActualPad_ =
162+ (wOutputActual_ - 1) * tilingData_.wStride + (tilingData_.wKernel - 1) * tilingData_.wDilation + 1;
163+ 
164+ wInputActualAlignedPad_ = CeilDivision(wInputActualPad_, blockSize_ / sizeof(T1)) * (blockSize_ / sizeof(T1));
165+ 
166+ int64_t inputPlaneSize = tilingData_.hInput * tilingData_.wInput;
167+ highInputAxisOffset_ = highAxisIndex_ * tilingData_.highAxisInner * inputPlaneSize;
168+ hInputAxisOffset_ = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.hStride * tilingData_.wInput;
169+ wInputAxisOffset_ = wAxisIndex_ * tilingData_.wOutputInner * tilingData_.wStride;
170+ if constexpr (IS_PAD == 1) {
171+ int64_t tRelBoundaryDistance =
172+ hAxisIndex_ * tilingData_.hOutputInner * tilingData_.hStride - tilingData_.padTop;
173+ 
174+ int64_t dRelBoundaryDistance = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.hStride +
175+ (hOutputActual_ - 1) * tilingData_.hStride + tilingData_.hKernel -
176+ tilingData_.hInput - tilingData_.padTop;
177+ 
178+ int64_t lRelBoundaryDistance =
179+ wAxisIndex_ * tilingData_.wOutputInner * tilingData_.wStride - tilingData_.padLeft;
180+ 
181+ int64_t rRelBoundaryDistance = wAxisIndex_ * tilingData_.wOutputInner * tilingData_.wStride +
182+ (wOutputActual_ - 1) * tilingData_.wStride + tilingData_.wKernel -
183+ tilingData_.wInput - tilingData_.padLeft;
184+ leftOffsetToInputLeft_ = lRelBoundaryDistance >= 0 ? 0 : -lRelBoundaryDistance;
185+ rightOffsetToInputRight_ = rRelBoundaryDistance >= 0 ? rRelBoundaryDistance : 0;
186+ topOffsetToInputTop_ = tRelBoundaryDistance >= 0 ? 0 : -tRelBoundaryDistance;
187+ downOffsetToInputDown_ = dRelBoundaryDistance >= 0 ? dRelBoundaryDistance : 0;
188+ 
189+ hInputActualNoPad_ = hInputActualPad_ - topOffsetToInputTop_ - downOffsetToInputDown_;
190+ wInputActualNoPad_ = wInputActualPad_ - leftOffsetToInputLeft_ - rightOffsetToInputRight_;
191+ hInputAxisOffset_ = topOffsetToInputTop_ == 0 ? hInputAxisOffset_ - tilingData_.padTop * tilingData_.wInput : 0;
192+ 
193+ wInputAxisOffset_ = leftOffsetToInputLeft_ == 0 ? wInputAxisOffset_ - tilingData_.padLeft : 0;
194+ }
195+}
196+template <typename T1, typename T2, const uint32_t IS_PAD>
197+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::ProcessPerLoop()
198+{
199+ CopyIn();
200+ Compute();
201+ CopyOut();
202+}
203+template <typename T1, typename T2, const uint32_t IS_PAD>
204+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::CopyIn()
205+{
206+ LocalTensor<T1> xLocal = inputQue_.AllocTensor<T1>();
207+ int64_t xGmOffset = highInputAxisOffset_ + hInputAxisOffset_ + wInputAxisOffset_;
208+ 
209+ LoopModeParams loopModeParamsT1;
210+ loopModeParamsT1.loop1Size = highAxisActual_;
211+ loopModeParamsT1.loop2Size = 1;
212+ loopModeParamsT1.loop1SrcStride = tilingData_.hInput * tilingData_.wInput * sizeof(T1);
213+ loopModeParamsT1.loop2SrcStride = 0;
214+ loopModeParamsT1.loop1DstStride = hInputActualPad_ * wInputActualAlignedPad_ * sizeof(T1);
215+ loopModeParamsT1.loop2DstStride = 0;
216+ 
217+ SetLoopModePara(loopModeParamsT1, DataCopyMVType::OUT_TO_UB);
218+ DataCopyPadExtParams<T1> paramsT1 = {false, 0, 0, 0};
219+ DataCopyExtParams copyOutParamT1;
220+ if constexpr (IS_PAD == 1) {
221+ copyOutParamT1.blockCount = static_cast<uint16_t>(hInputActualNoPad_);
222+ copyOutParamT1.blockLen = static_cast<uint32_t>(wInputActualNoPad_ * sizeof(T1));
223+ copyOutParamT1.srcStride = static_cast<uint32_t>((tilingData_.wInput - wInputActualNoPad_) * sizeof(T1));
224+ copyOutParamT1.dstStride = 0;
225+ copyOutParamT1.rsv = 0;
226+ } else {
227+ copyOutParamT1.blockCount = static_cast<uint16_t>(hInputActualPad_);
228+ copyOutParamT1.blockLen = static_cast<uint32_t>(wInputActualPad_ * sizeof(T1));
229+ copyOutParamT1.srcStride = static_cast<uint32_t>((tilingData_.wInput - wInputActualPad_) * sizeof(T1));
230+ copyOutParamT1.dstStride = 0;
231+ copyOutParamT1.rsv = 0;
232+ }
233+ DataCopyPad(xLocal, xGm_[xGmOffset], copyOutParamT1, paramsT1);
234+ inputQue_.EnQue(xLocal);
235+ ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
236+}
237+template <typename T1, typename T2, const uint32_t IS_PAD>
238+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::DupBufferNegInf(
239+ __local_mem__ T1* dstAddr, uint32_t repeatElm, uint16_t loop, uint32_t tail)
240+{
241+ MicroAPI::RegTensor<T1> v0;
242+ DuplicateNegInfReg<T1>(v0);
243+ MicroAPI::MaskReg preg = MicroAPI::CreateMask<T1, MicroAPI::MaskPattern::ALL>();
244+ uint32_t maskCount = tail;
245+ for (uint16_t i = 0; i < loop; i++) {
246+ MicroAPI::DataCopy<T1, MicroAPI::PostLiteral::POST_MODE_UPDATE>(dstAddr, v0, repeatElm, preg);
247+ }
248+ preg = MicroAPI::UpdateMask<T1>(maskCount);
249+ MicroAPI::DataCopy<T1, MicroAPI::PostLiteral::POST_MODE_UPDATE>(dstAddr, v0, repeatElm, preg);
250+}
251+template <typename T1, typename T2, const uint32_t IS_PAD>
252+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::CopyToCalcBuffer(
253+ __local_mem__ T1* dstAddr, __local_mem__ T1* srcAddr, uint16_t batch, uint16_t rows, uint16_t loopCols,
254+ uint16_t tailCols, uint32_t repeatElm, uint32_t srcBatchStride, uint32_t srcRowStride, uint32_t dstBatchStride,
255+ uint32_t dstRowStride, uint32_t dstRowOffset, uint32_t dstColOffset)
256+{
257+ MicroAPI::RegTensor<T1> v0;
258+ MicroAPI::UnalignReg u0;
259+ for (uint16_t i = 0; i < batch; i++) {
260+ for (uint16_t j = 0; j < rows; j++) {
261+ __local_mem__ T1* curSrcAddr = srcAddr + i * srcBatchStride + j * srcRowStride;
262+ __local_mem__ T1* curDstAddr =
263+ dstAddr + i * dstBatchStride + (j + dstRowOffset) * dstRowStride + dstColOffset;
264+ for (uint16_t k = 0; k < loopCols; k++) {
265+ MicroAPI::DataCopy<T1, MicroAPI::PostLiteral::POST_MODE_UPDATE>(v0, curSrcAddr, repeatElm);
266+ MicroAPI::DataCopyUnAlign(curDstAddr, v0, u0, repeatElm);
267+ }
268+ MicroAPI::DataCopy<T1, MicroAPI::PostLiteral::POST_MODE_UPDATE>(v0, curSrcAddr, repeatElm);
269+ MicroAPI::DataCopyUnAlign(curDstAddr, v0, u0, tailCols);
270+ MicroAPI::DataCopyUnAlignPost(curDstAddr, u0, 0);
271+ }
272+ }
273+}
274+template <typename T1, typename T2, const uint32_t IS_PAD>
275+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::DupAndCopyToCalcBuffer(
276+ __local_mem__ T1* dstAddr, __local_mem__ T1* srcAddr)
277+{
278+ uint16_t loopCols = wInputActualNoPad_ / vlT1_;
279+ uint16_t tailCols = wInputActualNoPad_ - loopCols * vlT1_;
280+ uint32_t wInputActualNoPadAlign =
281+ CeilDivision(wInputActualNoPad_, blockSize_ / sizeof(T1)) * blockSize_ / sizeof(T1);
282+ uint32_t dstBatchStride = hInputActualPad_ * wInputActualAlignedPad_;
283+ uint32_t totalInput = tilingData_.highAxisInner * hInputActualPad_ * wInputActualAlignedPad_;
284+ uint16_t loopDup = totalInput / vlT1_;
285+ uint32_t tailDup = totalInput - loopDup * vlT1_;
286+ uint32_t dstRowOffset = topOffsetToInputTop_;
287+ uint32_t dstColOffset = leftOffsetToInputLeft_;
288+ __VEC_SCOPE__
289+ {
290+ DupBufferNegInf(dstAddr, vlT1_, loopDup, tailDup);
291+ CopyToCalcBuffer(
292+ dstAddr, srcAddr, highAxisActual_, hInputActualNoPad_, loopCols, tailCols, vlT1_, dstBatchStride,
293+ wInputActualNoPadAlign, dstBatchStride, wInputActualAlignedPad_, dstRowOffset, dstColOffset);
294+ }
295+ return;
296+}
297+template <typename T1, typename T2, const uint32_t IS_PAD>
298+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::Compute()
299+{
300+ LocalTensor<T1> inputLocal = inputQue_.DeQue<T1>();
301+ LocalTensor<T1> caclBuffLocal;
302+ __local_mem__ T1* inputBuffAddr;
303+ __local_mem__ T1* inputQueAddr = (__local_mem__ T1*)inputLocal.GetPhyAddr();
304+ __local_mem__ T1* computeAddr = inputQueAddr;
305+ if constexpr (IS_PAD == 1) {
306+ caclBuffLocal = inputCalcBuff_.Get<T1>();
307+ inputBuffAddr = (__local_mem__ T1*)caclBuffLocal.GetPhyAddr();
308+ DupAndCopyToCalcBuffer(inputBuffAddr, inputQueAddr);
309+ computeAddr = inputBuffAddr;
310+ }
311+ LocalTensor<T1> maxValueLocal = maxValueQue_.AllocTensor<T1>();
312+ LocalTensor<T2> argmaxLocal = argmaxQue_.AllocTensor<T2>();
313+ __local_mem__ T1* maxValueAddr = (__local_mem__ T1*)maxValueLocal.GetPhyAddr();
314+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
315+ if (wOutputActual_ * RATIO > vlT2_) {
316+ SingleRowGather(computeAddr, maxValueAddr, argmaxAddr);
317+ } else if (hOutputActual_ * wOutputActual_ * RATIO > vlT2_) {
318+ MultiRowGather(computeAddr, maxValueAddr, argmaxAddr);
319+ } else {
320+ MultiNcGather(computeAddr, maxValueAddr, argmaxAddr);
321+ }
322+ 
323+ inputQue_.FreeTensor(inputLocal);
324+ maxValueQue_.EnQue(maxValueLocal);
325+ argmaxQue_.EnQue(argmaxLocal);
326+ return;
327+}
328+template <typename T1, typename T2, const uint32_t IS_PAD>
329+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::ConvertIndexWithoutPadAlignNc(
330+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, T2 left, T2 wInputActualNoPad, T2 hIndexBase,
331+ MicroAPI::RegTensor<T2>& dstReg, int32_t ncInputOffset, int32_t ncOutputCount, int32_t inputNcSize)
332+{
333+ MicroAPI::RegTensor<int32_t> ncIndexReg;
334+ MicroAPI::RegTensor<int32_t> divResultReg;
335+ MicroAPI::RegTensor<int32_t> constReg;
336+ MicroAPI::MaskReg allMaskB32 = MicroAPI::CreateMask<int32_t, MicroAPI::MaskPattern::ALL>();
337+ MicroAPI::Arange(ncIndexReg, static_cast<int32_t>(0));
338+ MicroAPI::Duplicate(constReg, static_cast<int32_t>(ncOutputCount));
339+ MicroAPI::Div(divResultReg, ncIndexReg, constReg, allMaskB32);
340+ MicroAPI::Muls(divResultReg, divResultReg, inputNcSize, allMaskB32);
341+ MicroAPI::Sub(srcReg, srcReg, divResultReg, allMaskB32);
342+ 
343+ ConvertIndexWithoutPadAlign(srcReg, wStrideOffset, left, wInputActualNoPad, hIndexBase, dstReg, ncInputOffset);
344+}
345+template <typename T1, typename T2, const uint32_t IS_PAD>
346+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::ConvertIndexWithoutPadAlign(
347+ MicroAPI::RegTensor<int32_t>& srcReg, uint32_t wStrideOffset, T2 left, T2 wInputActualNoPad, T2 hIndexBase,
348+ MicroAPI::RegTensor<T2>& dstReg, int32_t ncInputOffset)
349+{
350+ MicroAPI::RegTensor<T2> hIndexReg;
351+ MicroAPI::RegTensor<int32_t> constReg;
352+ MicroAPI::RegTensor<int32_t> divResultReg;
353+ MicroAPI::RegTensor<T2> divResultRegUnpack;
354+ MicroAPI::RegTensor<T2> wIndexReg;
355+ MicroAPI::RegTensor<int32_t> wIndexRegUnpack;
356+ MicroAPI::RegTensor<T2> zeroReg;
357+ MicroAPI::MaskReg negInfMask;
358+ MicroAPI::MaskReg allMaskB32 = MicroAPI::CreateMask<int32_t, MicroAPI::MaskPattern::ALL>();
359+ MicroAPI::MaskReg allMaskT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
360+ MicroAPI::Duplicate(constReg, static_cast<int32_t>(wStrideOffset));
361+ MicroAPI::Duplicate(zeroReg, static_cast<T2>(0));
362+ MicroAPI::Adds(srcReg, srcReg, -ncInputOffset, allMaskB32);
363+ MicroAPI::Div(divResultReg, srcReg, constReg, allMaskB32);
364+ if constexpr (std::is_same<T2, int64_t>::value) {
365+ MicroAPI::UnPack(divResultRegUnpack, divResultReg);
366+ MicroAPI::Adds(hIndexReg, divResultRegUnpack, hIndexBase, allMaskT2);
367+ } else {
368+ MicroAPI::Adds(hIndexReg, divResultReg, hIndexBase, allMaskB32);
369+ }
370+ if constexpr (IS_PAD == 1) {
371+ MicroAPI::Compare<T2, CMPMODE::LT>(negInfMask, hIndexReg, zeroReg, allMaskT2);
372+ MicroAPI::Select(hIndexReg, zeroReg, hIndexReg, negInfMask);
373+ }
374+ MicroAPI::Muls(hIndexReg, hIndexReg, wInputActualNoPad, allMaskT2);
375+ MicroAPI::Mul(divResultReg, divResultReg, constReg, allMaskB32);
376+ MicroAPI::Sub(wIndexRegUnpack, srcReg, divResultReg, allMaskB32);
377+ if constexpr (std::is_same<T2, int64_t>::value) {
378+ MicroAPI::UnPack(wIndexReg, wIndexRegUnpack);
379+ MicroAPI::Adds(wIndexReg, wIndexReg, left, allMaskT2);
380+ } else {
381+ MicroAPI::Adds(wIndexReg, wIndexRegUnpack, left, allMaskB32);
382+ }
383+ if constexpr (IS_PAD == 1) {
384+ MicroAPI::Compare<T2, CMPMODE::LT>(negInfMask, wIndexReg, zeroReg, allMaskT2);
385+ MicroAPI::Select(wIndexReg, zeroReg, wIndexReg, negInfMask);
386+ }
387+ MicroAPI::Add(dstReg, hIndexReg, wIndexReg, allMaskT2);
388+ return;
389+}
390+template <typename T1, typename T2, const uint32_t IS_PAD>
391+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::ProcessW(
392+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, int32_t hOffset, uint16_t wStrideOffset,
393+ MicroAPI::RegTensor<int32_t>& indexReg, uint16_t hKernel, uint16_t wKernel, uint16_t repeatElem,
394+ int32_t outputOffset, MicroAPI::RegTensor<int32_t>& maxIndexReg, uint32_t hDilation, uint32_t wDilation)
395+{
396+ MicroAPI::RegTensor<int32_t> indexWithOffset;
397+ MicroAPI::RegTensor<T1> calcReg;
398+ MicroAPI::RegTensor<int32_t> calcMaxIndexReg;
399+ uint32_t maskCount = repeatElem;
400+ MicroAPI::MaskReg allMaskU32 = MicroAPI::CreateMask<int32_t, MicroAPI::MaskPattern::ALL>();
401+ MicroAPI::MaskReg gatherMask = MicroAPI::UpdateMask<T1>(maskCount);
402+ MicroAPI::RegTensor<T1> maxReg;
403+ MicroAPI::MaskReg neMask;
404+ MicroAPI::MaskReg gtMask;
405+ MicroAPI::MaskReg tmpMask;
406+ MicroAPI::UnalignReg u0;
407+ 
408+ __local_mem__ T1* maxValueAddrLocal = maxValueAddr + outputOffset;
409+ DuplicateNegInfReg<T1>(maxReg);
410+ MicroAPI::Adds(maxIndexReg, indexReg, hOffset, allMaskU32);
411+ for (uint16_t i = 0; i < hKernel; i++) {
412+ for (uint16_t j = 0; j < wKernel; j++) {
413+ int32_t relIndex = i * wStrideOffset * hDilation + j * wDilation;
414+ int32_t offset = static_cast<int32_t>(hOffset + relIndex);
415+ MicroAPI::Adds(indexWithOffset, indexReg, offset, allMaskU32);
416+ if constexpr (std::is_same<T1, float>::value) {
417+ MicroAPI::DataCopyGather(
418+ calcReg, computeAddr, (MicroAPI::RegTensor<uint32_t>&)indexWithOffset, gatherMask);
419+ } else {
420+ MicroAPI::RegTensor<uint16_t> indexConvert;
421+ MicroAPI::Pack(indexConvert, indexWithOffset);
422+ MicroAPI::DataCopyGather(calcReg, computeAddr, indexConvert, gatherMask);
423+ }
424+ MicroAPI::Compare<T1, CMPMODE::GT>(gtMask, calcReg, maxReg, gatherMask);
425+ MicroAPI::Compare<T1, CMPMODE::NE>(neMask, calcReg, calcReg, gatherMask);
426+ MicroAPI::MaskOr(gtMask, gtMask, neMask, gatherMask);
427+ if constexpr (sizeof(int32_t) / sizeof(T1) == 1) {
428+ MicroAPI::Select(maxIndexReg, indexWithOffset, maxIndexReg, gtMask);
429+ } else {
430+ MicroAPI::MaskUnPack(tmpMask, gtMask);
431+ MicroAPI::Select(maxIndexReg, indexWithOffset, maxIndexReg, tmpMask);
432+ }
433+ MicroAPI::Max(maxReg, maxReg, calcReg, gatherMask);
434+ }
435+ }
436+ MicroAPI::DataCopyUnAlign(maxValueAddrLocal, maxReg, u0, repeatElem);
437+ MicroAPI::DataCopyUnAlignPost(maxValueAddrLocal, u0, 0);
438+ return;
439+}
440+template <typename T1, typename T2, const uint32_t IS_PAD>
441+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::SingleRowGather(
442+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr)
443+{
444+ uint16_t loopW = wOutputActual_ / vlT2_;
445+ uint16_t repeatsElem = vlT2_;
446+ uint16_t tailRepeatsElem = wOutputActual_ - loopW * vlT2_;
447+ if (tailRepeatsElem == 0) {
448+ loopW = loopW - 1;
449+ tailRepeatsElem = repeatsElem;
450+ }
451+ uint16_t hKernel = tilingData_.hKernel;
452+ uint16_t wKernel = tilingData_.wKernel;
453+ uint32_t hStride = tilingData_.hStride;
454+ uint32_t wStride = tilingData_.wStride;
455+ T2 left = wAxisIndex_ * tilingData_.wOutputInner * tilingData_.wStride - tilingData_.padLeft;
456+ T2 hIndexBase = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.hStride - tilingData_.padTop;
457+ T2 wInput = tilingData_.wInput;
458+ uint32_t highAxisActual = highAxisActual_;
459+ uint32_t hOutputActual = hOutputActual_;
460+ uint32_t wOutputActual = wOutputActual_;
461+ uint32_t hInputActualPad = hInputActualPad_;
462+ uint32_t wInputActualAlignedPad = wInputActualAlignedPad_;
463+ uint32_t wOutputActualAligned = wOutputActualAligned_;
464+ uint32_t hDilation = tilingData_.hDilation;
465+ uint32_t wDilation = tilingData_.wDilation;
466+ for (uint16_t nc = 0; nc < static_cast<uint16_t>(highAxisActual); nc++) {
467+ __VEC_SCOPE__
468+ {
469+ MicroAPI::RegTensor<int32_t> indexReg;
470+ MicroAPI::RegTensor<int32_t> maxIndexReg;
471+ MicroAPI::RegTensor<T2> maxIndexConvertReg;
472+ MicroAPI::UnalignReg u1;
473+ MicroAPI::Arange(indexReg, static_cast<int32_t>(0));
474+ MicroAPI::MaskReg preg = MicroAPI::CreateMask<T1, MicroAPI::MaskPattern::ALL>();
475+ MicroAPI::Muls(indexReg, indexReg, static_cast<int32_t>(wStride), preg);
476+ int32_t ncInputOffset = nc * hInputActualPad * wInputActualAlignedPad;
477+ int32_t ncOutputOffset = nc * hOutputActual * wOutputActual;
478+ __local_mem__ T2* argmaxAddrLocal = argmaxAddr + ncOutputOffset;
479+ for (uint16_t hLoop = 0; hLoop < static_cast<uint16_t>(hOutputActual); hLoop++) {
480+ for (uint16_t wLoop = 0; wLoop < loopW; wLoop++) {
481+ int32_t wOffset =
482+ ncInputOffset + hLoop * wInputActualAlignedPad * hStride + wLoop * repeatsElem * wStride;
483+ int32_t wOutputOffset = ncOutputOffset + hLoop * wOutputActual + wLoop * repeatsElem;
484+ ProcessW(
485+ computeAddr, maxValueAddr, wOffset, wInputActualAlignedPad, indexReg, hKernel, wKernel,
486+ repeatsElem, wOutputOffset, maxIndexReg, hDilation, wDilation);
487+ ConvertIndexWithoutPadAlign(
488+ maxIndexReg, wInputActualAlignedPad, left, wInput, hIndexBase, maxIndexConvertReg,
489+ ncInputOffset);
490+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, repeatsElem);
491+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
492+ }
493+ int32_t wOffsetTail =
494+ ncInputOffset + hLoop * wInputActualAlignedPad * hStride + loopW * repeatsElem * wStride;
495+ int32_t wOutputOffsetTail = ncOutputOffset + hLoop * wOutputActual + loopW * repeatsElem;
496+ ProcessW(
497+ computeAddr, maxValueAddr, wOffsetTail, wInputActualAlignedPad, indexReg, hKernel, wKernel,
498+ tailRepeatsElem, wOutputOffsetTail, maxIndexReg, hDilation, wDilation);
499+ ConvertIndexWithoutPadAlign(
500+ maxIndexReg, wInputActualAlignedPad, left, wInput, hIndexBase, maxIndexConvertReg, ncInputOffset);
501+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, tailRepeatsElem);
502+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
503+ }
504+ }
505+ }
506+ return;
507+}
508+template <typename T1, typename T2, const uint32_t IS_PAD>
509+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::MultiRowGather(
510+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr)
511+{
512+ uint32_t wOutputActual = wOutputActual_;
513+ uint16_t wKernel = tilingData_.wKernel;
514+ uint16_t hKernel = tilingData_.hKernel;
515+ uint32_t wStride = tilingData_.wStride;
516+ uint32_t rate2D = wInputActualAlignedPad_ * tilingData_.hStride;
517+ uint16_t hBatchCount = vlT2_ / wOutputActual_;
518+ uint16_t hLoopTimes = hOutputActual_ / hBatchCount;
519+ uint16_t hTail = hOutputActual_ - hLoopTimes * hBatchCount;
520+ if (hTail == 0) {
521+ hLoopTimes = hLoopTimes - 1;
522+ hTail = hBatchCount;
523+ }
524+ uint16_t repeatsElem = hBatchCount * wOutputActual_;
525+ uint16_t tailRepeatsElem = hTail * wOutputActual_;
526+ T2 left = wAxisIndex_ * tilingData_.wOutputInner * tilingData_.wStride - tilingData_.padLeft;
527+ T2 hIndexBase = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.hStride - tilingData_.padTop;
528+ T2 wInput = tilingData_.wInput;
529+ uint32_t highAxisActual = highAxisActual_;
530+ uint32_t hInputActualPad = hInputActualPad_;
531+ uint32_t wInputActualAlignedPad = wInputActualAlignedPad_;
532+ uint32_t wOutputActualAligned = wOutputActualAligned_;
533+ uint32_t hOutputActual = hOutputActual_;
534+ uint32_t hStride = tilingData_.hStride;
535+ uint32_t hDilation = tilingData_.hDilation;
536+ uint32_t wDilation = tilingData_.wDilation;
537+ __VEC_SCOPE__
538+ {
539+ MicroAPI::RegTensor<int32_t> indexReg;
540+ MicroAPI::RegTensor<int32_t> maxIndexReg;
541+ MicroAPI::RegTensor<T2> maxIndexConvertReg;
542+ MicroAPI::UnalignReg u1;
543+ __local_mem__ T2* argmaxAddrLocal = argmaxAddr;
544+ GenGatterIndex2D<int32_t>(indexReg, rate2D, wOutputActual, wStride);
545+ for (uint16_t nc = 0; nc < static_cast<uint16_t>(highAxisActual); nc++) {
546+ int32_t ncInputOffset = nc * hInputActualPad * wInputActualAlignedPad;
547+ for (uint16_t hLoop = 0; hLoop < hLoopTimes; hLoop++) {
548+ int32_t wOffset = ncInputOffset + hLoop * hBatchCount * hStride * wInputActualAlignedPad;
549+ int32_t wOutputOffset = nc * hOutputActual * wOutputActual + hLoop * hBatchCount * wOutputActual;
550+ ProcessW(
551+ computeAddr, maxValueAddr, wOffset, wInputActualAlignedPad, indexReg, hKernel, wKernel, repeatsElem,
552+ wOutputOffset, maxIndexReg, hDilation, wDilation);
553+ ConvertIndexWithoutPadAlign(
554+ maxIndexReg, wInputActualAlignedPad, left, wInput, hIndexBase, maxIndexConvertReg, ncInputOffset);
555+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, repeatsElem);
556+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
557+ }
558+ int32_t wOffsetTail = ncInputOffset + hLoopTimes * hBatchCount * hStride * wInputActualAlignedPad;
559+ int32_t wOutputOffsetTail = nc * hOutputActual * wOutputActual + hLoopTimes * hBatchCount * wOutputActual;
560+ ProcessW(
561+ computeAddr, maxValueAddr, wOffsetTail, wInputActualAlignedPad, indexReg, hKernel, wKernel,
562+ tailRepeatsElem, wOutputOffsetTail, maxIndexReg, hDilation, wDilation);
563+ ConvertIndexWithoutPadAlign(
564+ maxIndexReg, wInputActualAlignedPad, left, wInput, hIndexBase, maxIndexConvertReg, ncInputOffset);
565+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, tailRepeatsElem);
566+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
567+ }
568+ }
569+ return;
570+}
571+template <typename T1, typename T2, const uint32_t IS_PAD>
572+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::MultiNcGather(
573+ __local_mem__ T1* computeAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr)
574+{
575+ uint16_t wKernel = tilingData_.wKernel;
576+ uint16_t hKernel = tilingData_.hKernel;
577+ uint32_t wStride = tilingData_.wStride;
578+ uint16_t rate3D = hInputActualPad_ * wInputActualAlignedPad_;
579+ uint16_t num2D = hOutputActual_ * wOutputActual_;
580+ uint16_t rate2D = tilingData_.hStride * wInputActualAlignedPad_;
581+ uint16_t wOutputActual = wOutputActual_;
582+ uint16_t eachBatchCount = hOutputActual_ * wOutputActual_;
583+ uint16_t ncBatchCount = vlT2_ / eachBatchCount;
584+ uint16_t ncLoopTimes = highAxisActual_ / ncBatchCount;
585+ uint16_t ncTail = highAxisActual_ - ncLoopTimes * ncBatchCount;
586+ if (ncTail == 0) {
587+ ncLoopTimes = ncLoopTimes - 1;
588+ ncTail = ncBatchCount;
589+ }
590+ uint16_t repeatsElem = ncBatchCount * eachBatchCount;
591+ uint16_t tailRepeatsElem = ncTail * eachBatchCount;
592+ T2 left = wAxisIndex_ * tilingData_.wOutputInner * tilingData_.wStride - tilingData_.padLeft;
593+ T2 hIndexBase = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.hStride - tilingData_.padTop;
594+ T2 wInput = tilingData_.wInput;
595+ uint32_t hInputActualPad = hInputActualPad_;
596+ uint32_t wInputActualAlignedPad = wInputActualAlignedPad_;
597+ uint32_t hOutputActual = hOutputActual_;
598+ uint32_t wOutputActualAligned = wOutputActualAligned_;
599+ uint32_t hDilation = tilingData_.hDilation;
600+ uint32_t wDilation = tilingData_.wDilation;
601+ __VEC_SCOPE__
602+ {
603+ MicroAPI::RegTensor<int32_t> indexReg;
604+ MicroAPI::RegTensor<int32_t> maxIndexReg;
605+ MicroAPI::RegTensor<T2> maxIndexConvertReg;
606+ MicroAPI::UnalignReg u1;
607+ __local_mem__ T2* argmaxAddrLocal = argmaxAddr;
608+ GenGatterIndex3D<int32_t>(indexReg, rate3D, num2D, rate2D, wOutputActual, wStride);
609+ for (uint16_t nc = 0; nc < ncLoopTimes; nc++) {
610+ uint32_t wOffset = nc * ncBatchCount * hInputActualPad * wInputActualAlignedPad;
611+ uint32_t wOutputOffset = nc * ncBatchCount * hOutputActual * wOutputActual;
612+ ProcessW(
613+ computeAddr, maxValueAddr, wOffset, wInputActualAlignedPad, indexReg, hKernel, wKernel, repeatsElem,
614+ wOutputOffset, maxIndexReg, hDilation, wDilation);
615+ ConvertIndexWithoutPadAlignNc(
616+ maxIndexReg, wInputActualAlignedPad, left, wInput, hIndexBase, maxIndexConvertReg, wOffset, num2D,
617+ rate3D);
618+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, repeatsElem);
619+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
620+ }
621+ uint32_t wOffsetTail = ncLoopTimes * ncBatchCount * hInputActualPad * wInputActualAlignedPad;
622+ uint32_t wOutputOffsetTail = ncLoopTimes * ncBatchCount * hOutputActual * wOutputActual;
623+ ProcessW(
624+ computeAddr, maxValueAddr, wOffsetTail, wInputActualAlignedPad, indexReg, hKernel, wKernel, tailRepeatsElem,
625+ wOutputOffsetTail, maxIndexReg, hDilation, wDilation);
626+ ConvertIndexWithoutPadAlignNc(
627+ maxIndexReg, wInputActualAlignedPad, left, wInput, hIndexBase, maxIndexConvertReg, wOffsetTail, num2D,
628+ rate3D);
629+ MicroAPI::DataCopyUnAlign(argmaxAddrLocal, maxIndexConvertReg, u1, tailRepeatsElem);
630+ MicroAPI::DataCopyUnAlignPost(argmaxAddrLocal, u1, 0);
631+ }
632+ return;
633+}
634+template <typename T1, typename T2, const uint32_t IS_PAD>
635+__aicore__ inline void MaxPoolWithArgmaxV3GatherKernel<T1, T2, IS_PAD>::CopyOut()
636+{
637+ LocalTensor<T1> maxValueLocal = maxValueQue_.DeQue<T1>();
638+ LocalTensor<T2> argmaxLocal = argmaxQue_.DeQue<T2>();
639+ int64_t outputPlaneSize = tilingData_.hOutput * tilingData_.wOutput;
640+ int64_t highOutputAxisOffset = highAxisIndex_ * tilingData_.highAxisInner * outputPlaneSize;
641+ int64_t hOutputAxisOffset = hAxisIndex_ * tilingData_.hOutputInner * tilingData_.wOutput;
642+ int64_t wOutputAxisOffset = wAxisIndex_ * tilingData_.wOutputInner;
643+ int64_t outputGmOffset = highOutputAxisOffset + hOutputAxisOffset + wOutputAxisOffset;
644+ 
645+ DataCopyExtParams copyOutParamT1 = {
646+ static_cast<uint16_t>(1), static_cast<uint32_t>(highAxisActual_ * hOutputActual_ * wOutputActual_ * sizeof(T1)),
647+ static_cast<uint32_t>(0), static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
648+ 
649+ DataCopyPad(yGm_[outputGmOffset], maxValueLocal, copyOutParamT1);
650+ DataCopyExtParams copyOutParamT2 = {
651+ static_cast<uint16_t>(1), static_cast<uint32_t>(highAxisActual_ * hOutputActual_ * wOutputActual_ * sizeof(T2)),
652+ static_cast<uint32_t>(0), static_cast<uint32_t>(0), static_cast<uint32_t>(0)};
653+ DataCopyPad(argmaxGm_[outputGmOffset], argmaxLocal, copyOutParamT2);
654+ maxValueQue_.FreeTensor(maxValueLocal);
655+ argmaxQue_.FreeTensor(argmaxLocal);
656+ return;
657+}
658+} // namespace MaxPoolWithArgmaxV3GatherNameSpace
659+#endif
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_nhwc_big_c.h+798-0
@@ -0,0 +1,798 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_nhwc_big_c.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_WITH_ARGMAX_V3_NHWC_BIG_C_H_
17+#define MAX_POOL_WITH_ARGMAX_V3_NHWC_BIG_C_H_
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+#include "max_pool_with_argmax_v3_base.h"
23+ 
24+namespace MaxPoolWithArgmaxV3NHWC {
25+using namespace AscendC;
26+ 
27+constexpr uint32_t BUFFER_NUM = 2;
28+constexpr int64_t HELPER_BUFFER_SIZE = 1024;
29+constexpr int64_t HELPER_BUFFER_SIZE_512 = 512;
30+constexpr int64_t THREE_DIM = 3;
31+constexpr int64_t DIGIT_1 = 1;
32+constexpr int64_t DIGIT_2 = 2;
33+ 
34+template <typename T1, typename T2, const uint32_t IS_PAD = 0>
35+class MaxPoolWithArgmaxV3NhwCKernel {
36+public:
37+ __aicore__ inline MaxPoolWithArgmaxV3NhwCKernel(
38+ TPipe* pipe, const MaxPoolWithArgmaxV3NhwcTilingData* __restrict tiling)
39+ : pipe_(pipe), tilingData_(tiling){};
40+ 
41+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR argmax);
42+ __aicore__ inline void ParseTilingData(const MaxPoolWithArgmaxV3NhwcTilingData& tilingData);
43+ __aicore__ inline void Process();
44+ __aicore__ inline void ScalarCompute(int64_t loopNum);
45+ __aicore__ inline void CopyIn();
46+ __aicore__ inline void FillPadNegVF(__local_mem__ T1* xLocalAddr);
47+ __aicore__ inline void Compute(__local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal);
48+ 
49+ __aicore__ inline void InitHelpBuf();
50+ __aicore__ inline void CopyResultToUb(__local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal);
51+ 
52+ template <const bool IS_SPLIT_KERNEL>
53+ __aicore__ inline void MaxPoolAndArgmaxV3VF(
54+ __local_mem__ T1* xLocal, __local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal);
55+ template <const bool IS_SPLIT_KERNEL>
56+ __aicore__ inline void MaxPoolAndArgmaxV3VFForPad(
57+ __local_mem__ T1* xLocal, __local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal);
58+ __aicore__ inline void CopyOut();
59+ 
60+ TPipe* pipe_;
61+ TQue<QuePosition::VECIN, BUFFER_NUM> inputQue_;
62+ TQue<QuePosition::VECOUT, BUFFER_NUM> maxValueQue_;
63+ TQue<QuePosition::VECOUT, BUFFER_NUM> argmaxQue_;
64+ TBuf<TPosition::VECCALC> helperTBuf_;
65+ 
66+ GlobalTensor<T1> xGm_;
67+ GlobalTensor<T1> yGm_;
68+ GlobalTensor<T2> argmaxGm_;
69+ 
70+ const MaxPoolWithArgmaxV3NhwcTilingData* tilingData_;
71+ uint32_t blockIdx_ = 0;
72+ 
73+ constexpr static int32_t BLOCK_SIZE = platform::GetUbBlockSize();
74+ constexpr static int64_t MAX_DATA_NUM_IN_ONE_BLOCK =
75+ BLOCK_SIZE / sizeof(T1) >= BLOCK_SIZE / sizeof(T2) ? BLOCK_SIZE / sizeof(T1) : BLOCK_SIZE / sizeof(T2);
76+ constexpr static int64_t VREG_LENGTH_DATA_NUM_T2 = platform::GetVRegSize() / sizeof(T2);
77+ 
78+ // tilingdata
79+ int64_t cInput_ = 0;
80+ int64_t hInput_ = 0;
81+ int64_t wInput_ = 0;
82+ int64_t hOutput_ = 0;
83+ int64_t wOutput_ = 0;
84+ int64_t hKernel_ = 0;
85+ int64_t wKernel_ = 0;
86+ int64_t hStride_ = 0;
87+ int64_t wStride_ = 0;
88+ int64_t padLeft_ = 0;
89+ int64_t padTop_ = 0;
90+ int64_t hDilation_ = 0;
91+ int64_t wDilation_ = 0;
92+ int64_t nOutputInner_ = 0;
93+ int64_t nOutputTail_ = 0;
94+ int64_t nOutputOuter_ = 0;
95+ int64_t hOutputInner_ = 0;
96+ int64_t hOutputTail_ = 0;
97+ int64_t hOutputOuter_ = 0;
98+ int64_t wOutputInner_ = 0;
99+ int64_t wOutputTail_ = 0;
100+ int64_t wOutputOuter_ = 0;
101+ int64_t cOutputInner_ = 0;
102+ int64_t cOutputTail_ = 0;
103+ int64_t cOutputOuter_ = 0;
104+ int64_t normalCoreProcessNum_ = 0;
105+ int64_t tailCoreProcessNum_ = 0;
106+ int64_t usedCoreNum_ = 0;
107+ int64_t inputBufferSize_ = 0;
108+ int64_t maxValueBufferSize_ = 0;
109+ int64_t argmaxBufferSize_ = 0;
110+ int64_t isPad_ = 0;
111+ int64_t isSplitKernel_ = 0;
112+ int64_t hKernelInner_ = 0;
113+ int64_t hKernelTail_ = 0;
114+ int64_t hKernelOuter_ = 0;
115+ int64_t wKernelInner_ = 0;
116+ int64_t wKernelTail_ = 0;
117+ int64_t wKernelOuter_ = 0;
118+ 
119+ // 输出域大小
120+ int64_t nOutputActual_ = 1;
121+ int64_t hOutputActual_ = 1;
122+ int64_t wOutputActual_ = 1;
123+ int64_t cOutputActual_ = 1;
124+ 
125+ // c轴对齐到BlockSize
126+ int64_t cOutputActualAlign_ = 0;
127+ 
128+ // 输入域大小
129+ int64_t hInputActual_ = 1;
130+ int64_t wInputActual_ = 1;
131+ // 输入域大小包含前后pad大小
132+ int64_t hInputActualPad_ = 1;
133+ int64_t wInputActualPad_ = 1;
134+ 
135+ // 输入相对偏移
136+ int64_t nInputAxisOffset_ = 0;
137+ int64_t cInputAxisOffset_ = 0;
138+ int64_t hInputAxisOffset_ = 0;
139+ int64_t wInputAxisOffset_ = 0;
140+ // 输出相对偏移
141+ int64_t nOutputAxisOffset_ = 0;
142+ int64_t cOutputAxisOffset_ = 0;
143+ int64_t hOutputAxisOffset_ = 0;
144+ int64_t wOutputAxisOffset_ = 0;
145+ 
146+ // WH轴索引起始点
147+ int64_t indexWHplanOffset_ = 0;
148+ // 轴循环index
149+ int64_t nAxisIndex_ = 0;
150+ int64_t cAxisIndex_ = 0;
151+ int64_t hAxisIndex_ = 0;
152+ int64_t wAxisIndex_ = 0;
153+ 
154+ // 切kernel时kernel循环、大小
155+ int64_t hKernelIndex_ = 0;
156+ int64_t wKernelIndex_ = 0;
157+ int64_t hKernelActual_ = 0;
158+ int64_t wKernelActual_ = 0;
159+ 
160+ // 存在pad时，上下左右偏移
161+ int64_t baseBlockLeftOffsetInOcean_ = 0;
162+ int64_t baseBlockRightOffsetInOcean_ = 0;
163+ int64_t baseBlockTopOffsetInOcean_ = 0;
164+ int64_t baseBlockDownOffsetInOcean_ = 0;
165+ 
166+ // 存在pad时，上下左右偏移
167+ int64_t xLocalUbOffset = 0;
168+};
169+ 
170+template <typename T1, typename T2, const uint32_t IS_PAD>
171+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::ParseTilingData(
172+ const MaxPoolWithArgmaxV3NhwcTilingData& tilingData)
173+{
174+ cInput_ = tilingData.cInput;
175+ hInput_ = tilingData.hInput;
176+ wInput_ = tilingData.wInput;
177+ hOutput_ = tilingData.hOutput;
178+ wOutput_ = tilingData.wOutput;
179+ hKernel_ = tilingData.hKernel;
180+ wKernel_ = tilingData.wKernel;
181+ hStride_ = tilingData.hStride;
182+ wStride_ = tilingData.wStride;
183+ padLeft_ = tilingData.padLeft;
184+ padTop_ = tilingData.padTop;
185+ hDilation_ = tilingData.hDilation;
186+ wDilation_ = tilingData.wDilation;
187+ nOutputInner_ = tilingData.nOutputInner;
188+ nOutputTail_ = tilingData.nOutputTail;
189+ nOutputOuter_ = tilingData.nOutputOuter;
190+ hOutputInner_ = tilingData.hOutputInner;
191+ hOutputTail_ = tilingData.hOutputTail;
192+ hOutputOuter_ = tilingData.hOutputOuter;
193+ wOutputInner_ = tilingData.wOutputInner;
194+ wOutputTail_ = tilingData.wOutputTail;
195+ wOutputOuter_ = tilingData.wOutputOuter;
196+ cOutputInner_ = tilingData.cOutputInner;
197+ cOutputTail_ = tilingData.cOutputTail;
198+ cOutputOuter_ = tilingData.cOutputOuter;
199+ normalCoreProcessNum_ = tilingData.normalCoreProcessNum;
200+ tailCoreProcessNum_ = tilingData.tailCoreProcessNum;
201+ usedCoreNum_ = tilingData.usedCoreNum;
202+ inputBufferSize_ = tilingData.inputBufferSize;
203+ maxValueBufferSize_ = tilingData.maxValueBufferSize;
204+ argmaxBufferSize_ = tilingData.argmaxBufferSize;
205+ isPad_ = tilingData.isPad;
206+ isSplitKernel_ = tilingData.isSplitKernel;
207+ hKernelInner_ = tilingData.hKernelInner;
208+ hKernelTail_ = tilingData.hKernelTail;
209+ hKernelOuter_ = tilingData.hKernelOuter;
210+ wKernelInner_ = tilingData.wKernelInner;
211+ wKernelTail_ = tilingData.wKernelTail;
212+ wKernelOuter_ = tilingData.wKernelOuter;
213+}
214+ 
215+template <typename T1, typename T2, const uint32_t IS_PAD>
216+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR argmax)
217+{
218+ ParseTilingData(*tilingData_);
219+ blockIdx_ = GetBlockIdx();
220+ if (blockIdx_ >= usedCoreNum_) {
221+ return;
222+ }
223+ 
224+ xGm_.SetGlobalBuffer((__gm__ T1*)x);
225+ yGm_.SetGlobalBuffer((__gm__ T1*)y);
226+ argmaxGm_.SetGlobalBuffer((__gm__ T2*)argmax);
227+ 
228+ pipe_->InitBuffer(inputQue_, BUFFER_NUM, inputBufferSize_);
229+ pipe_->InitBuffer(maxValueQue_, BUFFER_NUM, maxValueBufferSize_);
230+ pipe_->InitBuffer(argmaxQue_, BUFFER_NUM, argmaxBufferSize_);
231+ pipe_->InitBuffer(helperTBuf_, HELPER_BUFFER_SIZE);
232+}
233+ 
234+template <typename T1, typename T2, const uint32_t IS_PAD>
235+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::ScalarCompute(int64_t loopNum)
236+{
237+ int64_t baseBlockIdx = blockIdx_ * normalCoreProcessNum_ + loopNum;
238+ int64_t hwc = hOutputOuter_ * wOutputOuter_ * cOutputOuter_;
239+ int64_t wc = wOutputOuter_ * cOutputOuter_;
240+ 
241+ nAxisIndex_ = baseBlockIdx / hwc;
242+ baseBlockIdx = baseBlockIdx % hwc;
243+ hAxisIndex_ = baseBlockIdx / wc;
244+ baseBlockIdx = baseBlockIdx % wc;
245+ wAxisIndex_ = baseBlockIdx / cOutputOuter_;
246+ cAxisIndex_ = baseBlockIdx % cOutputOuter_;
247+ 
248+ nOutputActual_ = nAxisIndex_ == (nOutputOuter_ - 1) ? nOutputTail_ : nOutputInner_;
249+ hOutputActual_ = hAxisIndex_ == (hOutputOuter_ - 1) ? hOutputTail_ : hOutputInner_;
250+ wOutputActual_ = wAxisIndex_ == (wOutputOuter_ - 1) ? wOutputTail_ : wOutputInner_;
251+ cOutputActual_ = cAxisIndex_ == (cOutputOuter_ - 1) ? cOutputTail_ : cOutputInner_;
252+ 
253+ cOutputActualAlign_ = ops::Aligned(cOutputActual_, int64_t(BLOCK_SIZE / sizeof(T1)));
254+ 
255+ hInputActual_ = (hOutputActual_ - 1) * hStride_ + hKernel_;
256+ wInputActual_ = (wOutputActual_ - 1) * wStride_ + wKernel_;
257+ // 输入相对偏移
258+ cInputAxisOffset_ = cAxisIndex_ * cOutputInner_;
259+ wInputAxisOffset_ = (wAxisIndex_ * wStride_ * wOutputInner_ + wKernelIndex_ * wKernelInner_) * cInput_;
260+ hInputAxisOffset_ = (hAxisIndex_ * hStride_ * hOutputInner_ + hKernelIndex_ * hKernelInner_) * wInput_ * cInput_;
261+ nInputAxisOffset_ = nAxisIndex_ * nOutputInner_ * hInput_ * wInput_ * cInput_;
262+ 
263+ // 输出相对偏移
264+ cOutputAxisOffset_ = cAxisIndex_ * cOutputInner_;
265+ wOutputAxisOffset_ = wAxisIndex_ * wOutputInner_ * cInput_;
266+ hOutputAxisOffset_ = hAxisIndex_ * hOutputInner_ * wOutput_ * cInput_;
267+ nOutputAxisOffset_ = nAxisIndex_ * nOutputInner_ * hOutput_ * wOutput_ * cInput_;
268+ 
269+ hKernelActual_ = hKernel_;
270+ wKernelActual_ = wKernel_;
271+ 
272+ // kernel切分
273+ if (isSplitKernel_ == 1) {
274+ wInputActual_ = wKernelIndex_ == (wKernelOuter_ - 1) ? wKernelTail_ : wKernelInner_;
275+ hInputActual_ = hKernelIndex_ == (hKernelOuter_ - 1) ? hKernelTail_ : hKernelInner_;
276+ hKernelActual_ = hInputActual_;
277+ wKernelActual_ = wInputActual_;
278+ }
279+ hInputActualPad_ = hInputActual_;
280+ wInputActualPad_ = wInputActual_;
281+ 
282+ if constexpr (IS_PAD == 1) {
283+ int64_t topOffset = hAxisIndex_ * hOutputInner_ * hStride_ + hKernelIndex_ * hKernelInner_ - padTop_;
284+ int64_t downOffset = hAxisIndex_ * hOutputInner_ * hStride_ + hKernelIndex_ * hKernelInner_ +
285+ (hOutputActual_ - 1) * hStride_ + hKernelActual_ - hInput_ - padTop_;
286+ int64_t leftOffset = wAxisIndex_ * wOutputInner_ * wStride_ + wKernelIndex_ * wKernelInner_ - padLeft_;
287+ int64_t rightOffset = wAxisIndex_ * wOutputInner_ * wStride_ + wKernelIndex_ * wKernelInner_ +
288+ (wOutputActual_ - 1) * wStride_ + wKernelActual_ - wInput_ - padLeft_;
289+ 
290+ baseBlockLeftOffsetInOcean_ = leftOffset >= 0 ? 0 : -leftOffset;
291+ baseBlockRightOffsetInOcean_ = rightOffset >= 0 ? rightOffset : 0;
292+ baseBlockTopOffsetInOcean_ = topOffset >= 0 ? 0 : -topOffset;
293+ baseBlockDownOffsetInOcean_ = downOffset >= 0 ? downOffset : 0;
294+ // PAD时输入偏移
295+ xLocalUbOffset = baseBlockTopOffsetInOcean_ * wInputActual_ * cOutputActualAlign_ +
296+ baseBlockLeftOffsetInOcean_ * cOutputActualAlign_;
297+ 
298+ hInputActual_ = hInputActual_ - baseBlockTopOffsetInOcean_ - baseBlockDownOffsetInOcean_;
299+ wInputActual_ = wInputActual_ - baseBlockLeftOffsetInOcean_ - baseBlockRightOffsetInOcean_;
300+ 
301+ hInputAxisOffset_ = baseBlockTopOffsetInOcean_ == 0 ? hInputAxisOffset_ - padTop_ * wInput_ * cInput_ : 0;
302+ wInputAxisOffset_ = baseBlockLeftOffsetInOcean_ == 0 ? wInputAxisOffset_ - padLeft_ * cInput_ : 0;
303+ }
304+ indexWHplanOffset_ = (hInputAxisOffset_ + wInputAxisOffset_) / cInput_;
305+}
306+ 
307+template <typename T1, typename T2, const uint32_t IS_PAD>
308+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::InitHelpBuf()
309+{
310+ __local_mem__ T1* maxValueHelp = (__local_mem__ T1*)helperTBuf_.Get<T1>().GetPhyAddr();
311+ __local_mem__ T2* argmaxHelp =
312+ (__local_mem__ T2*)helperTBuf_.Get<T2>().GetPhyAddr() + HELPER_BUFFER_SIZE_512 / sizeof(T1);
313+ 
314+ __VEC_SCOPE__
315+ {
316+ AscendC::MicroAPI::RegTensor<T1> negInfReg;
317+ AscendC::MicroAPI::RegTensor<T2> negOne;
318+ DuplicateNegInfReg(negInfReg);
319+ AscendC::MicroAPI::Duplicate(negOne, 0);
320+ AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask<T1, AscendC::MicroAPI::MaskPattern::ALL>();
321+ AscendC::MicroAPI::MaskReg pregAllT2 = AscendC::MicroAPI::CreateMask<T2, AscendC::MicroAPI::MaskPattern::ALL>();
322+ AscendC::MicroAPI::DataCopy(maxValueHelp, negInfReg, pregAll);
323+ AscendC::MicroAPI::DataCopy(argmaxHelp, negOne, pregAllT2);
324+ }
325+}
326+ 
327+template <typename T1, typename T2, const uint32_t IS_PAD>
328+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::CopyResultToUb(
329+ __local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal)
330+{
331+ __local_mem__ T1* maxValueHelp = (__local_mem__ T1*)helperTBuf_.Get<T1>().GetPhyAddr();
332+ __local_mem__ T2* argmaxHelp =
333+ (__local_mem__ T2*)helperTBuf_.Get<T2>().GetPhyAddr() + HELPER_BUFFER_SIZE_512 / sizeof(T1);
334+ 
335+ __VEC_SCOPE__
336+ {
337+ AscendC::MicroAPI::RegTensor<T1> vreg0;
338+ AscendC::MicroAPI::RegTensor<T2> argmaxUpdateVreg;
339+ AscendC::MicroAPI::MaskReg pregAllT1 = AscendC::MicroAPI::CreateMask<T1, AscendC::MicroAPI::MaskPattern::ALL>();
340+ AscendC::MicroAPI::MaskReg pregAllT2 = AscendC::MicroAPI::CreateMask<T2, AscendC::MicroAPI::MaskPattern::ALL>();
341+ AscendC::MicroAPI::DataCopy(vreg0, maxValueHelp);
342+ AscendC::MicroAPI::DataCopy(argmaxUpdateVreg, argmaxHelp);
343+ AscendC::MicroAPI::DataCopy(maxValueLocal, vreg0, pregAllT1);
344+ AscendC::MicroAPI::DataCopy(argmaxLocal, argmaxUpdateVreg, pregAllT2);
345+ }
346+}
347+ 
348+template <typename T1, typename T2, const uint32_t IS_PAD>
349+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::Process()
350+{
351+ if (blockIdx_ >= usedCoreNum_) {
352+ return;
353+ }
354+ 
355+ int64_t curCoreProcessNum = (blockIdx_ + 1 == usedCoreNum_) ? tailCoreProcessNum_ : normalCoreProcessNum_;
356+ 
357+ for (int64_t loopNum = 0; loopNum < curCoreProcessNum; loopNum++) {
358+ LocalTensor<T1> maxValueLocal = maxValueQue_.AllocTensor<T1>();
359+ LocalTensor<T2> argmaxLocal = argmaxQue_.AllocTensor<T2>();
360+ __local_mem__ T1* maxValueAddr = (__local_mem__ T1*)maxValueLocal.GetPhyAddr();
361+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
362+ if (isSplitKernel_ == 1) {
363+ InitHelpBuf();
364+ for (hKernelIndex_ = 0; hKernelIndex_ < tilingData_->hKernelOuter; ++hKernelIndex_) {
365+ for (wKernelIndex_ = 0; wKernelIndex_ < tilingData_->wKernelOuter; ++wKernelIndex_) {
366+ ScalarCompute(loopNum);
367+ if (hInputActual_ <= 0 || wInputActual_ <= 0) {
368+ continue;
369+ }
370+ CopyIn();
371+ Compute(maxValueAddr, argmaxAddr);
372+ }
373+ }
374+ CopyResultToUb(maxValueAddr, argmaxAddr);
375+ } else {
376+ ScalarCompute(loopNum);
377+ CopyIn();
378+ Compute(maxValueAddr, argmaxAddr);
379+ }
380+ 
381+ maxValueQue_.EnQue(maxValueLocal);
382+ argmaxQue_.EnQue(argmaxLocal);
383+ CopyOut();
384+ }
385+}
386+ 
387+template <typename T1, typename T2, const uint32_t IS_PAD>
388+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::Compute(
389+ __local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal)
390+{
391+ LocalTensor<T1> xLocal = inputQue_.DeQue<T1>();
392+ __local_mem__ T1* xAddr = (__local_mem__ T1*)xLocal.GetPhyAddr();
393+ 
394+ if constexpr (IS_PAD == 1) {
395+ if (isSplitKernel_ == 1) {
396+ MaxPoolAndArgmaxV3VFForPad<true>(xAddr, maxValueLocal, argmaxLocal);
397+ } else {
398+ MaxPoolAndArgmaxV3VFForPad<false>(xAddr, maxValueLocal, argmaxLocal);
399+ }
400+ } else {
401+ if (isSplitKernel_ == 1) {
402+ MaxPoolAndArgmaxV3VF<true>(xAddr, maxValueLocal, argmaxLocal);
403+ } else {
404+ MaxPoolAndArgmaxV3VF<false>(xAddr, maxValueLocal, argmaxLocal);
405+ }
406+ }
407+ inputQue_.FreeTensor(xLocal);
408+}
409+ 
410+template <typename T1, typename T2, const uint32_t IS_PAD>
411+template <const bool IS_SPLIT_KERNEL>
412+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::MaxPoolAndArgmaxV3VF(
413+ __local_mem__ T1* xLocal, __local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal)
414+{
415+ __local_mem__ T1* maxValueHelp = (__local_mem__ T1*)helperTBuf_.Get<T1>().GetPhyAddr();
416+ __local_mem__ T2* argmaxHelp =
417+ (__local_mem__ T2*)helperTBuf_.Get<T2>().GetPhyAddr() + HELPER_BUFFER_SIZE_512 / sizeof(T1);
418+ 
419+ int64_t nOutputActual = nOutputActual_;
420+ int64_t hOutputActual = hOutputActual_;
421+ int64_t wOutputActual = wOutputActual_;
422+ int64_t cLoop = ops::CeilDiv(cOutputActual_, int64_t(VREG_LENGTH_DATA_NUM_T2));
423+ int64_t cOutputActualAlign = cOutputActualAlign_;
424+ int64_t cOutputActual = cOutputActual_;
425+ int64_t wInputActual = wInputActual_;
426+ int64_t hInputActual = hInputActual_;
427+ 
428+ int64_t hStride = hStride_;
429+ int64_t wStride = wStride_;
430+ int64_t hKernel = hKernelActual_;
431+ int64_t wKernel = wKernelActual_;
432+ int64_t wInput = wInput_;
433+ 
434+ // wh平面偏移
435+ int64_t indexWHplanOffset = indexWHplanOffset_;
436+ int64_t kernelStartInBaseBlock = 0;
437+ 
438+ int64_t cOutputTail = (cOutputActual_ % VREG_LENGTH_DATA_NUM_T2) == 0 ? VREG_LENGTH_DATA_NUM_T2 :
439+ (cOutputActual_ % VREG_LENGTH_DATA_NUM_T2);
440+ 
441+ for (uint16_t nIndex = 0; nIndex < nOutputActual; ++nIndex) {
442+ for (uint16_t cIndex = 0; cIndex < cLoop; ++cIndex) {
443+ uint32_t computeLoopTmp = cIndex == (cLoop - 1) ? cOutputTail : VREG_LENGTH_DATA_NUM_T2;
444+ uint32_t computeLoopVL = computeLoopTmp;
445+ uint32_t computeLoopVLT2 = computeLoopTmp;
446+ __VEC_SCOPE__
447+ {
448+ AscendC::MicroAPI::RegTensor<T1> vreg0;
449+ AscendC::MicroAPI::RegTensor<T1> vreg1;
450+ AscendC::MicroAPI::RegTensor<T2> argmaxUpdateVreg;
451+ AscendC::MicroAPI::RegTensor<T2> argmaxResVreg;
452+ AscendC::MicroAPI::MaskReg neMask;
453+ AscendC::MicroAPI::MaskReg gtMask;
454+ AscendC::MicroAPI::MaskReg gtMaskT2;
455+ AscendC::MicroAPI::MaskReg gtMaskT4;
456+ 
457+ AscendC::MicroAPI::MaskReg computeMaskT1 = AscendC::MicroAPI::UpdateMask<T1>(computeLoopVL);
458+ AscendC::MicroAPI::MaskReg computeMaskT2 = AscendC::MicroAPI::UpdateMask<T2>(computeLoopVLT2);
459+ for (uint16_t hIndex = 0; hIndex < static_cast<uint16_t>(hOutputActual); ++hIndex) {
460+ for (uint16_t wIndex = 0; wIndex < static_cast<uint16_t>(wOutputActual); ++wIndex) {
461+ int64_t outputOffset = nIndex * hOutputActual * wOutputActual * cOutputActualAlign +
462+ hIndex * wOutputActual * cOutputActualAlign +
463+ wIndex * cOutputActualAlign + cIndex * VREG_LENGTH_DATA_NUM_T2;
464+ // UB内偏移
465+ int64_t offsetC = cIndex * VREG_LENGTH_DATA_NUM_T2;
466+ int64_t offsetW = wIndex * wStride * cOutputActualAlign;
467+ int64_t offsetH = hIndex * hStride * wInputActual * cOutputActualAlign;
468+ int64_t offsetN = nIndex * hInputActual * wInputActual * cOutputActualAlign;
469+ int64_t startInUb = offsetC + offsetW + offsetH + offsetN;
470+ 
471+ // 起始点hw面偏移
472+ int64_t scopeHWOffset = indexWHplanOffset + hIndex * hStride * wInput + wIndex * wStride;
473+ 
474+ if constexpr (IS_SPLIT_KERNEL == 1) {
475+ AscendC::MicroAPI::DataCopy(vreg0, maxValueHelp);
476+ AscendC::MicroAPI::DataCopy(argmaxResVreg, argmaxHelp);
477+ } else {
478+ AscendC::MicroAPI::DataCopy(vreg0, xLocal + startInUb);
479+ AscendC::MicroAPI::Duplicate(argmaxResVreg, scopeHWOffset);
480+ }
481+ 
482+ for (uint16_t hKernelIdx = 0; hKernelIdx < static_cast<uint16_t>(hKernel); ++hKernelIdx) {
483+ for (uint16_t wKernelIdx = 0; wKernelIdx < static_cast<uint16_t>(wKernel); wKernelIdx++) {
484+ AscendC::MicroAPI::DataCopy(
485+ vreg1,
486+ xLocal + startInUb + (hKernelIdx * wInputActual + wKernelIdx) * cOutputActualAlign);
487+ AscendC::MicroAPI::Compare<T1, CMPMODE::GT>(gtMask, vreg1, vreg0, computeMaskT1);
488+ AscendC::MicroAPI::Compare<T1, CMPMODE::NE>(neMask, vreg1, vreg1, computeMaskT1);
489+ AscendC::MicroAPI::MaskOr(gtMask, gtMask, neMask, computeMaskT1);
490+ 
491+ Duplicate(argmaxUpdateVreg, scopeHWOffset + hKernelIdx * wInput + wKernelIdx);
492+ if constexpr (sizeof(T2) / sizeof(T1) == DIGIT_1) {
493+ AscendC::MicroAPI::Select(argmaxResVreg, argmaxUpdateVreg, argmaxResVreg, gtMask);
494+ } else if constexpr (sizeof(T2) / sizeof(T1) == DIGIT_2) {
495+ AscendC::MicroAPI::MaskUnPack(gtMaskT2, gtMask);
496+ AscendC::MicroAPI::Select(argmaxResVreg, argmaxUpdateVreg, argmaxResVreg, gtMaskT2);
497+ } else {
498+ AscendC::MicroAPI::MaskUnPack(gtMaskT2, gtMask);
499+ AscendC::MicroAPI::MaskUnPack(gtMaskT4, gtMaskT2);
500+ AscendC::MicroAPI::Select(argmaxResVreg, argmaxUpdateVreg, argmaxResVreg, gtMaskT4);
501+ }
502+ 
503+ AscendC::MicroAPI::Max(vreg0, vreg0, vreg1, computeMaskT1);
504+ }
505+ }
506+ 
507+ if constexpr (IS_SPLIT_KERNEL == 1) {
508+ AscendC::MicroAPI::DataCopy(maxValueHelp, vreg0, computeMaskT1);
509+ AscendC::MicroAPI::DataCopy(argmaxHelp, argmaxResVreg, computeMaskT2);
510+ } else {
511+ AscendC::MicroAPI::DataCopy(maxValueLocal + outputOffset, vreg0, computeMaskT1);
512+ AscendC::MicroAPI::DataCopy(argmaxLocal + outputOffset, argmaxResVreg, computeMaskT2);
513+ }
514+ }
515+ }
516+ }
517+ }
518+ }
519+}
520+ 
521+template <typename T1, typename T2, const uint32_t IS_PAD>
522+template <const bool IS_SPLIT_KERNEL>
523+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::MaxPoolAndArgmaxV3VFForPad(
524+ __local_mem__ T1* xLocal, __local_mem__ T1* maxValueLocal, __local_mem__ T2* argmaxLocal)
525+{
526+ __local_mem__ T1* maxValueHelp = (__local_mem__ T1*)helperTBuf_.Get<T1>().GetPhyAddr();
527+ __local_mem__ T2* argmaxHelp =
528+ (__local_mem__ T2*)helperTBuf_.Get<T2>().GetPhyAddr() + HELPER_BUFFER_SIZE_512 / sizeof(T1);
529+ int64_t nOutputActual = nOutputActual_;
530+ int64_t hOutputActual = hOutputActual_;
531+ int64_t wOutputActual = wOutputActual_;
532+ int64_t cLoop = ops::CeilDiv(cOutputActual_, int64_t(VREG_LENGTH_DATA_NUM_T2));
533+ int64_t cOutputActualAlign = cOutputActualAlign_;
534+ int64_t cOutputActual = cOutputActual_;
535+ int64_t wInputActual = wInputActual_;
536+ int64_t hInputActual = hInputActual_;
537+ int64_t hStride = hStride_;
538+ int64_t wStride = wStride_;
539+ int64_t wInput = wInput_;
540+ int64_t wInputActualPad = wInputActualPad_;
541+ int64_t indexWHplanOffset = indexWHplanOffset_;
542+ int64_t kernelStartInBaseBlock = 0;
543+ int64_t cOutputTail = (cOutputActual_ % VREG_LENGTH_DATA_NUM_T2) == 0 ? VREG_LENGTH_DATA_NUM_T2 :
544+ (cOutputActual_ % VREG_LENGTH_DATA_NUM_T2);
545+ for (uint16_t nIndex = 0; nIndex < nOutputActual; ++nIndex) {
546+ for (uint16_t hIndex = 0; hIndex < hOutputActual; ++hIndex) {
547+ for (uint16_t wIndex = 0; wIndex < wOutputActual; ++wIndex) {
548+ for (uint16_t cIndex = 0; cIndex < cLoop; ++cIndex) {
549+ uint32_t computeLoopTmp = cIndex == (cLoop - 1) ? cOutputTail : VREG_LENGTH_DATA_NUM_T2;
550+ uint32_t computeLoopVL = computeLoopTmp;
551+ uint32_t computeLoopVLT2 = computeLoopTmp;
552+ uint32_t correctHKernel = hKernelActual_;
553+ uint32_t correctWKernel = wKernelActual_;
554+ int64_t topOffset = hAxisIndex_ * hOutputInner_ * hStride_ + hIndex * hStride_ +
555+ hKernelIndex_ * hKernelInner_ - padTop_;
556+ correctHKernel = topOffset >= 0 ? correctHKernel : correctHKernel + topOffset;
557+ int64_t downOffset = hAxisIndex_ * hOutputInner_ * hStride_ + hIndex * hStride_ +
558+ +hKernelIndex_ * hKernelInner_ + hKernelActual_ - padTop_ - hInput_;
559+ correctHKernel = downOffset >= 0 ? correctHKernel - downOffset : correctHKernel;
560+ 
561+ int64_t leftOffset = wAxisIndex_ * wOutputInner_ * wStride_ + wIndex * wStride_ +
562+ wKernelIndex_ * wKernelInner_ - padLeft_;
563+ correctWKernel = leftOffset >= 0 ? correctWKernel : correctWKernel + leftOffset;
564+ int64_t rightOffset = wAxisIndex_ * wOutputInner_ * wStride_ + wIndex * wStride_ +
565+ wKernelIndex_ * wKernelInner_ + wKernelActual_ - padLeft_ - wInput_;
566+ correctWKernel = rightOffset >= 0 ? correctWKernel - rightOffset : correctWKernel;
567+ 
568+ int64_t outputOffset = nIndex * hOutputActual * wOutputActual * cOutputActualAlign +
569+ hIndex * wOutputActual * cOutputActualAlign + wIndex * cOutputActualAlign +
570+ cIndex * VREG_LENGTH_DATA_NUM_T2;
571+ // UB内偏移
572+ int64_t kernelTopOffsetOnLand = hIndex * hStride_;
573+ kernelTopOffsetOnLand = kernelTopOffsetOnLand >= baseBlockTopOffsetInOcean_ ?
574+ kernelTopOffsetOnLand :
575+ baseBlockTopOffsetInOcean_ + kernelTopOffsetOnLand;
576+ int64_t kernelLeftOffsetOnLand = wIndex * wStride_;
577+ kernelLeftOffsetOnLand = kernelLeftOffsetOnLand >= baseBlockLeftOffsetInOcean_ ?
578+ kernelLeftOffsetOnLand :
579+ baseBlockLeftOffsetInOcean_ + kernelLeftOffsetOnLand;
580+ 
581+ int64_t offsetC = cIndex * VREG_LENGTH_DATA_NUM_T2;
582+ int64_t offsetW = kernelLeftOffsetOnLand * cOutputActualAlign;
583+ int64_t offsetH = kernelTopOffsetOnLand * wInputActualPad_ * cOutputActualAlign;
584+ int64_t offsetN = nIndex * hInputActualPad_ * wInputActualPad_ * cOutputActualAlign;
585+ int64_t startInUb = offsetC + offsetW + offsetH + offsetN;
586+ 
587+ // w,h 输入偏移.
588+ int64_t topOffsetCoast = topOffset >= 0 ? topOffset : 0;
589+ int64_t leftOffsetCoast = leftOffset >= 0 ? leftOffset : 0;
590+ int64_t kernelStartArgmaxOffset = topOffsetCoast * wInput_ + leftOffsetCoast;
591+ 
592+ __VEC_SCOPE__
593+ {
594+ AscendC::MicroAPI::RegTensor<T1> vreg0;
595+ AscendC::MicroAPI::RegTensor<T1> vreg1;
596+ 
597+ AscendC::MicroAPI::RegTensor<T2> argmaxUpdateVreg;
598+ AscendC::MicroAPI::RegTensor<T2> argmaxResVreg;
599+ 
600+ AscendC::MicroAPI::RegTensor<uint32_t> startOffsetRegU32;
601+ AscendC::MicroAPI::RegTensor<uint32_t> separateOffsetRegU32;
602+ AscendC::MicroAPI::RegTensor<uint16_t> separateOffsetRegU16;
603+ AscendC::MicroAPI::MaskReg computeMaskT1 = AscendC::MicroAPI::UpdateMask<T1>(computeLoopVL);
604+ AscendC::MicroAPI::MaskReg computeMaskT2 = AscendC::MicroAPI::UpdateMask<T2>(computeLoopVLT2);
605+ AscendC::MicroAPI::MaskReg neMask;
606+ AscendC::MicroAPI::MaskReg gtMask;
607+ AscendC::MicroAPI::MaskReg gtMaskT2;
608+ AscendC::MicroAPI::MaskReg gtMaskT4;
609+ 
610+ if constexpr (IS_SPLIT_KERNEL == 1) {
611+ AscendC::MicroAPI::DataCopy(vreg0, maxValueHelp);
612+ AscendC::MicroAPI::DataCopy(argmaxResVreg, argmaxHelp);
613+ } else {
614+ AscendC::MicroAPI::DataCopy(vreg0, xLocal + startInUb);
615+ AscendC::MicroAPI::Duplicate(argmaxResVreg, kernelStartArgmaxOffset);
616+ }
617+ 
618+ for (uint16_t hKernelIdx = 0; hKernelIdx < static_cast<uint16_t>(correctHKernel);
619+ ++hKernelIdx) {
620+ for (uint16_t wKernelIdx = 0; wKernelIdx < static_cast<uint16_t>(correctWKernel);
621+ ++wKernelIdx) {
622+ AscendC::MicroAPI::DataCopy(
623+ vreg1, xLocal + startInUb +
624+ (hKernelIdx * wInputActualPad + wKernelIdx) * cOutputActualAlign);
625+ 
626+ AscendC::MicroAPI::Compare<T1, CMPMODE::GT>(gtMask, vreg1, vreg0, computeMaskT1);
627+ AscendC::MicroAPI::Compare<T1, CMPMODE::NE>(neMask, vreg1, vreg1, computeMaskT1);
628+ AscendC::MicroAPI::MaskOr(gtMask, gtMask, neMask, computeMaskT1);
629+ Duplicate(argmaxUpdateVreg, hKernelIdx * wInput + wKernelIdx + kernelStartArgmaxOffset);
630+ if constexpr (sizeof(T2) / sizeof(T1) == DIGIT_1) {
631+ AscendC::MicroAPI::Select(argmaxResVreg, argmaxUpdateVreg, argmaxResVreg, gtMask);
632+ } else if constexpr (sizeof(T2) / sizeof(T1) == DIGIT_2) {
633+ AscendC::MicroAPI::MaskUnPack(gtMaskT2, gtMask);
634+ AscendC::MicroAPI::Select(argmaxResVreg, argmaxUpdateVreg, argmaxResVreg, gtMaskT2);
635+ } else {
636+ AscendC::MicroAPI::MaskUnPack(gtMaskT2, gtMask);
637+ AscendC::MicroAPI::MaskUnPack(gtMaskT4, gtMaskT2);
638+ AscendC::MicroAPI::Select(argmaxResVreg, argmaxUpdateVreg, argmaxResVreg, gtMaskT4);
639+ }
640+ 
641+ AscendC::MicroAPI::Max(vreg0, vreg0, vreg1, computeMaskT1);
642+ }
643+ }
644+ if constexpr (IS_SPLIT_KERNEL == 1) {
645+ AscendC::MicroAPI::DataCopy(maxValueHelp, vreg0, computeMaskT1);
646+ AscendC::MicroAPI::DataCopy(argmaxHelp, argmaxResVreg, computeMaskT2);
647+ } else {
648+ AscendC::MicroAPI::DataCopy(maxValueLocal + outputOffset, vreg0, computeMaskT1);
649+ AscendC::MicroAPI::DataCopy(argmaxLocal + outputOffset, argmaxResVreg, computeMaskT2);
650+ }
651+ }
652+ }
653+ }
654+ }
655+ }
656+}
657+ 
658+template <typename T1, typename T2, const uint32_t IS_PAD>
659+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::CopyOut()
660+{
661+ LocalTensor<T1> maxValueLocal = maxValueQue_.DeQue<T1>();
662+ LocalTensor<T2> argmaxLocal = argmaxQue_.DeQue<T2>();
663+ 
664+ {
665+ DataCopyExtParams copyOutParamT;
666+ copyOutParamT.blockCount = nOutputActual_ * hOutputActual_ * wOutputActual_;
667+ copyOutParamT.blockLen = cOutputActual_ * sizeof(T1);
668+ copyOutParamT.srcStride = (cOutputActualAlign_ - cOutputActual_) * sizeof(T1) / BLOCK_SIZE;
669+ copyOutParamT.dstStride = 0;
670+ 
671+ DataCopyPad(
672+ yGm_[nOutputAxisOffset_ + hOutputAxisOffset_ + wOutputAxisOffset_ + cOutputAxisOffset_], maxValueLocal,
673+ copyOutParamT);
674+ }
675+ 
676+ {
677+ DataCopyExtParams copyOutParamT;
678+ copyOutParamT.blockCount = nOutputActual_ * hOutputActual_ * wOutputActual_;
679+ copyOutParamT.blockLen = cOutputActual_ * sizeof(T2);
680+ copyOutParamT.srcStride = (cOutputActualAlign_ - cOutputActual_) * sizeof(T2) / BLOCK_SIZE;
681+ copyOutParamT.dstStride = 0;
682+ DataCopyPad(
683+ argmaxGm_[nOutputAxisOffset_ + hOutputAxisOffset_ + wOutputAxisOffset_ + cOutputAxisOffset_], argmaxLocal,
684+ copyOutParamT);
685+ }
686+ maxValueQue_.FreeTensor(maxValueLocal);
687+ argmaxQue_.FreeTensor(argmaxLocal);
688+}
689+ 
690+template <typename T1, typename T2, const uint32_t IS_PAD>
691+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::FillPadNegVF(__local_mem__ T1* xLocalAddr)
692+{
693+ int32_t top = baseBlockTopOffsetInOcean_;
694+ int32_t left = baseBlockLeftOffsetInOcean_;
695+ int32_t right = baseBlockRightOffsetInOcean_;
696+ int32_t down = baseBlockDownOffsetInOcean_;
697+ int64_t wInputActual = wInputActual_;
698+ int64_t hInputActual = hInputActual_;
699+ int32_t cOutputActualAlign = cOutputActualAlign_;
700+ int32_t hInputActualAmend = (hOutputActual_ - 1) * hStride_ + hKernel_;
701+ int32_t wInputActualAmend = (wOutputActual_ - 1) * wStride_ + wKernel_;
702+ uint32_t computeSize = platform::GetVRegSize() / sizeof(T1);
703+ 
704+ uint32_t topCount = top * wInputActualAmend * cOutputActualAlign;
705+ uint16_t topRepeatTimes = (topCount + computeSize - 1) / computeSize;
706+ 
707+ int32_t leftSingleRowCount = left * cOutputActualAlign;
708+ uint16_t leftSingleRowRepeatTimes = (leftSingleRowCount + computeSize - 1) / computeSize;
709+ int32_t leftStartOffset = topCount;
710+ 
711+ int32_t rightSingleRowCount = right * cOutputActualAlign;
712+ uint16_t rightSingleRowRepeatTimes = (rightSingleRowCount + computeSize - 1) / computeSize;
713+ int32_t rightStartOffset = topCount + (wInputActual + left) * cOutputActualAlign;
714+ 
715+ uint32_t downCount = down * wInputActualAmend * cOutputActualAlign;
716+ uint16_t downRepeatTimes = (downCount + computeSize - 1) / computeSize;
717+ int32_t downStartOffset = (hInputActual + top) * wInputActualAmend * cOutputActualAlign;
718+ uint16_t nOutputActual = nOutputActual_;
719+ int32_t nStartOffset = hInputActualAmend * wInputActualAmend * cOutputActualAlign;
720+ __VEC_SCOPE__
721+ 
722+ {
723+ AscendC::MicroAPI::RegTensor<T1> negInfReg;
724+ DuplicateNegInfReg(negInfReg);
725+ for (uint16_t n = 0; n < nOutputActual; n++) {
726+ int32_t nOffset = n * nStartOffset;
727+ // top
728+ for (uint16_t i = 0; i < topRepeatTimes; i++) {
729+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<T1>(topCount);
730+ AscendC::MicroAPI::DataCopy(xLocalAddr + nOffset + i * computeSize, negInfReg, preg);
731+ }
732+ 
733+ // left
734+ for (uint16_t hIndex = 0; hIndex < static_cast<uint16_t>(hInputActual); hIndex++) {
735+ int32_t leftOffset = hIndex * wInputActualAmend * cOutputActualAlign + leftStartOffset;
736+ uint32_t leftCount = leftSingleRowCount;
737+ for (uint16_t i = 0; i < leftSingleRowRepeatTimes; i++) {
738+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<T1>(leftCount);
739+ AscendC::MicroAPI::DataCopy(xLocalAddr + nOffset + leftOffset + i * computeSize, negInfReg, preg);
740+ }
741+ }
742+ 
743+ // right
744+ for (uint16_t hIndex = 0; hIndex < static_cast<uint16_t>(hInputActual); hIndex++) {
745+ int32_t rightOffset = hIndex * wInputActualAmend * cOutputActualAlign + rightStartOffset;
746+ uint32_t rightCount = rightSingleRowCount;
747+ for (uint16_t i = 0; i < rightSingleRowRepeatTimes; i++) {
748+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<T1>(rightCount);
749+ AscendC::MicroAPI::DataCopy(xLocalAddr + nOffset + rightOffset + i * computeSize, negInfReg, preg);
750+ }
751+ }
752+ 
753+ // down
754+ for (uint16_t i = 0; i < downRepeatTimes; i++) {
755+ AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<T1>(downCount);
756+ AscendC::MicroAPI::DataCopy(xLocalAddr + nOffset + downStartOffset + i * computeSize, negInfReg, preg);
757+ }
758+ }
759+ }
760+}
761+ 
762+template <typename T1, typename T2, const uint32_t IS_PAD>
763+__aicore__ inline void MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>::CopyIn()
764+{
765+ LocalTensor<T1> xLocal = inputQue_.AllocTensor<T1>();
766+ __local_mem__ T1* xLocalAddr = (__local_mem__ T1*)xLocal.GetPhyAddr();
767+ 
768+ int64_t nOutputActual = nOutputActual_;
769+ int64_t hInputWithPad = (hInputActual_ + baseBlockTopOffsetInOcean_ + baseBlockDownOffsetInOcean_);
770+ int64_t wInputWithPad = (wInputActual_ + baseBlockLeftOffsetInOcean_ + baseBlockRightOffsetInOcean_);
771+ int64_t cOutputActualAlign = cOutputActualAlign_;
772+ int64_t xGmOffset = nInputAxisOffset_ + hInputAxisOffset_ + wInputAxisOffset_ + cInputAxisOffset_;
773+ 
774+ LoopModeParams loopParams;
775+ loopParams.loop2Size = nOutputActual_;
776+ loopParams.loop2SrcStride = hInput_ * wInput_ * cInput_ * sizeof(T1);
777+ loopParams.loop2DstStride = hInputWithPad * wInputWithPad * cOutputActualAlign_ * sizeof(T1);
778+ 
779+ loopParams.loop1Size = hInputActual_;
780+ loopParams.loop1SrcStride = wInput_ * cInput_ * sizeof(T1);
781+ loopParams.loop1DstStride = wInputWithPad * cOutputActualAlign_ * sizeof(T1);
782+ 
783+ SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB);
784+ DataCopyExtParams copyExtParams;
785+ copyExtParams.blockCount = wInputActual_;
786+ copyExtParams.blockLen = cOutputActual_ * sizeof(T1);
787+ copyExtParams.srcStride = (cInput_ - cOutputActual_) * sizeof(T1);
788+ copyExtParams.dstStride = 0;
789+ DataCopyPadExtParams<T1> copyPadExtparams;
790+ copyPadExtparams.isPad = false;
791+ DataCopyPad(xLocal[xLocalUbOffset], xGm_[xGmOffset], copyExtParams, copyPadExtparams);
792+ ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
793+ 
794+ inputQue_.EnQue(xLocal);
795+}
796+ 
797+} // namespace MaxPoolWithArgmaxV3NHWC
798+#endif // MAX_POOL_WITH_ARGMAX_V3_NHWC_BIG_C_H_
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_nhwc_small_c.h+775-0
@@ -0,0 +1,775 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file max_pool_with_argmax_v3_nhwc_small_c_.h
13+ * \brief
14+ */
15+ 
16+#ifndef MAX_POOL_WITH_ARGMAX_V3_NHWC_SMALL_C__H_
17+#define MAX_POOL_WITH_ARGMAX_V3_NHWC_SMALL_C__H_
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+#include "../inc/platform.h"
22+#include "max_pool_with_argmax_v3_base.h"
23+#include "max_pool_with_argmax_v3_nhwc_big_c.h"
24+ 
25+namespace MaxPoolWithArgmaxV3SmallCNameSpace {
26+using namespace AscendC;
27+ 
28+constexpr int64_t THREE = 3;
29+constexpr int64_t DOUBLE = 2;
30+ 
31+constexpr AscendC::MicroAPI::CastTrait castTraitU32U16 = {
32+ AscendC::MicroAPI::RegLayout::ZERO,
33+ AscendC::MicroAPI::SatMode::NO_SAT,
34+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
35+ AscendC::RoundMode::CAST_RINT,
36+};
37+ 
38+template <typename T1, typename T2, const uint32_t IS_PAD = 0>
39+__aicore__ inline void MaxPoolWithArgMaxV3GatherImpl(
40+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr, uint16_t kH, uint16_t kW,
41+ uint32_t rowStrideInUb, uint16_t alignedC, int32_t gatterIndexOffset, MicroAPI::RegTensor<uint32_t>& gatterStartIdx,
42+ int32_t count, MicroAPI::RegTensor<T2>& argmaxHStart, MicroAPI::RegTensor<T2>& argmaxWStart, int32_t argmaxHOffset,
43+ int32_t argmaxWOffset, MicroAPI::RegTensor<uint32_t>& scatterStartIdx, int32_t scatterOffset, int32_t padH,
44+ int32_t padW, int32_t wInput)
45+{
46+ MicroAPI::RegTensor<T1> vd0;
47+ MicroAPI::RegTensor<T1> vd1;
48+ MicroAPI::RegTensor<uint32_t> gatterIndexReg;
49+ MicroAPI::RegTensor<uint32_t> scatterIndexReg;
50+ MicroAPI::RegTensor<uint16_t> gatterIdxU16Reg;
51+ MicroAPI::RegTensor<uint16_t> scatterIdxU16Reg;
52+ 
53+ AscendC::MicroAPI::RegTensor<T2> argmaxUpdateHVreg;
54+ AscendC::MicroAPI::RegTensor<T2> argmaxUpdateWVreg;
55+ AscendC::MicroAPI::RegTensor<T2> argmaxHRes;
56+ AscendC::MicroAPI::RegTensor<T2> argmaxWRes;
57+ 
58+ AscendC::MicroAPI::MaskReg neMask;
59+ AscendC::MicroAPI::MaskReg gtMask;
60+ AscendC::MicroAPI::MaskReg gtMaskT2;
61+ AscendC::MicroAPI::MaskReg gtMaskT4;
62+ 
63+ DuplicateNegInfReg<T1>(vd0);
64+ 
65+ uint32_t numU32 = count;
66+ uint32_t numT1 = count;
67+ uint32_t numT2 = count;
68+ 
69+ MicroAPI::MaskReg computeT1 = MicroAPI::UpdateMask<T1>(numT1);
70+ MicroAPI::MaskReg computeT2 = MicroAPI::UpdateMask<T2>(numT2);
71+ MicroAPI::MaskReg computeU32 = MicroAPI::UpdateMask<uint32_t>(numU32);
72+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
73+ MicroAPI::MaskReg maskAllT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
74+ 
75+ MicroAPI::Adds(argmaxHRes, argmaxHStart, argmaxHOffset, computeT2);
76+ MicroAPI::Adds(argmaxWRes, argmaxWStart, argmaxWOffset, computeT2);
77+ 
78+ for (uint16_t hIdx = 0; hIdx < kH; hIdx++) {
79+ int32_t hKernelOffset = hIdx * rowStrideInUb;
80+ int32_t argmaxHKernelOffset = hIdx + argmaxHOffset;
81+ 
82+ for (uint16_t wIdx = 0; wIdx < kW; wIdx++) {
83+ int32_t wKernelOffset = wIdx * alignedC;
84+ int32_t argmaxWKernelOffset = wIdx + argmaxWOffset;
85+ 
86+ int32_t gatterIndexOffsetTotal = gatterIndexOffset + hKernelOffset + wKernelOffset;
87+ MicroAPI::Adds(gatterIndexReg, gatterStartIdx, gatterIndexOffsetTotal, computeU32);
88+ 
89+ if constexpr (std::is_same<T1, float>::value) {
90+ MicroAPI::DataCopyGather(vd1, xAddr, gatterIndexReg, computeT1);
91+ } else {
92+ AscendC::MicroAPI::Cast<uint16_t, uint32_t, castTraitU32U16>(
93+ gatterIdxU16Reg, gatterIndexReg, computeU32);
94+ AscendC::MicroAPI::Pack(gatterIdxU16Reg, (AscendC::MicroAPI::RegTensor<uint32_t>&)gatterIdxU16Reg);
95+ AscendC::MicroAPI::DataCopyGather(vd1, xAddr, gatterIdxU16Reg, computeT1);
96+ }
97+ 
98+ AscendC::MicroAPI::Compare<T1, CMPMODE::GT>(gtMask, vd1, vd0, computeT1);
99+ AscendC::MicroAPI::Compare<T1, CMPMODE::NE>(neMask, vd1, vd1, computeT1);
100+ AscendC::MicroAPI::MaskOr(gtMask, gtMask, neMask, computeT1);
101+ 
102+ MicroAPI::Adds(argmaxUpdateHVreg, argmaxHStart, argmaxHKernelOffset, computeT2);
103+ MicroAPI::Adds(argmaxUpdateWVreg, argmaxWStart, argmaxWKernelOffset, computeT2);
104+ if constexpr (sizeof(T2) / sizeof(T1) == 1) {
105+ AscendC::MicroAPI::Select(argmaxHRes, argmaxUpdateHVreg, argmaxHRes, gtMask);
106+ AscendC::MicroAPI::Select(argmaxWRes, argmaxUpdateWVreg, argmaxWRes, gtMask);
107+ } else if constexpr (sizeof(T2) / sizeof(T1) == DOUBLE) {
108+ AscendC::MicroAPI::MaskUnPack(gtMaskT2, gtMask);
109+ AscendC::MicroAPI::Select(argmaxHRes, argmaxUpdateHVreg, argmaxHRes, gtMaskT2);
110+ AscendC::MicroAPI::Select(argmaxWRes, argmaxUpdateWVreg, argmaxWRes, gtMaskT2);
111+ } else {
112+ AscendC::MicroAPI::MaskUnPack(gtMaskT2, gtMask);
113+ AscendC::MicroAPI::MaskUnPack(gtMaskT4, gtMaskT2);
114+ AscendC::MicroAPI::Select(argmaxHRes, argmaxUpdateHVreg, argmaxHRes, gtMaskT4);
115+ AscendC::MicroAPI::Select(argmaxWRes, argmaxUpdateWVreg, argmaxWRes, gtMaskT4);
116+ }
117+ 
118+ MicroAPI::Max(vd0, vd1, vd0, computeT1);
119+ }
120+ }
121+ 
122+ if constexpr (IS_PAD == 1) {
123+ // 修正argmax
124+ MicroAPI::Adds(argmaxHRes, argmaxHRes, -padH, computeT2);
125+ MicroAPI::Adds(argmaxWRes, argmaxWRes, -padW, computeT2);
126+ 
127+ AscendC::MicroAPI::MaskReg hMask;
128+ AscendC::MicroAPI::MaskReg wMask;
129+ MicroAPI::RegTensor<T2> argmaxZero;
130+ AscendC::MicroAPI::Duplicate(argmaxZero, 0);
131+ 
132+ AscendC::MicroAPI::Compare<T2, CMPMODE::GE>(hMask, argmaxHRes, argmaxZero, computeT2);
133+ AscendC::MicroAPI::Select(argmaxHRes, argmaxHRes, argmaxZero, hMask);
134+ AscendC::MicroAPI::Compare<T2, CMPMODE::GE>(wMask, argmaxWRes, argmaxZero, computeT2);
135+ AscendC::MicroAPI::Select(argmaxWRes, argmaxWRes, argmaxZero, wMask);
136+ }
137+ 
138+ MicroAPI::RegTensor<T2> argmaxRes;
139+ MicroAPI::Muls(argmaxRes, argmaxHRes, wInput, computeT2);
140+ MicroAPI::Add(argmaxRes, argmaxRes, argmaxWRes, computeT2);
141+ 
142+ MicroAPI::Adds(scatterIndexReg, scatterStartIdx, scatterOffset, computeU32);
143+ 
144+ AscendC::MicroAPI::DataCopyScatter(argmaxAddr, argmaxRes, scatterIndexReg, computeT2);
145+ if constexpr (std::is_same<T1, float>::value) {
146+ AscendC::MicroAPI::DataCopyScatter(maxValueAddr, vd0, scatterIndexReg, computeT1);
147+ } else {
148+ AscendC::MicroAPI::Cast<uint16_t, uint32_t, castTraitU32U16>(scatterIdxU16Reg, scatterIndexReg, computeU32);
149+ AscendC::MicroAPI::Pack(scatterIdxU16Reg, (AscendC::MicroAPI::RegTensor<uint32_t>&)scatterIdxU16Reg);
150+ AscendC::MicroAPI::DataCopyScatter(maxValueAddr, vd0, scatterIdxU16Reg, computeT1);
151+ }
152+}
153+ 
154+template <typename T1, typename T2, const uint32_t IS_PAD = 0>
155+class MaxPoolWithArgmaxV3SmallC : public MaxPoolWithArgmaxV3NHWC::MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD> {
156+public:
157+ __aicore__ inline MaxPoolWithArgmaxV3SmallC(TPipe* pipe, const MaxPoolWithArgmaxV3NhwcTilingData* tiling)
158+ : MaxPoolWithArgmaxV3NHWC::MaxPoolWithArgmaxV3NhwCKernel<T1, T2, IS_PAD>(pipe, tiling){};
159+ __aicore__ inline void MaxPoolWithArgmaxV3SmallCProcess();
160+ __aicore__ inline void MaxPoolWithArgmaxV3SmallCCompute();
161+ __aicore__ inline void ComputeSingleRow(
162+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr);
163+ __aicore__ inline void ComputeMultiRow(
164+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr);
165+ __aicore__ inline void ComputeMultiRowForInt64(
166+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr,
167+ __local_mem__ uint32_t* helpAddr);
168+ __aicore__ inline void ComputeMultiBatch(
169+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr);
170+ __aicore__ inline void ComputeMultiBatchForInt64(
171+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr,
172+ __local_mem__ uint32_t* helpAddr);
173+ 
174+public:
175+ constexpr static uint32_t V_REG_SIZE = platform::GetVRegSize();
176+};
177+ 
178+template <typename T1, typename T2, const uint32_t IS_PAD>
179+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::MaxPoolWithArgmaxV3SmallCProcess()
180+{
181+ if (this->blockIdx_ >= this->usedCoreNum_) {
182+ return;
183+ }
184+ 
185+ int64_t curCoreProcessNum =
186+ (this->blockIdx_ + 1 == this->usedCoreNum_) ? this->tailCoreProcessNum_ : this->normalCoreProcessNum_;
187+ for (int64_t loopNum = 0; loopNum < curCoreProcessNum; loopNum++) {
188+ this->ScalarCompute(loopNum);
189+ this->CopyIn();
190+ MaxPoolWithArgmaxV3SmallCCompute();
191+ this->CopyOut();
192+ }
193+}
194+ 
195+template <typename T1, typename T2, const uint32_t IS_PAD>
196+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::MaxPoolWithArgmaxV3SmallCCompute()
197+{
198+ LocalTensor<T1> xLocal = this->inputQue_.template DeQue<T1>();
199+ LocalTensor<T1> maxValueLocal = this->maxValueQue_.template AllocTensor<T1>();
200+ LocalTensor<T2> argmaxLocal = this->argmaxQue_.template AllocTensor<T2>();
201+ LocalTensor<uint32_t> helpTensor = this->helperTBuf_.template Get<uint32_t>();
202+ 
203+ __local_mem__ T1* xAddr = (__local_mem__ T1*)xLocal.GetPhyAddr();
204+ __local_mem__ T1* maxValueAddr = (__local_mem__ T1*)maxValueLocal.GetPhyAddr();
205+ __local_mem__ T2* argmaxAddr = (__local_mem__ T2*)argmaxLocal.GetPhyAddr();
206+ __local_mem__ uint32_t* helpAddr = (__local_mem__ uint32_t*)helpTensor.GetPhyAddr();
207+ 
208+ if constexpr (IS_PAD == 1) {
209+ this->FillPadNegVF(xAddr);
210+ }
211+ 
212+ uint16_t repeatElm = platform::GetVRegSize() / sizeof(T2);
213+ if (repeatElm >= DOUBLE * this->hOutputActual_ * this->wOutputActual_ * this->cInput_) {
214+ if constexpr (std::is_same<T2, int64_t>::value) { // 拼nhw，并发nhwc
215+ ComputeMultiBatchForInt64(xAddr, maxValueAddr, argmaxAddr, helpAddr);
216+ } else {
217+ ComputeMultiBatch(xAddr, maxValueAddr, argmaxAddr);
218+ }
219+ } else if (repeatElm >= DOUBLE * this->wOutputActual_ * this->cInput_) {
220+ if constexpr (std::is_same<T2, int64_t>::value) { // 拼hw，并发hwc
221+ ComputeMultiRowForInt64(xAddr, maxValueAddr, argmaxAddr, helpAddr);
222+ } else {
223+ ComputeMultiRow(xAddr, maxValueAddr, argmaxAddr);
224+ }
225+ } else { // 拼w，并发wc
226+ ComputeSingleRow(xAddr, maxValueAddr, argmaxAddr);
227+ }
228+ 
229+ this->inputQue_.template FreeTensor(xLocal);
230+ this->maxValueQue_.template EnQue(maxValueLocal);
231+ this->argmaxQue_.template EnQue(argmaxLocal);
232+ this->helperTBuf_.template FreeTensor(helpTensor);
233+}
234+ 
235+template <typename T1, typename T2, const uint32_t IS_PAD>
236+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::ComputeMultiBatch(
237+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr)
238+{
239+ uint16_t kH = static_cast<uint16_t>(this->hKernel_);
240+ uint16_t kW = static_cast<uint16_t>(this->wKernel_);
241+ uint16_t hStride = static_cast<uint16_t>(this->hStride_);
242+ uint16_t padH = static_cast<uint16_t>(this->padTop_);
243+ uint16_t padW = static_cast<uint16_t>(this->padLeft_);
244+ int32_t wInput = static_cast<int32_t>(this->wInput_);
245+ uint16_t alignedC = static_cast<uint16_t>(this->cOutputActualAlign_);
246+ 
247+ constexpr uint16_t repeatElm = platform::GetVRegSize() / sizeof(T2);
248+ uint16_t nFactor = static_cast<uint16_t>(repeatElm / (this->hOutputActual_ * this->wOutputActual_ * this->cInput_));
249+ nFactor = nFactor > this->nOutputActual_ ? this->nOutputActual_ : nFactor;
250+ uint16_t loopN = static_cast<uint16_t>(this->nOutputActual_ / nFactor);
251+ uint16_t tailN = static_cast<uint16_t>(this->nOutputActual_ - loopN * nFactor);
252+ 
253+ int32_t hInputActualAmend = (this->hOutputActual_ - 1) * this->hStride_ + this->hKernel_;
254+ int32_t wInputActualAmend = (this->wOutputActual_ - 1) * this->wStride_ + this->wKernel_;
255+ int32_t ubNumHWC = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
256+ 
257+ int32_t wBlockArgmaxOffset = this->wAxisIndex_ * this->wStride_ * this->wOutputInner_;
258+ int32_t hBlockArgmaxOffset = this->hAxisIndex_ * this->hStride_ * this->hOutputInner_;
259+ 
260+ uint32_t oneLoopElements = static_cast<uint32_t>(
261+ nFactor * this->hOutputActual_ * this->wOutputActual_ * this->cInput_); // 一次循环处理的输出元素
262+ uint32_t tailLoopElements =
263+ static_cast<uint32_t>(tailN * this->hOutputActual_ * this->wOutputActual_ * this->cInput_); // 尾循环处理输出
264+ uint32_t rowStrideInUb = static_cast<uint32_t>(wInputActualAmend * this->cOutputActualAlign_);
265+ uint32_t oneNOutScatterElements =
266+ static_cast<uint32_t>(this->hOutputActual_ * this->wOutputActual_ * this->cOutputActualAlign_);
267+ 
268+ int32_t num1D = this->cInput_;
269+ int32_t rate2D = this->wStride_ * this->cOutputActualAlign_;
270+ int32_t num2D = this->wOutputActual_ * this->cInput_;
271+ int32_t rate3D = this->hStride_ * wInputActualAmend * this->cOutputActualAlign_;
272+ int32_t num3D = this->hOutputActual_ * this->wOutputActual_ * this->cInput_;
273+ int32_t rate4D = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
274+ 
275+ T2 argNum1D = this->cInput_;
276+ T2 argRate2D = this->wStride_;
277+ T2 argNum2D = this->wOutputActual_ * this->cInput_;
278+ T2 argNum3D = this->hOutputActual_ * this->wOutputActual_ * this->cInput_;
279+ int32_t scatterIdxNum1D = this->cInput_;
280+ int32_t scatterIdxRate2D = this->cOutputActualAlign_;
281+ 
282+ // 产生N的输出索引的索引
283+ __VEC_SCOPE__
284+ {
285+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
286+ MicroAPI::RegTensor<uint32_t> gatterNStartIdx;
287+ MicroAPI::RegTensor<T2> argmaxHStart;
288+ MicroAPI::RegTensor<T2> argmaxWStart;
289+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
290+ MicroAPI::RegTensor<uint32_t> scatterNStartIdx;
291+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
292+ MicroAPI::MaskReg maskAllT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
293+ 
294+ GenGatterIndex4D<int32_t>(
295+ (MicroAPI::RegTensor<int32_t>&)gatterStartIdx, rate4D, num3D, rate3D, num2D, rate2D, num1D);
296+ GenGatterIndex4D<T2>(argmaxWStart, 0, argNum3D, 0, argNum2D, argRate2D, argNum1D, 0);
297+ GenGatterIndex3D<T2>(argmaxHStart, 0, argNum3D, static_cast<T2>(hStride), argNum2D, 0);
298+ GenGatterIndex2D<int32_t>((MicroAPI::RegTensor<int32_t>&)scatterStartIdx, scatterIdxRate2D, scatterIdxNum1D);
299+ 
300+ for (uint16_t nIdex = 0; nIdex < loopN; nIdex++) {
301+ // 校正N
302+ MicroAPI::Adds(gatterNStartIdx, gatterStartIdx, nIdex * nFactor * ubNumHWC, maskAllU32);
303+ MicroAPI::Adds(scatterNStartIdx, scatterStartIdx, nIdex * nFactor * oneNOutScatterElements, maskAllU32);
304+ 
305+ int32_t gatterIndexOffset = 0;
306+ int32_t argmaxHOffset = hBlockArgmaxOffset;
307+ int32_t argmaxWOffset = wBlockArgmaxOffset;
308+ int32_t scatterOffset = 0;
309+ 
310+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
311+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterNStartIdx,
312+ oneLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterNStartIdx,
313+ scatterOffset, padH, padW, wInput);
314+ }
315+ 
316+ // tail N
317+ MicroAPI::Adds(gatterNStartIdx, gatterStartIdx, loopN * nFactor * ubNumHWC, maskAllU32);
318+ MicroAPI::Adds(scatterNStartIdx, scatterStartIdx, loopN * nFactor * oneNOutScatterElements, maskAllU32);
319+ 
320+ int32_t gatterIndexOffset = 0;
321+ int32_t argmaxHOffset = hBlockArgmaxOffset;
322+ int32_t argmaxWOffset = wBlockArgmaxOffset;
323+ int32_t scatterOffset = 0;
324+ 
325+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
326+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterNStartIdx,
327+ tailLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterNStartIdx, scatterOffset,
328+ padH, padW, wInput);
329+ }
330+}
331+ 
332+template <typename T1, typename T2, const uint32_t IS_PAD>
333+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::ComputeMultiBatchForInt64(
334+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr,
335+ __local_mem__ uint32_t* helpAddr)
336+{
337+ uint16_t kH = static_cast<uint16_t>(this->hKernel_);
338+ uint16_t kW = static_cast<uint16_t>(this->wKernel_);
339+ uint16_t hStride = static_cast<uint16_t>(this->hStride_);
340+ uint16_t padH = static_cast<uint16_t>(this->padTop_);
341+ uint16_t padW = static_cast<uint16_t>(this->padLeft_);
342+ int32_t wInput = static_cast<int32_t>(this->wInput_);
343+ uint16_t alignedC = static_cast<uint16_t>(this->cOutputActualAlign_);
344+ 
345+ constexpr uint16_t repeatElm = platform::GetVRegSize() / sizeof(T2);
346+ uint16_t nFactor = static_cast<uint16_t>(repeatElm / (this->hOutputActual_ * this->wOutputActual_ * this->cInput_));
347+ nFactor = nFactor > this->nOutputActual_ ? this->nOutputActual_ : nFactor;
348+ uint16_t loopN = static_cast<uint16_t>(this->nOutputActual_ / nFactor);
349+ uint16_t tailN = static_cast<uint16_t>(this->nOutputActual_ - loopN * nFactor);
350+ 
351+ int32_t hInputActualAmend = (this->hOutputActual_ - 1) * this->hStride_ + this->hKernel_;
352+ int32_t wInputActualAmend = (this->wOutputActual_ - 1) * this->wStride_ + this->wKernel_;
353+ int32_t ubNumHWC = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
354+ 
355+ int32_t wBlockArgmaxOffset = this->wAxisIndex_ * this->wStride_ * this->wOutputInner_;
356+ int32_t hBlockArgmaxOffset = this->hAxisIndex_ * this->hStride_ * this->hOutputInner_;
357+ 
358+ uint32_t oneLoopElements = static_cast<uint32_t>(
359+ nFactor * this->hOutputActual_ * this->wOutputActual_ * this->cInput_); // 一次循环处理的输出元素
360+ uint32_t tailLoopElements =
361+ static_cast<uint32_t>(tailN * this->hOutputActual_ * this->wOutputActual_ * this->cInput_); // 尾循环处理输出
362+ uint32_t rowStrideInUb = static_cast<uint32_t>(wInputActualAmend * this->cOutputActualAlign_);
363+ uint32_t oneNOutScatterElements =
364+ static_cast<uint32_t>(this->hOutputActual_ * this->wOutputActual_ * this->cOutputActualAlign_);
365+ 
366+ int32_t num1D = this->cInput_;
367+ int32_t rate2D = this->wStride_ * this->cOutputActualAlign_;
368+ int32_t num2D = this->wOutputActual_ * this->cInput_;
369+ int32_t rate3D = this->hStride_ * wInputActualAmend * this->cOutputActualAlign_;
370+ int32_t num3D = this->hOutputActual_ * this->wOutputActual_ * this->cInput_;
371+ int32_t rate4D = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
372+ 
373+ T2 argNum1D = this->cInput_;
374+ T2 argRate2D = this->wStride_;
375+ T2 argNum2D = this->wOutputActual_ * this->cInput_;
376+ T2 argNum3D = this->hOutputActual_ * this->wOutputActual_ * this->cInput_;
377+ int32_t scatterIdxNum1D = this->cInput_;
378+ int32_t scatterIdxRate2D = this->cOutputActualAlign_;
379+ 
380+ // 产生N的输出索引的索引
381+ __VEC_SCOPE__
382+ {
383+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
384+ MicroAPI::RegTensor<T2> argmaxHStart;
385+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
386+ 
387+ GenGatterIndex4D<int32_t>(
388+ (MicroAPI::RegTensor<int32_t>&)gatterStartIdx, rate4D, num3D, rate3D, num2D, rate2D, num1D);
389+ GenGatterIndex3D<T2>(argmaxHStart, 0, argNum3D, hStride, argNum2D, 0);
390+ 
391+ AscendC::MicroAPI::DataCopy(helpAddr, gatterStartIdx, maskAllU32);
392+ AscendC::MicroAPI::DataCopy(
393+ helpAddr + V_REG_SIZE / sizeof(uint32_t), (MicroAPI::RegTensor<uint32_t>&)argmaxHStart, maskAllU32);
394+ }
395+ 
396+ __VEC_SCOPE__
397+ {
398+ MicroAPI::RegTensor<T2> argmaxWStart;
399+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
400+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
401+ 
402+ GenGatterIndex4D<T2>(argmaxWStart, 0, argNum3D, 0, argNum2D, argRate2D, argNum1D, 0);
403+ GenGatterIndex2D<int32_t>((MicroAPI::RegTensor<int32_t>&)scatterStartIdx, scatterIdxRate2D, scatterIdxNum1D);
404+ 
405+ AscendC::MicroAPI::DataCopy(
406+ helpAddr + V_REG_SIZE / sizeof(uint32_t) * DOUBLE, (MicroAPI::RegTensor<uint32_t>&)argmaxWStart,
407+ maskAllU32);
408+ AscendC::MicroAPI::DataCopy(helpAddr + V_REG_SIZE / sizeof(uint32_t) * THREE, scatterStartIdx, maskAllU32);
409+ }
410+ 
411+ __VEC_SCOPE__
412+ {
413+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
414+ MicroAPI::RegTensor<uint32_t> gatterNStartIdx;
415+ MicroAPI::RegTensor<T2> argmaxHStart;
416+ MicroAPI::RegTensor<T2> argmaxWStart;
417+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
418+ MicroAPI::RegTensor<uint32_t> scatterNStartIdx;
419+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
420+ MicroAPI::MaskReg maskAllT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
421+ 
422+ AscendC::MicroAPI::DataCopy(gatterStartIdx, helpAddr);
423+ AscendC::MicroAPI::DataCopy(
424+ (MicroAPI::RegTensor<uint32_t>&)argmaxHStart, helpAddr + V_REG_SIZE / sizeof(uint32_t));
425+ AscendC::MicroAPI::DataCopy(
426+ (MicroAPI::RegTensor<uint32_t>&)argmaxWStart, helpAddr + V_REG_SIZE / sizeof(uint32_t) * DOUBLE);
427+ AscendC::MicroAPI::DataCopy(scatterStartIdx, helpAddr + V_REG_SIZE / sizeof(uint32_t) * THREE);
428+ 
429+ for (uint16_t nIdex = 0; nIdex < loopN; nIdex++) {
430+ // 校正N
431+ MicroAPI::Adds(gatterNStartIdx, gatterStartIdx, nIdex * nFactor * ubNumHWC, maskAllU32);
432+ MicroAPI::Adds(scatterNStartIdx, scatterStartIdx, nIdex * nFactor * oneNOutScatterElements, maskAllU32);
433+ 
434+ int32_t gatterIndexOffset = 0;
435+ int32_t argmaxHOffset = hBlockArgmaxOffset;
436+ int32_t argmaxWOffset = wBlockArgmaxOffset;
437+ int32_t scatterOffset = 0;
438+ 
439+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
440+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterNStartIdx,
441+ oneLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterNStartIdx,
442+ scatterOffset, padH, padW, wInput);
443+ }
444+ 
445+ // tail N
446+ MicroAPI::Adds(gatterNStartIdx, gatterStartIdx, loopN * nFactor * ubNumHWC, maskAllU32);
447+ MicroAPI::Adds(scatterNStartIdx, scatterStartIdx, loopN * nFactor * oneNOutScatterElements, maskAllU32);
448+ 
449+ int32_t gatterIndexOffset = 0;
450+ int32_t argmaxHOffset = hBlockArgmaxOffset;
451+ int32_t argmaxWOffset = wBlockArgmaxOffset;
452+ int32_t scatterOffset = 0;
453+ 
454+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
455+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterNStartIdx,
456+ tailLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterNStartIdx, scatterOffset,
457+ padH, padW, wInput);
458+ }
459+}
460+ 
461+template <typename T1, typename T2, const uint32_t IS_PAD>
462+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::ComputeMultiRow(
463+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr)
464+{
465+ uint16_t kH = static_cast<uint16_t>(this->hKernel_);
466+ uint16_t kW = static_cast<uint16_t>(this->wKernel_);
467+ uint16_t hStride = static_cast<uint16_t>(this->hStride_);
468+ uint16_t padH = static_cast<uint16_t>(this->padTop_);
469+ uint16_t padW = static_cast<uint16_t>(this->padLeft_);
470+ int32_t wInput = static_cast<int32_t>(this->wInput_);
471+ uint16_t wOutputActual = static_cast<uint16_t>(this->wOutputActual_);
472+ uint16_t alignedC = static_cast<uint16_t>(this->cOutputActualAlign_);
473+ 
474+ uint16_t loopN = static_cast<uint16_t>(this->nOutputActual_);
475+ constexpr uint32_t repeatElm = platform::GetVRegSize() / sizeof(T2);
476+ uint16_t hFactor = static_cast<uint16_t>(repeatElm / (this->wOutputActual_ * this->cInput_));
477+ hFactor = hFactor > this->hOutputActual_ ? this->hOutputActual_ : hFactor;
478+ uint16_t loopH = static_cast<uint16_t>(this->hOutputActual_ / hFactor);
479+ uint16_t tailH = static_cast<uint16_t>(this->hOutputActual_ - loopH * hFactor);
480+ 
481+ int32_t hInputActualAmend = (this->hOutputActual_ - 1) * this->hStride_ + this->hKernel_;
482+ int32_t wInputActualAmend = (this->wOutputActual_ - 1) * this->wStride_ + this->wKernel_;
483+ int32_t ubNumHWC = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
484+ 
485+ int32_t wBlockArgmaxOffset = this->wAxisIndex_ * this->wStride_ * this->wOutputInner_;
486+ int32_t hBlockArgmaxOffset = this->hAxisIndex_ * this->hStride_ * this->hOutputInner_;
487+ 
488+ uint32_t oneLoopStrideH =
489+ static_cast<uint32_t>(hFactor * this->hStride_ * wInputActualAmend * this->cOutputActualAlign_);
490+ uint32_t oneLoopElements = static_cast<uint32_t>(hFactor * this->wOutputActual_ * this->cInput_);
491+ uint32_t tailLoopElements = static_cast<uint32_t>(tailH * this->wOutputActual_ * this->cInput_);
492+ uint32_t rowStrideInUb = static_cast<uint32_t>(wInputActualAmend * this->cOutputActualAlign_);
493+ uint32_t oneNOutScatterElements =
494+ static_cast<uint32_t>(this->hOutputActual_ * this->wOutputActual_ * this->cOutputActualAlign_);
495+ 
496+ int32_t num1D = this->cInput_;
497+ int32_t rate2D = this->wStride_ * this->cOutputActualAlign_;
498+ int32_t num2D = this->wOutputActual_ * this->cInput_;
499+ int32_t rate3D = this->hStride_ * wInputActualAmend * this->cOutputActualAlign_;
500+ T2 argmaxNum1D = this->cInput_;
501+ T2 argMaxRate2D = this->wStride_;
502+ T2 argMaxNum2D = this->wOutputActual_ * this->cInput_;
503+ T2 argHRate3D = this->hStride_;
504+ int32_t scatterIdxNum1D = this->cInput_;
505+ int32_t scatterIdxRate2D = this->cOutputActualAlign_;
506+ 
507+ __VEC_SCOPE__
508+ {
509+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
510+ MicroAPI::RegTensor<uint32_t> gatterNStartIdx;
511+ MicroAPI::RegTensor<T2> argmaxHStart;
512+ MicroAPI::RegTensor<T2> argmaxWStart;
513+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
514+ MicroAPI::RegTensor<uint32_t> scatterNStartIdx;
515+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
516+ MicroAPI::MaskReg maskAllT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
517+ 
518+ GenGatterIndex3D<int32_t>((MicroAPI::RegTensor<int32_t>&)gatterStartIdx, rate3D, num2D, rate2D, num1D);
519+ GenGatterIndex3D<T2>(argmaxWStart, 0, argMaxNum2D, argMaxRate2D, argmaxNum1D, 0);
520+ GenGatterIndex2D<T2>(argmaxHStart, argHRate3D, argMaxNum2D, 0);
521+ GenGatterIndex2D<int32_t>((MicroAPI::RegTensor<int32_t>&)scatterStartIdx, scatterIdxRate2D, scatterIdxNum1D);
522+ 
523+ for (uint16_t nIdex = 0; nIdex < loopN; nIdex++) {
524+ // 校正N
525+ MicroAPI::Adds(gatterNStartIdx, gatterStartIdx, nIdex * ubNumHWC, maskAllU32);
526+ MicroAPI::Adds(scatterNStartIdx, scatterStartIdx, nIdex * oneNOutScatterElements, maskAllU32);
527+ 
528+ for (uint16_t j = 0; j < loopH; j++) {
529+ int32_t gatterIndexOffset = j * oneLoopStrideH;
530+ int32_t argmaxHOffset = j * hStride * hFactor + hBlockArgmaxOffset;
531+ int32_t argmaxWOffset = wBlockArgmaxOffset;
532+ int32_t scatterOffset = j * hFactor * wOutputActual * alignedC;
533+ 
534+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
535+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset,
536+ gatterNStartIdx, oneLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset,
537+ scatterNStartIdx, scatterOffset, padH, padW, wInput);
538+ }
539+ 
540+ // tail H
541+ int32_t gatterIndexOffset = loopH * oneLoopStrideH;
542+ int32_t argmaxHOffset = loopH * hStride * hFactor + hBlockArgmaxOffset;
543+ int32_t argmaxWOffset = wBlockArgmaxOffset;
544+ int32_t scatterOffset = loopH * hFactor * wOutputActual * alignedC;
545+ 
546+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
547+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterNStartIdx,
548+ tailLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterNStartIdx,
549+ scatterOffset, padH, padW, wInput);
550+ }
551+ }
552+}
553+ 
554+template <typename T1, typename T2, const uint32_t IS_PAD>
555+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::ComputeMultiRowForInt64(
556+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr,
557+ __local_mem__ uint32_t* helpAddr)
558+{
559+ uint16_t kH = static_cast<uint16_t>(this->hKernel_);
560+ uint16_t kW = static_cast<uint16_t>(this->wKernel_);
561+ uint16_t hStride = static_cast<uint16_t>(this->hStride_);
562+ uint16_t padH = static_cast<uint16_t>(this->padTop_);
563+ uint16_t padW = static_cast<uint16_t>(this->padLeft_);
564+ int32_t wInput = static_cast<int32_t>(this->wInput_);
565+ uint16_t wOutputActual = static_cast<uint16_t>(this->wOutputActual_);
566+ uint16_t alignedC = static_cast<uint16_t>(this->cOutputActualAlign_);
567+ 
568+ uint16_t loopN = static_cast<uint16_t>(this->nOutputActual_);
569+ constexpr uint32_t repeatElm = platform::GetVRegSize() / sizeof(T2);
570+ uint16_t hFactor = static_cast<uint16_t>(repeatElm / (this->wOutputActual_ * this->cInput_));
571+ hFactor = hFactor > this->hOutputActual_ ? this->hOutputActual_ : hFactor;
572+ uint16_t loopH = static_cast<uint16_t>(this->hOutputActual_ / hFactor);
573+ uint16_t tailH = static_cast<uint16_t>(this->hOutputActual_ - loopH * hFactor);
574+ 
575+ int32_t hInputActualAmend = (this->hOutputActual_ - 1) * this->hStride_ + this->hKernel_;
576+ int32_t wInputActualAmend = (this->wOutputActual_ - 1) * this->wStride_ + this->wKernel_;
577+ int32_t ubNumHWC = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
578+ 
579+ int32_t wBlockArgmaxOffset = this->wAxisIndex_ * this->wStride_ * this->wOutputInner_;
580+ int32_t hBlockArgmaxOffset = this->hAxisIndex_ * this->hStride_ * this->hOutputInner_;
581+ 
582+ uint32_t oneLoopStrideH =
583+ static_cast<uint32_t>(hFactor * this->hStride_ * wInputActualAmend * this->cOutputActualAlign_);
584+ uint32_t oneLoopElements = static_cast<uint32_t>(hFactor * this->wOutputActual_ * this->cInput_);
585+ uint32_t tailLoopElements = static_cast<uint32_t>(tailH * this->wOutputActual_ * this->cInput_);
586+ uint32_t rowStrideInUb = static_cast<uint32_t>(wInputActualAmend * this->cOutputActualAlign_);
587+ uint32_t oneNOutScatterElements =
588+ static_cast<uint32_t>(this->hOutputActual_ * this->wOutputActual_ * this->cOutputActualAlign_);
589+ 
590+ int32_t num1D = this->cInput_;
591+ int32_t rate2D = this->wStride_ * this->cOutputActualAlign_;
592+ int32_t num2D = this->wOutputActual_ * this->cInput_;
593+ int32_t rate3D = this->hStride_ * wInputActualAmend * this->cOutputActualAlign_;
594+ T2 argmaxNum1D = this->cInput_;
595+ T2 argMaxRate2D = this->wStride_;
596+ T2 argMaxNum2D = this->wOutputActual_ * this->cInput_;
597+ T2 argHRate3D = this->hStride_;
598+ int32_t scatterIdxNum1D = this->cInput_;
599+ int32_t scatterIdxRate2D = this->cOutputActualAlign_;
600+ 
601+ __VEC_SCOPE__
602+ {
603+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
604+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
605+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
606+ 
607+ GenGatterIndex3D<int32_t>((MicroAPI::RegTensor<int32_t>&)gatterStartIdx, rate3D, num2D, rate2D, num1D);
608+ GenGatterIndex2D<int32_t>((MicroAPI::RegTensor<int32_t>&)scatterStartIdx, scatterIdxRate2D, scatterIdxNum1D);
609+ 
610+ AscendC::MicroAPI::DataCopy(helpAddr, gatterStartIdx, maskAllU32);
611+ AscendC::MicroAPI::DataCopy(helpAddr + V_REG_SIZE / sizeof(uint32_t), scatterStartIdx, maskAllU32);
612+ }
613+ 
614+ __VEC_SCOPE__
615+ {
616+ MicroAPI::RegTensor<T2> argmaxHStart;
617+ MicroAPI::RegTensor<T2> argmaxWStart;
618+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
619+ 
620+ GenGatterIndex3D<T2>(argmaxWStart, 0, argMaxNum2D, argMaxRate2D, argmaxNum1D, 0);
621+ GenGatterIndex2D<T2>(argmaxHStart, argHRate3D, argMaxNum2D, 0);
622+ 
623+ AscendC::MicroAPI::DataCopy(
624+ helpAddr + V_REG_SIZE / sizeof(uint32_t) * DOUBLE, (MicroAPI::RegTensor<uint32_t>&)argmaxHStart,
625+ maskAllU32);
626+ AscendC::MicroAPI::DataCopy(
627+ helpAddr + V_REG_SIZE / sizeof(uint32_t) * THREE, (MicroAPI::RegTensor<uint32_t>&)argmaxWStart, maskAllU32);
628+ }
629+ 
630+ __VEC_SCOPE__
631+ {
632+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
633+ MicroAPI::RegTensor<uint32_t> gatterNStartIdx;
634+ MicroAPI::RegTensor<T2> argmaxHStart;
635+ MicroAPI::RegTensor<T2> argmaxWStart;
636+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
637+ MicroAPI::RegTensor<uint32_t> scatterNStartIdx;
638+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
639+ MicroAPI::MaskReg maskAllT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
640+ 
641+ AscendC::MicroAPI::DataCopy(gatterStartIdx, helpAddr);
642+ AscendC::MicroAPI::DataCopy(scatterStartIdx, helpAddr + V_REG_SIZE / sizeof(uint32_t));
643+ AscendC::MicroAPI::DataCopy(
644+ (MicroAPI::RegTensor<uint32_t>&)argmaxHStart, helpAddr + V_REG_SIZE / sizeof(uint32_t) * DOUBLE);
645+ AscendC::MicroAPI::DataCopy(
646+ (MicroAPI::RegTensor<uint32_t>&)argmaxWStart, helpAddr + V_REG_SIZE / sizeof(uint32_t) * THREE);
647+ 
648+ for (uint16_t nIdex = 0; nIdex < loopN; nIdex++) {
649+ // 校正N
650+ MicroAPI::Adds(gatterNStartIdx, gatterStartIdx, nIdex * ubNumHWC, maskAllU32);
651+ MicroAPI::Adds(scatterNStartIdx, scatterStartIdx, nIdex * oneNOutScatterElements, maskAllU32);
652+ 
653+ for (uint16_t j = 0; j < loopH; j++) {
654+ int32_t gatterIndexOffset = j * oneLoopStrideH;
655+ int32_t argmaxHOffset = j * hStride * hFactor + hBlockArgmaxOffset;
656+ int32_t argmaxWOffset = wBlockArgmaxOffset;
657+ int32_t scatterOffset = j * hFactor * wOutputActual * alignedC;
658+ 
659+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
660+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset,
661+ gatterNStartIdx, oneLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset,
662+ scatterNStartIdx, scatterOffset, padH, padW, wInput);
663+ }
664+ 
665+ // tail H
666+ int32_t gatterIndexOffset = loopH * oneLoopStrideH;
667+ int32_t argmaxHOffset = loopH * hStride * hFactor + hBlockArgmaxOffset;
668+ int32_t argmaxWOffset = wBlockArgmaxOffset;
669+ int32_t scatterOffset = loopH * hFactor * wOutputActual * alignedC;
670+ 
671+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
672+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterNStartIdx,
673+ tailLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterNStartIdx,
674+ scatterOffset, padH, padW, wInput);
675+ }
676+ }
677+}
678+ 
679+template <typename T1, typename T2, const uint32_t IS_PAD>
680+__aicore__ inline void MaxPoolWithArgmaxV3SmallC<T1, T2, IS_PAD>::ComputeSingleRow(
681+ __local_mem__ T1* xAddr, __local_mem__ T1* maxValueAddr, __local_mem__ T2* argmaxAddr)
682+{
683+ uint16_t kH = static_cast<uint16_t>(this->hKernel_);
684+ uint16_t kW = static_cast<uint16_t>(this->wKernel_);
685+ uint16_t hStride = static_cast<uint16_t>(this->hStride_);
686+ uint16_t wStride = static_cast<uint16_t>(this->hStride_);
687+ uint16_t padH = static_cast<uint16_t>(this->padTop_);
688+ uint16_t padW = static_cast<uint16_t>(this->padLeft_);
689+ int32_t wInput = static_cast<int32_t>(this->wInput_);
690+ uint16_t wOutputActual = static_cast<uint16_t>(this->wOutputActual_);
691+ uint16_t alignedC = static_cast<uint16_t>(this->cOutputActualAlign_);
692+ 
693+ uint16_t loopN = this->nOutputActual_;
694+ uint16_t loopH = this->hOutputActual_;
695+ 
696+ constexpr uint32_t repeatElm = platform::GetVRegSize() / sizeof(T2);
697+ uint16_t wFactor = repeatElm / this->cInput_;
698+ wFactor = wFactor > this->wOutputActual_ ? this->wOutputActual_ : wFactor;
699+ uint16_t loopW = static_cast<uint16_t>(this->wOutputActual_ / wFactor);
700+ uint16_t tailW = static_cast<uint16_t>(this->wOutputActual_ - loopW * wFactor);
701+ 
702+ int32_t hInputActualAmend = (this->hOutputActual_ - 1) * this->hStride_ + this->hKernel_;
703+ int32_t wInputActualAmend = (this->wOutputActual_ - 1) * this->wStride_ + this->wKernel_;
704+ int32_t ubNumHWC = hInputActualAmend * wInputActualAmend * this->cOutputActualAlign_;
705+ 
706+ int32_t wBlockArgmaxOffset = this->wAxisIndex_ * this->wStride_ * this->wOutputInner_;
707+ int32_t hBlockArgmaxOffset = this->hAxisIndex_ * this->hStride_ * this->hOutputInner_;
708+ 
709+ uint32_t oneLoopStrideH = static_cast<uint32_t>(this->hStride_ * wInputActualAmend * this->cOutputActualAlign_);
710+ uint32_t oneLoopStrideW = static_cast<uint32_t>(this->wStride_ * wFactor * this->cOutputActualAlign_);
711+ uint32_t oneLoopElements = static_cast<uint32_t>(wFactor * this->cInput_);
712+ uint32_t tailLoopElements = tailW * this->cInput_;
713+ 
714+ uint32_t oneNOutScatterElements =
715+ static_cast<uint32_t>(this->hOutputActual_ * this->wOutputActual_ * this->cOutputActualAlign_);
716+ uint32_t rowStrideInUb = static_cast<uint32_t>(wInputActualAmend * this->cOutputActualAlign_);
717+ 
718+ int32_t num1D = this->cInput_;
719+ int32_t rate2D = this->wStride_ * this->cOutputActualAlign_;
720+ int32_t argmaxNum1D = this->cInput_;
721+ T2 argmaxRate2D = this->wStride_;
722+ int32_t scatterIdxNum1D = this->cInput_;
723+ int32_t scatterIdxRate2D = this->cOutputActualAlign_;
724+ 
725+ for (uint16_t nIdex = 0; nIdex < loopN; nIdex++) {
726+ __VEC_SCOPE__
727+ {
728+ MicroAPI::RegTensor<uint32_t> gatterStartIdx;
729+ MicroAPI::RegTensor<T2> argmaxHStart;
730+ MicroAPI::RegTensor<T2> argmaxWStart;
731+ MicroAPI::RegTensor<uint32_t> scatterStartIdx;
732+ MicroAPI::MaskReg maskAllU32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
733+ MicroAPI::MaskReg maskAllT2 = MicroAPI::CreateMask<T2, MicroAPI::MaskPattern::ALL>();
734+ 
735+ GenGatterIndex2D<int32_t>((MicroAPI::RegTensor<int32_t>&)gatterStartIdx, rate2D, num1D);
736+ GenGatterIndex2D<T2>(argmaxWStart, argmaxRate2D, static_cast<T2>(argmaxNum1D), 0);
737+ AscendC::MicroAPI::Duplicate(argmaxHStart, 0);
738+ GenGatterIndex2D<int32_t>(
739+ (MicroAPI::RegTensor<int32_t>&)scatterStartIdx, scatterIdxRate2D, scatterIdxNum1D);
740+ 
741+ MicroAPI::Adds(gatterStartIdx, gatterStartIdx, nIdex * ubNumHWC, maskAllU32);
742+ MicroAPI::Adds(scatterStartIdx, scatterStartIdx, nIdex * oneNOutScatterElements, maskAllU32);
743+ 
744+ for (uint16_t i = 0; i < loopH; i++) {
745+ int32_t hOffset = i * oneLoopStrideH;
746+ int32_t argmaxHOffset = i * hStride + hBlockArgmaxOffset;
747+ 
748+ for (uint16_t j = 0; j < loopW; j++) {
749+ int32_t wOffset = j * oneLoopStrideW;
750+ int32_t argmaxWOffset = j * wStride * wFactor + wBlockArgmaxOffset;
751+ int32_t gatterIndexOffset = hOffset + wOffset;
752+ int32_t scatterOffset = (j * wFactor + i * wOutputActual) * alignedC;
753+ 
754+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
755+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset,
756+ gatterStartIdx, oneLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset,
757+ scatterStartIdx, scatterOffset, padH, padW, wInput);
758+ }
759+ 
760+ // tail w
761+ int32_t wOffset = loopW * oneLoopStrideW;
762+ int32_t argmaxWOffset = loopW * wStride * wFactor + wBlockArgmaxOffset;
763+ int32_t gatterIndexOffset = hOffset + wOffset;
764+ int32_t scatterOffset = (loopW * wFactor + i * wOutputActual) * alignedC;
765+ 
766+ MaxPoolWithArgMaxV3GatherImpl<T1, T2, IS_PAD>(
767+ xAddr, maxValueAddr, argmaxAddr, kH, kW, rowStrideInUb, alignedC, gatterIndexOffset, gatterStartIdx,
768+ tailLoopElements, argmaxHStart, argmaxWStart, argmaxHOffset, argmaxWOffset, scatterStartIdx,
769+ scatterOffset, padH, padW, wInput);
770+ }
771+ }
772+ }
773+}
774+} // namespace MaxPoolWithArgmaxV3SmallCNameSpace
775+#endif // MAX_POOL_WITH_ARGMAX_V3_NHWC_SMALL_C__H_
Apooling/max_pool_with_argmax_v3/op_kernel/arch35/max_pool_with_argmax_v3_simt.h+208-0
@@ -0,0 +1,208 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file max_pool_with_argmax_v3_simt.h
13+ * \brief max_pool_with_argmax_v3 implied by simt
14+ */
15+ 
16+#ifndef CANN_MAX_POOL_WITH_ARGMAX_V3_SIMT_H
17+#define CANN_MAX_POOL_WITH_ARGMAX_V3_SIMT_H
18+ 
19+#include "kernel_operator.h"
20+#include "kernel_tiling/kernel_tiling.h"
21+ 
22+#ifdef __CCE_KT_TEST__
23+#define LAUNCH_BOUND(threads)
24+#endif
25+ 
26+using namespace AscendC;
27+ 
28+namespace SimtProc {
29+constexpr static uint32_t THREAD_DIM = 256;
30+ 
31+template <typename idx_accscalar_t>
32+__aicore__ inline static void CycleUpdate(float val, idx_accscalar_t idxOffset, float* maxval, idx_accscalar_t* maxidx)
33+{
34+ if ((static_cast<float>(val) > *maxval) || Simt::IsNan(val)) {
35+ *maxidx = idxOffset;
36+ *maxval = val;
37+ }
38+}
39+} // namespace SimtProc
40+ 
41+template <typename VALUE_T, typename INDICES_T, int Format_T, bool useINT64Index>
42+class MaxPoolWithArgmaxV3 {
43+public:
44+ __aicore__ inline MaxPoolWithArgmaxV3(const MaxPoolWithArgmaxV3SimtTilingData* __restrict tilingData)
45+ : tilingData_(tilingData), blockIdx_(GetBlockIdx()), blockNum_(GetBlockNum())
46+ {}
47+ 
48+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR argmax);
49+ __aicore__ inline void Process();
50+ __aicore__ inline void Compute() const;
51+ 
52+private:
53+ __aicore__ inline static INDICES_T min(INDICES_T left, INDICES_T right)
54+ {
55+ if (left <= right) {
56+ return left;
57+ }
58+ return right;
59+ }
60+ 
61+private:
62+ AscendC::GlobalTensor<VALUE_T> x_;
63+ AscendC::GlobalTensor<VALUE_T> y_;
64+ AscendC::GlobalTensor<INDICES_T> argmax_;
65+ const MaxPoolWithArgmaxV3SimtTilingData* tilingData_;
66+ uint32_t blockIdx_ = 0;
67+ uint32_t blockNum_ = 1;
68+ const uint32_t F32_NEG_INF = 0xff800000;
69+};
70+ 
71+template <typename VALUE_T, typename INDICES_T, int Format_T, bool useINT64Index>
72+__aicore__ inline void MaxPoolWithArgmaxV3<VALUE_T, INDICES_T, Format_T, useINT64Index>::Init(
73+ GM_ADDR x, GM_ADDR y, GM_ADDR argmax)
74+{
75+ x_.SetGlobalBuffer((__gm__ VALUE_T*)(x));
76+ y_.SetGlobalBuffer((__gm__ VALUE_T*)(y));
77+ argmax_.SetGlobalBuffer((__gm__ INDICES_T*)(argmax));
78+}
79+ 
80+template <typename VALUE_T, typename INDICES_T, int Format_T, bool useINT64Index>
81+__aicore__ inline void MaxPoolWithArgmaxV3<VALUE_T, INDICES_T, Format_T, useINT64Index>::Process()
82+{
83+ Compute();
84+}
85+ 
86+template <typename scalar_t, typename idx_scalar_t, typename idx_accscalar_t>
87+__simt_vf__ __aicore__ LAUNCH_BOUND(SimtProc::THREAD_DIM) inline void MaxPoolForwardNchw(
88+ const int64_t count, const __gm__ scalar_t* bottomData, const int64_t height, const int64_t width,
89+ const int outputHeight, const int outputWidth, const int kernelH, const int kernelW, const int strideH,
90+ const int strideW, const int padH, const int padW, const int dilationH, const int dilationW,
91+ __gm__ scalar_t* topData, __gm__ idx_scalar_t* topMask, int blockIdx, int blockNum)
92+{
93+ for (idx_accscalar_t index = blockIdx * Simt::GetThreadNum() + Simt::GetThreadIdx(); index < count;
94+ index = index + blockNum * Simt::GetThreadNum()) {
95+ idx_accscalar_t pw = index % outputWidth;
96+ idx_accscalar_t ph = (index / outputWidth) % outputHeight;
97+ idx_accscalar_t nxc = index / outputWidth / outputHeight;
98+ idx_accscalar_t hstart = ph * strideH - padH;
99+ idx_accscalar_t wstart = pw * strideW - padW;
100+ idx_accscalar_t hend =
101+ (hstart + (kernelH - 1) * dilationH + 1) < height ? (hstart + (kernelH - 1) * dilationH + 1) : height;
102+ idx_accscalar_t wend =
103+ (wstart + (kernelW - 1) * dilationW + 1) < width ? (wstart + (kernelW - 1) * dilationW + 1) : width;
104+ while (hstart < 0)
105+ hstart += dilationH;
106+ while (wstart < 0)
107+ wstart += dilationW;
108+ float maxval = *reinterpret_cast<const float*>(&F32_NEG_INF); // -Infinity
109+ idx_accscalar_t maxidx = hstart * width + wstart;
110+ auto btmData = bottomData + nxc * height * width;
111+ for (idx_accscalar_t h = hstart; h < hend; h += dilationH) {
112+ for (idx_accscalar_t w = wstart; w < wend; w += dilationW) {
113+ idx_accscalar_t idxOffset = h * width + w;
114+ float val = static_cast<float>(btmData[idxOffset]);
115+ SimtProc::CycleUpdate<idx_accscalar_t>(val, idxOffset, &maxval, &maxidx);
116+ }
117+ }
118+ topData[index] = static_cast<scalar_t>(maxval);
119+ topMask[index] = static_cast<idx_scalar_t>(maxidx);
120+ }
121+}
122+ 
123+template <typename scalar_t, typename idx_scalar_t, typename idx_accscalar_t>
124+__simt_vf__ __aicore__ LAUNCH_BOUND(SimtProc::THREAD_DIM) inline void MaxPoolForwardNhwc(
125+ const int64_t count, const __gm__ scalar_t* bottomData, const int64_t channels, const int64_t height,
126+ const int64_t width, const int outputHeight, const int outputWidth, const int kernelH, const int kernelW,
127+ const int strideH, const int strideW, const int padH, const int padW, const int dilationH, const int dilationW,
128+ __gm__ scalar_t* topData, __gm__ idx_scalar_t* topMask, int blockIdx, int blockNum)
129+{
130+ for (idx_accscalar_t index = blockIdx * Simt::GetThreadNum() + Simt::GetThreadIdx(); index < count;
131+ index = index + blockNum * Simt::GetThreadNum()) {
132+ idx_accscalar_t c = index % channels;
133+ idx_accscalar_t pw = (index / channels) % outputWidth;
134+ idx_accscalar_t ph = (index / channels / outputWidth) % outputHeight;
135+ idx_accscalar_t n = index / channels / outputWidth / outputHeight;
136+ idx_accscalar_t hstart = ph * strideH - padH;
137+ idx_accscalar_t wstart = pw * strideW - padW;
138+ idx_accscalar_t hend =
139+ (hstart + (kernelH - 1) * dilationH + 1) < height ? (hstart + (kernelH - 1) * dilationH + 1) : height;
140+ idx_accscalar_t wend =
141+ (wstart + (kernelW - 1) * dilationW + 1) < width ? (wstart + (kernelW - 1) * dilationW + 1) : width;
142+ while (hstart < 0)
143+ hstart += dilationH;
144+ while (wstart < 0)
145+ wstart += dilationW;
146+ float maxval = *reinterpret_cast<const float*>(&F32_NEG_INF);
147+ idx_accscalar_t maxidx = hstart * width + wstart;
148+ auto btmData = bottomData + (n * height * width * channels);
149+ for (idx_accscalar_t h = hstart; h < hend; h += dilationH) {
150+ for (idx_accscalar_t w = wstart; w < wend; w += dilationW) {
151+ idx_accscalar_t idxOffset = h * width + w;
152+ scalar_t val = static_cast<float>(btmData[idxOffset * channels + c]);
153+ SimtProc::CycleUpdate<idx_accscalar_t>(val, idxOffset, &maxval, &maxidx);
154+ }
155+ }
156+ topData[index] = static_cast<scalar_t>(maxval);
157+ topMask[index] = static_cast<idx_scalar_t>(maxidx);
158+ }
159+}
160+ 
161+template <typename VALUE_T, typename INDICES_T, int Format_T, bool useINT64Index>
162+__aicore__ inline void MaxPoolWithArgmaxV3<VALUE_T, INDICES_T, Format_T, useINT64Index>::Compute() const
163+{
164+ const int kH = tilingData_->kSizeH;
165+ const int kW = tilingData_->kSizeW;
166+ 
167+ const int dH = tilingData_->stridesH;
168+ const int dW = tilingData_->stridesW;
169+ 
170+ const int padH = tilingData_->padH;
171+ const int padW = tilingData_->padW;
172+ 
173+ const int dilationH = tilingData_->dilationH;
174+ const int dilationW = tilingData_->dilationW;
175+ const int64_t nbatch = tilingData_->nDim;
176+ const int64_t inputChannel = tilingData_->cDim;
177+ const int64_t inputHeight = tilingData_->hInDim;
178+ const int64_t inputWidth = tilingData_->wInDim;
179+ 
180+ const int64_t outputHeight = tilingData_->hOutDim;
181+ const int64_t outputWidth = tilingData_->wOutDim;
182+ 
183+ auto inputData = (__gm__ VALUE_T*)x_.GetPhyAddr();
184+ auto outputData = (__gm__ VALUE_T*)y_.GetPhyAddr();
185+ auto indicesData = (__gm__ INDICES_T*)argmax_.GetPhyAddr();
186+ int64_t count = nbatch * inputChannel * outputHeight * outputWidth;
187+ if constexpr (Format_T == 0 && !useINT64Index) {
188+ Simt::VF_CALL<MaxPoolForwardNchw<VALUE_T, INDICES_T, int32_t>>(
189+ Simt::Dim3(SimtProc::THREAD_DIM), count, inputData, inputHeight, inputWidth, outputHeight, outputWidth, kH,
190+ kW, dH, dW, padH, padW, dilationH, dilationW, outputData, indicesData, blockIdx_, blockNum_);
191+ } else if constexpr (Format_T == 1 && !useINT64Index) {
192+ Simt::VF_CALL<MaxPoolForwardNhwc<VALUE_T, INDICES_T, int32_t>>(
193+ Simt::Dim3(SimtProc::THREAD_DIM), count, inputData, inputChannel, inputHeight, inputWidth, outputHeight,
194+ outputWidth, kH, kW, dH, dW, padH, padW, dilationH, dilationW, outputData, indicesData, blockIdx_,
195+ blockNum_);
196+ } else if constexpr (Format_T == 0 && useINT64Index) {
197+ Simt::VF_CALL<MaxPoolForwardNchw<VALUE_T, INDICES_T, int64_t>>(
198+ Simt::Dim3(SimtProc::THREAD_DIM), count, inputData, inputHeight, inputWidth, outputHeight, outputWidth, kH,
199+ kW, dH, dW, padH, padW, dilationH, dilationW, outputData, indicesData, blockIdx_, blockNum_);
200+ } else if constexpr (Format_T == 1 && useINT64Index) {
201+ Simt::VF_CALL<MaxPoolForwardNhwc<VALUE_T, INDICES_T, int64_t>>(
202+ Simt::Dim3(SimtProc::THREAD_DIM), count, inputData, inputChannel, inputHeight, inputWidth, outputHeight,
203+ outputWidth, kH, kW, dH, dW, padH, padW, dilationH, dilationW, outputData, indicesData, blockIdx_,
204+ blockNum_);
205+ }
206+}
207+ 
208+#endif // CANN_MAX_POOL_WITH_ARGMAX_V3_SIMT_H
Apooling/max_pool_with_argmax_v3/op_kernel/max_pool_with_argmax_v3.cpp+166-0
@@ -0,0 +1,166 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* !
12+ * \file max_pool_with_argmax_v3.cpp
13+ * \brief max_pool_with_argmax_v3 implied
14+ */
15+ 
16+#include <cstdint>
17+#include "kernel_operator.h"
18+#include "kernel_tiling/kernel_tiling.h"
19+#include "arch35/max_pool_with_argmax_v3_simt.h"
20+#include "arch35/max_pool_with_argmax_v3_big_kernel.h"
21+#include "arch35/max_pool_with_argmax_v3_gather_kernel.h"
22+#include "arch35/max_pool_with_argmax_v3_big_kernel_mul_core.h"
23+#include "arch35/max_pool_with_argmax_v3_nhwc_big_c.h"
24+#include "arch35/max_pool_with_argmax_v3_nhwc_small_c.h"
25+ 
26+#define BIG_KERNEL_FORMAT_NCHW 311110
27+#define NO_PADDING_TILING_KEY 300001
28+#define PADDING_TILING_KEY 300002
29+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT32 400001
30+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT64 400002
31+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT32 400003
32+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT64 400004
33+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT32 400005
34+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT64 400006
35+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NCHW 500001
36+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NHWC 500002
37+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NCHW_PAD 500011
38+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NHWC_PAD 500012
39+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_BIG_C 800001
40+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_BIG_C_PAD 800002
41+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_SMALL_C 700001
42+#define MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_SMALL_C_PAD 700002
43+ 
44+constexpr uint32_t PAD_DISABLE = 0;
45+constexpr uint32_t PAD_ENABLE = 1;
46+constexpr int NCHW = 0;
47+constexpr int NHWC = 1;
48+ 
49+extern "C" __global__ __aicore__ void max_pool_with_argmax_v3(
50+ GM_ADDR x, GM_ADDR y, GM_ADDR argmax, GM_ADDR workspace, GM_ADDR tiling)
51+{
52+ TPipe pipeBase;
53+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0);
54+ if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_SMALL_C)) {
55+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3NhwcTilingData, tilingDataIn, tiling);
56+ const MaxPoolWithArgmaxV3NhwcTilingData* __restrict tilingData = &tilingDataIn;
57+ MaxPoolWithArgmaxV3SmallCNameSpace::MaxPoolWithArgmaxV3SmallC<DTYPE_X, DTYPE_ARGMAX, 0> op(
58+ &pipeBase, tilingData);
59+ op.Init(x, y, argmax);
60+ op.MaxPoolWithArgmaxV3SmallCProcess();
61+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_SMALL_C_PAD)) {
62+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3NhwcTilingData, tilingDataIn, tiling);
63+ const MaxPoolWithArgmaxV3NhwcTilingData* __restrict tilingData = &tilingDataIn;
64+ MaxPoolWithArgmaxV3SmallCNameSpace::MaxPoolWithArgmaxV3SmallC<DTYPE_X, DTYPE_ARGMAX, 1> op(
65+ &pipeBase, tilingData);
66+ op.Init(x, y, argmax);
67+ op.MaxPoolWithArgmaxV3SmallCProcess();
68+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_BIG_C)) {
69+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3NhwcTilingData, tilingDataIn, tiling);
70+ const MaxPoolWithArgmaxV3NhwcTilingData* __restrict tilingData = &tilingDataIn;
71+ MaxPoolWithArgmaxV3NHWC::MaxPoolWithArgmaxV3NhwCKernel<DTYPE_X, DTYPE_ARGMAX, false> op(&pipeBase, tilingData);
72+ op.Init(x, y, argmax);
73+ op.Process();
74+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_NHWC_BIG_C_PAD)) {
75+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3NhwcTilingData, tilingDataIn, tiling);
76+ const MaxPoolWithArgmaxV3NhwcTilingData* __restrict tilingData = &tilingDataIn;
77+ MaxPoolWithArgmaxV3NHWC::MaxPoolWithArgmaxV3NhwCKernel<DTYPE_X, DTYPE_ARGMAX, true> op(&pipeBase, tilingData);
78+ op.Init(x, y, argmax);
79+ op.Process();
80+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NCHW)) {
81+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3SimtTilingData, tilingDataIn, tiling);
82+ const MaxPoolWithArgmaxV3SimtTilingData* __restrict tilingData = &tilingDataIn;
83+ MaxPoolWithArgmaxV3<DTYPE_X, DTYPE_ARGMAX, NCHW, false> op(tilingData);
84+ op.Init(x, y, argmax);
85+ op.Process();
86+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NHWC)) {
87+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3SimtTilingData, tilingDataIn, tiling);
88+ const MaxPoolWithArgmaxV3SimtTilingData* __restrict tilingData = &tilingDataIn;
89+ MaxPoolWithArgmaxV3<DTYPE_X, DTYPE_ARGMAX, NHWC, false> op(tilingData);
90+ op.Init(x, y, argmax);
91+ op.Process();
92+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NCHW_PAD)) {
93+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3SimtTilingData, tilingDataIn, tiling);
94+ const MaxPoolWithArgmaxV3SimtTilingData* __restrict tilingData = &tilingDataIn;
95+ MaxPoolWithArgmaxV3<DTYPE_X, DTYPE_ARGMAX, NCHW, true> op(tilingData);
96+ op.Init(x, y, argmax);
97+ op.Process();
98+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_SIMT_NHWC_PAD)) {
99+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3SimtTilingData, tilingDataIn, tiling);
100+ const MaxPoolWithArgmaxV3SimtTilingData* __restrict tilingData = &tilingDataIn;
101+ MaxPoolWithArgmaxV3<DTYPE_X, DTYPE_ARGMAX, NHWC, true> op(tilingData);
102+ op.Init(x, y, argmax);
103+ op.Process();
104+ } else if (TILING_KEY_IS(BIG_KERNEL_FORMAT_NCHW)) {
105+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelTilingData, tilingDataIn, tiling);
106+ const MaxPoolWithArgmaxV3BigKernelTilingData* __restrict tilingData = &tilingDataIn;
107+ MaxPoolWithArgMaxV3BigKernel::MaxPoolWithArgmaxV3BigKernel<DTYPE_X, float, DTYPE_ARGMAX> op(
108+ &pipeBase, tilingData);
109+ op.Init(x, y, argmax);
110+ op.Process();
111+ } else if (TILING_KEY_IS(NO_PADDING_TILING_KEY)) {
112+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3GatherTilingData, tilingDataIn, tiling);
113+ MaxPoolWithArgmaxV3GatherNameSpace::MaxPoolWithArgmaxV3GatherKernel<DTYPE_X, DTYPE_ARGMAX, PAD_DISABLE> op(
114+ pipeBase, tilingDataIn);
115+ op.Init(x, y, argmax);
116+ op.Process();
117+ } else if (TILING_KEY_IS(PADDING_TILING_KEY)) {
118+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3GatherTilingData, tilingDataIn, tiling);
119+ MaxPoolWithArgmaxV3GatherNameSpace::MaxPoolWithArgmaxV3GatherKernel<DTYPE_X, DTYPE_ARGMAX, PAD_ENABLE> op(
120+ pipeBase, tilingDataIn);
121+ op.Init(x, y, argmax);
122+ op.Process();
123+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT32)) {
124+ TPipe pipeBase;
125+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData, tilingDataIn, tiling);
126+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tilingData = &tilingDataIn;
127+ MaxPoolWithArgmaxV3BigKernelMulCore::MaxPoolWithArgmaxV3BigKernelMulCore<float, float, int32_t> op;
128+ op.Init(x, y, argmax, workspace, &pipeBase, tilingData);
129+ op.Process();
130+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_FP_INT64)) {
131+ TPipe pipeBase;
132+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData, tilingDataIn, tiling);
133+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tilingData = &tilingDataIn;
134+ MaxPoolWithArgmaxV3BigKernelMulCore::MaxPoolWithArgmaxV3BigKernelMulCore<float, float, int64_t> op;
135+ op.Init(x, y, argmax, workspace, &pipeBase, tilingData);
136+ op.Process();
137+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT32)) {
138+ TPipe pipeBase;
139+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData, tilingDataIn, tiling);
140+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tilingData = &tilingDataIn;
141+ MaxPoolWithArgmaxV3BigKernelMulCore::MaxPoolWithArgmaxV3BigKernelMulCore<bfloat16_t, float, int32_t> op;
142+ op.Init(x, y, argmax, workspace, &pipeBase, tilingData);
143+ op.Process();
144+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_BF16_INT64)) {
145+ TPipe pipeBase;
146+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData, tilingDataIn, tiling);
147+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tilingData = &tilingDataIn;
148+ MaxPoolWithArgmaxV3BigKernelMulCore::MaxPoolWithArgmaxV3BigKernelMulCore<bfloat16_t, float, int64_t> op;
149+ op.Init(x, y, argmax, workspace, &pipeBase, tilingData);
150+ op.Process();
151+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT32)) {
152+ TPipe pipeBase;
153+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData, tilingDataIn, tiling);
154+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tilingData = &tilingDataIn;
155+ MaxPoolWithArgmaxV3BigKernelMulCore::MaxPoolWithArgmaxV3BigKernelMulCore<half, half, int32_t> op;
156+ op.Init(x, y, argmax, workspace, &pipeBase, tilingData);
157+ op.Process();
158+ } else if (TILING_KEY_IS(MAX_POOL_WITH_ARGMAX_V3_TILING_KEY_BIG_KERNEL_SPLIT_HALF_INT64)) {
159+ TPipe pipeBase;
160+ GET_TILING_DATA_WITH_STRUCT(MaxPoolWithArgmaxV3BigKernelMulCoreTilingData, tilingDataIn, tiling);
161+ const MaxPoolWithArgmaxV3BigKernelMulCoreTilingData* __restrict tilingData = &tilingDataIn;
162+ MaxPoolWithArgmaxV3BigKernelMulCore::MaxPoolWithArgmaxV3BigKernelMulCore<half, half, int64_t> op;
163+ op.Init(x, y, argmax, workspace, &pipeBase, tilingData);
164+ op.Process();
165+ }
166+}
Apooling/max_pool_with_argmax_v3/tests/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
1+# This program is free software, you can redistribute it and/or modify.
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+# This file is a part of the CANN Open Software.
4+# Licensed under 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, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+# See LICENSE in the root of the software repository for the full text of the License.
8+#/
9+ 
10+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
11+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
12+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
13+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Apooling/max_pool_with_argmax_v3/tests/ut/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
1+# This program is free software, you can redistribute it and/or modify.
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+# This file is a part of the CANN Open Software.
4+# Licensed under 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, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+# See LICENSE in the root of the software repository for the full text of the License.
8+#/
9+ 
10+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
11+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
12+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
13+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Apooling/max_pool_with_argmax_v3/tests/ut/op_host/CMakeLists.txt+14-0
@@ -0,0 +1,14 @@
1+# This program is free software, you can redistribute it and/or modify.
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+# This file is a part of the CANN Open Software.
4+# Licensed under 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, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+# See LICENSE in the root of the software repository for the full text of the License.
8+#/
9+ 
10+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
11+if(UT_TEST_ALL OR OP_HOST_UT)
12+ add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13+ add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14+endif()
Apooling/max_pool_with_argmax_v3/tests/ut/op_host/test_max_pool_with_argmax_v3_infershape.cpp+251-0
@@ -0,0 +1,251 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include "exe_graph/runtime/storage_format.h"
13+#include "exe_graph/runtime/storage_shape.h"
14+#include <gtest/gtest.h>
15+#include "kernel_run_context_facker.h"
16+#include "register/op_impl_registry.h"
17+#include "log/log.h"
18+#include "platform/platform_info.h"
19+#include "../../../op_graph/max_pool_with_argmax_v3_proto.h"
20+ 
21+namespace {
22+template <typename T>
23+std::string Shape2String(const T& shape)
24+{
25+ std::ostringstream oss;
26+ oss << "[";
27+ if (shape.GetDimNum() > 0) {
28+ for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) {
29+ oss << shape.GetDim(i) << ", ";
30+ }
31+ oss << shape.GetDim(shape.GetDimNum() - 1);
32+ }
33+ oss << "]";
34+ return oss.str();
35+}
36+ 
37+class MaxPoolWithArgmaxV3Infer : public testing::Test {
38+protected:
39+ static void SetUpTestCase()
40+ {
41+ std::cout << "MaxPoolWithArgmaxV3InferTest SetUp" << std::endl;
42+ }
43+ 
44+ static void TearDownTestCase()
45+ {
46+ std::cout << "MaxPoolWithArgmaxV3InferTest TearDown" << std::endl;
47+ }
48+};
49+ 
50+TEST_F(MaxPoolWithArgmaxV3Infer, maxpool_with_argmax_v3_infershape_test_1)
51+{
52+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolWithArgmaxV3")->infer_shape;
53+ 
54+ gert::StorageShape xShape = {{4, 512, 16, 16}, {}};
55+ gert::StorageShape yShape = {{}, {}};
56+ gert::StorageShape indicesShape = {{}, {}};
57+ auto holder = gert::InferShapeContextFaker()
58+ .NodeIoNum(1, 2)
59+ .IrInstanceNum({1, 2})
60+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
61+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
62+ .NodeOutputTd(1, ge::DT_INT32, ge::Format::FORMAT_NCHW, ge::Format::FORMAT_RESERVED)
63+ .NodeAttrs(
64+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
65+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 2})},
66+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
67+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
68+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
69+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(true)},
70+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
71+ .InputShapes({&xShape})
72+ .OutputShapes({&yShape, &indicesShape})
73+ .Build();
74+ 
75+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
76+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
77+ gert::Shape* indices = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);
78+ ASSERT_EQ(Shape2String(*output), "[4, 512, 8, 8]");
79+ ASSERT_EQ(Shape2String(*indices), "[4, 512, 8, 8]");
80+}
81+ 
82+TEST_F(MaxPoolWithArgmaxV3Infer, maxpool_with_argmax_v3_infershape_test_2)
83+{
84+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolWithArgmaxV3")->infer_shape;
85+ 
86+ gert::StorageShape xShape = {{5, 256, 144, 589}, {}};
87+ gert::StorageShape yShape = {{}, {}};
88+ gert::StorageShape indicesShape = {{}, {}};
89+ auto holder = gert::InferShapeContextFaker()
90+ .NodeIoNum(1, 2)
91+ .IrInstanceNum({1, 2})
92+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
93+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
94+ .NodeOutputTd(1, ge::DT_INT32, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
95+ .NodeAttrs(
96+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 3})},
97+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 3})},
98+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
99+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
100+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({9, 5})},
101+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
102+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
103+ .InputShapes({&xShape})
104+ .OutputShapes({&yShape, &indicesShape})
105+ .Build();
106+ 
107+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
108+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
109+ gert::Shape* indices = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);
110+ ASSERT_EQ(Shape2String(*output), "[5, 125, 46, 589]");
111+ ASSERT_EQ(Shape2String(*indices), "[5, 125, 46, 589]");
112+}
113+ 
114+TEST_F(MaxPoolWithArgmaxV3Infer, maxpool_with_argmax_v3_infershape_test_3)
115+{
116+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolWithArgmaxV3")->infer_shape;
117+ 
118+ gert::StorageShape xShape = {{256, 144, 589}, {}};
119+ gert::StorageShape yShape = {{}, {}};
120+ gert::StorageShape indicesShape = {{}, {}};
121+ auto holder = gert::InferShapeContextFaker()
122+ .NodeIoNum(1, 2)
123+ .IrInstanceNum({1, 2})
124+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_ND, ge::Format::FORMAT_RESERVED)
125+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_ND, ge::Format::FORMAT_RESERVED)
126+ .NodeOutputTd(1, ge::DT_INT32, ge::Format::FORMAT_ND, ge::Format::FORMAT_RESERVED)
127+ .NodeAttrs(
128+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 3})},
129+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 3})},
130+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
131+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
132+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({9, 5})},
133+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
134+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
135+ .InputShapes({&xShape})
136+ .OutputShapes({&yShape, &indicesShape})
137+ .Build();
138+ 
139+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
140+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
141+ gert::Shape* indices = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);
142+ ASSERT_EQ(Shape2String(*output), "[125, 46, 589]");
143+ ASSERT_EQ(Shape2String(*indices), "[125, 46, 589]");
144+}
145+ 
146+TEST_F(MaxPoolWithArgmaxV3Infer, maxpool_with_argmax_v3_infershape_test_4)
147+{
148+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolWithArgmaxV3")->infer_shape;
149+ 
150+ gert::StorageShape xShape = {{1, 3, -1, -1}, {}};
151+ gert::StorageShape yShape = {{}, {}};
152+ gert::StorageShape indicesShape = {{}, {}};
153+ auto holder = gert::InferShapeContextFaker()
154+ .NodeIoNum(1, 2)
155+ .IrInstanceNum({1, 2})
156+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
157+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
158+ .NodeOutputTd(1, ge::DT_INT32, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
159+ .NodeAttrs(
160+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
161+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 2})},
162+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
163+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
164+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
165+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(true)},
166+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
167+ .InputShapes({&xShape})
168+ .OutputShapes({&yShape, &indicesShape})
169+ .Build();
170+ 
171+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
172+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
173+ gert::Shape* indices = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);
174+ ASSERT_EQ(Shape2String(*output), "[1, 3, -1, -1]");
175+ ASSERT_EQ(Shape2String(*indices), "[1, 3, -1, -1]");
176+}
177+ 
178+TEST_F(MaxPoolWithArgmaxV3Infer, maxpool_with_argmax_v3_infershape_test_5)
179+{
180+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolWithArgmaxV3")->infer_shape;
181+ 
182+ gert::StorageShape xShape = {{-2}, {}};
183+ gert::StorageShape yShape = {{}, {}};
184+ gert::StorageShape indicesShape = {{}, {}};
185+ auto holder = gert::InferShapeContextFaker()
186+ .NodeIoNum(1, 2)
187+ .IrInstanceNum({1, 2})
188+ .NodeInputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
189+ .NodeOutputTd(0, ge::DT_FLOAT, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
190+ .NodeOutputTd(1, ge::DT_INT32, ge::Format::FORMAT_NHWC, ge::Format::FORMAT_RESERVED)
191+ .NodeAttrs(
192+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
193+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({2, 2})},
194+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
195+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
196+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
197+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(true)},
198+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
199+ .InputShapes({&xShape})
200+ .OutputShapes({&yShape, &indicesShape})
201+ .Build();
202+ 
203+ ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
204+ gert::Shape* output = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
205+ gert::Shape* indices = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);
206+ ASSERT_EQ(Shape2String(*output), "[-2]");
207+ ASSERT_EQ(Shape2String(*indices), "[-2]");
208+}
209+ 
210+TEST_F(MaxPoolWithArgmaxV3Infer, max_pool_with_argmax_v3_inferdtype_success_01)
211+{
212+ fe::PlatformInfo platformInfo;
213+ fe::OptionalInfo optiCompilationInfo;
214+ platformInfo.soc_info.ai_core_cnt = 64;
215+ platformInfo.str_info.short_soc_version = "Ascend910_95";
216+ optiCompilationInfo.soc_version = "Ascend910_9589";
217+ fe::PlatformInfoManager::Instance().platform_info_map_["Ascend910_9589"] = platformInfo;
218+ fe::PlatformInfoManager::Instance().SetOptionalCompilationInfo(optiCompilationInfo);
219+ 
220+ auto inferDtypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MaxPoolWithArgmaxV3")->infer_datatype;
221+ 
222+ ge::DataType x_dtype = ge::DT_FLOAT16;
223+ ge::DataType y_dtype = ge::DT_FLOAT16;
224+ ge::DataType argmax_dtype = ge::DT_INT32;
225+ ge::DataType expect_output_dtype = ge::DT_FLOAT16;
226+ 
227+ auto holder = gert::InferDataTypeContextFaker()
228+ .NodeIoNum(1, 2)
229+ .IrInstanceNum({
230+ 1,
231+ })
232+ .InputDataTypes({&x_dtype})
233+ .OutputDataTypes({&y_dtype, &argmax_dtype})
234+ .NodeAttrs(
235+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({4, 4})},
236+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({4, 4})},
237+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
238+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
239+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
240+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
241+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
242+ .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
243+ .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
244+ .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
245+ .Build();
246+ auto context = holder.GetContext<gert::InferDataTypeContext>();
247+ ASSERT_EQ(inferDtypeFunc(context), ge::GRAPH_SUCCESS);
248+ ASSERT_NE(context, nullptr);
249+ EXPECT_EQ(context->GetOutputDataType(0), expect_output_dtype);
250+}
251+} // namespace
Apooling/max_pool_with_argmax_v3/tests/ut/op_host/test_max_pool_with_argmax_v3_tiling.cpp+1172-0
@@ -0,0 +1,1172 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include <fstream>
13+#include <vector>
14+#include <gtest/gtest.h>
15+ 
16+#include "log/log.h"
17+#include "kernel_run_context_facker.h"
18+#include "test_cube_util.h"
19+#include "exe_graph/runtime/storage_format.h"
20+#include "exe_graph/runtime/storage_shape.h"
21+#include "platform/platform_infos_def.h"
22+#include "ut_op_util.h"
23+#include "../../../op_host/arch35/max_pool_with_argmax_v3_tiling.h"
24+ 
25+using namespace ut_util;
26+using namespace std;
27+using namespace ge;
28+ 
29+class MaxPoolWithArgmaxV3Tiling : public testing::Test {
30+protected:
31+ static void SetUpTestCase()
32+ {
33+ std::cout << "MaxPoolWithArgmaxV3Tiling SetUp" << std::endl;
34+ }
35+ 
36+ static void TearDownTestCase()
37+ {
38+ std::cout << "MaxPoolWithArgmaxV3Tiling TearDown" << std::endl;
39+ }
40+};
41+ 
42+static void ExecuteTestCase(
43+ gert::StorageShape xShape, gert::StorageShape yShape, gert::StorageShape argmaxShape, std::vector<int64_t> ksize,
44+ std::vector<int64_t> strides, std::vector<int64_t> pads, std::vector<int64_t> dilation, ge::DataType dtype,
45+ int64_t index_dtype, bool ceil_mode, std::string data_format, uint64_t except_tilingkey, std::string expect)
46+{
47+ dlog_setlevel(0, 0, 0);
48+ 
49+ string compile_info_string = R"({
50+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
51+ "Intrinsic_fix_pipe_l0c2out": false,
52+ "Intrinsic_data_move_l12ub": true,
53+ "Intrinsic_data_move_l0c2ub": true,
54+ "Intrinsic_data_move_out2l1_nd2nz": false,
55+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
56+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
57+ "CORE_NUM": 64}
58+ })";
59+ map<string, string> soc_infos;
60+ map<string, string> aicore_spec;
61+ map<string, string> intrinsics;
62+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
63+ std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend910_95"}};
64+ 
65+ // platform info
66+ fe::PlatFormInfos platform_info;
67+ platform_info.Init();
68+ // compile info
69+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
70+ 
71+ std::string op_type("MaxPoolWithArgmaxV3");
72+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
73+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
74+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
75+ 
76+ // tilingParseFunc simulate
77+ auto kernel_holder =
78+ gert::KernelRunContextFaker()
79+ .KernelIONum(2, 1)
80+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
81+ .Outputs({&compile_info})
82+ .Build();
83+ 
84+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
85+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
86+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
87+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
88+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
89+ "AICoreintrinsicDtypeMap", intrinsics);
90+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
91+ "version", soc_version_infos);
92+ 
93+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
94+ 
95+ // tilingFunc simulate
96+ auto param = gert::TilingData::CreateCap(4096);
97+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
98+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
99+ ASSERT_NE(param, nullptr);
100+ auto holder = gert::TilingContextFaker()
101+ .SetOpType(op_type)
102+ .NodeIoNum(1, 2)
103+ .IrInstanceNum({1})
104+ .InputShapes({&xShape})
105+ .OutputShapes({&yShape, &argmaxShape})
106+ .CompileInfo(&compile_info)
107+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
108+ .NodeInputTd(0, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
109+ .NodeOutputTd(0, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
110+ .NodeOutputTd(1, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
111+ .NodeAttrs(
112+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(ksize)},
113+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(strides)},
114+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(pads)},
115+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(index_dtype)},
116+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>(dilation)},
117+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(ceil_mode)},
118+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>(data_format)}})
119+ .TilingData(param.get())
120+ .Workspace(ws_size)
121+ .Build();
122+ 
123+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
124+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
125+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
126+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
127+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
128+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
129+ 
130+ // workspaces nullptr return failed
131+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
132+ auto tiling_key = tiling_context->GetTilingKey();
133+ ASSERT_EQ(tiling_key, except_tilingkey);
134+ // auto tilingData = tiling_context->GetRawTilingData();
135+ // ASSERT_NE(tilingData, nullptr);
136+ // dlog_setlevel(0, 3, 0);
137+ // auto tiling_data_result = to_string<int64_t>(tilingData->GetData(), tilingData->GetDataSize());
138+ // std::cout << tiling_data_result << std::endl;
139+ // EXPECT_EQ(tiling_data_result, expect);
140+}
141+ 
142+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_Test1)
143+{
144+ dlog_setlevel(0, 0, 0);
145+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
146+ gert::StorageShape yShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
147+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
148+ 
149+ string compile_info_string = R"({
150+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
151+ "Intrinsic_fix_pipe_l0c2out": false,
152+ "Intrinsic_data_move_l12ub": true,
153+ "Intrinsic_data_move_l0c2ub": true,
154+ "Intrinsic_data_move_out2l1_nd2nz": false,
155+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
156+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
157+ "CORE_NUM": 64}
158+ })";
159+ map<string, string> soc_infos;
160+ map<string, string> aicore_spec;
161+ map<string, string> intrinsics;
162+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
163+ 
164+ // platform info
165+ fe::PlatFormInfos platform_info;
166+ platform_info.Init();
167+ // compile info
168+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
169+ 
170+ std::string op_type("MaxPoolWithArgmaxV3");
171+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
172+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
173+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
174+ 
175+ // tilingParseFunc simulate
176+ auto kernel_holder =
177+ gert::KernelRunContextFaker()
178+ .KernelIONum(2, 1)
179+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
180+ .Outputs({&compile_info})
181+ .Build();
182+ 
183+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
184+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
185+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
186+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
187+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
188+ "AICoreintrinsicDtypeMap", intrinsics);
189+ 
190+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
191+ 
192+ // tilingFunc simulate
193+ auto param = gert::TilingData::CreateCap(4096);
194+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
195+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
196+ ASSERT_NE(param, nullptr);
197+ auto holder = gert::TilingContextFaker()
198+ .SetOpType(op_type)
199+ .NodeIoNum(1, 2)
200+ .IrInstanceNum({1})
201+ .InputShapes({&xShape})
202+ .OutputShapes({&yShape, &argmaxShape})
203+ .CompileInfo(&compile_info)
204+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
205+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
206+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
207+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
208+ .NodeAttrs(
209+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
210+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
211+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
212+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
213+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
214+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
215+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NCHW")}})
216+ .TilingData(param.get())
217+ .Workspace(ws_size)
218+ .Build();
219+ 
220+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
221+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
222+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
223+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
224+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
225+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
226+ 
227+ // workspaces nullptr return failed
228+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
229+ auto tiling_key = tiling_context->GetTilingKey();
230+ ASSERT_EQ(tiling_key, 400001);
231+ // auto tilingData = tiling_context->GetRawTilingData();
232+ // ASSERT_NE(tilingData, nullptr);
233+ // dlog_setlevel(0, 3, 0);
234+ // EXPECT_EQ(to_string<int64_t>(tilingData->GetData(), tilingData->GetDataSize()),
235+ // "64 64 1 1 64 64 64 64 0 0 1 1 6 10 7 1 0 0 0 0 30016 256 512 ");
236+}
237+ 
238+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test2)
239+{
240+ dlog_setlevel(0, 0, 0);
241+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
242+ gert::StorageShape yShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
243+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
244+ 
245+ string compile_info_string = R"({
246+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
247+ "Intrinsic_fix_pipe_l0c2out": false,
248+ "Intrinsic_data_move_l12ub": true,
249+ "Intrinsic_data_move_l0c2ub": true,
250+ "Intrinsic_data_move_out2l1_nd2nz": false,
251+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
252+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
253+ "CORE_NUM": 64}
254+ })";
255+ map<string, string> soc_infos;
256+ map<string, string> aicore_spec;
257+ map<string, string> intrinsics;
258+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
259+ 
260+ // platform info
261+ fe::PlatFormInfos platform_info;
262+ platform_info.Init();
263+ // compile info
264+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
265+ 
266+ std::string op_type("MaxPoolWithArgmaxV3");
267+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
268+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
269+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
270+ 
271+ // tilingParseFunc simulate
272+ auto kernel_holder =
273+ gert::KernelRunContextFaker()
274+ .KernelIONum(2, 1)
275+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
276+ .Outputs({&compile_info})
277+ .Build();
278+ 
279+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
280+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
281+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
282+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
283+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
284+ "AICoreintrinsicDtypeMap", intrinsics);
285+ 
286+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
287+ 
288+ // tilingFunc simulate
289+ auto param = gert::TilingData::CreateCap(4096);
290+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
291+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
292+ ASSERT_NE(param, nullptr);
293+ auto holder = gert::TilingContextFaker()
294+ .SetOpType(op_type)
295+ .NodeIoNum(1, 2)
296+ .IrInstanceNum({1})
297+ .InputShapes({&xShape})
298+ .OutputShapes({&yShape, &argmaxShape})
299+ .CompileInfo(&compile_info)
300+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
301+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
302+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
303+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
304+ .NodeAttrs(
305+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
306+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
307+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
308+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
309+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
310+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
311+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
312+ .TilingData(param.get())
313+ .Workspace(ws_size)
314+ .Build();
315+ 
316+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
317+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
318+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
319+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
320+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
321+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
322+ 
323+ // workspaces nullptr return failed
324+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
325+ auto tiling_key = tiling_context->GetTilingKey();
326+ ASSERT_EQ(tiling_key, 800001);
327+ // auto tilingData = tiling_context->GetRawTilingData();
328+ // ASSERT_NE(tilingData, nullptr);
329+ // EXPECT_EQ(to_string<int64_t>(tilingData->GetData(), tilingData->GetDataSize()),
330+ // "64 3 64 3 1 64 64 64 64 0 0 1 1 1 1 2 1 1 3 1 1 1 64 64 1 1 1 6 114688 256 256 0 1 7 1 10 64 64 1
331+ // 800001 ");
332+ // dlog_setlevel(0, 3, 0);
333+}
334+ 
335+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test3)
336+{
337+ dlog_setlevel(0, 0, 0);
338+ gert::StorageShape xShape = {{2, 3, 64}, {2, 3, 64}};
339+ gert::StorageShape yShape = {{2, 3, 1}, {2, 3, 1}};
340+ gert::StorageShape argmaxShape = {{2, 3, 1}, {2, 3, 1}};
341+ 
342+ string compile_info_string = R"({
343+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
344+ "Intrinsic_fix_pipe_l0c2out": false,
345+ "Intrinsic_data_move_l12ub": true,
346+ "Intrinsic_data_move_l0c2ub": true,
347+ "Intrinsic_data_move_out2l1_nd2nz": false,
348+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
349+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
350+ "CORE_NUM": 64}
351+ })";
352+ map<string, string> soc_infos;
353+ map<string, string> aicore_spec;
354+ map<string, string> intrinsics;
355+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
356+ 
357+ // platform info
358+ fe::PlatFormInfos platform_info;
359+ platform_info.Init();
360+ // compile info
361+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
362+ 
363+ std::string op_type("MaxPoolWithArgmaxV3");
364+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
365+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
366+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
367+ 
368+ // tilingParseFunc simulate
369+ auto kernel_holder =
370+ gert::KernelRunContextFaker()
371+ .KernelIONum(2, 1)
372+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
373+ .Outputs({&compile_info})
374+ .Build();
375+ 
376+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
377+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
378+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
379+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
380+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
381+ "AICoreintrinsicDtypeMap", intrinsics);
382+ 
383+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
384+ 
385+ // tilingFunc simulate
386+ auto param = gert::TilingData::CreateCap(4096);
387+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
388+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
389+ ASSERT_NE(param, nullptr);
390+ auto holder = gert::TilingContextFaker()
391+ .SetOpType(op_type)
392+ .NodeIoNum(1, 2)
393+ .IrInstanceNum({1})
394+ .InputShapes({&xShape})
395+ .OutputShapes({&yShape, &argmaxShape})
396+ .CompileInfo(&compile_info)
397+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
398+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
399+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
400+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
401+ .NodeAttrs(
402+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
403+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
404+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
405+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
406+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
407+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
408+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
409+ .TilingData(param.get())
410+ .Workspace(ws_size)
411+ .Build();
412+ 
413+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
414+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
415+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
416+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
417+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
418+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
419+ 
420+ // workspaces nullptr return failed
421+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
422+ dlog_setlevel(0, 3, 0);
423+}
424+ 
425+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test4)
426+{
427+ dlog_setlevel(0, 0, 0);
428+ gert::StorageShape xShape = {{2, 3, 0, 64}, {2, 3, 0, 64}};
429+ gert::StorageShape yShape = {{2, 3, 0, 1}, {2, 3, 0, 1}};
430+ gert::StorageShape argmaxShape = {{2, 3, 0, 1}, {2, 3, 0, 1}};
431+ 
432+ string compile_info_string = R"({
433+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
434+ "Intrinsic_fix_pipe_l0c2out": false,
435+ "Intrinsic_data_move_l12ub": true,
436+ "Intrinsic_data_move_l0c2ub": true,
437+ "Intrinsic_data_move_out2l1_nd2nz": false,
438+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
439+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
440+ "CORE_NUM": 64}
441+ })";
442+ map<string, string> soc_infos;
443+ map<string, string> aicore_spec;
444+ map<string, string> intrinsics;
445+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
446+ 
447+ // platform info
448+ fe::PlatFormInfos platform_info;
449+ platform_info.Init();
450+ // compile info
451+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
452+ 
453+ std::string op_type("MaxPoolWithArgmaxV3");
454+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
455+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
456+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
457+ 
458+ // tilingParseFunc simulate
459+ auto kernel_holder =
460+ gert::KernelRunContextFaker()
461+ .KernelIONum(2, 1)
462+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
463+ .Outputs({&compile_info})
464+ .Build();
465+ 
466+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
467+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
468+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
469+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
470+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
471+ "AICoreintrinsicDtypeMap", intrinsics);
472+ 
473+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
474+ 
475+ // tilingFunc simulate
476+ auto param = gert::TilingData::CreateCap(4096);
477+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
478+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
479+ ASSERT_NE(param, nullptr);
480+ auto holder = gert::TilingContextFaker()
481+ .SetOpType(op_type)
482+ .NodeIoNum(1, 2)
483+ .IrInstanceNum({1})
484+ .InputShapes({&xShape})
485+ .OutputShapes({&yShape, &argmaxShape})
486+ .CompileInfo(&compile_info)
487+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
488+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
489+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
490+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
491+ .NodeAttrs(
492+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
493+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
494+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
495+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
496+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
497+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
498+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
499+ .TilingData(param.get())
500+ .Workspace(ws_size)
501+ .Build();
502+ 
503+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
504+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
505+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
506+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
507+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
508+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
509+ 
510+ // workspaces nullptr return failed
511+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
512+ dlog_setlevel(0, 3, 0);
513+}
514+ 
515+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test5)
516+{
517+ dlog_setlevel(0, 0, 0);
518+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
519+ gert::StorageShape yShape = {{2, 3, 1, 0}, {2, 3, 1, 0}};
520+ gert::StorageShape argmaxShape = {{2, 3, 1, 0}, {2, 3, 1, 0}};
521+ 
522+ string compile_info_string = R"({
523+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
524+ "Intrinsic_fix_pipe_l0c2out": false,
525+ "Intrinsic_data_move_l12ub": true,
526+ "Intrinsic_data_move_l0c2ub": true,
527+ "Intrinsic_data_move_out2l1_nd2nz": false,
528+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
529+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
530+ "CORE_NUM": 64}
531+ })";
532+ map<string, string> soc_infos;
533+ map<string, string> aicore_spec;
534+ map<string, string> intrinsics;
535+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
536+ 
537+ // platform info
538+ fe::PlatFormInfos platform_info;
539+ platform_info.Init();
540+ // compile info
541+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
542+ 
543+ std::string op_type("MaxPoolWithArgmaxV3");
544+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
545+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
546+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
547+ 
548+ // tilingParseFunc simulate
549+ auto kernel_holder =
550+ gert::KernelRunContextFaker()
551+ .KernelIONum(2, 1)
552+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
553+ .Outputs({&compile_info})
554+ .Build();
555+ 
556+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
557+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
558+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
559+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
560+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
561+ "AICoreintrinsicDtypeMap", intrinsics);
562+ 
563+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
564+ 
565+ // tilingFunc simulate
566+ auto param = gert::TilingData::CreateCap(4096);
567+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
568+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
569+ ASSERT_NE(param, nullptr);
570+ auto holder = gert::TilingContextFaker()
571+ .SetOpType(op_type)
572+ .NodeIoNum(1, 2)
573+ .IrInstanceNum({1})
574+ .InputShapes({&xShape})
575+ .OutputShapes({&yShape, &argmaxShape})
576+ .CompileInfo(&compile_info)
577+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
578+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
579+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
580+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
581+ .NodeAttrs(
582+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
583+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
584+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
585+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
586+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
587+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
588+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("ND")}})
589+ .TilingData(param.get())
590+ .Workspace(ws_size)
591+ .Build();
592+ 
593+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
594+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
595+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
596+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
597+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
598+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
599+ 
600+ // workspaces nullptr return failed
601+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
602+ dlog_setlevel(0, 3, 0);
603+}
604+ 
605+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test6)
606+{
607+ dlog_setlevel(0, 0, 0);
608+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
609+ gert::StorageShape yShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
610+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
611+ 
612+ string compile_info_string = R"({
613+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
614+ "Intrinsic_fix_pipe_l0c2out": false,
615+ "Intrinsic_data_move_l12ub": true,
616+ "Intrinsic_data_move_l0c2ub": true,
617+ "Intrinsic_data_move_out2l1_nd2nz": false,
618+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
619+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
620+ "CORE_NUM": 64}
621+ })";
622+ map<string, string> soc_infos;
623+ map<string, string> aicore_spec;
624+ map<string, string> intrinsics;
625+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
626+ 
627+ // platform info
628+ fe::PlatFormInfos platform_info;
629+ platform_info.Init();
630+ // compile info
631+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
632+ 
633+ std::string op_type("MaxPoolWithArgmaxV3");
634+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
635+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
636+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
637+ 
638+ // tilingParseFunc simulate
639+ auto kernel_holder =
640+ gert::KernelRunContextFaker()
641+ .KernelIONum(2, 1)
642+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
643+ .Outputs({&compile_info})
644+ .Build();
645+ 
646+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
647+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
648+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
649+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
650+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
651+ "AICoreintrinsicDtypeMap", intrinsics);
652+ 
653+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
654+ 
655+ // tilingFunc simulate
656+ auto param = gert::TilingData::CreateCap(4096);
657+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
658+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
659+ ASSERT_NE(param, nullptr);
660+ auto holder = gert::TilingContextFaker()
661+ .SetOpType(op_type)
662+ .NodeIoNum(1, 2)
663+ .IrInstanceNum({1})
664+ .InputShapes({&xShape})
665+ .OutputShapes({&yShape, &argmaxShape})
666+ .CompileInfo(&compile_info)
667+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
668+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
669+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
670+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
671+ .NodeAttrs(
672+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
673+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
674+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
675+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
676+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
677+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
678+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
679+ .TilingData(param.get())
680+ .Workspace(ws_size)
681+ .Build();
682+ 
683+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
684+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
685+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
686+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
687+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
688+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
689+ 
690+ // workspaces nullptr return failed
691+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
692+ dlog_setlevel(0, 3, 0);
693+}
694+ 
695+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test7)
696+{
697+ dlog_setlevel(0, 0, 0);
698+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
699+ gert::StorageShape yShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
700+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
701+ 
702+ string compile_info_string = R"({
703+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
704+ "Intrinsic_fix_pipe_l0c2out": false,
705+ "Intrinsic_data_move_l12ub": true,
706+ "Intrinsic_data_move_l0c2ub": true,
707+ "Intrinsic_data_move_out2l1_nd2nz": false,
708+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
709+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
710+ "CORE_NUM": 64}
711+ })";
712+ map<string, string> soc_infos;
713+ map<string, string> aicore_spec;
714+ map<string, string> intrinsics;
715+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
716+ 
717+ // platform info
718+ fe::PlatFormInfos platform_info;
719+ platform_info.Init();
720+ // compile info
721+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
722+ 
723+ std::string op_type("MaxPoolWithArgmaxV3");
724+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
725+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
726+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
727+ 
728+ // tilingParseFunc simulate
729+ auto kernel_holder =
730+ gert::KernelRunContextFaker()
731+ .KernelIONum(2, 1)
732+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
733+ .Outputs({&compile_info})
734+ .Build();
735+ 
736+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
737+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
738+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
739+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
740+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
741+ "AICoreintrinsicDtypeMap", intrinsics);
742+ 
743+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
744+ 
745+ // tilingFunc simulate
746+ auto param = gert::TilingData::CreateCap(4096);
747+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
748+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
749+ ASSERT_NE(param, nullptr);
750+ auto holder = gert::TilingContextFaker()
751+ .SetOpType(op_type)
752+ .NodeIoNum(1, 2)
753+ .IrInstanceNum({1})
754+ .InputShapes({&xShape})
755+ .OutputShapes({&yShape, &argmaxShape})
756+ .CompileInfo(&compile_info)
757+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
758+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
759+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
760+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
761+ .NodeAttrs(
762+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
763+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
764+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
765+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
766+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({1, 1})},
767+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
768+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
769+ .TilingData(param.get())
770+ .Workspace(ws_size)
771+ .Build();
772+ 
773+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
774+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
775+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
776+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
777+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
778+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
779+ 
780+ // workspaces nullptr return failed
781+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
782+ dlog_setlevel(0, 3, 0);
783+}
784+ 
785+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NHWC_Test8)
786+{
787+ dlog_setlevel(0, 0, 0);
788+ gert::StorageShape xShape = {{2, 3, 64, 64}, {2, 3, 64, 64}};
789+ gert::StorageShape yShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
790+ gert::StorageShape argmaxShape = {{2, 3, 1, 1}, {2, 3, 1, 1}};
791+ 
792+ string compile_info_string = R"({
793+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
794+ "Intrinsic_fix_pipe_l0c2out": false,
795+ "Intrinsic_data_move_l12ub": true,
796+ "Intrinsic_data_move_l0c2ub": true,
797+ "Intrinsic_data_move_out2l1_nd2nz": false,
798+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
799+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
800+ "CORE_NUM": 64}
801+ })";
802+ map<string, string> soc_infos;
803+ map<string, string> aicore_spec;
804+ map<string, string> intrinsics;
805+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
806+ 
807+ // platform info
808+ fe::PlatFormInfos platform_info;
809+ platform_info.Init();
810+ // compile info
811+ optiling::MaxPoolWithArgmaxV3CompileInfo compile_info;
812+ 
813+ std::string op_type("MaxPoolWithArgmaxV3");
814+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
815+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
816+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
817+ 
818+ // tilingParseFunc simulate
819+ auto kernel_holder =
820+ gert::KernelRunContextFaker()
821+ .KernelIONum(2, 1)
822+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
823+ .Outputs({&compile_info})
824+ .Build();
825+ 
826+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
827+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
828+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
829+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
830+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
831+ "AICoreintrinsicDtypeMap", intrinsics);
832+ 
833+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
834+ 
835+ // tilingFunc simulate
836+ auto param = gert::TilingData::CreateCap(4096);
837+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
838+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
839+ ASSERT_NE(param, nullptr);
840+ auto holder = gert::TilingContextFaker()
841+ .SetOpType(op_type)
842+ .NodeIoNum(1, 2)
843+ .IrInstanceNum({1})
844+ .InputShapes({&xShape})
845+ .OutputShapes({&yShape, &argmaxShape})
846+ .CompileInfo(&compile_info)
847+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
848+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
849+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
850+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
851+ .NodeAttrs(
852+ {{"ksize", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
853+ {"strides", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({64, 64})},
854+ {"pads", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
855+ {"dtype", Ops::NN::AnyValue::CreateFrom<int64_t>(3)},
856+ {"dilation", Ops::NN::AnyValue::CreateFrom<std::vector<int64_t>>({0, 0})},
857+ {"ceil_mode", Ops::NN::AnyValue::CreateFrom<bool>(false)},
858+ {"data_format", Ops::NN::AnyValue::CreateFrom<std::string>("NHWC")}})
859+ .TilingData(param.get())
860+ .Workspace(ws_size)
861+ .Build();
862+ 
863+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
864+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
865+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
866+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
867+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
868+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
869+ 
870+ // workspaces nullptr return failed
871+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
872+ dlog_setlevel(0, 3, 0);
873+}
874+ 
875+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_Test9)
876+{
877+ gert::StorageShape xShape = {{4, 163, 1024, 600}, {4, 163, 1024, 600}};
878+ gert::StorageShape yShape = {{4, 163, 2, 2}, {4, 163, 2, 2}};
879+ gert::StorageShape argmaxShape = {{4, 163, 2, 2}, {4, 163, 2, 2}};
880+ std::vector<int64_t> ksize = {324, 457};
881+ std::vector<int64_t> strides = {858, 457};
882+ std::vector<int64_t> pads = {30, 132};
883+ std::vector<int64_t> dilation = {1, 1};
884+ ge::DataType dtype = ge::DT_FLOAT;
885+ int64_t index_dtype = 3;
886+ bool ceil_mode = false;
887+ std::string data_format = "NCHW";
888+ uint64_t except_tilingkey = 311110;
889+ std::string expect = "1024 600 2 2 457 324 457 858 132 30 1 1 40 48 2608 64 29184 0 ";
890+ ExecuteTestCase(
891+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
892+ except_tilingkey, expect);
893+}
894+ 
895+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_Test10)
896+{
897+ gert::StorageShape xShape = {{4, 833, 1024, 640}, {4, 833, 1024, 640}};
898+ gert::StorageShape yShape = {{4, 833, 2, 2}, {4, 833, 2, 2}};
899+ gert::StorageShape argmaxShape = {{4, 833, 2, 2}, {4, 833, 2, 2}};
900+ std::vector<int64_t> ksize = {455, 513};
901+ std::vector<int64_t> strides = {455, 256};
902+ std::vector<int64_t> pads = {106, 163};
903+ std::vector<int64_t> dilation = {1, 1};
904+ ge::DataType dtype = ge::DT_FLOAT;
905+ int64_t index_dtype = 9;
906+ bool ceil_mode = true;
907+ std::string data_format = "NCHW";
908+ uint64_t except_tilingkey = 311110;
909+ std::string expect = "1024 640 2 2 513 455 256 455 163 106 1 1 208 16 13328 64 29184 0 ";
910+ ExecuteTestCase(
911+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
912+ except_tilingkey, expect);
913+}
914+ 
915+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_Gather_Test0)
916+{
917+ gert::StorageShape xShape = {{4, 16, 32, 32}, {4, 16, 32, 32}};
918+ gert::StorageShape yShape = {{4, 16, 16, 16}, {4, 16, 16, 16}};
919+ gert::StorageShape argmaxShape = {{4, 16, 16, 16}, {4, 16, 16, 16}};
920+ std::vector<int64_t> ksize = {2, 2};
921+ std::vector<int64_t> strides = {2, 2};
922+ std::vector<int64_t> pads = {0, 0};
923+ std::vector<int64_t> dilation = {1, 1};
924+ ge::DataType dtype = ge::DT_FLOAT;
925+ int64_t index_dtype = 3;
926+ bool ceil_mode = true;
927+ std::string data_format = "NCHW";
928+ uint64_t except_tilingkey = 300001;
929+ std::string expect = "32 32 16 16 2 2 2 2 0 0 2 2 32 16 16 1 16 16 1 1 1 32 8192 2048 2048 0 1 1 ";
930+ ExecuteTestCase(
931+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
932+ except_tilingkey, expect);
933+}
934+ 
935+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_Gather_Test1)
936+{
937+ gert::StorageShape xShape = {{4, 16, 30, 30}, {4, 16, 30, 30}};
938+ gert::StorageShape yShape = {{4, 16, 16, 16}, {4, 16, 16, 16}};
939+ gert::StorageShape argmaxShape = {{4, 16, 16, 16}, {4, 16, 16, 16}};
940+ std::vector<int64_t> ksize = {2, 2};
941+ std::vector<int64_t> strides = {2, 2};
942+ std::vector<int64_t> pads = {1, 1};
943+ std::vector<int64_t> dilation = {1, 1};
944+ ge::DataType dtype = ge::DT_FLOAT;
945+ int64_t index_dtype = 3;
946+ bool ceil_mode = false;
947+ std::string data_format = "NCHW";
948+ uint64_t except_tilingkey = 300002;
949+ std::string expect = "30 30 16 16 2 2 2 2 1 1 2 2 32 16 16 1 16 16 1 1 1 32 8192 2048 2048 1 1 1 ";
950+ ExecuteTestCase(
951+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
952+ except_tilingkey, expect);
953+}
954+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Float_Test01)
955+{
956+ gert::StorageShape xShape = {{1, 1, 68, 22}, {1, 1, 68, 22}};
957+ gert::StorageShape yShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
958+ gert::StorageShape argmaxShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
959+ std::vector<int64_t> ksize = {64, 16};
960+ std::vector<int64_t> strides = {2, 2};
961+ std::vector<int64_t> pads = {0, 0};
962+ std::vector<int64_t> dilation = {1, 1};
963+ ge::DataType dtype = ge::DT_FLOAT;
964+ int64_t index_dtype = 3;
965+ bool ceil_mode = true;
966+ std::string data_format = "NCHW";
967+ uint64_t except_tilingkey = 400001;
968+ std::string expect = "68 22 3 4 16 64 2 2 0 0 1 1 12 5 13 12 0 0 0 0 30016 256 512 ";
969+ ExecuteTestCase(
970+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
971+ except_tilingkey, expect);
972+}
973+ 
974+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Float_Test02)
975+{
976+ gert::StorageShape xShape = {{1, 1, 4, 522}, {1, 1, 4, 522}};
977+ gert::StorageShape yShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
978+ gert::StorageShape argmaxShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
979+ std::vector<int64_t> ksize = {2, 516};
980+ std::vector<int64_t> strides = {2, 2};
981+ std::vector<int64_t> pads = {0, 0};
982+ std::vector<int64_t> dilation = {1, 1};
983+ ge::DataType dtype = ge::DT_FLOAT;
984+ int64_t index_dtype = 3;
985+ bool ceil_mode = true;
986+ std::string data_format = "NCHW";
987+ uint64_t except_tilingkey = 400001;
988+ std::string expect = "4 522 2 4 516 2 2 2 0 0 1 1 8 8 0 0 1 129 129 4 30016 256 512 ";
989+ ExecuteTestCase(
990+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
991+ except_tilingkey, expect);
992+}
993+ 
994+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Float_Test03)
995+{
996+ gert::StorageShape xShape = {{1, 1, 68, 22}, {1, 1, 68, 22}};
997+ gert::StorageShape yShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
998+ gert::StorageShape argmaxShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
999+ std::vector<int64_t> ksize = {64, 16};
1000+ std::vector<int64_t> strides = {2, 2};
1001+ std::vector<int64_t> pads = {0, 0};
1002+ std::vector<int64_t> dilation = {1, 1};
1003+ ge::DataType dtype = ge::DT_FLOAT;
1004+ int64_t index_dtype = 9;
1005+ bool ceil_mode = true;
1006+ std::string data_format = "NCHW";
1007+ uint64_t except_tilingkey = 400002;
1008+ std::string expect = "68 22 3 4 16 64 2 2 0 0 1 1 12 5 13 12 0 0 0 0 30016 256 512 ";
1009+ ExecuteTestCase(
1010+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1011+ except_tilingkey, expect);
1012+}
1013+ 
1014+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Float_Test04)
1015+{
1016+ gert::StorageShape xShape = {{1, 1, 4, 522}, {1, 1, 4, 522}};
1017+ gert::StorageShape yShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1018+ gert::StorageShape argmaxShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1019+ std::vector<int64_t> ksize = {2, 516};
1020+ std::vector<int64_t> strides = {2, 2};
1021+ std::vector<int64_t> pads = {0, 0};
1022+ std::vector<int64_t> dilation = {1, 1};
1023+ ge::DataType dtype = ge::DT_FLOAT;
1024+ int64_t index_dtype = 9;
1025+ bool ceil_mode = true;
1026+ std::string data_format = "NCHW";
1027+ uint64_t except_tilingkey = 400002;
1028+ std::string expect = "4 522 2 4 516 2 2 2 0 0 1 1 8 8 0 0 1 129 129 4 30016 256 512 ";
1029+ ExecuteTestCase(
1030+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1031+ except_tilingkey, expect);
1032+}
1033+ 
1034+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Bfloat16_Test01)
1035+{
1036+ gert::StorageShape xShape = {{1, 1, 68, 22}, {1, 1, 68, 22}};
1037+ gert::StorageShape yShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
1038+ gert::StorageShape argmaxShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
1039+ std::vector<int64_t> ksize = {64, 16};
1040+ std::vector<int64_t> strides = {2, 2};
1041+ std::vector<int64_t> pads = {0, 0};
1042+ std::vector<int64_t> dilation = {1, 1};
1043+ ge::DataType dtype = ge::DT_BF16;
1044+ int64_t index_dtype = 3;
1045+ bool ceil_mode = true;
1046+ std::string data_format = "NCHW";
1047+ uint64_t except_tilingkey = 300001;
1048+ std::string expect = "68 22 3 4 64 16 2 2 0 0 1 1 1 1 1 3 1 1 4 1 1 12 2048 32 64 0 1 1 ";
1049+ ExecuteTestCase(
1050+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1051+ except_tilingkey, expect);
1052+}
1053+ 
1054+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Bfloat16_Test02)
1055+{
1056+ gert::StorageShape xShape = {{1, 1, 4, 522}, {1, 1, 4, 522}};
1057+ gert::StorageShape yShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1058+ gert::StorageShape argmaxShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1059+ std::vector<int64_t> ksize = {2, 516};
1060+ std::vector<int64_t> strides = {2, 2};
1061+ std::vector<int64_t> pads = {0, 0};
1062+ std::vector<int64_t> dilation = {1, 1};
1063+ ge::DataType dtype = ge::DT_BF16;
1064+ int64_t index_dtype = 3;
1065+ bool ceil_mode = true;
1066+ std::string data_format = "NCHW";
1067+ uint64_t except_tilingkey = 400003;
1068+ std::string expect = "4 522 2 4 516 2 2 2 0 0 1 1 8 8 0 0 1 129 129 4 30016 256 512 ";
1069+ ExecuteTestCase(
1070+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1071+ except_tilingkey, expect);
1072+}
1073+ 
1074+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Bfloat16_Test03)
1075+{
1076+ gert::StorageShape xShape = {{1, 1, 4, 522}, {1, 1, 4, 522}};
1077+ gert::StorageShape yShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1078+ gert::StorageShape argmaxShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1079+ std::vector<int64_t> ksize = {2, 516};
1080+ std::vector<int64_t> strides = {2, 2};
1081+ std::vector<int64_t> pads = {0, 0};
1082+ std::vector<int64_t> dilation = {1, 1};
1083+ ge::DataType dtype = ge::DT_BF16;
1084+ int64_t index_dtype = 9;
1085+ bool ceil_mode = true;
1086+ std::string data_format = "NCHW";
1087+ uint64_t except_tilingkey = 400004;
1088+ std::string expect = "4 522 2 4 516 2 2 2 0 0 1 1 8 8 0 0 1 129 129 4 30016 256 512 ";
1089+ ExecuteTestCase(
1090+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1091+ except_tilingkey, expect);
1092+}
1093+ 
1094+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Half_Test01)
1095+{
1096+ gert::StorageShape xShape = {{1, 1, 68, 22}, {1, 1, 68, 22}};
1097+ gert::StorageShape yShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
1098+ gert::StorageShape argmaxShape = {{1, 1, 3, 4}, {1, 1, 3, 4}};
1099+ std::vector<int64_t> ksize = {64, 16};
1100+ std::vector<int64_t> strides = {2, 2};
1101+ std::vector<int64_t> pads = {0, 0};
1102+ std::vector<int64_t> dilation = {1, 1};
1103+ ge::DataType dtype = ge::DT_FLOAT16;
1104+ int64_t index_dtype = 3;
1105+ bool ceil_mode = true;
1106+ std::string data_format = "NCHW";
1107+ uint64_t except_tilingkey = 300001;
1108+ std::string expect = "68 22 3 4 64 16 2 2 0 0 1 1 1 1 1 3 1 1 4 1 1 12 2048 32 64 0 1 1 ";
1109+ ExecuteTestCase(
1110+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1111+ except_tilingkey, expect);
1112+}
1113+ 
1114+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Half_Test02)
1115+{
1116+ gert::StorageShape xShape = {{1, 1, 4, 522}, {1, 1, 4, 522}};
1117+ gert::StorageShape yShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1118+ gert::StorageShape argmaxShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1119+ std::vector<int64_t> ksize = {2, 516};
1120+ std::vector<int64_t> strides = {2, 2};
1121+ std::vector<int64_t> pads = {0, 0};
1122+ std::vector<int64_t> dilation = {1, 1};
1123+ ge::DataType dtype = ge::DT_FLOAT16;
1124+ int64_t index_dtype = 3;
1125+ bool ceil_mode = true;
1126+ std::string data_format = "NCHW";
1127+ uint64_t except_tilingkey = 400005;
1128+ std::string expect = "4 522 2 4 516 2 2 2 0 0 1 1 8 8 0 0 1 129 129 4 30016 256 512 ";
1129+ ExecuteTestCase(
1130+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1131+ except_tilingkey, expect);
1132+}
1133+ 
1134+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_NCHW_MulCore_Half_Test03)
1135+{
1136+ gert::StorageShape xShape = {{1, 1, 4, 522}, {1, 1, 4, 522}};
1137+ gert::StorageShape yShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1138+ gert::StorageShape argmaxShape = {{1, 1, 2, 4}, {1, 1, 2, 4}};
1139+ std::vector<int64_t> ksize = {2, 516};
1140+ std::vector<int64_t> strides = {2, 2};
1141+ std::vector<int64_t> pads = {0, 0};
1142+ std::vector<int64_t> dilation = {1, 1};
1143+ ge::DataType dtype = ge::DT_FLOAT16;
1144+ int64_t index_dtype = 9;
1145+ bool ceil_mode = true;
1146+ std::string data_format = "NCHW";
1147+ uint64_t except_tilingkey = 400006;
1148+ std::string expect = "4 522 2 4 516 2 2 2 0 0 1 1 8 8 0 0 1 129 129 4 30016 256 512 ";
1149+ ExecuteTestCase(
1150+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1151+ except_tilingkey, expect);
1152+}
1153+ 
1154+TEST_F(MaxPoolWithArgmaxV3Tiling, MaxPoolWithArgmaxV3Tiling_SIMT_Test01)
1155+{
1156+ gert::StorageShape xShape = {{4, 4, 2665, 4}, {4, 4, 2665, 4}};
1157+ gert::StorageShape yShape = {{4, 4, 738, 2}, {4, 4, 738, 2}};
1158+ gert::StorageShape argmaxShape = {{4, 4, 738, 2}, {4, 4, 738, 2}};
1159+ std::vector<int64_t> ksize = {456, 4};
1160+ std::vector<int64_t> strides = {3, 3};
1161+ std::vector<int64_t> pads = {1, 1};
1162+ std::vector<int64_t> dilation = {1, 1};
1163+ ge::DataType dtype = ge::DT_BF16;
1164+ int64_t index_dtype = 3;
1165+ bool ceil_mode = true;
1166+ std::string data_format = "NCHW";
1167+ uint64_t except_tilingkey = 500001;
1168+ std::string expect = "256 64 4 4 2665 4 738 2 456 4 3 3 1 1 1 1 1 ";
1169+ ExecuteTestCase(
1170+ xShape, yShape, argmaxShape, ksize, strides, pads, dilation, dtype, index_dtype, ceil_mode, data_format,
1171+ except_tilingkey, expect);
1172+}
Mquant/dynamic_mx_quant/README.md+1-1
@@ -100,4 +100,4 @@
100 100 
101| 调用方式 | 调用样例 | 说明 |101| 调用方式 | 调用样例 | 说明 |
102|--------------|------------------------------------------------------------------------|--------------------------------------------------------------|102|--------------|------------------------------------------------------------------------|--------------------------------------------------------------|
103-| aclnn调用 | [test_aclnn_dynamic_mx_quant](./examples/test_aclnn_dynamic_mx_quant.cpp) | 通过[aclnnDynamicMxQuant](./docs/aclnnDynamicMxQuant.md)接口方式调用Inxdex算子。 |103+| aclnn调用 | [test_aclnn_dynamic_mx_quant](./examples/test_aclnn_dynamic_mx_quant.cpp) | 通过[aclnnDynamicMxQuant](./docs/aclnnDynamicMxQuant.md)接口方式调用DynamicMxQuant算子。 |