Fix sharing scalar bug

onnxruntime/core/optimizer/constant_sharing.cc

@@ -129,7 +129,9 @@ struct GetOrAddValueInConstantStoreDispatcher {
 }  // namespace
 
 Status ConstantSharing::ApplyImpl(Graph& graph, bool& modified, int /*graph_level*/,
-                                  const logging::Logger& /*logger*/) const {
+                                  const logging::Logger& logger) const {
+  int shared_count = 0;
+
   // Accumulated map from type/value/rank to initializer:
   // > The key is a string representation of initializer's data type, value and rank.
   // > The value is newly created initializer NodeArg* to be shared.
@@ -138,9 +140,11 @@ Status ConstantSharing::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve
   InlinedVector<std::string> original_initializer_names;
   original_initializer_names.reserve(initialized_tensor_set.size());
   for (const auto& entry : initialized_tensor_set) {
-    // Ignore if the initializer already handled, or not a constant initializer.
+    // Ignore if the initializer exists in graph output, already handled,
+    // or not a constant initializer (implicitly excludes the graph input).
     if (IsSharedInitializer(entry.first) ||
         !graph_utils::IsConstantInitializer(graph, entry.first) ||
+        graph.IsOutput(graph.GetNodeArg(entry.first)) ||
         excluded_initializers_.find(entry.first) != excluded_initializers_.end()) {
       continue;
     }
@@ -191,6 +195,8 @@ Status ConstantSharing::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve
       NodeArg& shared_scalar_initializer_node_arg = graph_utils::AddInitializer(graph,
                                                                                 constant_tensor_proto_as_replacement);
       pattern_key_to_shared_arg_map[pattern_key] = &shared_scalar_initializer_node_arg;
+    } else {
+      shared_count += 1;
     }
 
     ReplaceInputsToUseSharedInitializer(graph, consumer_node_to_input_ports_map, origin_initializer_node_arg,
@@ -199,6 +205,8 @@ Status ConstantSharing::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve
     modified = true;
   }
 
+  LOGS(logger, INFO) << "Total shared scalar initializer count: " << shared_count;
+
   return Status::OK();
 }
 

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -4339,7 +4339,7 @@ TEST_F(GraphTransformationTests, BitmaskDropoutFusionTest) {
 
 /*
 This test build a graph like:
-	     input0  input1
+             input0  input1
                   \ /
                   Add
  -----------------|
@@ -4359,7 +4359,7 @@ This test build a graph like:
 
 After fusion, the graph become:
              input0  input1
-	          \ /
+                  \ /
                   Add     (Constant Initializer)
                    \       /
                     Reshape
@@ -4436,15 +4436,16 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) {
         builder.AddNode("Unsqueeze", {gather_out_1}, {unsqueeze_out_1}).AddAttribute("axes", std::vector<int64_t>{0});
       }
       builder.AddNode("ConcatTraining", {unsqueeze_out_0, unsqueeze_out_1, single_value_1d_int_16, single_value_1d_int_64},
-                      {concattraining1_out, concattraining1_length}, "com.microsoft").AddAttribute("axis", static_cast<int64_t>(0));
+                      {concattraining1_out, concattraining1_length}, "com.microsoft")
+          .AddAttribute("axis", static_cast<int64_t>(0));
       builder.AddNode("Reshape", {add_out, concattraining1_out}, {out});
     };
 
     std::unique_ptr<GraphTransformer> transformer = std::make_unique<ReshapeFusion>();
     if (opset_version == 15 && shape_test_for_opset15) {
       ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset_version, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
-                                                pre_graph_checker, pre_graph_checker));
-    } else{
+                                            pre_graph_checker, pre_graph_checker));
+    } else {
       ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset_version, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
                                             pre_graph_checker, post_graph_checker));
     }
@@ -6243,6 +6244,69 @@ TEST_F(GraphTransformationTests, ConstantSharing_ShareIntMaxOrFloatInfinityIniti
   }
 }
 
+/*
+Test graph as below.
+           graph input [2] (float)                Constant (1.0float)                Constant (1.0uint8)
+                  \_______________  ________________/        |                             |
+                                  \/                         |                             |
+                                 Add                         |                             |
+                                  |                          |                             |
+                            graph output [2](float)   graph output [](float)       graph output [](int8)
+
+Be noted: expected result graph should maintain original graph outputs,
+   both float and unin8 constant values are not shared.
+*/
+TEST_F(GraphTransformationTests, ConstantSharing_ShouldNotShareForGraphOutput) {
+  constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "scalar_const_not_share.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  {
+    std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+    ASSERT_TRUE(op_to_count["Add"] == 1);
+    // Be noted, constant nodes are converted to initialized already.
+    ASSERT_TRUE(graph.GetAllInitializedTensors().size() == 2U);
+  }
+
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  std::unique_ptr<GraphTransformer> transformer = std::make_unique<ConstantSharing>();
+  ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(transformer), TransformerLevel::Level1));
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+  {
+    const InitializedTensorSet& initialized_tensor_set = graph.GetAllInitializedTensors();
+    ASSERT_TRUE(initialized_tensor_set.size() == 2U);
+    const NodeArg* add_initializer = nullptr;
+    for (auto& node : graph.Nodes()) {
+      if (node.OpType().compare("Add") == 0) {
+        add_initializer = node.InputDefs()[1];
+        ASSERT_TRUE(add_initializer->Shape()->dim_size() == 0);
+        ASSERT_TRUE(add_initializer->Name().compare("y_scale") == 0);
+      }
+    }
+    ASSERT_TRUE(add_initializer != nullptr);
+    for (const auto& entry : initialized_tensor_set) {
+      if (entry.first.compare(add_initializer->Name()) == 0) {
+        const ONNX_NAMESPACE::TensorProto* tensor_proto = entry.second;
+        onnxruntime::Initializer int64_const{*tensor_proto, graph.ModelPath()};
+        ASSERT_TRUE(int64_const.size() == 1);
+        float float_const_value = *(int64_const.data<float>());
+        ASSERT_TRUE(float_const_value == 1);
+      } else {
+        const ONNX_NAMESPACE::TensorProto* tensor_proto = entry.second;
+        onnxruntime::Initializer uint8_const{*tensor_proto, graph.ModelPath()};
+        ASSERT_TRUE(uint8_const.size() == 1);
+        uint8_t uint8_const_value = *(uint8_const.data<uint8_t>());
+        ASSERT_TRUE(uint8_const_value == static_cast<uint8_t>(1));
+      }
+    }
+
+    auto op_count = CountOpsInGraph(graph);
+    ASSERT_TRUE(op_count.size() == 1U);
+    ASSERT_TRUE(op_count["Add"] == 1);
+  }
+}
+
 TEST_F(GraphTransformationTests, GatherToSplitFusion) {
   auto build_test_case = [&](ModelTestBuilder& builder) {
     auto* data_arg = builder.MakeInput<float>({{54}});

onnxruntime/test/testdata/transform/scalar_const_not_share.py

@@ -0,0 +1,19 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import numpy as np
+import onnx
+import onnxscript
+from onnx import numpy_helper
+from onnxscript import opset17 as op
+
+
+@onnxscript.script()
+def build_model(x: onnxscript.FLOAT):
+    y_scale = op.Constant(value_float=1.0)
+    y_zero_point = op.Constant(value=numpy_helper.from_array(np.array(1, dtype=np.uint8)))
+    return op.Add(x, y_scale), y_scale, y_zero_point
+
+
+model: onnx.ModelProto = build_model.to_model_proto()
+onnx.save(model, "scalar_const_not_share.onnx")