Skip to content

Commit 0272ddd

Browse files
mvafinmvafin
authored
[ONNX FE] Support Mean and InvStdDev outputs of LayerNormalization (#36086)
### Details: - Extend ONNX LayerNormalization translator to emit Mean and InvStdDev (computed via the spec's reference decomposition); MVN remains the fast path for Y. - Drop stale skip_segfault block for 19 LayerNormalization conformance tests; failure is now resolved at translation time. - Dedupe scale/bias reshape via a shared lambda. ### Tickets: - CVS-90649 ### AI Assistance: - AI assistance used: yes - Used to investigate the skip_segfault block, design and implement the Mean/InvStdDev decomposition, and simplify the function.
1 parent 3d614db commit 0272ddd

2 files changed

Lines changed: 69 additions & 54 deletions

File tree

‎src/frontends/onnx/frontend/src/op/layer_normalization.cpp‎

Lines changed: 69 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
// SPDX-License-Identifier: Apache-2.0
33
//
44

5+
#include <memory>
6+
7+
#include "core/null_node.hpp"
58
#include "core/operator_set.hpp"
69
#include "exceptions.hpp"
710
#include "openvino/core/validation_util.hpp"
@@ -10,13 +13,17 @@
1013
#include "openvino/op/constant.hpp"
1114
#include "openvino/op/convert.hpp"
1215
#include "openvino/op/convert_like.hpp"
16+
#include "openvino/op/divide.hpp"
1317
#include "openvino/op/multiply.hpp"
1418
#include "openvino/op/mvn.hpp"
1519
#include "openvino/op/range.hpp"
20+
#include "openvino/op/reduce_mean.hpp"
1621
#include "openvino/op/reshape.hpp"
1722
#include "openvino/op/shape_of.hpp"
1823
#include "openvino/op/slice.hpp"
24+
#include "openvino/op/sqrt.hpp"
1925
#include "openvino/op/squeeze.hpp"
26+
#include "openvino/op/subtract.hpp"
2027
#include "utils/common.hpp"
2128
using namespace ov::op;
2229
using namespace ov::op::v0;
@@ -43,12 +50,11 @@ ov::OutputVector layer_normalization(const ov::frontend::onnx::Node& node) {
4350
num_inputs == 2 || num_inputs == 3,
4451
"LayerNormalization expects 2 or 3 input tensors. Got: ",
4552
num_inputs);
53+
const auto num_outputs = node.get_outputs_size();
4654
CHECK_VALID_NODE(node,
47-
node.get_outputs_size() == 1,
48-
"LayerNormalization expects 1 output tensor to be used in a model, other configurations are used "
49-
"for training and are not supported. Got: ",
50-
node.get_outputs_size(),
51-
" outputs.");
55+
num_outputs >= 1 && num_outputs <= 3,
56+
"LayerNormalization expects 1, 2 or 3 output tensors. Got: ",
57+
num_outputs);
5258

