diff --git a/onnxruntime/core/optimizer/layer_norm_fusion.cc b/onnxruntime/core/optimizer/layer_norm_fusion.cc index e4716b5a0d366..59a69411dcb30 100644 --- a/onnxruntime/core/optimizer/layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc @@ -15,8 +15,6 @@ namespace onnxruntime { // LayerNorm supports limited data types. static constexpr std::array supported_data_types{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}; -// Default epsilon -static constexpr float DEFAULT_LAYERNORM_EPSILON = 1e-5f; static bool IsSupportedDataType(const Node& node, int first_n_inputs = -1) { int input_index = 0; @@ -79,7 +77,8 @@ static std::vector GetAxesFromReduceMeanNode(Node& reduce_mean_node, co return axes_values; }; -static bool TryGetScalarInitializerAsDouble(const Graph& graph, const NodeArg& node_arg, double& value) { +template +static bool TryGetScalarInitializer(const Graph& graph, const NodeArg& node_arg, T& value) { const auto* tensor_proto = graph_utils::GetConstantInitializer(graph, node_arg.Name()); if (tensor_proto == nullptr) { return false; @@ -92,46 +91,70 @@ static bool TryGetScalarInitializerAsDouble(const Graph& graph, const NodeArg& n switch (tensor_proto->data_type()) { case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: - value = initializer.data()[0]; + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT8: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT32: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT64: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT8: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT32: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT64: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; default: return false; } } +static bool TryGetScalarInitializerOrFloatCastAsDouble(const Graph& graph, const NodeArg& node_arg, double& value) { + if (TryGetScalarInitializer(graph, node_arg, value)) { + return true; + } + + const Node* cast_node = graph.GetProducerNode(node_arg.Name()); + if (cast_node == nullptr || + !graph_utils::IsSupportedOptypeVersionAndDomain(*cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) || + cast_node->InputDefs().empty() || cast_node->InputDefs()[0] == nullptr) { + return false; + } + + const auto to_it = cast_node->GetAttributes().find("to"); + float cast_value = 0.0f; + if (to_it == cast_node->GetAttributes().end() || + to_it->second.i() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT || + !TryGetScalarInitializer(graph, *cast_node->InputDefs()[0], cast_value)) { + return false; + } + + value = static_cast(cast_value); + return true; +} + static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) { const auto& pow_inputs = pow_node.InputDefs(); if (pow_inputs.size() < 2 || pow_inputs[1] == nullptr) { @@ -139,7 +162,7 @@ static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) { } double exponent_value = 0.0; - if (TryGetScalarInitializerAsDouble(graph, *pow_inputs[1], exponent_value)) { + if (TryGetScalarInitializer(graph, *pow_inputs[1], exponent_value)) { return exponent_value == 2.0; } @@ -150,7 +173,7 @@ static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) { return false; } - return TryGetScalarInitializerAsDouble(graph, *exponent_input_node->InputDefs()[0], exponent_value) && + return TryGetScalarInitializer(graph, *exponent_input_node->InputDefs()[0], exponent_value) && exponent_value == 2.0; } @@ -562,20 +585,21 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, continue; } + const NodeArg* epsilon_input = GetOtherAddInput(add2_node, *reduce_mean2_node.MutableOutputDefs()[0]); + double epsilon = 0.0; + if (epsilon_input == nullptr || + !TryGetScalarInitializerOrFloatCastAsDouble(graph, *epsilon_input, epsilon) || + static_cast(static_cast(epsilon)) != epsilon) { + continue; + } + InlinedVector layer_norm_input_defs{x_input, scale, bias}; Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName(mul_node.Name() + "/LayerNormFusion/"), "LayerNormalization", "fused LayerNorm subgraphs ", layer_norm_input_defs, {}, mul_node, nullptr, kOnnxDomain); - - // Get constant "epsilon" from "Add2" node if available. Else, default value will be used. - double epsilon = 0.0; - if (TryGetScalarInitializerAsDouble(graph, *add2_node.MutableInputDefs()[1], epsilon)) { - layer_norm_node.AddAttribute("epsilon", static_cast(epsilon)); - } else { - layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); - } + layer_norm_node.AddAttribute("epsilon", static_cast(epsilon)); // The axis definition of layer_norm is ranging from axis to the last dim layer_norm_node.AddAttribute("axis", static_cast(axes_values[0])); @@ -814,25 +838,17 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr continue; } + double epsilon = 0.0; + if (!TryGetScalarInitializerOrFloatCastAsDouble(graph, *epsilon_input, epsilon) || + static_cast(static_cast(epsilon)) != epsilon) { + continue; + } + InlinedVector layer_norm_input_defs{x_input, scale}; Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName(mul_node.Name() + "/SimplifiedLayerNormFusion/"), "SimplifiedLayerNormalization", "fused LayerNorm subgraphs ", layer_norm_input_defs, {}, mul_node, nullptr, kOnnxDomain); - - // Get constant "epsilon" from "Add" node if available. Else, default value will be used. - const ONNX_NAMESPACE::TensorProto* tensor_proto = - graph_utils::GetConstantInitializer(graph, epsilon_input->Name()); - if (tensor_proto != nullptr && tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { - Initializer initializer{graph, *tensor_proto, graph.ModelPath()}; - // epsilon must be a scalar/1-element tensor; fall back to default otherwise. - if (initializer.size() == 1) { - layer_norm_node.AddAttribute("epsilon", initializer.data()[0]); - } else { - layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); - } - } else { - layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); - } + layer_norm_node.AddAttribute("epsilon", static_cast(epsilon)); // Set stash_type to double if any input is double, default value if float. if (x_input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE || diff --git a/onnxruntime/core/optimizer/matmul_scale_fusion.cc b/onnxruntime/core/optimizer/matmul_scale_fusion.cc index 151cd19df63c8..71ddb2dd28230 100644 --- a/onnxruntime/core/optimizer/matmul_scale_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_scale_fusion.cc @@ -3,6 +3,9 @@ #include "core/optimizer/matmul_scale_fusion.h" +#include +#include + #include "onnx/defs/attr_proto_util.h" #include "core/common/optional.h" @@ -16,12 +19,54 @@ namespace onnxruntime { namespace { template -struct ExtractScalarAsFloatDispatchTarget { +bool IsIntegralExactlyRepresentableAsFloat(T value) { + static_assert(std::is_integral_v); + using U = std::make_unsigned_t; + const U unsigned_value = static_cast(value); + const U magnitude = [&]() { + if constexpr (std::is_signed_v) { + return value < 0 ? U{0} - unsigned_value : unsigned_value; + } else { + return unsigned_value; + } + }(); + + size_t significant_bits = 0; + for (U remaining = magnitude; remaining != 0; remaining >>= 1) { + ++significant_bits; + } + + constexpr size_t float_digits = std::numeric_limits::digits; + if (significant_bits <= float_digits) { + return true; + } + + const size_t discarded_bits = significant_bits - float_digits; + return (magnitude & ((U{1} << discarded_bits) - 1)) == 0; +} + +struct ScalarScale { + double value; + bool is_double; +}; + +template +struct ExtractScalarDispatchTarget { Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const std::filesystem::path& model_path, - float& scalar_float) { + ScalarScale& scale, bool& lossless) { T scalar; ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, model_path, &scalar, 1)); - scalar_float = static_cast(scalar); + const float scalar_float = static_cast(scalar); + if constexpr (std::is_same_v) { + lossless = static_cast(scalar_float) == scalar; + } else if constexpr (std::is_integral_v) { + lossless = IsIntegralExactlyRepresentableAsFloat(scalar); + } else { + lossless = true; + } + scale = {std::is_same_v ? static_cast(scalar) + : static_cast(scalar_float), + std::is_same_v}; return Status::OK(); } }; @@ -41,7 +86,7 @@ std::optional GetTensorShape(const NodeArg& node_arg) { // if it does not have any broadcasting effect on the Mul or Div output shape. // Because the dimension lengths can only be 1, we only need to consider additional leading dimensions being prepended. // `max_rank` should be set to the rank of the other Mul or Div input to avoid that. -std::optional GetScalarConstantInitializer(const Graph& graph, const NodeArg& node_arg, size_t max_rank) { +std::optional GetScalarConstantInitializer(const Graph& graph, const NodeArg& node_arg, size_t max_rank) { const auto* initializer = graph_utils::GetConstantInitializer(graph, node_arg.Name()); if (!initializer) { @@ -59,18 +104,20 @@ std::optional GetScalarConstantInitializer(const Graph& graph, const Node return {}; } - float scalar{}; + ScalarScale scale{}; + bool lossless = false; utils::MLTypeCallDispatcher< uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double, BFloat16> dispatcher{initializer->data_type()}; ORT_THROW_IF_ERROR( - (dispatcher.InvokeRet(*initializer, graph.ModelPath(), scalar))); + (dispatcher.InvokeRet( + *initializer, graph.ModelPath(), scale, lossless))); - return {scalar}; + return lossless ? std::optional{scale} : std::nullopt; } // gets the scale value and its input index if node is a fusable scale (Mul or Div by scalar constant) -std::optional> GetScaleFromNode( +std::optional> GetScaleFromNode( const Graph& graph, const Node& scale_node, const InlinedHashSet& excluded_initializer_names) { const auto is_excluded_initializer = @@ -95,7 +142,11 @@ std::optional> GetScaleFromNode( if (!divisor.has_value()) return std::nullopt; - return {std::make_pair(1.0f / *divisor, scale_reciprocal_arg_index)}; + ScalarScale scale = *divisor; + scale.value = scale.is_double + ? 1.0 / scale.value + : static_cast(1.0f / static_cast(scale.value)); + return {std::make_pair(scale, scale_reciprocal_arg_index)}; } if (graph_utils::IsSupportedOptypeVersionAndDomain(scale_node, "Mul", {7, 13, 14})) { @@ -128,7 +179,7 @@ struct ScaleMergeInfo { // the edge from the base node to the original node Node::EdgeConstIterator node_to_merge_edge; // the scale of the original node - float scale; + ScalarScale scale; // the index of the input or output def on the original node // this def is moved to the fused node // for a leading scale (scale -> MatMul), it will be the unscaled input @@ -236,15 +287,34 @@ Status ProcessNode( { ONNX_NAMESPACE::AttributeProto& alpha_attr = fused_node_attrs["alpha"]; - float total_scale = utils::HasFloat(alpha_attr) ? alpha_attr.f() : 1.0f; - auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) { - total_scale *= fusion.scale; - }; - - std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale); - std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale); - - alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", total_scale); + const bool uses_double_scale = + std::any_of(input_node_merges.begin(), input_node_merges.end(), + [](const ScaleMergeInfo& fusion) { return fusion.scale.is_double; }) || + std::any_of(output_node_merges.begin(), output_node_merges.end(), + [](const ScaleMergeInfo& fusion) { return fusion.scale.is_double; }); + + if (uses_double_scale) { + double total_scale = utils::HasFloat(alpha_attr) ? static_cast(alpha_attr.f()) : 1.0; + auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) { + total_scale *= fusion.scale.value; + }; + std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale); + std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale); + + const float fused_scale = static_cast(total_scale); + if (static_cast(fused_scale) != total_scale) { + return Status::OK(); + } + alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", fused_scale); + } else { + float total_scale = utils::HasFloat(alpha_attr) ? alpha_attr.f() : 1.0f; + auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) { + total_scale *= static_cast(fusion.scale.value); + }; + std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale); + std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale); + alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", total_scale); + } } auto get_mutable_node_to_merge = [&graph](const ScaleMergeInfo& merge) -> Node& { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 1f4fd0aa541c1..fa2d33a9ac7be 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5,6 +5,7 @@ #pragma warning(disable : 4244) #endif +#include #include #include #include @@ -11219,6 +11220,133 @@ TEST_F(GraphTransformationTests, MatMulScaleFusionDoubleType) { {kCpuExecutionProvider}); } +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyDoubleScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + auto* scale = builder.MakeScalarInitializer(std::nextafter(1.0, 2.0)); + auto* scaled_a = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 1); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyDoubleDivScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + auto* divisor = builder.MakeScalarInitializer(3.0); + auto* scaled_a = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Div", {input_a, divisor}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Div") == 1); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyCombinedDoubleScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + const double scale_value = static_cast(std::nextafter(1.0f, 2.0f)); + auto* scale1 = builder.MakeScalarInitializer(scale_value); + auto* scale2 = builder.MakeScalarInitializer(scale_value); + auto* scaled_a = builder.MakeIntermediate(); + auto* matmul_output = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale1}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {matmul_output}); + builder.AddNode("Mul", {matmul_output, scale2}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 2); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsCancellingLossyDoubleScales) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{1, 1}}); + auto* input_b = builder.MakeInput({{1, 1}}); + auto* scale1 = builder.MakeScalarInitializer(1e300); + auto* scale2 = builder.MakeScalarInitializer(1e-300); + auto* scaled_a = builder.MakeIntermediate(); + auto* matmul_output = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale1}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {matmul_output}); + builder.AddNode("Mul", {matmul_output, scale2}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 2); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyIntegerScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + auto* scale = builder.MakeScalarInitializer((int64_t{1} << 24) + 1); + auto* scaled_a = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 1); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + TEST_F(GraphTransformationTests, MatMulScaleFusionWithScaleInput) { TestMatMulScaleFusion( MODEL_FOLDER "fusion/matmul_scale_with_scale_input.onnx", *logger_, diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index d3092e0561b03..2a6544d6c95fc 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc @@ -243,11 +243,12 @@ TEST_F(GraphTransformationTests, LayerNormWithSubDupFusionTest) { } void BuildLayerNorm(ModelTestBuilder& builder, std::vector reduce1_axes = {-1}, - std::vector reduce2_axes = {-1}) { + std::vector reduce2_axes = {-1}, bool epsilon_input_first = false, + float epsilon = 1e-5f) { std::vector input_shape = {2, 3, 3, 3}; auto* data_arg = builder.MakeInput(input_shape); auto* pow_initializer = builder.MakeInitializer({}, {2.0f}); - auto* add_initializer = builder.MakeInitializer({}, {1e-5f}); + auto* add_initializer = builder.MakeInitializer({}, {epsilon}); std::vector normalized_shape = {}; int64_t normalized_shape_size = 1; auto raxes = reduce1_axes; @@ -295,7 +296,11 @@ void BuildLayerNorm(ModelTestBuilder& builder, std::vector reduce1_axes } else { builder.AddNode("ReduceMean", {pow_out}, {reduce_mean_out_2}).AddAttribute("axes", reduce2_axes); } - builder.AddNode("Add", {reduce_mean_out_2, add_initializer}, {add_out_1}); + if (epsilon_input_first) { + builder.AddNode("Add", {add_initializer, reduce_mean_out_2}, {add_out_1}); + } else { + builder.AddNode("Add", {reduce_mean_out_2, add_initializer}, {add_out_1}); + } builder.AddNode("Sqrt", {add_out_1}, {sqrt_out}); builder.AddNode("Div", {cast_out_1, sqrt_out}, {div_out}); builder.AddNode("Cast", {div_out}, {cast_out_2}) @@ -344,6 +349,31 @@ TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest_5) { TransformerLevel::Level2, 1, pre_graph_checker, post_graph_checker)); } +TEST_F(GraphTransformationTests, LayerNormFusionAddEpsilonInput0) { + auto build_test_case = [](ModelTestBuilder& builder) { + BuildLayerNorm(builder, {-1}, {-1}, true, 1e-4f); + }; + + auto post_graph_checker = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("LayerNormalization") == 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "LayerNormalization") { + const auto epsilon_it = node.GetAttributes().find("epsilon"); + TEST_RETURN_IF_NOT(epsilon_it != node.GetAttributes().end()); + TEST_RETURN_IF_NOT(epsilon_it->second.f() == 1e-4f); + } + } + return Status::OK(); + }; + + const InlinedHashSet compatible_execution_providers; + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 14, *logger_, + std::make_unique(compatible_execution_providers, TransformerLevel::Level2), + TransformerLevel::Level2, 1, nullptr, post_graph_checker)); +} + TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest_6) { auto build_test_case = [](ModelTestBuilder& builder) { BuildLayerNorm(builder, {-2}, {-1}); @@ -498,6 +528,187 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionTest) { } } +TEST_F(GraphTransformationTests, LayerNormFusionsRejectLossyDoubleEpsilon) { + auto build_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* bias = builder.MakeInitializer({4}, {0.0, 0.0, 0.0, 0.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(1e-50); + auto* mean1 = builder.MakeIntermediate(); + auto* sub = builder.MakeIntermediate(); + auto* squared = builder.MakeIntermediate(); + auto* mean2 = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* scaled = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("ReduceMean", {input}, {mean1}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Sub", {input, mean1}, {sub}); + builder.AddNode("Pow", {sub, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean2}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean2, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {sub, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {scaled}); + builder.AddNode("Add", {scaled, bias}, {output}); + }; + + auto build_simplified_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(1e-50); + auto* squared = builder.MakeIntermediate(); + auto* mean = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Pow", {input, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {input, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF(op_to_count.count("LayerNormalization") != 0); + TEST_RETURN_IF(op_to_count.count("SimplifiedLayerNormalization") != 0); + TEST_RETURN_IF_NOT(op_to_count.at("Pow") == 1); + return Status::OK(); + }; + + const InlinedHashSet compatible_execution_providers; + ASSERT_STATUS_OK(TestGraphTransformer( + build_layer_norm, 14, *logger_, + std::make_unique(compatible_execution_providers, TransformerLevel::Level2), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); + ASSERT_STATUS_OK(TestGraphTransformer( + build_simplified_layer_norm, 14, *logger_, + std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, LayerNormFusionsAllowExactDoubleEpsilon) { + auto build_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* bias = builder.MakeInitializer({4}, {0.0, 0.0, 0.0, 0.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(0.125); + auto* mean1 = builder.MakeIntermediate(); + auto* sub = builder.MakeIntermediate(); + auto* squared = builder.MakeIntermediate(); + auto* mean2 = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* scaled = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("ReduceMean", {input}, {mean1}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Sub", {input, mean1}, {sub}); + builder.AddNode("Pow", {sub, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean2}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean2, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {sub, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {scaled}); + builder.AddNode("Add", {scaled, bias}, {output}); + }; + + auto build_simplified_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(0.125); + auto* squared = builder.MakeIntermediate(); + auto* mean = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Pow", {input, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {input, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {output}); + }; + + auto make_checker = [](std::string_view fused_op) { + return [fused_op](Graph& graph) { + TEST_RETURN_IF_NOT(OpCount(CountOpsInGraph(graph), std::string{fused_op}) == 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == fused_op) { + const auto epsilon_it = node.GetAttributes().find("epsilon"); + TEST_RETURN_IF_NOT(epsilon_it != node.GetAttributes().end()); + TEST_RETURN_IF_NOT(epsilon_it->second.f() == 0.125f); + } + } + return Status::OK(); + }; + }; + + const InlinedHashSet compatible_execution_providers; + ASSERT_STATUS_OK(TestGraphTransformer( + build_layer_norm, 14, *logger_, + std::make_unique(compatible_execution_providers, TransformerLevel::Level2), + TransformerLevel::Level2, 1, nullptr, make_checker("LayerNormalization"))); + ASSERT_STATUS_OK(TestGraphTransformer( + build_simplified_layer_norm, 14, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, nullptr, make_checker("SimplifiedLayerNormalization"))); +} + +TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionPreservesIntegerToFloatCastEpsilon) { + constexpr uint64_t kEpsilonSource = (uint64_t{1} << 63) + (uint64_t{1} << 39) + 1; + auto build_test_case = [=](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0f, 1.0f, 1.0f, 1.0f}); + auto* exponent = builder.MakeScalarInitializer(2.0f); + auto* epsilon_source = builder.MakeScalarInitializer(kEpsilonSource); + auto* epsilon = builder.MakeIntermediate(); + auto* squared = builder.MakeIntermediate(); + auto* mean = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Cast", {epsilon_source}, {epsilon}) + .AddAttribute("to", static_cast(TensorProto_DataType_FLOAT)); + builder.AddNode("Pow", {input, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {input, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {output}); + }; + + auto check_after = [=](Graph& graph) { + TEST_RETURN_IF_NOT(OpCount(CountOpsInGraph(graph), "SimplifiedLayerNormalization") == 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "SimplifiedLayerNormalization") { + const auto epsilon_it = node.GetAttributes().find("epsilon"); + TEST_RETURN_IF_NOT(epsilon_it != node.GetAttributes().end()); + TEST_RETURN_IF_NOT(epsilon_it->second.f() == static_cast(kEpsilonSource)); + } + } + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 17, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, nullptr, check_after)); +} + TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionSharedCastPowExponent) { for (const bool has_leading_cast : {false, true}) { auto build_test_case = [has_leading_cast](ModelTestBuilder& builder) { @@ -2598,26 +2809,16 @@ TEST_F(GraphTransformationTests, LayerNormFusion_ZeroElementEpsilon) { // Must handle zero-element epsilon gracefully. ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - // Verify the fusion fired and produced a LayerNormalization with the - // default epsilon (1e-5f), per the zero-element fallback in - // layer_norm_fusion.cc. + // A non-scalar epsilon cannot be represented by the fused float attribute, + // so the original graph must remain unchanged. std::map op_to_count = CountOpsInGraph(graph); - EXPECT_EQ(op_to_count["LayerNormalization"], 1); - EXPECT_EQ(op_to_count["ReduceMean"], 0); - EXPECT_EQ(op_to_count["Sub"], 0); - EXPECT_EQ(op_to_count["Pow"], 0); - EXPECT_EQ(op_to_count["Sqrt"], 0); - EXPECT_EQ(op_to_count["Div"], 0); - EXPECT_EQ(op_to_count["Mul"], 0); - - for (const auto& node : graph.Nodes()) { - if (node.OpType() == "LayerNormalization") { - const auto& attrs = node.GetAttributes(); - auto it = attrs.find("epsilon"); - ASSERT_NE(it, attrs.end()); - EXPECT_FLOAT_EQ(it->second.f(), 1e-5f); - } - } + EXPECT_EQ(op_to_count["LayerNormalization"], 0); + EXPECT_EQ(op_to_count["ReduceMean"], 2); + EXPECT_EQ(op_to_count["Sub"], 1); + EXPECT_EQ(op_to_count["Pow"], 1); + EXPECT_EQ(op_to_count["Sqrt"], 1); + EXPECT_EQ(op_to_count["Div"], 1); + EXPECT_EQ(op_to_count["Mul"], 1); } // Test: SimplifiedLayerNormFusion gracefully handles a zero-element epsilon initializer. @@ -2707,25 +2908,15 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusion_ZeroElementEpsilon) { ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - // Verify the fusion fired and produced a SimplifiedLayerNormalization with - // the default epsilon (1e-5f), per the zero-element fallback in - // layer_norm_fusion.cc. + // A non-scalar epsilon cannot be represented by the fused float attribute, + // so the original graph must remain unchanged. std::map op_to_count = CountOpsInGraph(graph); - EXPECT_EQ(op_to_count["SimplifiedLayerNormalization"], 1); - EXPECT_EQ(op_to_count["Pow"], 0); - EXPECT_EQ(op_to_count["ReduceMean"], 0); - EXPECT_EQ(op_to_count["Sqrt"], 0); - EXPECT_EQ(op_to_count["Div"], 0); - EXPECT_EQ(op_to_count["Mul"], 0); - - for (const auto& node : graph.Nodes()) { - if (node.OpType() == "SimplifiedLayerNormalization") { - const auto& attrs = node.GetAttributes(); - auto it = attrs.find("epsilon"); - ASSERT_NE(it, attrs.end()); - EXPECT_FLOAT_EQ(it->second.f(), 1e-5f); - } - } + EXPECT_EQ(op_to_count["SimplifiedLayerNormalization"], 0); + EXPECT_EQ(op_to_count["Pow"], 1); + EXPECT_EQ(op_to_count["ReduceMean"], 1); + EXPECT_EQ(op_to_count["Sqrt"], 1); + EXPECT_EQ(op_to_count["Div"], 1); + EXPECT_EQ(op_to_count["Mul"], 1); } #endif