Fix: LayerNormFusion should reject shape-expanding scale and bias patterns - #32660
SIDDARTHA REDDY (SIDDARTHAREDDY8) wants to merge 1 commit into
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
🟡 Changes recommended
Dynamic dimensions can still permit a semantics-changing fusion, and symbolic comparison lacks regression coverage.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Fixes unsafe LayerNorm fusion patterns involving operand order and shape-expanding broadcasts.
Changes:
- Selects scale and bias using graph connectivity.
- Adds broadcast and symbolic-dimension validation.
- Adds positive and negative regression tests.
File summaries
| File | Description |
|---|---|
onnxruntime/core/optimizer/layer_norm_fusion.cc |
Strengthens LayerNorm fusion validation. |
onnxruntime/test/optimizer/graph_transform_test_layernorm.cc |
Adds operand-order and expanding-broadcast tests. |
Review details
Suppressed comments (1)
onnxruntime/core/optimizer/layer_norm_fusion.cc:97
- Allowing an unproven mismatch is still semantics-changing for dynamic shapes. For example, with
Xshaped[N]and scale/bias shaped[2],N=1is a valid execution of the original graph andMulexpands the result to[2]; after fusion,LayerNormalizationpreservesX's shape (and ORT's kernel rejects scale[2]againstX[1]). Treat any aligned dimension that is neither1nor provably equal as unsafe and skip fusion.
// Symbolic or unknown dims that cannot be proven different may still match at runtime.
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Balanced (auto)
Note
Copilot is running an experiment and ran this review at Balanced.
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| static bool AreDimsProvablyEqual(const TensorShapeProto::Dimension& a, const TensorShapeProto::Dimension& b) { | ||
| if (a.has_dim_value() && b.has_dim_value()) { | ||
| return a.dim_value() == b.dim_value(); | ||
| } | ||
| return a.has_dim_param() && b.has_dim_param() && !a.dim_param().empty() && a.dim_param() == b.dim_param(); |
b47894c to
f38538c
Compare
|
Addressed the Copilot review: any aligned scale/bias dim that is neither 1 nor provably equal to the input dim now blocks fusion, which closes the symbolic-dim gap (e.g. scale [2] against input [N], where N == 1 at runtime would let Mul expand the result while LayerNormalization preserves the input shape). Also added three regression tests: matching dim_params still fuse, an unproven symbolic mismatch does not fuse, and the normal dynamic-batch [batch, 64] pattern still fuses. |
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
f5a011c to
f38538c
Compare
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
Reopening: this PR was auto-closed by GitHub (not by a reviewer) when an automated push momentarily reset the branch to upstream
|
Description
Fixes #32603.
Implemented the fix for microsoft/onnxruntime issue #32603 in LayerNormFusion::ApplyImpl (onnxruntime/core/optimizer/layer_norm_fusion.cc), pushed to branch oss-bhai-2026-09-16b on SIDDARTHAREDDY8/onnxruntime at upstream main HEAD. Changes: (1) scale/bias are now selected by graph connectivity (the Mul/Add input that is not the normalized data path, tracked through an optional Cast after Div) rather than rank alone, with the rank-vs-axes check kept as validation; (2) new IsBroadcastKnownToExpandShape helper rejects fusion when the trailing Mul/Add broadcast provably expands the normalized input shape (e.g. Mul([1,1],[2]) -> [1,2]) while allowing unprovable symbolic/unknown dims; (3) new AreDimsProvablyEqual compares dims via dim_param when dim_value()==0, fixing both the scale-vs-bias same-shape check and the expansion check. Added two regression tests in onnxruntime/test/optimizer/graph_transform_test_layernorm.cc (opset 17, Level1): LayerNormFusionFusesScaleFirstOperandOrder (rank-1 input, Mul(scale,div_out)/Add(bias,mul_out) still fuses AND uses the correct scale/bias initializer operands, verified by input name) and LayerNormFusionRejectsExpandingScaleBiasBroadcast (Mul([1,1],[2])+[2] must not fuse). CONTRIBUTING base branch 'main' confirmed and used. AI-disclosure: this contribution is AI-assisted.
Motivation and Context
Fixes microsoft/onnxruntime issue #32603: LayerNormFusion could fuse patterns where the scale/bias Mul/Add broadcast expands the normalized input shape, producing a wrong LayerNormalization node, and could pick the wrong operands when the scale/bias appears as the first Mul/Add operand instead of the second.
Testing: full onnxruntime C++ test suite is not runnable in this environment (build requires hours plus the full dep toolchain). Verification done: (a) standalone g++ harness compiled the two new helpers copied verbatim from the patch with stubbed protobuf types - 20/20 logic checks passed covering the issue's Mul([1,1],[2]) example, normal/1-d/symbolic/unknown/multi-dim scale shapes, and dim_param equality (different params correctly unequal, which the old dim_value() comparison got wrong); (b) manual trace of both new tests through ApplyImpl confirms old code fuses the negative case and picks wrong operands in the positive case, new code rejects the former and picks correct operands in the latter; (c) line-length/style checked against surrounding code; std::optional usage matches existing patterns in the test file.