5359
auto default_stash_type_i = static_cast<int64_t>(TensorProto_DataType::TensorProto_DataType_FLOAT);
5460
int64_t stash_type_i = node.get_attribute_value<int64_t>("stash_type", default_stash_type_i);
@@ -81,36 +87,68 @@ ov::OutputVector layer_normalization(const ov::frontend::onnx::Node& node) {
8187
if (needs_type_casting)
8288
normalized = std::make_shared<ConvertLike>(normalized, inputs.at(0));
8389

84-
ov::Output<ov::Node> normalized_shape = std::make_shared<v0::ShapeOf>(normalized);
85-
ov::Output<ov::Node> sub_shape = std::make_shared<v8::Slice>(normalized_shape,
86-
Constant::create(element::i64, {1}, {axis}),
87-
Constant::create(element::i64, {1}, {INT_MAX}),
88-
Constant::create(element::i64, {1}, {1}));
89-
auto normalized_rank = normalized.get_partial_shape().rank();
90-
91-
auto scale = inputs.at(1);
92-
auto scale_rank = scale.get_partial_shape().rank();
93-
if ((scale_rank.is_dynamic() && normalized_rank.is_dynamic()) ||
94-
((scale_rank.is_static() && normalized_rank.is_static()) &&
95-
scale_rank.get_length() + normalize_axis(axis, normalized_rank.get_length()) !=
96-
static_cast<size_t>(normalized_rank.get_length()))) {
97-
scale = std::make_shared<v1::Reshape>(scale, sub_shape, false);
98-
}
99-
auto scaled = std::make_shared<Multiply>(normalized, scale);
90+
// Use int32 max as the slice stop value; int64 max is not supported by all plugins (WA).
91+
constexpr auto slice_stop = std::numeric_limits<std::int32_t>::max();
92+
auto sub_shape = std::make_shared<v8::Slice>(std::make_shared<v0::ShapeOf>(normalized),
93+
Constant::create(element::i64, {1}, {axis}),
94+
Constant::create(element::i64, {1}, {slice_stop}),
95+
Constant::create(element::i64, {1}, {1}));
96+
const auto normalized_rank = normalized.get_partial_shape().rank();
97+
const auto reshape_to_sub_shape = [&](ov::Output<ov::Node> param) -> ov::Output<ov::Node> {
98+
const auto param_rank = param.get_partial_shape().rank();
99+
const bool both_dynamic = param_rank.is_dynamic() && normalized_rank.is_dynamic();
100+
const bool size_mismatch = param_rank.is_static() && normalized_rank.is_static() &&
101+
param_rank.get_length() + normalize_axis(axis, normalized_rank.get_length()) !=
102+
static_cast<size_t>(normalized_rank.get_length());
103+
if (both_dynamic || size_mismatch) {
104+
return std::make_shared<v1::Reshape>(param, sub_shape, false);
105+
}
106+
return param;
107+
};
100108

109+
ov::Output<ov::Node> y = std::make_shared<Multiply>(normalized, reshape_to_sub_shape(inputs.at(1)));
101110
if (common::is_input_valid(node, 2)) {
102-
auto bias = inputs.at(2);
103-
auto bias_rank = bias.get_partial_shape().rank();
104-
if ((bias_rank.is_dynamic() && normalized_rank.is_dynamic()) ||
105-
((bias_rank.is_static() && normalized_rank.is_static()) &&
106-
bias_rank.get_length() + normalize_axis(axis, normalized_rank.get_length()) !=
107-
static_cast<size_t>(normalized_rank.get_length()))) {
108-
bias = std::make_shared<v1::Reshape>(bias, sub_shape, false);
111+
y = std::make_shared<Add>(y, reshape_to_sub_shape(inputs.at(2)));
112+
}
113+
114+
ov::OutputVector results{y};
115+
if (num_outputs == 1) {
116+
return results;
117+
}
118+
119+
// Mean and InvStdDev are emitted in stash_type. MVN doesn't expose them, so they're recomputed via the
120+
// spec's reference decomposition (reduce over the same axes with keep_dims=true).
121+
const auto& output_names = node.get_output_names();
122+
const auto wanted = [&](size_t i) {
123+
return num_outputs > i && output_names.size() > i && !output_names[i].get().empty();
124+
};
125+
const auto null_output = []() {
126+
return std::make_shared<NullNode>()->output(0);
127+
};
128+
129+
// Only build the reference decomposition when Mean and/or InvStdDev are actually requested, so inference-only
130+
// models that keep the extra outputs but leave them empty don't get redundant ReduceMean nodes.
131+
constexpr auto keep_dims = true;
132+
std::shared_ptr<ov::Node> mean;
133+
if (wanted(1) || wanted(2)) {
134+
mean = std::make_shared<v1::ReduceMean>(data, axes, keep_dims);
135+
}
136+
if (num_outputs >= 2) {
137+
results.push_back(wanted(1) ? mean->output(0) : null_output());
138+
}
139+
if (num_outputs >= 3) {
140+
if (wanted(2)) {
141+
auto deviation = std::make_shared<v1::Subtract>(data, mean);
142+
auto variance =
143+
std::make_shared<v1::ReduceMean>(std::make_shared<Multiply>(deviation, deviation), axes, keep_dims);
144+
auto std_dev = std::make_shared<v0::Sqrt>(
145+
std::make_shared<v1::Add>(variance, Constant::create(stash_type, {}, {epsilon})));
146+
results.push_back(std::make_shared<v1::Divide>(Constant::create(stash_type, {}, {1}), std_dev)->output(0));
147+
} else {
148+
results.push_back(null_output());
109149
}
110-
return {std::make_shared<Add>(scaled, bias)->output(0)};
111-
} else {
112-
return {scaled->output(0)};
113150
}
151+
return results;
114152
}
115153

116154
ONNX_OP("LayerNormalization", OPSET_SINCE(1), ai_onnx::opset_1::layer_normalization);

‎src/frontends/onnx/tests/tests_python/test_backend.py‎

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
xfail_issue_63043,
3434
xfail_issue_63137,
3535
xfail_issue_69444,
36-
skip_segfault,
3736
xfail_issue_82039,
3837
xfail_issue_90649,
3938
skip_bitwise_ui64,
@@ -313,28 +312,6 @@ def expect_fail(test_case_path, xfail): # type: (str) -> None
313312
"OnnxBackendNodeModelTest.test_resize_downsample_scales_cubic_A_n0p5_exclude_outside_cpu",
314313
"OnnxBackendNodeModelTest.test_resize_upsample_scales_cubic_A_n0p5_exclude_outside_cpu",
315314
),
316-
(
317-
skip_segfault,
318-
"OnnxBackendNodeModelTest.test_layer_normalization_2d_axis0_cpu", # ticket: 90649
319-
"OnnxBackendNodeModelTest.test_layer_normalization_2d_axis1_cpu", # ticket: 90649
320-
"OnnxBackendNodeModelTest.test_layer_normalization_2d_axis_negative_1_cpu", # ticket: 90649
321-
"OnnxBackendNodeModelTest.test_layer_normalization_2d_axis_negative_2_cpu", # ticket: 90649
322-
"OnnxBackendNodeModelTest.test_layer_normalization_3d_axis0_epsilon_cpu", # ticket: 90649
323-
"OnnxBackendNodeModelTest.test_layer_normalization_3d_axis1_epsilon_cpu", # ticket: 90649
324-
"OnnxBackendNodeModelTest.test_layer_normalization_3d_axis2_epsilon_cpu", # ticket: 90649
325-
"OnnxBackendNodeModelTest.test_layer_normalization_3d_axis_negative_1_epsilon_cpu", # ticket: 90649
326-
"OnnxBackendNodeModelTest.test_layer_normalization_3d_axis_negative_2_epsilon_cpu", # ticket: 90649
327-
"OnnxBackendNodeModelTest.test_layer_normalization_3d_axis_negative_3_epsilon_cpu", # ticket: 90649
328-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis0_cpu", # ticket: 90649
329-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis1_cpu", # ticket: 90649
330-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis2_cpu", # ticket: 90649
331-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis3_cpu", # ticket: 90649
332-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis_negative_1_cpu", # ticket: 90649
333-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis_negative_2_cpu", # ticket: 90649
334-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis_negative_3_cpu", # ticket: 90649
335-
"OnnxBackendNodeModelTest.test_layer_normalization_4d_axis_negative_4_cpu", # ticket: 90649
336-
"OnnxBackendNodeModelTest.test_layer_normalization_default_axis_cpu", # ticket: 90649
337-
),
338315
(
339316
xfail_issue_82039,
340317
"OnnxBackendNodeModelTest.test_identity_opt_cpu",

0 commit comments

Comments
 (0)