Star964
1903
代码
代码
Issues243
Pull Requests371
Actions
项目讨论
Wiki
分析
Star964
1903
已合并
InplaceAddLayerNormFusionPass 融合规则迁移 #9018
rk创建于 28 天前
InplaceAddLayerNormFusionPass 融合规则迁移 #9018
已合并
rk创建于 28 天前
从wangpengbo26/ops-nninplace_aln_pass合入到cann/ops-nnmaster
共 4 个文件变更+849-0
Anorm/inplace_add_layer_norm/op_graph/fusion_pass/inplace_add_layer_norm_fusion_pass.cpp+482-0
@@ -0,0 +1,482 @@
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+/*!
12+ * \file inplace_add_layer_norm_fusion_pass.cpp
13+ * \brief AddLayerNorm --> InplaceAddLayerNorm (graph base route)
14+ *
15+ * 使用 graph_metadef 的原生 Graph/GNode 接口改图。
16+ */
17+#include "inplace_add_layer_norm_fusion_pass.h"
18+ 
19+#include <dlfcn.h>
20+ 
21+#include <cstdlib>
22+#include <set>
23+#include <string>
24+#include <utility>
25+#include <vector>
26+ 
27+#include "acl/acl_rt.h"
28+#include "common/inc/error_util.h"
29+#include "graph/operator_factory.h"
30+#include "platform/platform_info.h"
31+#include "version/ge-compiler_version.h"
32+#include "ge/fusion/pass/pattern_fusion_pass.h"
33+ 
34+namespace ops {
35+namespace {
36+// GetOptionValue @since 9.0.0，按 D4 走 dlsym。
37+const char* const kGetOptionValueSymbol = "_ZNK2ge17CustomPassContext14GetOptionValueERKNS_12AscendStringERS1_";
38+using GetOptionValueFn = graphStatus (*)(const void*, const AscendString&, AscendString&);
39+ 
40+// GE 库以 RTLD_GLOBAL 加载，故查全局符号表；
41+GetOptionValueFn ResolveGetOptionValue()
42+{
43+ static GetOptionValueFn fn = reinterpret_cast<GetOptionValueFn>(dlsym(RTLD_DEFAULT, kGetOptionValueSymbol));
44+ return fn;
45+}
46+ 
47+const std::string kPassName = "ZInplaceAddLayerNormFusionPass";
48+ 
49+const char* const kAddLayerNormType = "AddLayerNorm";
50+const char* const kInplaceAddLayerNormType = "InplaceAddLayerNorm";
51+const char* const kNewNodeNameSuffix = "_inplace";
52+ 
53+// IR 位序：输入 x1 x2 gamma beta [bias]；输出 y/x1' mean rstd x/x2'
54+const int32_t kInputIdxX1 = 0;
55+const int32_t kInputIdxX2 = 1;
56+const int32_t kInputIdxGamma = 2;
57+const size_t kInputNumWithoutBias = 4U;
58+const size_t kInputNumWithBias = 5U;
59+const size_t kOutputNum = 4U;
60+ 
61+// 只有 x1/x2 会被原地写回，独占性只需检查这两个输入
62+const int32_t kInputRefNums = 2;
63+ 
64+const int64_t kMainPartsNums = 4L;
65+const int64_t kL2UpperFactor = 2L;
66+const int64_t kL2LowerFactor = 2L;
67+ 
68+const int64_t kUnknownDim = -1L;
69+const int64_t kUnknownRank = -2L;
70+ 
71+const int64_t kReduceAxisValue = 5120L;
72+ 
73+const char* const kAttrEpsilon = "epsilon";
74+const char* const kAttrAdditionalOutput = "additional_output";
75+const char* const kAttrContinuousOutput = "continuous_output";
76+const char* const kOptionGraphRunMode = "ge.graphRunMode";
77+const int32_t kGraphRunModeTrain = 1;
78+const int32_t kDecimalBase = 10;
79+ 
80+const int32_t kMinGeCompilerVersion = 90000000;
81+ 
82+// 含未知维或未知秩即动态 shape。GE 未提供作用于对外 Shape 的等价接口，自备。
83+bool IsDynamicShape(const Shape& shape)
84+{
85+ const size_t dimNum = shape.GetDimNum();
86+ for (size_t i = 0U; i < dimNum; ++i) {
87+ const int64_t dim = shape.GetDim(i);
88+ if ((dim == kUnknownDim) || (dim == kUnknownRank)) {
89+ return true;
90+ }
91+ }
92+ return false;
93+}
94+ 
95+enum class SceneCheckResult {
96+ INFERENCE,
97+ TRAIN,
98+ API_UNAVAILABLE // 运行时 < 9.0.0
99+};
100+ 
101+// 守卫 G1：训练场景不做原地改写；读不到 option 则返回 API_UNAVAILABLE 保持静默。
102+SceneCheckResult CheckScene(CustomPassContext& passContext)
103+{
104+ // 运行期check
105+ int32_t runtimeVersion = 0;
106+ char geCompilerName[] = "ge_compiler";
107+ (void)aclsysGetVersionNum(geCompilerName, &runtimeVersion);
108+ if (runtimeVersion > 0 && runtimeVersion < kMinGeCompilerVersion) {
109+ return SceneCheckResult::API_UNAVAILABLE;
110+ }
111+ 
112+ // 符号可达性
113+ const GetOptionValueFn getOptionValue = ResolveGetOptionValue();
114+ if (getOptionValue == nullptr) {
115+ return SceneCheckResult::API_UNAVAILABLE;
116+ }
117+ 
118+ AscendString value;
119+ if (getOptionValue(&passContext, AscendString(kOptionGraphRunMode), value) != GRAPH_SUCCESS) {
120+ OPS_LOG_D(kPassName.c_str(), "Option %s is not set, treat as inference scene.", kOptionGraphRunMode);
121+ return SceneCheckResult::INFERENCE;
122+ }
123+ const char* modeStr = value.GetString();
124+ if (modeStr == nullptr) {
125+ return SceneCheckResult::INFERENCE;
126+ }
127+ if (static_cast<int32_t>(std::strtol(modeStr, nullptr, kDecimalBase)) == kGraphRunModeTrain) {
128+ return SceneCheckResult::TRAIN;
129+ }
130+ return SceneCheckResult::INFERENCE;
131+}
132+ 
133+// 守卫 G2：平台check。
134+bool IsSupportedPlatform(int64_t& l2Size)
135+{
136+ fe::PlatformInfo platformInfo;
137+ fe::OptionalInfo optionalInfo;
138+ if (fe::PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo) !=
139+ ge::SUCCESS) {
140+ OPS_LOG_D(kPassName.c_str(), "Get platform info failed, skip.");
141+ return false;
142+ }
143+ 
144+ const std::string curSoc = platformInfo.str_info.short_soc_version;
145+ static const std::set<std::string> kSupportSoc = {"Ascend910B", "Ascend910_93", "Ascend950"};
146+ if (kSupportSoc.count(curSoc) == 0U) {
147+ OPS_LOG_D(kPassName.c_str(), "Platform %s is not supported, skip.", curSoc.c_str());
148+ return false;
149+ }
150+ 
151+ l2Size = static_cast<int64_t>(platformInfo.soc_info.l2_size);
152+ return true;
153+}
154+ 
155+// 守卫 G4/G5：x1、x2 必须被本节点独占消费，否则原地写回会踩踏别人的数据。
156+bool IsInplaceSafe(const GNode& node)
157+{
158+ bool isContinuousOutput = false;
159+ (void)node.GetAttr(AscendString(kAttrContinuousOutput), isContinuousOutput);
160+ 
161+ for (int32_t inputIdx = 0; inputIdx < kInputRefNums; ++inputIdx) {
162+ const std::pair<GNodePtr, int32_t> peer = node.GetInDataNodesAndPortIndexs(inputIdx);
163+ const GNodePtr producer = peer.first;
164+ if (producer == nullptr) {
165+ OPS_LOG_D(kPassName.c_str(), "Input %d has no producer, skip.", inputIdx);
166+ return false;
167+ }
168+ 
169+ if (producer->GetOutDataNodesAndPortIndexs(peer.second).size() > 1U) {
170+ OPS_LOG_D(kPassName.c_str(), "Input %d is shared by other consumers, skip.", inputIdx);
171+ return false;
172+ }
173+ 
174+ if (producer->GetOutputsSize() > 1U && isContinuousOutput) {
175+ OPS_LOG_D(kPassName.c_str(), "Producer of input %d has continuous multi-outputs, skip.", inputIdx);
176+ return false;
177+ }
178+ }
179+ return true;
180+}
181+ 
182+// 守卫 G6：原地收益与 L2 容量相关。太大装不下，太小不值得。
183+bool IsShapeSupport(const GNode& node, int64_t l2Size)
184+{
185+ TensorDesc x1Desc;
186+ if (node.GetInputDesc(kInputIdxX1, x1Desc) != GRAPH_SUCCESS) {
187+ OPS_LOG_D(kPassName.c_str(), "Get x1 input desc failed, skip.");
188+ return false;
189+ }
190+ const Shape x1Shape = x1Desc.GetShape();
191+ 
192+ if (IsDynamicShape(x1Shape)) {
193+ const size_t dimNum = x1Shape.GetDimNum();
194+ if (dimNum == 0U) {
195+ OPS_LOG_D(kPassName.c_str(), "Dynamic shape with zero dim, skip.");
196+ return false;
197+ }
198+ if (x1Shape.GetDim(dimNum - 1U) != kReduceAxisValue) {
199+ OPS_LOG_D(kPassName.c_str(), "Dynamic shape last dim is not %ld, skip.", kReduceAxisValue);
200+ return false;
201+ }
202+ return true;
203+ }
204+ 
205+ if (l2Size <= 0L) {
206+ OPS_LOG_D(kPassName.c_str(), "Invalid l2_size %ld, skip.", l2Size);
207+ return false;
208+ }
209+ 
210+ const int64_t x1Size = x1Shape.GetShapeSize() * static_cast<int64_t>(GetSizeByDataType(x1Desc.GetDataType()));
211+ if (x1Size > l2Size * kL2UpperFactor) {
212+ OPS_LOG_D(kPassName.c_str(), "Input x1 size %ld far exceeds l2 %ld, skip.", x1Size, l2Size);
213+ return false;
214+ }
215+ if (x1Size * kMainPartsNums < l2Size / kL2LowerFactor) {
216+ OPS_LOG_D(kPassName.c_str(), "Input x1 size %ld far below l2 %ld, skip.", x1Size, l2Size);
217+ return false;
218+ }
219+ return true;
220+}
221+ 
222+// 守卫 G7：x1/x2/gamma 必须同 dtype，混合类型下原地语义不成立。
223+bool IsSameInputDataType(const GNode& node)
224+{
225+ TensorDesc x1Desc;
226+ TensorDesc x2Desc;
227+ TensorDesc gammaDesc;
228+ if (node.GetInputDesc(kInputIdxX1, x1Desc) != GRAPH_SUCCESS ||
229+ node.GetInputDesc(kInputIdxX2, x2Desc) != GRAPH_SUCCESS ||
230+ node.GetInputDesc(kInputIdxGamma, gammaDesc) != GRAPH_SUCCESS) {
231+ OPS_LOG_D(kPassName.c_str(), "Get input desc for dtype check failed, skip.");
232+ return false;
233+ }
234+ 
235+ const DataType x1Dtype = x1Desc.GetDataType();
236+ if (x1Dtype != x2Desc.GetDataType() || x1Dtype != gammaDesc.GetDataType()) {
237+ OPS_LOG_D(kPassName.c_str(), "Inputs x1/x2/gamma have different dtypes, skip.");
238+ return false;
239+ }
240+ return true;
241+}
242+ 
243+// 逐节点守卫链。全局守卫（场景、平台）在 Run 入口只判一次。
244+bool MeetGuards(const GNode& node, int64_t l2Size)
245+{
246+ const size_t inputNum = node.GetInputsSize();
247+ if (inputNum != kInputNumWithoutBias && inputNum != kInputNumWithBias) {
248+ OPS_LOG_D(kPassName.c_str(), "Unexpected input num %zu, skip.", inputNum);
249+ return false;
250+ }
251+ if (node.GetOutputsSize() != kOutputNum) {
252+ OPS_LOG_D(kPassName.c_str(), "Unexpected output num %zu, skip.", node.GetOutputsSize());
253+ return false;
254+ }
255+ if (!IsInplaceSafe(node)) {
256+ return false;
257+ }
258+ if (!IsShapeSupport(node, l2Size)) {
259+ return false;
260+ }
261+ if (!IsSameInputDataType(node)) {
262+ return false;
263+ }
264+ 
265+ OPS_LOG_D(kPassName.c_str(), "All guards passed, input num %zu.", inputNum);
266+ return true;
267+}
268+ 
269+bool CreateInplaceNode(Graph& graph, const GNode& oldNode, GNode& newNode)
270+{
271+ AscendString oldName;
272+ if (oldNode.GetName(oldName) != GRAPH_SUCCESS || oldName.GetString() == nullptr) {
273+ return false;
274+ }
275+ const std::string newName = std::string(oldName.GetString()) + kNewNodeNameSuffix;
276+ 
277+ Operator op = OperatorFactory::CreateOperator(newName.c_str(), kInplaceAddLayerNormType);
278+ newNode = graph.AddNodeByOp(op);
279+ 
280+ AscendString newType;
281+ if (newNode.GetType(newType) != GRAPH_SUCCESS || newType != AscendString(kInplaceAddLayerNormType)) {
282+ OPS_LOG_E(kPassName.c_str(), "Create %s node failed.", kInplaceAddLayerNormType);
283+ return false;
284+ }
285+ 
286+ // attr 原样透传，读不到则保留 IR 默认值。
287+ float32_t epsilon = 1e-5F;
288+ if (oldNode.GetAttr(AscendString(kAttrEpsilon), epsilon) == GRAPH_SUCCESS) {
289+ (void)newNode.SetAttr(AscendString(kAttrEpsilon), epsilon);
290+ }
291+ bool additionalOutput = false;
292+ if (oldNode.GetAttr(AscendString(kAttrAdditionalOutput), additionalOutput) == GRAPH_SUCCESS) {
293+ (void)newNode.SetAttr(AscendString(kAttrAdditionalOutput), additionalOutput);
294+ }
295+ return true;
296+}
297+ 
298+// 输入逐位接线并拷贝 TensorDesc。bias 缺省（4 输入）时第 5 个端口保持未接。
299+bool RewireInputs(Graph& graph, const GNode& oldNode, GNode& newNode, size_t inputNum)
300+{
301+ for (size_t i = 0U; i < inputNum; ++i) {
302+ const int32_t idx = static_cast<int32_t>(i);
303+ const std::pair<GNodePtr, int32_t> peer = oldNode.GetInDataNodesAndPortIndexs(idx);
304+ if (peer.first == nullptr) {
305+ continue;
306+ }
307+ if (graph.AddDataEdge(*peer.first, peer.second, newNode, idx) != GRAPH_SUCCESS) {
308+ OPS_LOG_E(kPassName.c_str(), "Add input edge %zu failed.", i);
309+ return false;
310+ }
311+ TensorDesc desc;
312+ if (oldNode.GetInputDesc(idx, desc) == GRAPH_SUCCESS) {
313+ (void)newNode.UpdateInputDesc(idx, desc);
314+ }
315+ }
316+ return true;
317+}
318+ 
319+// 输出逐位改接并拷贝 TensorDesc。替换前后 shape/dtype 完全相同，直接拷贝，无需 InferShape。
320+bool RewireOutputs(Graph& graph, GNode& oldNode, GNode& newNode)
321+{
322+ for (size_t j = 0U; j < kOutputNum; ++j) {
323+ const int32_t idx = static_cast<int32_t>(j);
324+ TensorDesc desc;
325+ if (oldNode.GetOutputDesc(idx, desc) == GRAPH_SUCCESS) {
326+ (void)newNode.UpdateOutputDesc(idx, desc);
327+ }
328+ 
329+ // 该输出的所有下游都要改接，漏一个就断图
330+ const std::vector<std::pair<GNodePtr, int32_t>> consumers = oldNode.GetOutDataNodesAndPortIndexs(idx);
331+ for (const std::pair<GNodePtr, int32_t>& consumer : consumers) {
332+ if (consumer.first == nullptr) {
333+ continue;
334+ }
335+ if (graph.RemoveEdge(oldNode, idx, *consumer.first, consumer.second) != GRAPH_SUCCESS) {
336+ OPS_LOG_E(kPassName.c_str(), "Remove output edge %zu failed.", j);
337+ return false;
338+ }
339+ if (graph.AddDataEdge(newNode, idx, *consumer.first, consumer.second) != GRAPH_SUCCESS) {
340+ OPS_LOG_E(kPassName.c_str(), "Add output edge %zu failed.", j);
341+ return false;
342+ }
343+ }
344+ }
345+ return true;
346+}
347+ 
348+// 控制边搬运。
349+bool RewireControlEdges(Graph& graph, const GNode& oldNode, GNode& newNode)
350+{
351+ for (const GNodePtr& src : oldNode.GetInControlNodes()) {
352+ if (src == nullptr) {
353+ continue;
354+ }
355+ if (graph.AddControlEdge(*src, newNode) != GRAPH_SUCCESS) {
356+ OPS_LOG_E(kPassName.c_str(), "Add in-control edge failed.");
357+ return false;
358+ }
359+ }
360+ for (const GNodePtr& dst : oldNode.GetOutControlNodes()) {
361+ if (dst == nullptr) {
362+ continue;
363+ }
364+ if (graph.AddControlEdge(newNode, *dst) != GRAPH_SUCCESS) {
365+ OPS_LOG_E(kPassName.c_str(), "Add out-control edge failed.");
366+ return false;
367+ }
368+ }
369+ return true;
370+}
371+ 
372+bool RemoveOldNode(Graph& graph, GNode& oldNode, size_t inputNum)
373+{
374+ for (size_t i = 0U; i < inputNum; ++i) {
375+ const int32_t idx = static_cast<int32_t>(i);
376+ const std::pair<GNodePtr, int32_t> peer = oldNode.GetInDataNodesAndPortIndexs(idx);
377+ if (peer.first == nullptr) {
378+ continue;
379+ }
380+ if (graph.RemoveEdge(*peer.first, peer.second, oldNode, idx) != GRAPH_SUCCESS) {
381+ OPS_LOG_E(kPassName.c_str(), "Remove input edge %zu failed.", i);
382+ return false;
383+ }
384+ }
385+ if (graph.RemoveNode(oldNode) != GRAPH_SUCCESS) {
386+ OPS_LOG_E(kPassName.c_str(), "Remove old node failed.");
387+ return false;
388+ }
389+ return true;
390+}
391+ 
392+// 单节点替换：建点 -> 接输入 -> 改接输出 -> 搬控制边 -> 断旧边删旧点
393+bool ReplaceOneNode(Graph& graph, GNode& oldNode)
394+{
395+ const size_t inputNum = oldNode.GetInputsSize();
396+ 
397+ GNode newNode;
398+ if (!CreateInplaceNode(graph, oldNode, newNode)) {
399+ return false;
400+ }
401+ if (!RewireInputs(graph, oldNode, newNode, inputNum)) {
402+ return false;
403+ }
404+ if (!RewireOutputs(graph, oldNode, newNode)) {
405+ return false;
406+ }
407+ if (!RewireControlEdges(graph, oldNode, newNode)) {
408+ return false;
409+ }
410+ if (!RemoveOldNode(graph, oldNode, inputNum)) {
411+ return false;
412+ }
413+ 
414+ OPS_LOG_D(kPassName.c_str(), "Replaced, with_bias=%d.", static_cast<int32_t>(inputNum == kInputNumWithBias));
415+ return true;
416+}
417+} // namespace
418+ 
419+Status ZInplaceAddLayerNormFusionPass::Run(GraphPtr& graph, CustomPassContext& passContext)
420+{
421+ if (graph == nullptr) {
422+ return GRAPH_NOT_CHANGED;
423+ }
424+ 
425+ // 全局守卫：与具体节点无关，只判一次
426+ const SceneCheckResult scene = CheckScene(passContext);
427+ if (scene == SceneCheckResult::API_UNAVAILABLE) {
428+ // 兼容，保持静默
429+ OPS_LOG_D(kPassName.c_str(), "GetOptionValue is unavailable below CANN 9.0.0, stay silent.");
430+ return GRAPH_NOT_CHANGED;
431+ }
432+ if (scene == SceneCheckResult::TRAIN) {
433+ OPS_LOG_D(kPassName.c_str(), "Train mode is not supported, skip.");
434+ return GRAPH_NOT_CHANGED;
435+ }
436+ 
437+ int64_t l2Size = 0L;
438+ if (!IsSupportedPlatform(l2Size)) {
439+ return GRAPH_NOT_CHANGED;
440+ }
441+ 
442+ // 扫图取候选。
443+ std::vector<GNode> candidates;
444+ for (auto& node : graph->GetDirectNode()) {
445+ AscendString nodeType;
446+ if (node.GetType(nodeType) != GRAPH_SUCCESS) {
447+ continue;
448+ }
449+ if (nodeType != AscendString(kAddLayerNormType)) {
450+ continue;
451+ }
452+ candidates.emplace_back(node);
453+ }
454+ if (candidates.empty()) {
455+ OPS_LOG_D(kPassName.c_str(), "No %s node found.", kAddLayerNormType);
456+ return GRAPH_NOT_CHANGED;
457+ }
458+ 
459+ // 替换失败时整图回退，避免留下半改的图
460+ Graph originGraph = *graph;
461+ bool changed = false;
462+ for (auto& node : candidates) {
463+ if (!MeetGuards(node, l2Size)) {
464+ continue;
465+ }
466+ 
467+ if (!ReplaceOneNode(*graph, node)) {
468+ OPS_LOG_E(kPassName.c_str(), "Replacement failed, rollback whole graph.");
469+ passContext.SetErrorMessage(AscendString("ZInplaceAddLayerNormFusionPass replacement failed."));
470+ *graph = originGraph;
471+ return FAILED;
472+ }
473+ changed = true;
474+ }
475+ 
476+ return changed ? SUCCESS : GRAPH_NOT_CHANGED;
477+}
478+ 
479+#if defined(GE_COMPILER_VERSION_NUM) && (GE_COMPILER_VERSION_NUM >= 90000000)
480+REG_FUSION_PASS(ZInplaceAddLayerNormFusionPass).Stage(CustomPassStage::kAfterBuiltinFusionPass);
481+#endif
482+} // namespace ops
Anorm/inplace_add_layer_norm/op_graph/fusion_pass/inplace_add_layer_norm_fusion_pass.h+37-0
@@ -0,0 +1,37 @@
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+/*!
12+ * \file inplace_add_layer_norm_fusion_pass.h
13+ * \brief AddLayerNorm --> InplaceAddLayerNorm (graph base route)
14+ *
15+ * x1 x2 gamma beta [bias] x1 x2 gamma beta [bias]
16+ * \ \ | / / \ \ | / /
17+ * AddLayerNorm --> InplaceAddLayerNorm
18+ * / | | \ / | | \
19+ * y mean rstd x x1' mean rstd x2'
20+ *
21+ */
22+#ifndef OPS_NORM_INPLACE_ADD_LAYER_NORM_OP_GRAPH_FUSION_PASS_INPLACE_ADD_LAYER_NORM_FUSION_PASS_H_
23+#define OPS_NORM_INPLACE_ADD_LAYER_NORM_OP_GRAPH_FUSION_PASS_INPLACE_ADD_LAYER_NORM_FUSION_PASS_H_
24+ 
25+#include "ge/fusion/pass/fusion_base_pass.h"
26+ 
27+namespace ops {
28+using namespace ge;
29+using namespace ge::fusion;
30+ 
31+class __attribute__((visibility("default"))) ZInplaceAddLayerNormFusionPass : public FusionBasePass {
32+public:
33+ Status Run(GraphPtr& graph, CustomPassContext& passContext) override;
34+};
35+} // namespace ops
36+ 
37+#endif // OPS_NORM_INPLACE_ADD_LAYER_NORM_OP_GRAPH_FUSION_PASS_INPLACE_ADD_LAYER_NORM_FUSION_PASS_H_
Anorm/inplace_add_layer_norm/tests/ut/op_graph/CMakeLists.txt+12-0
@@ -0,0 +1,12 @@
1+# Copyright (c) 2026 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+# 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+# 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+# See LICENSE in the root of the software repository for the full text of the License.
8+#/
9+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
10+if(UT_TEST_ALL OR OP_GRAPH_UT)
11+ add_modules_ut_sources(HOSTNAME ${OP_GRAPH_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
12+endif()
Anorm/inplace_add_layer_norm/tests/ut/op_graph/test_inplace_add_layer_norm_fusion_pass.cpp+318-0
@@ -0,0 +1,318 @@
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+#include <map>
11+#include <string>
12+#include <vector>
13+ 
14+#include <gtest/gtest.h>
15+ 
16+#include "es_nn_ops.h"
17+#include "ge/es_graph_builder.h"
18+#include "platform/platform_info.h"
19+#include "register/register_custom_pass.h"
20+#include "../../../op_graph/inplace_add_layer_norm_proto.h"
21+#include "../../../op_graph/fusion_pass/inplace_add_layer_norm_fusion_pass.h"
22+ 
23+using namespace fe;
24+using namespace ge;
25+using namespace ops;
26+ 
27+namespace {
28+constexpr float kEpsilon = 0.00001f;
29+constexpr int64_t kX1Idx = 0;
30+constexpr int64_t kX2Idx = 1;
31+constexpr int64_t kGammaIdx = 2;
32+constexpr int64_t kBetaIdx = 3;
33+constexpr int64_t kBiasIdx = 4;
34+ 
35+// 守卫 G6 要求 l2/8 <= x1_bytes <= 2*l2。
36+constexpr int64_t kL2Size = 134217728L;
37+constexpr int64_t kRowsInBand = 16384L; // fp16 [16384,1024] = 32MB，落在 [16MB,256MB]
38+constexpr int64_t kRowsTooSmall = 256L; // 0.5MB，低于下界
39+constexpr int64_t kRowsTooLarge = 262144L; // 512MB，高于上界
40+constexpr int64_t kCols = 1024L;
41+constexpr int64_t kDynamicReduceAxis = 5120L; // 动态 shape 下的经验归约轴
42+ 
43+const char* const kAddLayerNorm = "AddLayerNorm";
44+const char* const kInplaceAddLayerNorm = "InplaceAddLayerNorm";
45+ 
46+class ZInplaceAddLayerNormFusionPassTest : public testing::Test {
47+protected:
48+ void SetUp() override { SetPlatform("Ascend950"); }
49+ 
50+ static void SetPlatform(const std::string& soc, int64_t l2Size = kL2Size)
51+ {
52+ PlatformInfo platformInfo;
53+ OptionalInfo optionalInfo;
54+ platformInfo.soc_info.ai_core_cnt = 64;
55+ platformInfo.soc_info.l2_size = static_cast<uint64_t>(l2Size);
56+ platformInfo.str_info.short_soc_version = soc;
57+ optionalInfo.soc_version = soc;
58+ PlatformInfoManager::Instance().platform_info_map_[soc] = platformInfo;
59+ PlatformInfoManager::Instance().SetOptionalCompilationInfo(optionalInfo);
60+ }
61+ 
62+ // 构造单个 AddLayerNorm 节点的图。
63+ static std::shared_ptr<Graph> BuildGraph(bool withBias, int64_t rows = kRowsInBand, int64_t cols = kCols,
64+ DataType xDtype = DT_FLOAT16, DataType gammaDtype = DT_FLOAT16,
65+ bool extraConsumerOnX1 = false, bool extraConsumerOnY = false)
66+ {
67+ const std::vector<int64_t> xShape = {rows, cols};
68+ const std::vector<int64_t> pShape = {cols};
69+ 
70+ auto builder = es::EsGraphBuilder("inplace_add_layer_norm_fusion_test");
71+ auto x1 = builder.CreateInput(kX1Idx, "x1", xDtype, FORMAT_ND, xShape);
72+ auto x2 = builder.CreateInput(kX2Idx, "x2", xDtype, FORMAT_ND, xShape);
73+ auto gamma = builder.CreateInput(kGammaIdx, "gamma", gammaDtype, FORMAT_ND, pShape);
74+ auto beta = builder.CreateInput(kBetaIdx, "beta", gammaDtype, FORMAT_ND, pShape);
75+ 
76+ es::AddLayerNormOutput out;
77+ if (withBias) {
78+ auto bias = builder.CreateInput(kBiasIdx, "bias", gammaDtype, FORMAT_ND, pShape);
79+ out = es::AddLayerNorm(x1, x2, gamma, beta, bias, kEpsilon, false);
80+ } else {
81+ out = es::AddLayerNorm(x1, x2, gamma, beta, nullptr, kEpsilon, false);
82+ }
83+ 
84+ UpdateInputDesc(out.y, 0, xDtype, xShape);
85+ UpdateInputDesc(out.y, 1, xDtype, xShape);
86+ UpdateInputDesc(out.y, 2, gammaDtype, pShape);
87+ UpdateInputDesc(out.y, 3, gammaDtype, pShape);
88+ if (withBias) {
89+ UpdateInputDesc(out.y, 4, gammaDtype, pShape);
90+ }
91+ UpdateOutputDesc(out.y, 0, xDtype, xShape);
92+ UpdateOutputDesc(out.mean, 1, DT_FLOAT, {rows, 1});
93+ UpdateOutputDesc(out.rstd, 2, DT_FLOAT, {rows, 1});
94+ UpdateOutputDesc(out.x, 3, xDtype, xShape);
95+ 
96+ std::vector<es::EsTensorHolder> outputs = {out.y, out.mean, out.rstd, out.x};
97+ if (extraConsumerOnX1) {
98+ outputs.emplace_back(es::Relu(x1));
99+ }
100+ if (extraConsumerOnY) {
101+ outputs.emplace_back(es::Relu(out.y));
102+ }
103+ return builder.BuildAndReset(outputs);
104+ }
105+ 
106+ static void UpdateInputDesc(const es::EsTensorHolder& tensor, int32_t index, DataType dtype,
107+ const std::vector<int64_t>& shape)
108+ {
109+ TensorDesc desc;
110+ tensor.GetProducer()->GetInputDesc(index, desc);
111+ desc.SetDataType(dtype);
112+ desc.SetFormat(FORMAT_ND);
113+ desc.SetShape(Shape(shape));
114+ tensor.GetProducer()->UpdateInputDesc(index, desc);
115+ }
116+ 
117+ static void UpdateOutputDesc(const es::EsTensorHolder& tensor, int32_t index, DataType dtype,
118+ const std::vector<int64_t>& shape)
119+ {
120+ TensorDesc desc;
121+ tensor.GetProducer()->GetOutputDesc(index, desc);
122+ desc.SetDataType(dtype);
123+ desc.SetFormat(FORMAT_ND);
124+ desc.SetShape(Shape(shape));
125+ tensor.GetProducer()->UpdateOutputDesc(index, desc);
126+ }
127+ 
128+ static Status RunPass(std::shared_ptr<Graph>& graph)
129+ {
130+ CustomPassContext passContext;
131+ ZInplaceAddLayerNormFusionPass pass;
132+ return pass.Run(graph, passContext);
133+ }
134+ 
135+ static int CountOpType(const std::shared_ptr<Graph>& graph, const char* opType)
136+ {
137+ int count = 0;
138+ for (auto node : graph->GetAllNodes()) {
139+ AscendString type;
140+ node.GetType(type);
141+ if (type == AscendString(opType)) {
142+ ++count;
143+ }
144+ }
145+ return count;
146+ }
147+ 
148+ static bool FindNodeByType(const std::shared_ptr<Graph>& graph, const char* opType, GNode& found)
149+ {
150+ for (auto node : graph->GetAllNodes()) {
151+ AscendString type;
152+ node.GetType(type);
153+ if (type == AscendString(opType)) {
154+ found = node;
155+ return true;
156+ }
157+ }
158+ return false;
159+ }
160+};
161+ 
162+TEST_F(ZInplaceAddLayerNormFusionPassTest, control_edges_transferred_to_new_node)
163+{
164+ auto graph = BuildGraph(false, kRowsInBand, kCols, DT_FLOAT16, DT_FLOAT16, false, true);
165+ 
166+ GNode addLn;
167+ ASSERT_TRUE(FindNodeByType(graph, kAddLayerNorm, addLn));
168+ GNode relu;
169+ ASSERT_TRUE(FindNodeByType(graph, "Relu", relu));
170+ 
171+ // Relu --ctrl--> AddLayerNorm --ctrl--> Relu 都挂上，覆盖入/出两个方向
172+ ASSERT_EQ(graph->AddControlEdge(relu, addLn), GRAPH_SUCCESS);
173+ ASSERT_EQ(graph->AddControlEdge(addLn, relu), GRAPH_SUCCESS);
174+ ASSERT_EQ(addLn.GetInControlNodes().size(), 1U);
175+ ASSERT_EQ(addLn.GetOutControlNodes().size(), 1U);
176+ 
177+ ASSERT_EQ(RunPass(graph), SUCCESS);
178+ 
179+ GNode inplaceNode;
180+ ASSERT_TRUE(FindNodeByType(graph, kInplaceAddLayerNorm, inplaceNode));
181+ EXPECT_EQ(inplaceNode.GetInControlNodes().size(), 1U);
182+ EXPECT_EQ(inplaceNode.GetOutControlNodes().size(), 1U);
183+}
184+ 
185+// ---------- 正向：两种 bias 形态 ----------
186+TEST_F(ZInplaceAddLayerNormFusionPassTest, fusion_success_without_bias)
187+{
188+ auto graph = BuildGraph(false);
189+ EXPECT_EQ(RunPass(graph), SUCCESS);
190+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 1);
191+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 0);
192+}
193+ 
194+TEST_F(ZInplaceAddLayerNormFusionPassTest, fusion_success_with_bias)
195+{
196+ auto graph = BuildGraph(true);
197+ EXPECT_EQ(RunPass(graph), SUCCESS);
198+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 1);
199+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 0);
200+}
201+ 
202+TEST_F(ZInplaceAddLayerNormFusionPassTest, fusion_success_bf16)
203+{
204+ auto graph = BuildGraph(false, kRowsInBand, kCols, DT_BF16, DT_BF16);
205+ EXPECT_EQ(RunPass(graph), SUCCESS);
206+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 1);
207+}
208+ 
209+// ---------- graph 路线特有：接线与 attr 必须逐位保住 ----------
210+TEST_F(ZInplaceAddLayerNormFusionPassTest, inputs_and_outputs_rewired)
211+{
212+ auto graph = BuildGraph(true);
213+ ASSERT_EQ(RunPass(graph), SUCCESS);
214+ 
215+ GNode newNode;
216+ ASSERT_TRUE(FindNodeByType(graph, kInplaceAddLayerNorm, newNode));
217+ EXPECT_EQ(newNode.GetOutputsSize(), 4U);
218+ for (int32_t i = 0; i < 5; ++i) {
219+ EXPECT_NE(newNode.GetInDataNodesAndPortIndexs(i).first, nullptr) << "input " << i << " 未接线";
220+ }
221+}
222+ 
223+TEST_F(ZInplaceAddLayerNormFusionPassTest, downstream_consumer_rewired_to_new_node)
224+{
225+ // y 除了作为图输出，还被一个 Relu 消费；替换后该 Relu 的生产者必须是新节点
226+ auto graph = BuildGraph(false, kRowsInBand, kCols, DT_FLOAT16, DT_FLOAT16, false, true);
227+ ASSERT_EQ(RunPass(graph), SUCCESS);
228+ 
229+ GNode relu;
230+ ASSERT_TRUE(FindNodeByType(graph, "Relu", relu));
231+ auto producer = relu.GetInDataNodesAndPortIndexs(0).first;
232+ ASSERT_NE(producer, nullptr);
233+ AscendString producerType;
234+ producer->GetType(producerType);
235+ EXPECT_EQ(producerType, AscendString(kInplaceAddLayerNorm));
236+}
237+ 
238+TEST_F(ZInplaceAddLayerNormFusionPassTest, attrs_transferred_to_new_node)
239+{
240+ auto graph = BuildGraph(false);
241+ ASSERT_EQ(RunPass(graph), SUCCESS);
242+ 
243+ GNode newNode;
244+ ASSERT_TRUE(FindNodeByType(graph, kInplaceAddLayerNorm, newNode));
245+ float32_t epsilon = 0.0F;
246+ EXPECT_EQ(newNode.GetAttr(AscendString("epsilon"), epsilon), GRAPH_SUCCESS);
247+ EXPECT_FLOAT_EQ(epsilon, kEpsilon);
248+ bool additionalOutput = true;
249+ EXPECT_EQ(newNode.GetAttr(AscendString("additional_output"), additionalOutput), GRAPH_SUCCESS);
250+ EXPECT_FALSE(additionalOutput);
251+}
252+ 
253+// ---------- 守卫 G2：平台校验 ----------
254+TEST_F(ZInplaceAddLayerNormFusionPassTest, guard_reject_unsupported_platform)
255+{
256+ SetPlatform("Ascend310P");
257+ auto graph = BuildGraph(false);
258+ EXPECT_EQ(RunPass(graph), GRAPH_NOT_CHANGED);
259+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 1);
260+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 0);
261+}
262+ 
263+TEST_F(ZInplaceAddLayerNormFusionPassTest, platform_ascend910b_supported)
264+{
265+ SetPlatform("Ascend910B");
266+ auto graph = BuildGraph(false);
267+ EXPECT_EQ(RunPass(graph), SUCCESS);
268+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 1);
269+}
270+ 
271+// ---------- 守卫 G4：x1 被其它节点消费时不可原地 ----------
272+TEST_F(ZInplaceAddLayerNormFusionPassTest, guard_reject_input_shared_by_other_consumers)
273+{
274+ auto graph = BuildGraph(false, kRowsInBand, kCols, DT_FLOAT16, DT_FLOAT16, true);
275+ EXPECT_EQ(RunPass(graph), GRAPH_NOT_CHANGED);
276+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 1);
277+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 0);
278+}
279+ 
280+// ---------- 守卫 G6：shape 必须落在 L2 区间 ----------
281+TEST_F(ZInplaceAddLayerNormFusionPassTest, guard_reject_shape_below_l2_band)
282+{
283+ auto graph = BuildGraph(false, kRowsTooSmall);
284+ EXPECT_EQ(RunPass(graph), GRAPH_NOT_CHANGED);
285+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 1);
286+}
287+ 
288+TEST_F(ZInplaceAddLayerNormFusionPassTest, guard_reject_shape_above_l2_band)
289+{
290+ auto graph = BuildGraph(false, kRowsTooLarge);
291+ EXPECT_EQ(RunPass(graph), GRAPH_NOT_CHANGED);
292+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 1);
293+}
294+ 
295+// ---------- 守卫 G6：动态 shape 走末维判定 ----------
296+TEST_F(ZInplaceAddLayerNormFusionPassTest, fusion_success_dynamic_shape_on_reduce_axis)
297+{
298+ auto graph = BuildGraph(false, -1L, kDynamicReduceAxis);
299+ EXPECT_EQ(RunPass(graph), SUCCESS);
300+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 1);
301+}
302+ 
303+TEST_F(ZInplaceAddLayerNormFusionPassTest, guard_reject_dynamic_shape_other_axis)
304+{
305+ auto graph = BuildGraph(false, -1L, kCols);
306+ EXPECT_EQ(RunPass(graph), GRAPH_NOT_CHANGED);
307+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 1);
308+}
309+ 
310+// ---------- 守卫 G7：x1/x2/gamma dtype 必须一致 ----------
311+TEST_F(ZInplaceAddLayerNormFusionPassTest, guard_reject_mixed_input_dtype)
312+{
313+ auto graph = BuildGraph(false, kRowsInBand, kCols, DT_FLOAT16, DT_FLOAT);
314+ EXPECT_EQ(RunPass(graph), GRAPH_NOT_CHANGED);
315+ EXPECT_EQ(CountOpType(graph, kAddLayerNorm), 1);
316+ EXPECT_EQ(CountOpType(graph, kInplaceAddLayerNorm), 0);
317+}
318+} // namespace