Open
[feat][ops] norm_multiply_dropout反向算子op_host实现与torch适配层 #233
penghuiyangAdd on 5 days ago
[feat][ops] norm_multiply_dropout反向算子op_host实现与torch适配层 #233
Open
Total: 7 files changed+828-1
Aascendc/activation/norm_multiply_dropout/norm_multiply_dropout_backward/op_host/CMakeLists.txt+85-0
| @@ -0,0 +1,85 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# norm_multiply_dropout_backward/op_host：算子 host 侧实现（OpDef + tiling + infershape） | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 使用 CANN 的 npu_op_code_gen + npu_op_library 编译。 | ||
| 5 | +# npu_op_package 已在上层 CMakeLists.txt 调用。 | ||
| 6 | +set(MY_COMMON_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../../../common) | ||
| 7 | + | ||
| 8 | +# OPS_REC_SOC_CXX_DEFS 由根 CMakeLists.txt 统一设置（SUPPORT_950 / SUPPORT_A2 / SUPPORT_A3） | ||
| 9 | + | ||
| 10 | +file(GLOB ops_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp) | ||
| 11 | +file(GLOB ops_cc_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cc) | ||
| 12 | +list(APPEND ops_srcs ${ops_cc_srcs}) | ||
| 13 | + | ||
| 14 | +# 收集 SOC 架构特化源码（op_host/atlas950/ 或 op_host/atlasA2/ 目录下的额外实现） | ||
| 15 | +ops_rec_collect_arch_sources(ops_srcs) | ||
| 16 | + | ||
| 17 | +npu_op_code_gen( | ||
| 18 | + SRC ${ops_srcs} | ||
| 19 | + PACKAGE ${VENDOR_NAME} | ||
| 20 | + OUT_DIR ${ASCEND_AUTOGEN_PATH} | ||
| 21 | + COMPILE_OPTIONS | ||
| 22 | + ${OPS_REC_SOC_CXX_DEFS} | ||
| 23 | + -I${MY_COMMON_DIR}/log | ||
| 24 | + -I${MY_COMMON_DIR}/utils | ||
| 25 | +) | ||
| 26 | + | ||
| 27 | +file(GLOB autogen_aclnn_src ${ASCEND_AUTOGEN_PATH}/aclnn_*.cpp) | ||
| 28 | +set_source_files_properties(${autogen_aclnn_src} PROPERTIES GENERATED TRUE) | ||
| 29 | +npu_op_library(cust_opapi ACLNN | ||
| 30 | + ${autogen_aclnn_src} | ||
| 31 | +) | ||
| 32 | + | ||
| 33 | +target_include_directories(cust_opapi PRIVATE | ||
| 34 | + ${MY_COMMON_DIR}/log | ||
| 35 | + ${MY_COMMON_DIR}/utils | ||
| 36 | +) | ||
| 37 | + | ||
| 38 | +target_compile_options(cust_opapi PRIVATE | ||
| 39 | + -fvisibility=hidden | ||
| 40 | +) | ||
| 41 | + | ||
| 42 | +file(GLOB group_proto_src ${ASCEND_AUTOGEN_PATH}/group_op_proto/*.cc) | ||
| 43 | +file(GLOB proto_src ${ASCEND_AUTOGEN_PATH}/op_proto.cc) | ||
| 44 | +set_source_files_properties(${group_proto_src} PROPERTIES GENERATED TRUE) | ||
| 45 | +set_source_files_properties(${proto_src} PROPERTIES GENERATED TRUE) | ||
| 46 | +npu_op_library(cust_op_proto GRAPH | ||
| 47 | + ${ops_srcs} | ||
| 48 | + ${group_proto_src} | ||
| 49 | + ${proto_src} | ||
| 50 | +) | ||
| 51 | + | ||
| 52 | +target_include_directories(cust_op_proto PRIVATE | ||
| 53 | + ${MY_COMMON_DIR}/log | ||
| 54 | + ${MY_COMMON_DIR}/utils | ||
| 55 | +) | ||
| 56 | + | ||
| 57 | +target_compile_options(cust_op_proto PRIVATE | ||
| 58 | + -fvisibility=hidden | ||
| 59 | + # npu_op_code_gen 的 COMPILE_OPTIONS 不传播到 npu_op_library 目标， | ||
| 60 | + # op_host 源码中的 #ifdef SUPPORT_950 等宏需在此显式传递 | ||
| 61 | + ${OPS_REC_SOC_CXX_DEFS} | ||
| 62 | +) | ||
| 63 | + | ||
| 64 | +file(GLOB fallback_src ${ASCEND_AUTOGEN_PATH}/fallback_*.cpp) | ||
| 65 | +set_source_files_properties(${fallback_src} PROPERTIES GENERATED TRUE) | ||
| 66 | +npu_op_library(cust_optiling TILING | ||
| 67 | + ${ops_srcs} | ||
| 68 | + ${fallback_src} | ||
| 69 | +) | ||
| 70 | + | ||
| 71 | +target_include_directories(cust_optiling PRIVATE | ||
| 72 | + ${MY_COMMON_DIR}/log | ||
| 73 | + ${MY_COMMON_DIR}/utils | ||
| 74 | +) | ||
| 75 | + | ||
| 76 | +target_compile_options(cust_optiling PRIVATE | ||
| 77 | + ${OPS_REC_SOC_CXX_DEFS} | ||
| 78 | +) | ||
| 79 | + | ||
| 80 | +npu_op_package_add(${VENDOR_NAME} | ||
| 81 | + LIBRARY | ||
| 82 | + cust_optiling | ||
| 83 | + cust_opapi | ||
| 84 | + cust_op_proto | ||
| 85 | +) | ||
| @@ -0,0 +1,326 @@ | |||
| 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 <cmath> | ||
| 16 | + | ||
| 17 | +#include "register/op_def_registry.h" | ||
| 18 | +#include "tiling/platform/platform_ascendc.h" | ||
| 19 | +#include "ops_log.h" | ||
| 20 | +#include "norm_multiply_dropout_backward_tiling.h" | ||
| 21 | +#include "../op_kernel/norm_multiply_dropout_backward_tilingKey.h" | ||
| 22 | + | ||
| 23 | +#ifdef SUPPORT_950 | ||
| 24 | +constexpr uint64_t UB_RESERVED_LENGTH = 2048 * 5; | ||
| 25 | +#else | ||
| 26 | +constexpr uint64_t UB_RESERVED_LENGTH = 2048; | ||
| 27 | +#endif | ||
| 28 | +constexpr int32_t GRAD_TENSOR_NUM_PER_ROW = 12; | ||
| 29 | +constexpr int MASK_ELE_RATIO = 8; | ||
| 30 | +constexpr int REDUCE_ELEMENTS = 64; | ||
| 31 | +constexpr int DROP_BUF = 20 * 1024; | ||
| 32 | +constexpr float DROPOUT_RATIO_ZERO_EPS = 1e-10; | ||
| 33 | +// 参数范围约束，与前向 op_host check_params 保持一致 | ||
| 34 | +constexpr int64_t MAX_DIM0 = 1000000; | ||
| 35 | +constexpr int64_t BASE_DIM1 = 512; | ||
| 36 | +constexpr int64_t MAX_DIM1 = 1024; | ||
| 37 | +constexpr float EPS_MIN = 1e-10; | ||
| 38 | +constexpr float EPS_MAX = 1e-4; | ||
| 39 | +constexpr float DROPOUT_RATIO_MIN = 0.0; | ||
| 40 | +constexpr float DROPOUT_RATIO_MAX = 1.0; | ||
| 41 | +constexpr int INPUT_INDEX_DOUT = 0; | ||
| 42 | +constexpr int INPUT_INDEX_X = 1; | ||
| 43 | +constexpr int INPUT_INDEX_U = 2; | ||
| 44 | +constexpr int INPUT_INDEX_WEIGHT = 3; | ||
| 45 | +constexpr int INPUT_INDEX_BIAS = 4; | ||
| 46 | +constexpr int INPUT_INDEX_MEAN = 5; | ||
| 47 | +constexpr int INPUT_INDEX_VAR = 6; | ||
| 48 | +constexpr int INPUT_INDEX_MASK = 7; | ||
| 49 | +constexpr int ATTR_INDEX_EPS = 0; | ||
| 50 | +constexpr int ATTR_INDEX_DROPOUT = 1; | ||
| 51 | + | ||
| 52 | +namespace optiling { | ||
| 53 | +static ge::graphStatus TilingFunc(gert::TilingContext* context) | ||
| 54 | +{ | ||
| 55 | + NormMultiplyDropoutBackwardTilingData tiling; | ||
| 56 | + // ===== 入参校验（与前向 op_host check_params 对称）===== | ||
| 57 | + for (uint32_t idx = 0; idx < 8; idx++) { // d_out/x/u/weight/bias/mean/var/mask shape 指针非空 | ||
| 58 | + OPS_LOG_E_IF_NULL("input shape", context->GetInputShape(idx), return ge::GRAPH_FAILED); | ||
| 59 | + } | ||
| 60 | + const gert::StorageShape* dOutShape = context->GetInputShape(INPUT_INDEX_DOUT); | ||
| 61 | + const gert::StorageShape* xShape = context->GetInputShape(INPUT_INDEX_X); | ||
| 62 | + const gert::StorageShape* uShape = context->GetInputShape(INPUT_INDEX_U); | ||
| 63 | + const gert::StorageShape* wShape = context->GetInputShape(INPUT_INDEX_WEIGHT); | ||
| 64 | + const gert::StorageShape* bShape = context->GetInputShape(INPUT_INDEX_BIAS); | ||
| 65 | + const gert::StorageShape* meanShape = context->GetInputShape(INPUT_INDEX_MEAN); | ||
| 66 | + const gert::StorageShape* varShape = context->GetInputShape(INPUT_INDEX_VAR); | ||
| 67 | + const gert::StorageShape* maskShape = context->GetInputShape(INPUT_INDEX_MASK); | ||
| 68 | + auto xStorageShape = xShape->GetStorageShape(); | ||
| 69 | + OPS_CHECK(xStorageShape.GetDimNum() != 2, | ||
| 70 | + OPS_LOG_E("", "Input x dim num must be 2, but get %ld", xStorageShape.GetDimNum()), | ||
| 71 | + return ge::GRAPH_FAILED); | ||
| 72 | + OPS_CHECK(dOutShape->GetStorageShape().GetDimNum() != 2, | ||
| 73 | + OPS_LOG_E("", "Input d_out dim num must be 2, but get %ld", dOutShape->GetStorageShape().GetDimNum()), | ||
| 74 | + return ge::GRAPH_FAILED); | ||
| 75 | + OPS_CHECK(uShape->GetStorageShape().GetDimNum() != 2, | ||
| 76 | + OPS_LOG_E("", "Input u dim num must be 2, but get %ld", uShape->GetStorageShape().GetDimNum()), | ||
| 77 | + return ge::GRAPH_FAILED); | ||
| 78 | + OPS_CHECK(wShape->GetStorageShape().GetDimNum() != 1, | ||
| 79 | + OPS_LOG_E("", "Input weight dim num must be 1, but get %ld", wShape->GetStorageShape().GetDimNum()), | ||
| 80 | + return ge::GRAPH_FAILED); | ||
| 81 | + OPS_CHECK(bShape->GetStorageShape().GetDimNum() != 1, | ||
| 82 | + OPS_LOG_E("", "Input bias dim num must be 1, but get %ld", bShape->GetStorageShape().GetDimNum()), | ||
| 83 | + return ge::GRAPH_FAILED); | ||
| 84 | + OPS_CHECK(meanShape->GetStorageShape().GetDimNum() != 1, | ||
| 85 | + OPS_LOG_E("", "Input mean dim num must be 1, but get %ld", meanShape->GetStorageShape().GetDimNum()), | ||
| 86 | + return ge::GRAPH_FAILED); | ||
| 87 | + OPS_CHECK(varShape->GetStorageShape().GetDimNum() != 1, | ||
| 88 | + OPS_LOG_E("", "Input var dim num must be 1, but get %ld", varShape->GetStorageShape().GetDimNum()), | ||
| 89 | + return ge::GRAPH_FAILED); | ||
| 90 | + OPS_CHECK(maskShape->GetStorageShape().GetDimNum() != 1, | ||
| 91 | + OPS_LOG_E("", "Input mask dim num must be 1, but get %ld", maskShape->GetStorageShape().GetDimNum()), | ||
| 92 | + return ge::GRAPH_FAILED); | ||
| 93 | + int64_t xRowCount64 = xStorageShape.GetDim(0); | ||
| 94 | + int64_t xColCount64 = xStorageShape.GetDim(1); | ||
| 95 | + OPS_CHECK( | ||
| 96 | + xRowCount64 != dOutShape->GetStorageShape().GetDim(0) || xColCount64 != dOutShape->GetStorageShape().GetDim(1), | ||
| 97 | + OPS_LOG_E("", "Input d_out shape [%ld,%ld] is not equal to x shape [%ld,%ld].", | ||
| 98 | + dOutShape->GetStorageShape().GetDim(0), dOutShape->GetStorageShape().GetDim(1), xRowCount64, | ||
| 99 | + xColCount64), | ||
| 100 | + return ge::GRAPH_FAILED); | ||
| 101 | + OPS_CHECK(xRowCount64 != uShape->GetStorageShape().GetDim(0) || xColCount64 != uShape->GetStorageShape().GetDim(1), | ||
| 102 | + OPS_LOG_E("", "Input x shape [%ld,%ld] is not equal to u shape [%ld,%ld].", xRowCount64, xColCount64, | ||
| 103 | + uShape->GetStorageShape().GetDim(0), uShape->GetStorageShape().GetDim(1)), | ||
| 104 | + return ge::GRAPH_FAILED); | ||
| 105 | + OPS_CHECK(xColCount64 != wShape->GetStorageShape().GetDim(0), | ||
| 106 | + OPS_LOG_E("", "Input x dim1: %ld, weight dim0: %ld must be same.", xColCount64, | ||
| 107 | + wShape->GetStorageShape().GetDim(0)), | ||
| 108 | + return ge::GRAPH_FAILED); | ||
| 109 | + OPS_CHECK(wShape->GetStorageShape().GetDim(0) != bShape->GetStorageShape().GetDim(0), | ||
| 110 | + OPS_LOG_E("", "Input weight dim0: %ld and bias dim0: %ld must be same.", | ||
| 111 | + wShape->GetStorageShape().GetDim(0), bShape->GetStorageShape().GetDim(0)), | ||
| 112 | + return ge::GRAPH_FAILED); | ||
| 113 | + OPS_CHECK(xRowCount64 != meanShape->GetStorageShape().GetDim(0), | ||
| 114 | + OPS_LOG_E("", "Input x dim0: %ld, mean dim0: %ld must be same.", xRowCount64, | ||
| 115 | + meanShape->GetStorageShape().GetDim(0)), | ||
| 116 | + return ge::GRAPH_FAILED); | ||
| 117 | + OPS_CHECK(xRowCount64 != varShape->GetStorageShape().GetDim(0), | ||
| 118 | + OPS_LOG_E("", "Input x dim0: %ld, var dim0: %ld must be same.", xRowCount64, | ||
| 119 | + varShape->GetStorageShape().GetDim(0)), | ||
| 120 | + return ge::GRAPH_FAILED); | ||
| 121 | + OPS_CHECK(xRowCount64 < 1 || xRowCount64 > MAX_DIM0, OPS_LOG_E("", "Input x dim0 must in range:[1,%ld].", MAX_DIM0), | ||
| 122 | + return ge::GRAPH_FAILED); | ||
| 123 | + OPS_CHECK(xColCount64 != BASE_DIM1 && xColCount64 != MAX_DIM1, | ||
| 124 | + OPS_LOG_E("", "Input x dim1 must equal with %ld or %ld.", BASE_DIM1, MAX_DIM1), return ge::GRAPH_FAILED); | ||
| 125 | + // dtype 校验：仅支持 fp16/bf16/fp32 | ||
| 126 | + OPS_LOG_E_IF_NULL("xInputDesc", context->GetInputDesc(INPUT_INDEX_X), return ge::GRAPH_FAILED); | ||
| 127 | + ge::DataType dataType = context->GetInputDesc(INPUT_INDEX_X)->GetDataType(); | ||
| 128 | + OPS_CHECK(dataType != ge::DT_FLOAT16 && dataType != ge::DT_BF16 && dataType != ge::DT_FLOAT, | ||
| 129 | + OPS_LOG_E("", "Input x dtype must be fp16/bf16/fp32."), return ge::GRAPH_FAILED); | ||
| 130 | +#ifdef SUPPORT_950 | ||
| 131 | + OPS_CHECK(dataType == ge::DT_FLOAT, OPS_LOG_E("", "Input x not support fp32 in ascend 950 series device."), | ||
| 132 | + return ge::GRAPH_FAILED); | ||
| 133 | +#endif | ||
| 134 | + // dtype 一致性校验：d_out/u/weight/bias dtype 必须与 x 一致 | ||
| 135 | + OPS_LOG_E_IF_NULL("dOutInputDesc", context->GetInputDesc(INPUT_INDEX_DOUT), return ge::GRAPH_FAILED); | ||
| 136 | + OPS_LOG_E_IF_NULL("uInputDesc", context->GetInputDesc(INPUT_INDEX_U), return ge::GRAPH_FAILED); | ||
| 137 | + OPS_LOG_E_IF_NULL("weightInputDesc", context->GetInputDesc(INPUT_INDEX_WEIGHT), return ge::GRAPH_FAILED); | ||
| 138 | + OPS_LOG_E_IF_NULL("biasInputDesc", context->GetInputDesc(INPUT_INDEX_BIAS), return ge::GRAPH_FAILED); | ||
| 139 | + OPS_CHECK(context->GetInputDesc(INPUT_INDEX_DOUT)->GetDataType() != dataType, | ||
| 140 | + OPS_LOG_E("", "Input d_out dtype must be same as x dtype."), return ge::GRAPH_FAILED); | ||
| 141 | + OPS_CHECK(context->GetInputDesc(INPUT_INDEX_U)->GetDataType() != dataType, | ||
| 142 | + OPS_LOG_E("", "Input u dtype must be same as x dtype."), return ge::GRAPH_FAILED); | ||
| 143 | + OPS_CHECK(context->GetInputDesc(INPUT_INDEX_WEIGHT)->GetDataType() != dataType, | ||
| 144 | + OPS_LOG_E("", "Input weight dtype must be same as x dtype."), return ge::GRAPH_FAILED); | ||
| 145 | + OPS_CHECK(context->GetInputDesc(INPUT_INDEX_BIAS)->GetDataType() != dataType, | ||
| 146 | + OPS_LOG_E("", "Input bias dtype must be same as x dtype."), return ge::GRAPH_FAILED); | ||
| 147 | + // mask dtype 校验：必须为 uint8 位图 | ||
| 148 | + OPS_LOG_E_IF_NULL("maskInputDesc", context->GetInputDesc(INPUT_INDEX_MASK), return ge::GRAPH_FAILED); | ||
| 149 | + OPS_CHECK(context->GetInputDesc(INPUT_INDEX_MASK)->GetDataType() != ge::DT_UINT8, | ||
| 150 | + OPS_LOG_E("", "Input mask dtype must be uint8."), return ge::GRAPH_FAILED); | ||
| 151 | + int32_t xRowCount = xStorageShape.GetDim(0); | ||
| 152 | + int32_t xColCount = xStorageShape.GetDim(1); | ||
| 153 | + | ||
| 154 | + OPS_LOG_E_IF_NULL("GetAttrs ret", context->GetAttrs(), return ge::GRAPH_FAILED); | ||
| 155 | + OPS_LOG_E_IF_NULL("Attr eps ptr", context->GetAttrs()->GetAttrPointer<float>(ATTR_INDEX_EPS), | ||
| 156 | + return ge::GRAPH_FAILED); | ||
| 157 | + OPS_LOG_E_IF_NULL("Attr dropout ptr", context->GetAttrs()->GetAttrPointer<float>(ATTR_INDEX_DROPOUT), | ||
| 158 | + return ge::GRAPH_FAILED); | ||
| 159 | + float eps = *(context->GetAttrs()->GetAttrPointer<float>(ATTR_INDEX_EPS)); | ||
| 160 | + float dropoutRatio = *(context->GetAttrs()->GetAttrPointer<float>(ATTR_INDEX_DROPOUT)); | ||
| 161 | + OPS_CHECK(eps < EPS_MIN || eps > EPS_MAX, OPS_LOG_E("", "eps must in range:[%e,%e].", EPS_MIN, EPS_MAX), | ||
| 162 | + return ge::GRAPH_FAILED); | ||
| 163 | + OPS_CHECK(dropoutRatio < DROPOUT_RATIO_MIN || dropoutRatio > DROPOUT_RATIO_MAX, | ||
| 164 | + OPS_LOG_E("", "dropout_ratio must in range:[%e,%e].", DROPOUT_RATIO_MIN, DROPOUT_RATIO_MAX), | ||
| 165 | + return ge::GRAPH_FAILED); | ||
| 166 | + // mask 元素数：无 dropout 时允许 {1} 占位，有 dropout 时须为 x 元素数的 1/8 | ||
| 167 | + int64_t maskNumel = maskShape->GetStorageShape().GetDim(0); | ||
| 168 | + bool isNeedDropCheck = std::fabs(dropoutRatio) > DROPOUT_RATIO_ZERO_EPS; | ||
| 169 | + int64_t expectMaskNumel = (xRowCount64 * xColCount64) / MASK_ELE_RATIO; | ||
| 170 | + OPS_CHECK(maskNumel != expectMaskNumel && !(isNeedDropCheck == false && maskNumel == 1), | ||
| 171 | + OPS_LOG_E("", "Input mask numel must be %ld (or 1 when dropout_ratio is 0), but get %ld.", | ||
| 172 | + expectMaskNumel, maskNumel), | ||
| 173 | + return ge::GRAPH_FAILED); | ||
| 174 | + tiling.set_eps(eps); | ||
| 175 | + tiling.set_dropoutRatio(dropoutRatio); | ||
| 176 | + bool isNeedDrop = std::fabs(dropoutRatio) > DROPOUT_RATIO_ZERO_EPS; | ||
| 177 | + | ||
| 178 | + auto platformInfo = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); | ||
| 179 | + uint32_t coreNum = platformInfo.GetCoreNumAiv(); | ||
| 180 | + | ||
| 181 | + uint64_t ubSize = 0; | ||
| 182 | + platformInfo.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); | ||
| 183 | + ubSize -= UB_RESERVED_LENGTH; | ||
| 184 | + | ||
| 185 | + int32_t useCoreNum = coreNum > xRowCount ? xRowCount : coreNum; | ||
| 186 | + | ||
| 187 | + // 每行内存系数（不含常量） | ||
| 188 | + int32_t coeffPerRow = xColCount * sizeof(float) * GRAD_TENSOR_NUM_PER_ROW + // x, u, output, norm, ln_out (5个) | ||
| 189 | + xColCount / MASK_ELE_RATIO + // mask (uint8) | ||
| 190 | + sizeof(float) * 2 + // mean, rstd (2个) | ||
| 191 | + (xColCount / REDUCE_ELEMENTS + 1) * sizeof(float) * 2; // meanBuf, squareBuf | ||
| 192 | + | ||
| 193 | + // 常量内存 | ||
| 194 | + int32_t constMemory = xColCount * sizeof(float) * 2 + // weight, bias | ||
| 195 | + DROP_BUF; // dropBuf (20KB) | ||
| 196 | + | ||
| 197 | + // 计算最大行数 | ||
| 198 | + int32_t singleBlockRows = (ubSize - constMemory) / coeffPerRow; | ||
| 199 | + | ||
| 200 | + // 向下对齐到2的幂次或保持整数 | ||
| 201 | + singleBlockRows = std::max(1, singleBlockRows); // 至少1行 | ||
| 202 | + | ||
| 203 | + int32_t bigCoreNum = xRowCount % useCoreNum; | ||
| 204 | + int32_t littleCoreNum = useCoreNum - bigCoreNum; | ||
| 205 | + int32_t avgCoreCalcRows = xRowCount / useCoreNum; | ||
| 206 | + | ||
| 207 | + context->SetBlockDim(useCoreNum); | ||
| 208 | + | ||
| 209 | + tiling.set_bigCoreNum(bigCoreNum); | ||
| 210 | + tiling.set_littleCoreNum(littleCoreNum); | ||
| 211 | + tiling.set_avgCoreCalcRows(avgCoreCalcRows); | ||
| 212 | + tiling.set_xRowCount(xRowCount); | ||
| 213 | + tiling.set_xColCount(xColCount); | ||
| 214 | + tiling.set_singleBlockRows(singleBlockRows); | ||
| 215 | + tiling.set_useCoreNum(useCoreNum); | ||
| 216 | + | ||
| 217 | + const uint64_t tilingKey = GET_TPL_TILING_KEY(isNeedDrop); | ||
| 218 | + context->SetTilingKey(tilingKey); | ||
| 219 | + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); | ||
| 220 | + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); | ||
| 221 | + | ||
| 222 | + return ge::GRAPH_SUCCESS; | ||
| 223 | +} | ||
| 224 | +} // namespace optiling | ||
| 225 | + | ||
| 226 | +namespace ge { | ||
| 227 | +static ge::graphStatus InferShape(gert::InferShapeContext* context) | ||
| 228 | +{ | ||
| 229 | + OPS_LOG_E_IF_NULL("context", context, return ge::GRAPH_FAILED); | ||
| 230 | + const gert::Shape* xShape = context->GetInputShape(0); | ||
| 231 | + gert::Shape* outputShape = context->GetOutputShape(0); | ||
| 232 | + OPS_LOG_E_IF_NULL("xShape", xShape, return ge::GRAPH_FAILED); | ||
| 233 | + OPS_LOG_E_IF_NULL("outputShape", outputShape, return ge::GRAPH_FAILED); | ||
| 234 | + *outputShape = *xShape; | ||
A | |||
| 235 | + return GRAPH_SUCCESS; | ||
| 236 | +} | ||
| 237 | +static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) | ||
| 238 | +{ | ||
| 239 | + OPS_LOG_E_IF_NULL("context", context, return ge::GRAPH_FAILED); | ||
| 240 | + const auto inputDataType = context->GetInputDataType(0); | ||
| 241 | + context->SetOutputDataType(0, inputDataType); | ||
| 242 | + return ge::GRAPH_SUCCESS; | ||
| 243 | +} | ||
| 244 | +} // namespace ge | ||
| 245 | + | ||
| 246 | +namespace ops { | ||
| 247 | +class NormMultiplyDropoutBackward : public OpDef { | ||
| 248 | +public: | ||
| 249 | + explicit NormMultiplyDropoutBackward(const char* name) : OpDef(name) | ||
| 250 | + { | ||
| 251 | + this->Input("d_out") | ||
| 252 | + .ParamType(REQUIRED) | ||
| 253 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 254 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 255 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 256 | + this->Input("x") | ||
| 257 | + .ParamType(REQUIRED) | ||
| 258 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 259 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 260 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 261 | + this->Input("u") | ||
| 262 | + .ParamType(REQUIRED) | ||
| 263 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 264 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 265 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 266 | + this->Input("weight") | ||
| 267 | + .ParamType(REQUIRED) | ||
| 268 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 269 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 270 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 271 | + this->Input("bias") | ||
| 272 | + .ParamType(REQUIRED) | ||
| 273 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 274 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 275 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 276 | + this->Input("mean") | ||
| 277 | + .ParamType(REQUIRED) | ||
| 278 | + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) | ||
| 279 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 280 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 281 | + this->Input("var") | ||
| 282 | + .ParamType(REQUIRED) | ||
| 283 | + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) | ||
| 284 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 285 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 286 | + this->Input("mask") | ||
| 287 | + .ParamType(REQUIRED) | ||
| 288 | + .DataType({ge::DT_UINT8, ge::DT_UINT8, ge::DT_UINT8}) | ||
| 289 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 290 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 291 | + this->Output("d_x") | ||
| 292 | + .ParamType(REQUIRED) | ||
| 293 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 294 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 295 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 296 | + this->Output("d_u") | ||
| 297 | + .ParamType(REQUIRED) | ||
| 298 | + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) | ||
| 299 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 300 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 301 | + this->Output("d_weight") | ||
| 302 | + .ParamType(REQUIRED) | ||
| 303 | + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) | ||
| 304 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 305 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 306 | + this->Output("d_bias") | ||
| 307 | + .ParamType(REQUIRED) | ||
| 308 | + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) | ||
| 309 | + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 310 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 311 | + this->Attr("eps").Float(); | ||
| 312 | + this->Attr("dropout_ratio").Float(); | ||
| 313 | + | ||
| 314 | + this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType); | ||
| 315 | + | ||
| 316 | + this->AICore().SetTiling(optiling::TilingFunc); | ||
| 317 | + this->AICore().AddConfig("ascend910b"); | ||
| 318 | + this->AICore().AddConfig("ascend910_93"); | ||
| 319 | +#ifdef SUPPORT_950 | ||
| 320 | + this->AICore().AddConfig("ascend950"); | ||
| 321 | +#endif | ||
| 322 | + } | ||
| 323 | +}; | ||
| 324 | + | ||
| 325 | +OP_ADD(NormMultiplyDropoutBackward); | ||
| 326 | +} // namespace ops | ||
| @@ -0,0 +1,35 @@ | |||
| 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 NORM_MULTIPLY_DROPOUT_BACKWARD_TILING | ||
| 16 | +#define NORM_MULTIPLY_DROPOUT_BACKWARD_TILING | ||
| 17 | + | ||
| 18 | +#include "register/tilingdata_base.h" | ||
| 19 | + | ||
| 20 | +namespace optiling { | ||
| 21 | +BEGIN_TILING_DATA_DEF(NormMultiplyDropoutBackwardTilingData) | ||
| 22 | +TILING_DATA_FIELD_DEF(uint32_t, bigCoreNum); | ||
| 23 | +TILING_DATA_FIELD_DEF(uint32_t, littleCoreNum); | ||
| 24 | +TILING_DATA_FIELD_DEF(uint32_t, avgCoreCalcRows); | ||
| 25 | +TILING_DATA_FIELD_DEF(uint32_t, xRowCount); | ||
| 26 | +TILING_DATA_FIELD_DEF(uint32_t, xColCount); | ||
| 27 | +TILING_DATA_FIELD_DEF(uint32_t, singleBlockRows); | ||
| 28 | +TILING_DATA_FIELD_DEF(uint32_t, useCoreNum); | ||
| 29 | +TILING_DATA_FIELD_DEF(float, eps); | ||
| 30 | +TILING_DATA_FIELD_DEF(float, dropoutRatio); | ||
| 31 | +END_TILING_DATA_DEF; | ||
| 32 | + | ||
| 33 | +REGISTER_TILING_DATA_CLASS(NormMultiplyDropoutBackward, NormMultiplyDropoutBackwardTilingData) | ||
| 34 | +} // namespace optiling | ||
| 35 | +#endif // NORM_MULTIPLY_DROPOUT_BACKWARD_TILING | ||
| @@ -0,0 +1,127 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | +# ============================================================================== | ||
| 16 | +"""norm_multiply_dropout 算子的 autograd 注册（inductor 兼容）。 | ||
| 17 | + | ||
| 18 | +前向算子 ops_rec::norm_multiply_dropout_forward 由 plugin 层内部调用 | ||
| 19 | +aclnnDropoutGenMaskV2 生成 dropout 掩码（与 RecSDK 原始实现语义一致）， | ||
| 20 | +返回 (output, mean, var, mask)。 | ||
| 21 | +反向算子 ops_rec::norm_multiply_dropout_backward 接收上游梯度与前向保存的 | ||
| 22 | +中间结果，输出 (d_x, d_u, d_weight, d_bias)。 | ||
| 23 | + | ||
| 24 | +通过 torch.library.register_autograd 注册，使前向算子支持 loss.backward() | ||
| 25 | +并在 torch.compile (inductor) 图模式下正确工作。 | ||
| 26 | +""" | ||
| 27 | + | ||
| 28 | +import torch | ||
| 29 | + | ||
| 30 | +# 前向算子全名：ops_rec::norm_multiply_dropout_forward | ||
| 31 | +# 反向算子全名：ops_rec::norm_multiply_dropout_backward | ||
| 32 | +# 使用字符串而非 torch.ops.ops_rec.xxx 对象：延迟解析，避免本模块被 import 时 | ||
| 33 | +# C++ 插件库尚未加载导致 AttributeError；register_autograd 也仅接受字符串形式 | ||
| 34 | +_OP_NAMESPACE = "ops_rec" | ||
| 35 | +_FORWARD_OP = f"{_OP_NAMESPACE}::norm_multiply_dropout_forward" | ||
| 36 | +_BACKWARD_OP = f"{_OP_NAMESPACE}::norm_multiply_dropout_backward" | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +def norm_multiply_dropout_wrapper(x, u, weight, bias, eps=1e-5, dropout_ratio=0.0): | ||
| 40 | + """norm_multiply_dropout 前向接口（支持自动求导）。 | ||
| 41 | + | ||
| 42 | + mask 由 plugin 层自动生成并经第 4 输出返回，本接口仅对外返回 | ||
| 43 | + 前向输出的第一个元素 output；mean/var/mask 保存在 autograd 上下文 | ||
| 44 | + 中供反向使用，对用户不可见。 | ||
| 45 | + | ||
| 46 | + Args: | ||
| 47 | + x: 输入张量 [B, C] | ||
| 48 | + u: 乘法因子 [B, C] | ||
| 49 | + weight: LayerNorm 权重 [C] | ||
| 50 | + bias: LayerNorm 偏置 [C] | ||
| 51 | + eps: 防除 0 极小值，取值范围 [1e-10, 1e-4] | ||
| 52 | + dropout_ratio: dropout 概率，取值范围 [0.0, 1.0] | ||
| 53 | + | ||
| 54 | + Returns: | ||
| 55 | + output | ||
| 56 | + """ | ||
| 57 | + return torch.ops.ops_rec.norm_multiply_dropout_forward(x, u, weight, bias, eps, dropout_ratio)[0] | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +def _setup_context(ctx, inputs, output): | ||
| 61 | + """前向后保存反向所需张量。 | ||
| 62 | + | ||
| 63 | + Args: | ||
| 64 | + ctx: autograd context | ||
| 65 | + inputs: 前向输入 (x, u, weight, bias, eps, dropout_ratio) | ||
| 66 | + output: 前向输出 (output, mean, var, mask) | ||
| 67 | + """ | ||
| 68 | + x, u, weight, bias, eps, dropout_ratio = inputs | ||
| 69 | + # output[1]/output[2] 为 mean/var，output[3] 为 mask，反向计算必需 | ||
| 70 | + ctx.save_for_backward(x, u, weight, bias, output[1], output[2], output[3]) | ||
| 71 | + ctx.eps = eps | ||
| 72 | + ctx.dropout_ratio = dropout_ratio | ||
| 73 | + | ||
| 74 | + | ||
| 75 | +def _backward(ctx, out_grad, mean_grad, var_grad, mask_grad): | ||
| 76 | + """反向传播。 | ||
| 77 | + | ||
| 78 | + Args: | ||
| 79 | + ctx: autograd context | ||
| 80 | + out_grad: 前向 output 的上游梯度 | ||
| 81 | + mean_grad/var_grad/mask_grad: 反向中间结果，不参与梯度回传（忽略） | ||
| 82 | + | ||
| 83 | + Returns: | ||
| 84 | + 与前向输入 (x, u, weight, bias, eps, dropout_ratio) 一一对应的梯度 | ||
| 85 | + """ | ||
| 86 | + x, u, weight, bias, mean, var, mask = ctx.saved_tensors | ||
| 87 | + # 反向算子输出顺序：(d_x, d_u, d_weight, d_bias) | ||
| 88 | + d_x, d_u, d_weight, d_bias = torch.ops.ops_rec.norm_multiply_dropout_backward( | ||
| 89 | + out_grad, | ||
| 90 | + x, | ||
| 91 | + u, | ||
| 92 | + weight, | ||
| 93 | + bias, | ||
| 94 | + mean, | ||
| 95 | + var, | ||
| 96 | + mask, | ||
| 97 | + ctx.eps, | ||
| 98 | + ctx.dropout_ratio, | ||
| 99 | + ) | ||
| 100 | + # 返回顺序对应前向输入：x, u, weight, bias, eps, dropout_ratio | ||
| 101 | + return d_x, d_u, d_weight, d_bias, None, None | ||
| 102 | + | ||
| 103 | + | ||
| 104 | +def register_autograd(): | ||
| 105 | + """注册 norm_multiply_dropout 的 autograd 公式。 | ||
| 106 | + | ||
| 107 | + 在 ops_rec.autograd 插件包自动加载时被调用，也可手动调用。 | ||
| 108 | + 重复注册会抛异常，因此做一次幂等保护。 | ||
| 109 | + """ | ||
| 110 | + # 先绑定对外接口，再注册 autograd（幂等 early-return 路径也能保证接口存在）： | ||
| 111 | + # plugin 注册名已改为 norm_multiply_dropout_forward，activation 模块的 | ||
| 112 | + # __getattr__ 代理不再能命中 norm_multiply_dropout 原名，此处覆盖为 | ||
| 113 | + # 仅返回 output、支持 autograd 的高级封装。 | ||
| 114 | + from ops_rec.ascendc import activation as _activation | ||
| 115 | + | ||
| 116 | + _activation.norm_multiply_dropout = norm_multiply_dropout_wrapper | ||
| 117 | + | ||
| 118 | + try: | ||
| 119 | + torch.library.register_autograd(_FORWARD_OP, _backward, setup_context=_setup_context) | ||
| 120 | + except RuntimeError as exc: | ||
| 121 | + if "already" in str(exc).lower() or "registered" in str(exc).lower(): | ||
| 122 | + return | ||
| 123 | + raise | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +if __name__ == "__main__": | ||
| 127 | + register_autograd() | ||
Aascendc/activation/norm_multiply_dropout/torch_extension/norm_multiply_dropout_backward_plugin.cpp+103-0
| @@ -0,0 +1,103 @@ | |||||||||||||||||||||
| 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 | +#define OPS_REC_CUST_VENDORS "norm_multiply_dropout_backward" | ||||||||||||||||||||
| 16 | +#include <torch/torch.h> | ||||||||||||||||||||
| 17 | + | ||||||||||||||||||||
| 18 | +#include <stdexcept> | ||||||||||||||||||||
| 19 | +#include <torch/csrc/autograd/custom_function.h> | ||||||||||||||||||||
| 20 | +#include <torch/library.h> | ||||||||||||||||||||
| 21 | + | ||||||||||||||||||||
| 22 | +#include "pytorch_npu_helper.hpp" | ||||||||||||||||||||
| 23 | +#include "common_utils.h" | ||||||||||||||||||||
| 24 | +using namespace at; | ||||||||||||||||||||
| 25 | + | ||||||||||||||||||||
| 26 | +void CheckNormMultiplyDropoutBackwardInput(const at::Tensor& dOut, const at::Tensor& x, const at::Tensor& u, | ||||||||||||||||||||
| 27 | + const at::Tensor& weight, const at::Tensor& bias, const at::Tensor& mean, | ||||||||||||||||||||
| 28 | + const at::Tensor& var, const at::Tensor& mask) | ||||||||||||||||||||
| 29 | +{ | ||||||||||||||||||||
| 30 | + check_tensor_non_empty(dOut, "d_out"); | ||||||||||||||||||||
| 31 | + check_tensor_non_empty(x, "x"); | ||||||||||||||||||||
| 32 | + check_tensor_non_empty(u, "u"); | ||||||||||||||||||||
| 33 | + check_tensor_non_empty(weight, "weight"); | ||||||||||||||||||||
| 34 | + check_tensor_non_empty(bias, "bias"); | ||||||||||||||||||||
| 35 | + | ||||||||||||||||||||
| 36 | + // 检查NPU设备且设备ID一致 | ||||||||||||||||||||
| 37 | + std::vector<at::Tensor> tensors = {dOut, x, u, weight, bias, mean, var, mask}; | ||||||||||||||||||||
| 38 | + std::vector<std::string> names = {"d_out", "x", "u", "weight", "bias", "mean", "var", "mask"}; | ||||||||||||||||||||
| 39 | + check_tensor_npu_device(tensors, names); | ||||||||||||||||||||
| 40 | +} | ||||||||||||||||||||
| 41 | + | ||||||||||||||||||||
| 42 | +// 反向 NPU 实现，返回 (d_x, d_u, d_weight, d_bias) | ||||||||||||||||||||
| 43 | +std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> NormMultiplyDropoutBackwardInter( | ||||||||||||||||||||
| 44 | + const at::Tensor& dOut, const at::Tensor& x, const at::Tensor& u, const at::Tensor& weight, const at::Tensor& bias, | ||||||||||||||||||||
| 45 | + const at::Tensor& mean, const at::Tensor& var, const at::Tensor& mask, double eps, double dropoutRatio) | ||||||||||||||||||||
| 46 | +{ | ||||||||||||||||||||
| 47 | + CheckNormMultiplyDropoutBackwardInput(dOut, x, u, weight, bias, mean, var, mask); | ||||||||||||||||||||
| 48 | + TORCH_CHECK(eps > 0, "eps must be greater than 0"); | ||||||||||||||||||||
| 49 | + TORCH_CHECK(dropoutRatio >= 0 && dropoutRatio <= 1, "dropout_ratio must be in range [0.0, 1.0]"); | ||||||||||||||||||||
| 50 | + | ||||||||||||||||||||
| 51 | + auto dOutCont = dOut.contiguous(); | ||||||||||||||||||||
| 52 | + auto xCont = x.contiguous(); | ||||||||||||||||||||
| 53 | + auto uCont = u.contiguous(); | ||||||||||||||||||||
| 54 | + auto weightCont = weight.contiguous(); | ||||||||||||||||||||
| 55 | + auto biasCont = bias.contiguous(); | ||||||||||||||||||||
| 56 | + auto meanCont = mean.contiguous(); | ||||||||||||||||||||
| 57 | + auto varCont = var.contiguous(); | ||||||||||||||||||||
| 58 | + auto maskCont = mask.contiguous(); | ||||||||||||||||||||
| 59 | + | ||||||||||||||||||||
| 60 | + at::TensorOptions options = xCont.options(); | ||||||||||||||||||||
| 61 | + auto dX = at::empty(xCont.sizes(), options); | ||||||||||||||||||||
| 62 | + auto dU = at::empty(xCont.sizes(), options); | ||||||||||||||||||||
| 63 | + // dWeight/dBias 由行内规约累加得到，须初始化为全 0 | ||||||||||||||||||||
| 64 | + auto dWeight = at::zeros({xCont.size(1)}, options.dtype(at::kFloat)); | ||||||||||||||||||||
| 65 | + auto dBias = at::zeros({xCont.size(1)}, options.dtype(at::kFloat)); | ||||||||||||||||||||
| 66 | + | ||||||||||||||||||||
| 67 | + // aclnn 生成的接口 Float 属性为 double，直接透传 | ||||||||||||||||||||
| 68 | + EXEC_NPU_CMD(aclnnNormMultiplyDropoutBackward, dOutCont, xCont, uCont, weightCont, biasCont, meanCont, varCont, | ||||||||||||||||||||
| 69 | + maskCont, eps, dropoutRatio, dX, dU, dWeight, dBias); | ||||||||||||||||||||
| 70 | + return {dX, dU, dWeight, dBias}; | ||||||||||||||||||||
 🟡 Medium Priority 改动行： 受影响契约：前向/后向文档声明 weight/bias 的 dtype “同 x”，而 x 支持 float16/bfloat16（ 失败模式： 仓库同类实现的做法（同契约证据）： 建议：在 return 前把 dWeight/dBias 转为 weight/bias 的 options（dtype 与设备），使梯度 dtype 与 autograd 输入一致；dX/dU 已使用 x 的 options，保持一致。 Suggestions
  Incorrect? | |||||||||||||||||||||
| 71 | +} | ||||||||||||||||||||
| 72 | + | ||||||||||||||||||||
| 73 | +namespace { | ||||||||||||||||||||
| 74 | +// Meta 实现：供 torch.compile (inductor) trace 阶段做 shape/dtype 推断 | ||||||||||||||||||||
| 75 | +std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> NormMultiplyDropoutBackwardInterMeta( | ||||||||||||||||||||
| 76 | + const at::Tensor& dOut, const at::Tensor& x, const at::Tensor&, const at::Tensor& weight, const at::Tensor&, | ||||||||||||||||||||
| 77 | + const at::Tensor&, const at::Tensor&, const at::Tensor&, double, double) | ||||||||||||||||||||
| 78 | +{ | ||||||||||||||||||||
| 79 | + at::TensorOptions options = x.options(); | ||||||||||||||||||||
| 80 | + auto dX = at::empty(x.sizes(), options); | ||||||||||||||||||||
| 81 | + auto dU = at::empty(dOut.sizes(), options); | ||||||||||||||||||||
| 82 | + auto dWeight = at::zeros({x.size(1)}, options.dtype(at::kFloat)); | ||||||||||||||||||||
| 83 | + auto dBias = at::zeros({x.size(1)}, options.dtype(at::kFloat)); | ||||||||||||||||||||
| 84 | + return {dX, dU, dWeight, dBias}; | ||||||||||||||||||||
| 85 | +} | ||||||||||||||||||||
| 86 | +} // namespace | ||||||||||||||||||||
| 87 | + | ||||||||||||||||||||
| 88 | +TORCH_LIBRARY_FRAGMENT(ops_rec, m) | ||||||||||||||||||||
| 89 | +{ | ||||||||||||||||||||
| 90 | + m.def("norm_multiply_dropout_backward(Tensor d_out, Tensor x, Tensor u, Tensor weight, Tensor bias, " | ||||||||||||||||||||
| 91 | + "Tensor mean, Tensor var, Tensor mask, float eps=1e-05, float dropout_ratio=0.0) " | ||||||||||||||||||||
| 92 | + "-> (Tensor, Tensor, Tensor, Tensor)"); | ||||||||||||||||||||
| 93 | +} | ||||||||||||||||||||
| 94 | + | ||||||||||||||||||||
| 95 | +TORCH_LIBRARY_IMPL(ops_rec, PrivateUse1, m) | ||||||||||||||||||||
| 96 | +{ | ||||||||||||||||||||
| 97 | + m.impl("norm_multiply_dropout_backward", TORCH_FN(NormMultiplyDropoutBackwardInter)); | ||||||||||||||||||||
| 98 | +} | ||||||||||||||||||||
| 99 | + | ||||||||||||||||||||
| 100 | +TORCH_LIBRARY_IMPL(ops_rec, Meta, m) | ||||||||||||||||||||
| 101 | +{ | ||||||||||||||||||||
| 102 | + m.impl("norm_multiply_dropout_backward", &NormMultiplyDropoutBackwardInterMeta); | ||||||||||||||||||||
| 103 | +} | ||||||||||||||||||||
| @@ -0,0 +1,145 @@ | |||
| 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", | ||
| 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 | +#define OPS_REC_CUST_VENDORS "norm_multiply_dropout_forward" | ||
| 16 | +#include <torch/torch.h> | ||
| 17 | + | ||
| 18 | +#include <cmath> | ||
| 19 | +#include <random> | ||
| 20 | +#include <stdexcept> | ||
| 21 | +#include <torch/csrc/autograd/custom_function.h> | ||
| 22 | +#include <torch/library.h> | ||
| 23 | + | ||
| 24 | +#include "pytorch_npu_helper.hpp" | ||
| 25 | +#include "common_utils.h" | ||
| 26 | +using namespace at; | ||
| 27 | + | ||
| 28 | +// dropout 概率小于该阈值时不执行 dropout，mask 无效 | ||
| 29 | +constexpr float NO_DROP_THRESHOLD = 1e-10; | ||
| 30 | +constexpr int64_t MASK_BITS = 8; | ||
| 31 | +constexpr int64_t MASK_ALIGN = 128; | ||
| 32 | + | ||
| 33 | +void CheckNormMultiplyDropoutInput(const at::Tensor& x, const at::Tensor& u, const at::Tensor& weight, | ||
| 34 | + const at::Tensor& bias) | ||
| 35 | +{ | ||
| 36 | + check_tensor_non_empty(x, "x"); | ||
| 37 | + check_tensor_non_empty(u, "u"); | ||
| 38 | + check_tensor_non_empty(weight, "weight"); | ||
| 39 | + check_tensor_non_empty(bias, "bias"); | ||
| 40 | + | ||
| 41 | + // 检查NPU设备且设备ID一致 | ||
| 42 | + std::vector<at::Tensor> tensors = {x, u, weight, bias}; | ||
| 43 | + std::vector<std::string> names = {"x", "u", "weight", "bias"}; | ||
| 44 | + check_tensor_npu_device(tensors, names); | ||
| 45 | +} | ||
| 46 | + | ||
| 47 | +// 生成 dropout 掩码（迁移自 RecSDK dropout_gen_mask_impl，语义保持一致）： | ||
| 48 | +// kernel 内 dropout 使用 bit 模式，mask 为 uint8 tensor， | ||
| 49 | +// 元素数为 numels 向上对齐 128 再除以 8，bit 保留概率 = 1 - dropout_ratio | ||
| 50 | +at::Tensor DropoutGenMask(const at::Tensor& query, double keep_prob, const int64_t numels) | ||
| 51 | +{ | ||
| 52 | + int64_t length = (numels + MASK_ALIGN - 1) / MASK_ALIGN * MASK_ALIGN / MASK_BITS; | ||
| 53 | + c10::TensorOptions options = query.options(); | ||
| 54 | + at::Tensor mask = at::zeros(at::IntArrayRef{length}, options.dtype(at::kByte)); | ||
| 55 | + c10::SmallVector<int64_t> shape_size = {numels}; | ||
| 56 | + at::IntArrayRef shapeArray = at::IntArrayRef(shape_size); | ||
| 57 | + double prob; | ||
| 58 | + at::Scalar prob_scalar; | ||
| 59 | + if (query.scalar_type() == at::kHalf) { | ||
| 60 | + prob_scalar = at::Scalar(at::Half(1.0) - at::Half(keep_prob)); | ||
| 61 | + } else if (query.scalar_type() == at::kBFloat16) { | ||
| 62 | + prob_scalar = at::Scalar(at::BFloat16(1.0) - at::BFloat16(keep_prob)); | ||
| 63 | + } else { | ||
| 64 | + prob_scalar = at::Scalar(static_cast<float>(1.0) - static_cast<float>(keep_prob)); | ||
| 65 | + } | ||
| 66 | + prob = prob_scalar.toDouble(); | ||
| 67 | + aclDataType probDataType = aclDataType::ACL_FLOAT; | ||
| 68 | + std::mt19937 rng(std::random_device{}()); | ||
| 69 | + int64_t seed = static_cast<int64_t>(rng()); | ||
| 70 | + int64_t offset = 0; | ||
| 71 | + EXEC_NPU_CMD(aclnnDropoutGenMaskV2, shapeArray, prob, seed, offset, probDataType, mask); | ||
| 72 | + return mask; | ||
| 73 | +} | ||
| 74 | + | ||
| 75 | +// 前向 NPU 实现，返回 (output, mean, var, mask) | ||
| 76 | +// mask 由 plugin 内部调用 aclnnDropoutGenMaskV2 生成，并通过第 4 输出 | ||
| 77 | +// 返回给调用方，供反向算子使用 | ||
| 78 | +std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> NormMultiplyDropoutForwardInter( | ||
| 79 | + const at::Tensor& x, const at::Tensor& u, const at::Tensor& weight, const at::Tensor& bias, double eps, | ||
| 80 | + double dropoutRatio) | ||
| 81 | +{ | ||
| 82 | + CheckNormMultiplyDropoutInput(x, u, weight, bias); | ||
| 83 | + TORCH_CHECK(eps > 0, "eps must be greater than 0"); | ||
| 84 | + TORCH_CHECK(dropoutRatio >= 0 && dropoutRatio <= 1, "dropout_ratio must be in range [0.0, 1.0]"); | ||
| 85 | + | ||
| 86 | + auto xCont = x.contiguous(); | ||
| 87 | + auto uCont = u.contiguous(); | ||
| 88 | + auto weightCont = weight.contiguous(); | ||
| 89 | + auto biasCont = bias.contiguous(); | ||
| 90 | + | ||
| 91 | + at::TensorOptions options = xCont.options(); | ||
| 92 | + auto output = at::empty(xCont.sizes(), options); | ||
| 93 | + auto mean = at::empty({xCont.size(0)}, options.dtype(at::kFloat)); | ||
| 94 | + auto var = at::empty({xCont.size(0)}, options.dtype(at::kFloat)); | ||
| 95 | + | ||
| 96 | + // dropout 概率小于阈值时不执行 dropout，返回 1 元素占位 mask（kernel 不消费） | ||
| 97 | + at::Tensor dropMask; | ||
| 98 | + if (std::fabs(dropoutRatio) > NO_DROP_THRESHOLD) { | ||
| 99 | + dropMask = DropoutGenMask(xCont, 1.0 - dropoutRatio, xCont.numel()); | ||
| 100 | + } else { | ||
| 101 | + dropMask = at::empty({1}, options.dtype(at::kByte)); | ||
| 102 | + } | ||
| 103 | + | ||
| 104 | + // aclnn 生成的接口 Float 属性为 double，直接透传 | ||
| 105 | + EXEC_NPU_CMD(aclnnNormMultiplyDropout, xCont, uCont, weightCont, biasCont, dropMask, eps, dropoutRatio, output, | ||
| 106 | + mean, var); | ||
| 107 | + return {output, mean, var, dropMask}; | ||
| 108 | +} | ||
| 109 | + | ||
| 110 | +namespace { | ||
| 111 | +// Meta 实现：供 torch.compile (inductor) trace 阶段做 shape/dtype 推断 | ||
| 112 | +std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> NormMultiplyDropoutForwardInterMeta( | ||
| 113 | + const at::Tensor& x, const at::Tensor&, const at::Tensor&, const at::Tensor&, double, double dropoutRatio) | ||
| 114 | +{ | ||
| 115 | + at::TensorOptions options = x.options(); | ||
| 116 | + auto output = at::empty(x.sizes(), options); | ||
| 117 | + auto mean = at::empty({x.size(0)}, options.dtype(at::kFloat)); | ||
| 118 | + auto var = at::empty({x.size(0)}, options.dtype(at::kFloat)); | ||
| 119 | + at::Tensor dropMask; | ||
| 120 | + if (std::fabs(dropoutRatio) > NO_DROP_THRESHOLD) { | ||
| 121 | + int64_t numels = x.numel(); | ||
| 122 | + int64_t maskNumel = (numels + MASK_ALIGN - 1) / MASK_ALIGN * MASK_ALIGN / MASK_BITS; | ||
| 123 | + dropMask = at::empty({maskNumel}, options.dtype(at::kByte)); | ||
| 124 | + } else { | ||
| 125 | + dropMask = at::empty({1}, options.dtype(at::kByte)); | ||
| 126 | + } | ||
| 127 | + return {output, mean, var, dropMask}; | ||
| 128 | +} | ||
| 129 | +} // namespace | ||
| 130 | + | ||
| 131 | +TORCH_LIBRARY_FRAGMENT(ops_rec, m) | ||
| 132 | +{ | ||
| 133 | + m.def("norm_multiply_dropout_forward(Tensor x, Tensor u, Tensor weight, Tensor bias, " | ||
| 134 | + "float eps=1e-05, float dropout_ratio=0.0) -> (Tensor, Tensor, Tensor, Tensor)"); | ||
| 135 | +} | ||
| 136 | + | ||
| 137 | +TORCH_LIBRARY_IMPL(ops_rec, PrivateUse1, m) | ||
| 138 | +{ | ||
| 139 | + m.impl("norm_multiply_dropout_forward", TORCH_FN(NormMultiplyDropoutForwardInter)); | ||
| 140 | +} | ||
| 141 | + | ||
| 142 | +TORCH_LIBRARY_IMPL(ops_rec, Meta, m) | ||
| 143 | +{ | ||
| 144 | + m.impl("norm_multiply_dropout_forward", &NormMultiplyDropoutForwardInterMeta); | ||
| 145 | +} | ||
| @@ -275,7 +275,13 @@ inline void* GetOpApiFuncAddrForVendorList(const char* vendors, const char* apiN | |||
| 275 | } | 275 | } |
| 276 | start = pos + 1; | 276 | start = pos + 1; |
| 277 | } | 277 | } |
| 278 | - return nullptr; | 278 | + // CANN 内置算子（如 aclnnDropoutGenMaskV2）的符号在系统 libopapi.so， |
| 279 | + // 不在 vendor 的 libcust_opapi.so，回退到系统库查找 | ||
| 280 | + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); | ||
| 281 | + if (opApiHandler == nullptr) { | ||
| 282 | + return nullptr; | ||
| 283 | + } | ||
| 284 | + return GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); | ||
| 279 | } | 285 | } |
| 280 | 286 | ||
| 281 | inline c10::Scalar ConvertTensorToScalar(const at::Tensor& tensor) | 287 | inline c10::Scalar ConvertTensorToScalar(const at::Tensor& tensor) |
InferShape/InferDataType只处理了第 0 个输出，图模式下 d_u/d_weight/d_bias 的形状和 dtype 都是空的问题描述：OpDef 注册了 4 个输出
d_x/d_u/d_weight/d_bias（L291-L310），但InferShape只取了GetOutputShape(0)并赋值（本行），InferDataType只执行了SetOutputDataType(0, inputDataType)（L241）。d_u应该和 x 同 shape，d_weight/d_bias应该是[C]且固定 FP32，现在这三路既没有 shape 也没有 dtype。aclnn 单算子路径不受影响，因为输出 tensor 是 plugin 自己按正确 shape/dtype 创建后传进去的；但走 GE/ACL 图模式时，图编译器拿不到d_u/d_weight/d_bias的形状和类型——d_weight/d_bias在 OpDef 里是DT_FLOAT，而InferDataType只把 output 0 设成输入 dtype，另外三路可能被当成 fp16/bf16，kernel 却按float*写，结果是类型对不上或编译期直接报「输出类型未推断」。PR 描述里写的是「inductor 兼容」，而 torch.compile 走的是 Meta 实现，所以现有 graph 用例发现不了这个问题。建议：
InferShape里补*context->GetOutputShape(1) = *xShape（d_u与 u 同 shape），GetOutputShape(2)/[3]设为 1 维长度xShape->GetDim(1)；InferDataType里把 output 1 设成输入 dtype、output 2/3 设成DT_FLOAT。