diff --git a/onnxruntime/core/optimizer/constant_sharing.cc b/onnxruntime/core/optimizer/constant_sharing.cc index 96c60bfd14..fa9a309098 100644 --- a/onnxruntime/core/optimizer/constant_sharing.cc +++ b/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 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(); } diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 3b4b793ffc..fde8392d94 100755 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/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{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(0)); + {concattraining1_out, concattraining1_length}, "com.microsoft") + .AddAttribute("axis", static_cast(0)); builder.AddNode("Reshape", {add_out, concattraining1_out}, {out}); }; std::unique_ptr transformer = std::make_unique(); 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; + ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); + Graph& graph = model->MainGraph(); + { + std::map 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 transformer = std::make_unique(); + 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("y_scale") == 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()); + 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()); + ASSERT_TRUE(uint8_const_value == static_cast(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({{54}}); diff --git a/onnxruntime/test/testdata/transform/scalar_const_not_share.onnx b/onnxruntime/test/testdata/transform/scalar_const_not_share.onnx new file mode 100644 index 0000000000..13e572950f Binary files /dev/null and b/onnxruntime/test/testdata/transform/scalar_const_not_share.onnx differ diff --git a/onnxruntime/test/testdata/transform/scalar_const_not_share.py b/onnxruntime/test/testdata/transform/scalar_const_not_share.py new file mode 100644 index 0000000000..16fa1e825b --- /dev/null +++ b/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")