Open
[feat][ops] norm_multiply_dropout反向算子op_kernel实现 #232
penghuiyangAdd on 5 days ago
[feat][ops] norm_multiply_dropout反向算子op_kernel实现 #232
Open
Total: 3 files changed+601-0
Aascendc/activation/norm_multiply_dropout/norm_multiply_dropout_backward/op_kernel/CMakeLists.txt+24-0
| @@ -0,0 +1,24 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# norm_multiply_dropout_backward/op_kernel：AscendC kernel 编译 | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 使用 CANN 的 npu_op_kernel_* 函数编译 kernel。 | ||
| 5 | +# npu_op_package 已在上层 CMakeLists.txt 调用。 | ||
| 6 | +# OPS_REC_SOC_CXX_DEFS 由根 CMakeLists.txt 统一设置（SUPPORT_950 / SUPPORT_A2 / SUPPORT_A3） | ||
| 7 | +set(_KERNEL_TARGET "${VENDOR_NAME}_kernels") | ||
| 8 | + | ||
| 9 | +npu_op_kernel_options(${_KERNEL_TARGET} ALL | ||
| 10 | + OPTIONS "--cce-long-call=true" ${OPS_REC_SOC_CXX_DEFS} | ||
| 11 | +) | ||
| 12 | + | ||
| 13 | +npu_op_kernel_sources(${_KERNEL_TARGET} | ||
| 14 | + KERNEL_DIR ./ | ||
| 15 | +) | ||
| 16 | + | ||
| 17 | +npu_op_kernel_library(${_KERNEL_TARGET} | ||
| 18 | + SRC_BASE ${CMAKE_CURRENT_SOURCE_DIR}/ | ||
| 19 | + TILING_LIBRARY cust_optiling | ||
| 20 | +) | ||
| 21 | + | ||
| 22 | +npu_op_package_add(${VENDOR_NAME} | ||
| 23 | + LIBRARY ${_KERNEL_TARGET} | ||
| 24 | +) | ||
| @@ -0,0 +1,553 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | +#include "kernel_operator.h" | ||
| 16 | + | ||
| 17 | +#include "norm_multiply_dropout_backward_tilingKey.h" | ||
| 18 | + | ||
| 19 | +using namespace AscendC; | ||
| 20 | + | ||
| 21 | +namespace kernels { | ||
| 22 | +constexpr uint32_t VEC_BLOCK_BYTES = 32; | ||
| 23 | +constexpr int32_t BUFFER_NUM = 1; | ||
| 24 | +constexpr int32_t NO_DOUBLE = 1; | ||
| 25 | +constexpr int32_t REDUCE_CALC_BYTE = 256; | ||
| 26 | +constexpr int32_t MASK_ELEMENTS_BITS = 8; | ||
| 27 | +constexpr float NEGATIVE_ONE_F = -1.0; | ||
| 28 | +constexpr float POSITIVE_ONE_F = 1.0; | ||
| 29 | +constexpr int32_t DROP_TEMP_BUF = 2048; | ||
| 30 | + | ||
| 31 | +template <typename T, bool isNeedDrop> | ||
| 32 | +class NormMultiplyDropoutBackward { | ||
| 33 | +public: | ||
| 34 | + __aicore__ NormMultiplyDropoutBackward() {}; | ||
| 35 | + | ||
| 36 | + __aicore__ inline void init(GM_ADDR dOut, GM_ADDR x, GM_ADDR u, GM_ADDR weight, GM_ADDR bias, GM_ADDR mean, | ||
| 37 | + GM_ADDR var, GM_ADDR mask, GM_ADDR dX, GM_ADDR dU, GM_ADDR dWeight, GM_ADDR dBias, | ||
| 38 | + GM_ADDR workspace, const NormMultiplyDropoutBackwardTilingData& tiling) | ||
| 39 | + { | ||
| 40 | + ASSERT(GetBlockNum() != 0 && "useful core num can not be zero!!!"); | ||
| 41 | + this->tiling = tiling; | ||
| 42 | + int32_t coreId = GetBlockIdx(); | ||
| 43 | + | ||
| 44 | + bigCoreNum = tiling.bigCoreNum; | ||
| 45 | + avgCoreCalcRows = tiling.avgCoreCalcRows; | ||
| 46 | + xRowCount = tiling.xRowCount; | ||
| 47 | + xColCount = tiling.xColCount; | ||
| 48 | + singleBlockRows = tiling.singleBlockRows; | ||
| 49 | + eps = tiling.eps; | ||
| 50 | + dropoutRatio = tiling.dropoutRatio; | ||
| 51 | + | ||
| 52 | + if (bigCoreNum == 0) { | ||
| 53 | + coreOffset = coreId * avgCoreCalcRows; | ||
| 54 | + coreCalcRows = avgCoreCalcRows; | ||
| 55 | + } else if (coreId < bigCoreNum) { | ||
| 56 | + coreOffset = coreId * (avgCoreCalcRows + 1); | ||
| 57 | + coreCalcRows = avgCoreCalcRows + 1; | ||
| 58 | + } else { | ||
| 59 | + coreOffset = bigCoreNum * (avgCoreCalcRows + 1) + (coreId - bigCoreNum) * avgCoreCalcRows; | ||
| 60 | + coreCalcRows = avgCoreCalcRows; | ||
| 61 | + } | ||
 | |||
| 62 | + | ||
| 63 | + tail = coreCalcRows % singleBlockRows; | ||
| 64 | + loop = coreCalcRows / singleBlockRows; | ||
| 65 | + | ||
| 66 | + dOutGM.SetGlobalBuffer((__gm__ T*)dOut); | ||
| 67 | + xGM.SetGlobalBuffer((__gm__ T*)x); | ||
| 68 | + uGM.SetGlobalBuffer((__gm__ T*)u); | ||
| 69 | + weightGM.SetGlobalBuffer((__gm__ T*)weight); | ||
| 70 | + biasGM.SetGlobalBuffer((__gm__ T*)bias); | ||
| 71 | + meanGM.SetGlobalBuffer((__gm__ float*)mean); | ||
| 72 | + varGM.SetGlobalBuffer((__gm__ float*)var); | ||
| 73 | + maskGM.SetGlobalBuffer((__gm__ uint8_t*)mask); | ||
| 74 | + dUGM.SetGlobalBuffer((__gm__ T*)dU); | ||
| 75 | + dXGM.SetGlobalBuffer((__gm__ T*)dX); | ||
| 76 | + dWeightGM.SetGlobalBuffer((__gm__ float*)dWeight); | ||
| 77 | + dBiasGM.SetGlobalBuffer((__gm__ float*)dBias); | ||
| 78 | + | ||
| 79 | + pipe.InitBuffer(dOutQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(float)); // n * C * 4 | ||
| 80 | + pipe.InitBuffer(xQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(float)); // n * C * 4 | ||
| 81 | + pipe.InitBuffer(uQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(float)); // n * C * 4 | ||
| 82 | + pipe.InitBuffer(weightQue, NO_DOUBLE, xColCount * sizeof(float)); // C * 4 | ||
| 83 | + pipe.InitBuffer(biasQue, NO_DOUBLE, xColCount * sizeof(float)); // C * 4 | ||
| 84 | + pipe.InitBuffer(meanQue, NO_DOUBLE, singleBlockRows * sizeof(float)); // n * 4 | ||
| 85 | + pipe.InitBuffer(varQue, NO_DOUBLE, singleBlockRows * sizeof(float)); // n * 4 | ||
| 86 | + if constexpr (isNeedDrop) { | ||
| 87 | + pipe.InitBuffer(maskQue, BUFFER_NUM, | ||
| 88 | + singleBlockRows * xColCount / MASK_ELEMENTS_BITS * sizeof(uint8_t)); // n * C * 2 | ||
| 89 | + } | ||
| 90 | + pipe.InitBuffer(duQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(float)); // n * C * 4 | ||
| 91 | + pipe.InitBuffer(dxQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(float)); // n * C * 4 | ||
| 92 | + | ||
| 93 | + if constexpr (!std::is_same<T, float>::value) { | ||
| 94 | + pipe.InitBuffer(dOutcastQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(T)); // n * C * 2 | ||
| 95 | + pipe.InitBuffer(xCastQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(T)); // n * C * 2 | ||
| 96 | + pipe.InitBuffer(uCastQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(T)); // n * C * 2 | ||
| 97 | + pipe.InitBuffer(weightCastQue, NO_DOUBLE, xColCount * sizeof(T)); // C * 2 | ||
| 98 | + pipe.InitBuffer(biasCastQue, NO_DOUBLE, xColCount * sizeof(T)); // C * 2 | ||
| 99 | + pipe.InitBuffer(dxCastQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(T)); // n * C * 2 | ||
| 100 | + pipe.InitBuffer(duCastQue, BUFFER_NUM, singleBlockRows * xColCount * sizeof(T)); // n * C * 2 | ||
| @@ -0,0 +1,24 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | +#ifndef MXREC_NORM_MULTIPLY_DROPOUT_BACKWARD_TILINGKEY_H | ||
| 16 | +#define MXREC_NORM_MULTIPLY_DROPOUT_BACKWARD_TILINGKEY_H | ||
| 17 | + | ||
| 18 | +#include "ascendc/host_api/tiling/template_argument.h" | ||
| 19 | + | ||
| 20 | +ASCENDC_TPL_ARGS_DECL(NormMultiplyDropoutBackward, ASCENDC_TPL_BOOL_DECL(isNeedDrop, 0, 1)); | ||
| 21 | + | ||
| 22 | +ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_BOOL_DECL(isNeedDrop, 0, 1))); | ||
| 23 | + | ||
| 24 | +#endif // MXREC_NORM_MULTIPLY_DROPOUT_BACKWARD_TILINGKEY_H | ||
🟡 Medium Priority
变更行：第 44-63 行读取
useCoreNum = tiling.useCoreNum;（第 50 行）后从未使用，littleCoreNum（第 45 行）同样只赋值不参与任何判断；核间行切分完全按 coreId 与 bigCoreNum 计算：失败模式：一旦实际启动的 block 数（GetBlockNum()/SetBlockDim）大于 tiling 规划的 useCoreNum，coreId ∈ [useCoreNum, blockNum) 的核不会得到"本核 0 行"的结论，而是继续按 avg 行做计算，
coreOffset + coreCalcRows越过 xRowCount，导致：1) 对 dOut/x/u/weight/mean/var/mask 的 GM 越界读；2) 对 dX/dU 的 GM 越界写（内存踩踏）；3) 即使不越界，同一批行被两个核重复计算，SetAtomicAdd会把 dWeight/dBias 重复累加两遍，结果放大。建议：在 init() 中用 useCoreNum 做上界保护：coreId >= useCoreNum 时把 coreCalcRows/loop/tail 置 0 并让 process() 直接返回（同时跳过 dWeight/dBias 的原子累加），避免多余核对 GM 的越界读写和重复累加。