Preserve double constants in graph fusions

onnxruntime/core/optimizer/layer_norm_fusion.cc

@@ -562,20 +562,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 ||
+        !TryGetScalarInitializerAsDouble(graph, *epsilon_input, epsilon) ||
+        static_cast<double>(static_cast<float>(epsilon)) != epsilon) {
+      continue;
+    }
+
     InlinedVector<NodeArg*> 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<float>(epsilon));
-    } else {
-      layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON);
-    }
+    layer_norm_node.AddAttribute("epsilon", static_cast<float>(epsilon));
 
     // The axis definition of layer_norm is ranging from axis to the last dim
     layer_norm_node.AddAttribute("axis", static_cast<int64_t>(axes_values[0]));
@@ -814,25 +815,17 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr
       continue;
     }
 
+    double epsilon = 0.0;
+    if (!TryGetScalarInitializerAsDouble(graph, *epsilon_input, epsilon) ||
+        static_cast<double>(static_cast<float>(epsilon)) != epsilon) {
+      continue;
+    }
+
     InlinedVector<NodeArg*> 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<float>()[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<float>(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 ||

onnxruntime/core/optimizer/matmul_scale_fusion.cc

@@ -3,6 +3,9 @@
 
 #include "core/optimizer/matmul_scale_fusion.h"
 
+#include <limits>
+#include <type_traits>
+
 #include "onnx/defs/attr_proto_util.h"
 
 #include "core/common/optional.h"
@@ -15,13 +18,46 @@
 namespace onnxruntime {
 
 namespace {
+template <typename T>
+bool IsExactlyRepresentableAsFloat(T value, float converted) {
+  if constexpr (std::is_integral_v<T>) {
+    using U = std::make_unsigned_t<T>;
+    const U unsigned_value = static_cast<U>(value);
+    const U magnitude = [&]() {
+      if constexpr (std::is_signed_v<T>) {
+        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<float>::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;
+  } else if constexpr (std::is_same_v<T, double>) {
+    return static_cast<double>(converted) == value;
+  } else {
+    return true;
+  }
+}
+
 template <typename T>
 struct ExtractScalarAsFloatDispatchTarget {
   Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const std::filesystem::path& model_path,
-                    float& scalar_float) {
+                    float& scalar_float, bool& lossless) {
     T scalar;
     ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, model_path, &scalar, 1));
     scalar_float = static_cast<float>(scalar);
+    lossless = IsExactlyRepresentableAsFloat(scalar, scalar_float);
     return Status::OK();
   }
 };
@@ -60,13 +96,15 @@ std::optional<float> GetScalarConstantInitializer(const Graph& graph, const Node
   }
 
   float scalar{};
+  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<Status, ExtractScalarAsFloatDispatchTarget>(*initializer, graph.ModelPath(), scalar)));
+      (dispatcher.InvokeRet<Status, ExtractScalarAsFloatDispatchTarget>(
+          *initializer, graph.ModelPath(), scalar, lossless)));
 
-  return {scalar};
+  return lossless ? std::optional<float>{scalar} : std::nullopt;
 }
 
 // gets the scale value and its input index if node is a fusable scale (Mul or Div by scalar constant)

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -5,6 +5,7 @@
 #pragma warning(disable : 4244)
 #endif
 
+#include <cmath>
 #include <functional>
 #include <fstream>
 #include <random>
@@ -11219,6 +11220,54 @@ TEST_F(GraphTransformationTests, MatMulScaleFusionDoubleType) {
       {kCpuExecutionProvider});
 }
 
+TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyDoubleScale) {
+  auto build_test_case = [](ModelTestBuilder& builder) {
+    auto* input_a = builder.MakeInput<double>({{2, 2}});
+    auto* input_b = builder.MakeInput<double>({{2, 2}});
+    auto* scale = builder.MakeScalarInitializer<double>(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<MatMulScaleFusion>(),
+      TransformerLevel::Level2, 1, check_not_fused, check_not_fused));
+}
+
+TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyIntegerScale) {
+  auto build_test_case = [](ModelTestBuilder& builder) {
+    auto* input_a = builder.MakeInput<int64_t>({{2, 2}});
+    auto* input_b = builder.MakeInput<int64_t>({{2, 2}});
+    auto* scale = builder.MakeScalarInitializer<int64_t>((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<MatMulScaleFusion>(),
+      TransformerLevel::Level2, 1, check_not_fused, check_not_fused));
+}
+
 TEST_F(GraphTransformationTests, MatMulScaleFusionWithScaleInput) {
   TestMatMulScaleFusion(
       MODEL_FOLDER "fusion/matmul_scale_with_scale_input.onnx", *logger_,

onnxruntime/test/optimizer/graph_transform_test_layernorm.cc

@@ -243,7 +243,7 @@ TEST_F(GraphTransformationTests, LayerNormWithSubDupFusionTest) {
 }
 
 void BuildLayerNorm(ModelTestBuilder& builder, std::vector<int64_t> reduce1_axes = {-1},
-                    std::vector<int64_t> reduce2_axes = {-1}) {
+                    std::vector<int64_t> reduce2_axes = {-1}, bool epsilon_input_first = false) {
   std::vector<int64_t> input_shape = {2, 3, 3, 3};
   auto* data_arg = builder.MakeInput<MLFloat16>(input_shape);
   auto* pow_initializer = builder.MakeInitializer<float>({}, {2.0f});
@@ -295,7 +295,11 @@ void BuildLayerNorm(ModelTestBuilder& builder, std::vector<int64_t> 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 +348,24 @@ 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);
+  };
+
+  auto post_graph_checker = [](Graph& graph) {
+    const auto op_to_count = CountOpsInGraph(graph);
+    TEST_RETURN_IF_NOT(op_to_count.at("LayerNormalization") == 1);
+    return Status::OK();
+  };
+
+  const InlinedHashSet<std::string_view> compatible_execution_providers;
+  ASSERT_STATUS_OK(TestGraphTransformer(
+      build_test_case, 14, *logger_,
+      std::make_unique<LayerNormFusion>(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 +520,73 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionTest) {
   }
 }
 
+TEST_F(GraphTransformationTests, LayerNormFusionsRejectLossyDoubleEpsilon) {
+  auto build_layer_norm = [](ModelTestBuilder& builder) {
+    auto* input = builder.MakeInput<double>({{2, 4}});
+    auto* scale = builder.MakeInitializer<double>({4}, {1.0, 1.0, 1.0, 1.0});
+    auto* bias = builder.MakeInitializer<double>({4}, {0.0, 0.0, 0.0, 0.0});
+    auto* exponent = builder.MakeScalarInitializer<double>(2.0);
+    auto* epsilon = builder.MakeScalarInitializer<double>(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<int64_t>{-1});
+    builder.AddNode("Sub", {input, mean1}, {sub});
+    builder.AddNode("Pow", {sub, exponent}, {squared});
+    builder.AddNode("ReduceMean", {squared}, {mean2}).AddAttribute("axes", std::vector<int64_t>{-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<double>({{2, 4}});
+    auto* scale = builder.MakeInitializer<double>({4}, {1.0, 1.0, 1.0, 1.0});
+    auto* exponent = builder.MakeScalarInitializer<double>(2.0);
+    auto* epsilon = builder.MakeScalarInitializer<double>(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<int64_t>{-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<std::string_view> compatible_execution_providers;
+  ASSERT_STATUS_OK(TestGraphTransformer(
+      build_layer_norm, 14, *logger_,
+      std::make_unique<LayerNormFusion>(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<SimplifiedLayerNormFusion>(),
+      TransformerLevel::Level2, 1, check_not_fused, check_not_fused));
+}
+
 TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionSharedCastPowExponent) {
   for (const bool has_leading_cast : {false, true}) {
     auto build_test_case = [has_leading_cast](ModelTestBuilder& builder) {
@@ -2598,26 +2687,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<std::string, int> 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 +2786,15 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusion_ZeroElementEpsilon) {
   ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique<SimplifiedLayerNormFusion>(), 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<std::string, int> 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