已合并
relu6 算子 dtype 分发机制修改为由def驱动,补全 InferShape/InferDataType 规范化实现。 #8721
relu6 算子 dtype 分发机制修改为由def驱动,补全 InferShape/InferDataType 规范化实现。 #8721
已合并
张伟伟创建于 24 天前
7 个文件变更+59-48
@@ -0,0 +1,35 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/* Generated By CANNBot */
12+ 
13+/*!
14+ * \file relu6_graph_infer.cpp
15+ * \brief Relu6 InferDataType: output dtype = input x dtype
16+ */
17+ 
18+#include "register/op_impl_registry.h"
19+#include "log/log.h"
20+ 
21+using namespace ge;
22+ 
23+namespace ops {
24+ 
25+static ge::graphStatus InferDataTypeForRelu6(gert::InferDataTypeContext* context)
26+{
27+ OP_LOGD(context->GetNodeName(), "Begin to do InferDataTypeForRelu6");
28+ const ge::DataType xDataType = context->GetInputDataType(0);
29+ context->SetOutputDataType(0, xDataType);
30+ return ge::GRAPH_SUCCESS;
31+}
32+ 
33+IMPL_OP(Relu6).InferDataType(InferDataTypeForRelu6);
34+ 
35+} // namespace ops
@@ -20,7 +20,7 @@
20 * 3. 多核切分:blockFactor = ceil(totalNum / coreNum)20 * 3. 多核切分:blockFactor = ceil(totalNum / coreNum)
21 * 4. UB 切分:ubFactor = floor_align(floor_div(ubSize / typeSize / 4), ubBlockSize)21 * 4. UB 切分:ubFactor = floor_align(floor_div(ubSize / typeSize / 4), ubBlockSize)
22 * Relu6 需要 4 个 LocalTensor: inputLocal x2(双缓冲), tmpLocal, outputLocal22 * Relu6 需要 4 个 LocalTensor: inputLocal x2(双缓冲), tmpLocal, outputLocal
23- * 5. 设置 TilingKey(由模板参数选择 dtype)23+ * 5. dtype 由 def 文件驱动,无需设置 TilingKey
24 */24 */
25 25 
26#include "register/op_def_registry.h"26#include "register/op_def_registry.h"
@@ -28,7 +28,6 @@
28#include "op_common/op_host/util/math_util.h"28#include "op_common/op_host/util/math_util.h"
29#include "op_common/op_host/util/platform_util.h"29#include "op_common/op_host/util/platform_util.h"
30#include "../../op_kernel/arch35/relu6_tiling_data.h"30#include "../../op_kernel/arch35/relu6_tiling_data.h"
31-#include "../../op_kernel/arch35/relu6_tiling_key.h"
32 31 
33namespace optiling {32namespace optiling {
34 33 
@@ -153,14 +152,7 @@ static ge::graphStatus Relu6TilingFunc(gert::TilingContext* context)
153 tiling->ubFactor = FloorAlign(FloorDiv(FloorDiv(static_cast<int64_t>(ubSize), typeSize), BUFFER_NUM), ubBlockSize);152 tiling->ubFactor = FloorAlign(FloorDiv(FloorDiv(static_cast<int64_t>(ubSize), typeSize), BUFFER_NUM), ubBlockSize);
154 OP_CHECK_IF(tiling->ubFactor == 0, OP_LOGE(context, "ubFactor is 0, UB too small"), return ge::GRAPH_FAILED);153 OP_CHECK_IF(tiling->ubFactor == 0, OP_LOGE(context, "ubFactor is 0, UB too small"), return ge::GRAPH_FAILED);
155 154 
156- // 设置 dataType 字段:保存原始 ge::DataType 枚举值,供 Kernel 侧模板参数实例化使用
157- tiling->dataType = static_cast<int32_t>(dataType);
158- 
159 context->SetBlockDim(usedCoreNum);155 context->SetBlockDim(usedCoreNum);
160- 
161- // 5. 设置 TilingKey(dtype 模板参数选择)
162- uint32_t dType = static_cast<uint32_t>(dataType);
163- ASCENDC_TPL_SEL_PARAM(context, dType);
164 return ge::GRAPH_SUCCESS;156 return ge::GRAPH_SUCCESS;
165}157}
166 158 
@@ -4,10 +4,17 @@
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+/* Generated By CANNBot */
12+ 
13+/*!
14+ * \file relu6_infershape.cpp
15+ * \brief Relu6 shape inference: output shape = input x shape
16+ */
17+ 
11#include "register/op_impl_registry.h"18#include "register/op_impl_registry.h"
12#include "exe_graph/runtime/infer_shape_context.h"19#include "exe_graph/runtime/infer_shape_context.h"
13 20 
@@ -17,21 +24,20 @@ namespace ops {
17 24 
18static ge::graphStatus InferShape4Relu6(gert::InferShapeContext* context)25static ge::graphStatus InferShape4Relu6(gert::InferShapeContext* context)
19{26{
20- const gert::Shape* input_shape = context->GetInputShape(0);27+ const gert::Shape* inputShape = context->GetInputShape(0);
21- if (input_shape == nullptr) {28+ if (inputShape == nullptr) {
22 return ge::GRAPH_FAILED;29 return ge::GRAPH_FAILED;
23 }30 }
24 31 
25- gert::Shape* output_shape = context->GetOutputShape(0);32+ gert::Shape* outputShape = context->GetOutputShape(0);
26- if (output_shape == nullptr) {33+ if (outputShape == nullptr) {
27 return ge::GRAPH_FAILED;34 return ge::GRAPH_FAILED;
28 }35 }
29 36 
30- *output_shape = *input_shape;37+ *outputShape = *inputShape;
31- 
32 return ge::GRAPH_SUCCESS;38 return ge::GRAPH_SUCCESS;
33}39}
34 40 
35IMPL_OP_INFERSHAPE(Relu6).InferShape(InferShape4Relu6);41IMPL_OP_INFERSHAPE(Relu6).InferShape(InferShape4Relu6);
36 42 
37-} // namespace ops43+} // namespace ops
@@ -14,13 +14,11 @@
14 * \file relu6_tiling_data.h14 * \file relu6_tiling_data.h
15 * \brief Relu6 TilingData 结构体定义(arch35 = Ascend950)15 * \brief Relu6 TilingData 结构体定义(arch35 = Ascend950)
16 *16 *
17- * 迭代一:仅 float16 dtype 骨架17+ * dtype def 文件驱动(构建系统注入 DTYPE_X 宏),TilingData 无需存储 dataType。
18- * 迭代二:新增 dataType 字段,支持多 dtype 分发
19 * 字段说明:18 * 字段说明:
20 * - totalNum: 输入张量展平后的总元素数量19 * - totalNum: 输入张量展平后的总元素数量
21 * - blockFactor: 每个 AI Core 处理的元素数量20 * - blockFactor: 每个 AI Core 处理的元素数量
22 * - ubFactor: UB 单次循环处理的元素数量21 * - ubFactor: UB 单次循环处理的元素数量
23- * - dataType: 数据类型标识(0=float16, 1=float, 2=int32, 3=bfloat16)
24 */22 */
25 23 
26#ifndef _RELU6_TILING_DATA_H_24#ifndef _RELU6_TILING_DATA_H_
@@ -32,7 +30,6 @@ struct Relu6TilingData {
32 int64_t totalNum = 0; // 总元素数量30 int64_t totalNum = 0; // 总元素数量
33 int64_t blockFactor = 0; // 每个核处理的元素数量31 int64_t blockFactor = 0; // 每个核处理的元素数量
34 int64_t ubFactor = 0; // 每次 UB 循环处理的元素数量32 int64_t ubFactor = 0; // 每次 UB 循环处理的元素数量
35- int32_t dataType = 0; // 数据类型:0=float16, 1=float, 2=int32, 3=bfloat16
36};33};
37 34 
38#endif // _RELU6_TILING_DATA_H_35#endif // _RELU6_TILING_DATA_H_
@@ -12,26 +12,14 @@
12 12 
13/*!13/*!
14 * \file relu6_tiling_key.h14 * \file relu6_tiling_key.h
15- * \brief Relu6 TilingKey 模板参数定义(arch35 = Ascend950)15+ * \brief Relu6 TilingKey 定义(arch35 = Ascend950)
16 *16 *
17- * 迭代一:仅 float16 dtype 骨架预留 float/int32/bfloat16 扩展位置17+ * dtype def 文件(DataType 列表)驱动构建系统注入 DTYPE_X 宏,
18- * 模板参数类型参考:18+ * Kernel 直接使用 DTYPE_X 获取类型,TilingKey 无需编码 dtype。
19- * - DATATYPE: 原生数据类型(C_DT_FLOAT16, C_DT_FLOAT, C_DT_INT32, C_DT_BF16)19+ * Relu6 无算法分支/调度模式变体,故无 ASCENDC_TPL 声明。
20- * 参考:ascendc/host_api/tiling/template_argument.h
21 */20 */
22 21 
23#ifndef __RELU6_TILING_KEY_H__22#ifndef __RELU6_TILING_KEY_H__
24#define __RELU6_TILING_KEY_H__23#define __RELU6_TILING_KEY_H__
25 24 
26-#include "ascendc/host_api/tiling/template_argument.h"
27- 
28-// 迭代一:仅 float16;迭代二扩展全部 4 种 dtype
29-ASCENDC_TPL_ARGS_DECL(Relu6, ASCENDC_TPL_DATATYPE_DECL(D_T, C_DT_FLOAT16, C_DT_FLOAT, C_DT_INT32, C_DT_BF16,
30- ASCENDC_TPL_INPUT(0)));
31- 
32-ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DATATYPE_SEL(D_T, C_DT_FLOAT16)),
33- ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DATATYPE_SEL(D_T, C_DT_FLOAT)),
34- ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DATATYPE_SEL(D_T, C_DT_INT32)),
35- ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DATATYPE_SEL(D_T, C_DT_BF16)), );
36- 
37#endif // __RELU6_TILING_KEY_H__25#endif // __RELU6_TILING_KEY_H__
@@ -14,8 +14,8 @@
14 * \file relu6.cpp14 * \file relu6.cpp
15 * \brief Relu6 Kernel 入口(arch35 = Ascend950)15 * \brief Relu6 Kernel 入口(arch35 = Ascend950)
16 *16 *
17- * 模板参数说(与 relu6_tiling_key.h ASCENDC_TPL_ARGS_DECL 定义对应):17+ * dtype 驱动方式:def 文件声DataType 列表,构建系统注入 DTYPE_X 宏,
18- * - D_T: 数据类型, ASCENDC_TPL_DATATYPE_DECL 定义18+ * Kernel 直接使用 DTYPE_X 获取输入 x 的实际 C++ 类型,无需 TilingKey 编码 dtype。
19 *19 *
20 * 核函数参数顺序(固定):20 * 核函数参数顺序(固定):
21 * GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling21 * GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling
@@ -23,12 +23,11 @@
23 23 
24#include "arch35/relu6.h"24#include "arch35/relu6.h"
25 25 
26-template <typename D_T>
27__global__ __aicore__ void relu6(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)26__global__ __aicore__ void relu6(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
28{27{
29 REGISTER_TILING_DEFAULT(Relu6TilingData);28 REGISTER_TILING_DEFAULT(Relu6TilingData);
30 GET_TILING_DATA_WITH_STRUCT(Relu6TilingData, tilingData, tiling);29 GET_TILING_DATA_WITH_STRUCT(Relu6TilingData, tilingData, tiling);
31- NsRelu6::Relu6<D_T> op;30+ NsRelu6::Relu6<DTYPE_X> op;
32 op.Init(x, y, &tilingData);31 op.Init(x, y, &tilingData);
33 op.Process();32 op.Process();
34}33}
@@ -65,7 +65,6 @@ TEST_F(relu6_test, test_case_fp32_small)
65 tilingData->totalNum = static_cast<int64_t>(dataNum);65 tilingData->totalNum = static_cast<int64_t>(dataNum);
66 tilingData->blockFactor = static_cast<int64_t>(dataNum);66 tilingData->blockFactor = static_cast<int64_t>(dataNum);
67 tilingData->ubFactor = static_cast<int64_t>(dataNum);67 tilingData->ubFactor = static_cast<int64_t>(dataNum);
68- tilingData->dataType = 1;
69 68 
70 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);69 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);
71 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {70 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
@@ -103,7 +102,6 @@ TEST_F(relu6_test, test_case_fp32_multi_loop)
103 tilingData->totalNum = static_cast<int64_t>(dataNum);102 tilingData->totalNum = static_cast<int64_t>(dataNum);
104 tilingData->blockFactor = static_cast<int64_t>(dataNum);103 tilingData->blockFactor = static_cast<int64_t>(dataNum);
105 tilingData->ubFactor = static_cast<int64_t>(128);104 tilingData->ubFactor = static_cast<int64_t>(128);
106- tilingData->dataType = 1;
107 105 
108 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);106 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);
109 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {107 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
@@ -141,7 +139,6 @@ TEST_F(relu6_test, test_case_fp16_small)
141 tilingData->totalNum = static_cast<int64_t>(dataNum);139 tilingData->totalNum = static_cast<int64_t>(dataNum);
142 tilingData->blockFactor = static_cast<int64_t>(dataNum);140 tilingData->blockFactor = static_cast<int64_t>(dataNum);
143 tilingData->ubFactor = static_cast<int64_t>(dataNum);141 tilingData->ubFactor = static_cast<int64_t>(dataNum);
144- tilingData->dataType = 0;
145 142 
146 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);143 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);
147 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {144 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
@@ -179,7 +176,6 @@ TEST_F(relu6_test, test_case_int32_small)
179 tilingData->totalNum = static_cast<int64_t>(dataNum);176 tilingData->totalNum = static_cast<int64_t>(dataNum);
180 tilingData->blockFactor = static_cast<int64_t>(dataNum);177 tilingData->blockFactor = static_cast<int64_t>(dataNum);
181 tilingData->ubFactor = static_cast<int64_t>(dataNum);178 tilingData->ubFactor = static_cast<int64_t>(dataNum);
182- tilingData->dataType = 2;
183 179 
184 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);180 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);
185 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {181 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
@@ -217,7 +213,6 @@ TEST_F(relu6_test, test_case_bf16_small)
217 tilingData->totalNum = static_cast<int64_t>(dataNum);213 tilingData->totalNum = static_cast<int64_t>(dataNum);
218 tilingData->blockFactor = static_cast<int64_t>(dataNum);214 tilingData->blockFactor = static_cast<int64_t>(dataNum);
219 tilingData->ubFactor = static_cast<int64_t>(dataNum);215 tilingData->ubFactor = static_cast<int64_t>(dataNum);
220- tilingData->dataType = 3;
221 216 
222 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);217 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);
223 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {218 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
@@ -255,7 +250,6 @@ TEST_F(relu6_test, test_case_fp16_multi_loop)
255 tilingData->totalNum = static_cast<int64_t>(dataNum);250 tilingData->totalNum = static_cast<int64_t>(dataNum);
256 tilingData->blockFactor = static_cast<int64_t>(dataNum);251 tilingData->blockFactor = static_cast<int64_t>(dataNum);
257 tilingData->ubFactor = static_cast<int64_t>(128);252 tilingData->ubFactor = static_cast<int64_t>(128);
258- tilingData->dataType = 0;
259 253 
260 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);254 ReadFile(path + "/relu6_data/input_x.bin", xByteSize, x, xByteSize);
261 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {255 auto KernelRelu6 = [](GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
@@ -270,4 +264,4 @@ TEST_F(relu6_test, test_case_fp16_multi_loop)
270 AscendC::GmFree(y);264 AscendC::GmFree(y);
271 AscendC::GmFree(workspace);265 AscendC::GmFree(workspace);
272 AscendC::GmFree(tiling);266 AscendC::GmFree(tiling);
273-}267+}