diff --git a/include/onnxruntime/core/graph/constants.h b/include/onnxruntime/core/graph/constants.h index d213874038..097665862d 100644 --- a/include/onnxruntime/core/graph/constants.h +++ b/include/onnxruntime/core/graph/constants.h @@ -3,8 +3,6 @@ #pragma once -#include -#include #include #include "core/common/common.h" diff --git a/onnxruntime/core/optimizer/common_subexpression_elimination.cc b/onnxruntime/core/optimizer/common_subexpression_elimination.cc new file mode 100644 index 0000000000..20c5093ea7 --- /dev/null +++ b/onnxruntime/core/optimizer/common_subexpression_elimination.cc @@ -0,0 +1,448 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "common_subexpression_elimination.h" +#include "core/optimizer/utils.h" +#include "core/graph/graph_utils.h" + +#include +#include +#include +#include +#include + +// This optimization pass will collapse expressions that always evaluate to the same value +// into one node. As an example, consider the following abstract function where x1, x2,... +// denote graph inputs, and F1, F2,... - operations. +// +// return F3(F2(F1(x1, x2), x3)) + F4(F2(F1(x1, x2), x3)) +// +// Because F1 operations are given the same inputs, they can be merged into one node: +// +// y1 = F1(x1, x2) +// return F3(F2(y1, x3)) + F4(F2(y1, x3)) +// +// Now we can see that F2 operations are given the same inputs, so they can be merged too: +// +// y1 = F1(x1, x2) +// y2 = F2(y1, x3) +// return F3(y2) + F4(y2) +// +// This is implemented using value numbering (https://en.wikipedia.org/wiki/Value_numbering): +// first every graph input, constant initializer and graph node output are assigned +// an equivalence class, and then nodes that have the same operation and equivalent inputs +// are collapsed. + +namespace onnxruntime { + +namespace { +struct DeepPointerEquality { + template + bool operator()(const Ptr& lhs, const Ptr& rhs) const { + if (lhs == nullptr || rhs == nullptr) { + return lhs == rhs; + } + return *lhs == *rhs; + } +}; + +struct DeepPointerHash { + template + std::size_t operator()(const Ptr& value) const { + if (value == nullptr) { + return 0; + } + return std::hash::type>{}(*value); + } +}; + +template +void UpdateHash(const T& x, Hasher hasher, std::size_t& hash) { + constexpr std::size_t kPrime = 31013; + hash = hash * kPrime + hasher(x); +} + +template +void UpdateHash(const T& x, std::size_t& hash) { + UpdateHash(x, std::hash{}, hash); +} + +template +void UpdateHashWithContainer(const Container& c, std::size_t& hash) { + for (const auto& element : c) { + UpdateHash(element, hash); + } +} + +using OutputIndex = int; + +constexpr OutputIndex kInvalidOutputIndex = -1; + +const NodeArg* Normalize(const NodeArg* node_arg) { + return node_arg == nullptr || !node_arg->Exists() ? nullptr : node_arg; +} + +// Represents an equivalence class of expressions (inputs, constant initializers and node outputs) +// that will always evaluate to the same runtime value. +class EquivalenceClass { + public: + bool operator==(const EquivalenceClass& other) const; + + friend struct ::std::hash; + friend std::vector> Normalize(const Node& node, const std::vector& inputs); + + explicit EquivalenceClass(const NodeArg* non_op_value) + : attributes_(nullptr), + output_index_(kInvalidOutputIndex), + non_op_value_(Normalize(non_op_value)), + discriminator_(0), + hash_(CalculateHash()) { + } + + EquivalenceClass(const Node& node, const std::vector& explicit_inputs, + OutputIndex output_index, int discriminator) + : op_type_(node.OpType()), + domain_(node.Domain()), + inputs_(Normalize(node, explicit_inputs)), + attributes_(&node.GetAttributes()), + output_index_(output_index), + non_op_value_(nullptr), + discriminator_(discriminator), + hash_(CalculateHash()) { + } + + private: + std::size_t CalculateHash() const; + + // Operation and domain of the node that produces this value. + const std::string op_type_; + const std::string domain_; + + // Explicit inputs to the operation, sequence of inputs for each formal parameter. + const std::vector> inputs_; + + // Attributes of the operation. + const NodeAttributes* attributes_; + + // Index of this value in the output list of the operation. + const OutputIndex output_index_; + + // When the value is not an output of an operation, (i.e., a constant initializer or an input), + // non_op_value is set to the corresponding NodeArg, and other fields are empty. + // Currently, different inputs/initializers are always considered different values, although we + // could merge equal initializers. + const NodeArg* non_op_value_; + + // When an operation is not supported by the CSE optimization pass, we consider its + // outputs unique (not equal to other values). For this purpose we assign a unique + // discriminator for such values. + const int discriminator_; + + const std::size_t hash_; +}; + +std::vector> Normalize(const Node& node, const std::vector& inputs) { + const auto& arg_count = node.InputArgCount(); + auto input_iter = inputs.begin(); + std::vector> result(arg_count.size()); + + for (std::size_t arg_index = 0; arg_index < arg_count.size(); ++arg_index) { + auto& arg = result[arg_index]; + for (int i = 0; i < arg_count[arg_index]; ++i) { + if (input_iter != inputs.end()) { + arg.push_back(*input_iter); + ++input_iter; + } + } + + // Remove missing optional inputs from the back + while (!arg.empty() && arg.back()->output_index_ == kInvalidOutputIndex && arg.back()->non_op_value_ == nullptr) { + arg.pop_back(); + } + } + + return result; +} + +template +bool AreRangesEqual(const Range& lhs, const Range& rhs) { + return lhs.size() == rhs.size() && + std::equal(lhs.begin(), lhs.end(), rhs.begin()); +} + +bool AreEqual(const ONNX_NAMESPACE::AttributeProto& lhs, const ONNX_NAMESPACE::AttributeProto& rhs) { + if (&lhs == &rhs) { + return true; + } + + if (lhs.type() != rhs.type() || lhs.name() != rhs.name()) { + return false; + } + + switch (lhs.type()) { + case onnx::AttributeProto_AttributeType_FLOAT: + return lhs.f() == rhs.f(); + case onnx::AttributeProto_AttributeType_INT: + return lhs.i() == rhs.i(); + case onnx::AttributeProto_AttributeType_STRING: + return lhs.s() == rhs.s(); + case onnx::AttributeProto_AttributeType_FLOATS: + return AreRangesEqual(lhs.floats(), rhs.floats()); + case onnx::AttributeProto_AttributeType_INTS: + return AreRangesEqual(lhs.ints(), rhs.ints()); + case onnx::AttributeProto_AttributeType_STRINGS: + return AreRangesEqual(lhs.strings(), rhs.strings()); + case onnx::AttributeProto_AttributeType_TENSOR: + case onnx::AttributeProto_AttributeType_GRAPH: + case onnx::AttributeProto_AttributeType_SPARSE_TENSOR: + case onnx::AttributeProto_AttributeType_TENSORS: + case onnx::AttributeProto_AttributeType_GRAPHS: + case onnx::AttributeProto_AttributeType_SPARSE_TENSORS: + case onnx::AttributeProto_AttributeType_UNDEFINED: + return false; // Don't support these attributes for now; corresponding nodes will be considered distinct. + } + + return false; +} + +std::size_t GetAttributeHash(const ONNX_NAMESPACE::AttributeProto& attr) { + std::size_t hash = 0; + UpdateHash( + static_cast::type>(attr.type()), + hash); + UpdateHash(attr.name(), hash); + switch (attr.type()) { + case onnx::AttributeProto_AttributeType_FLOAT: + UpdateHash(attr.f(), hash); + break; + case onnx::AttributeProto_AttributeType_INT: + UpdateHash(attr.i(), hash); + break; + case onnx::AttributeProto_AttributeType_STRING: + UpdateHash(attr.s(), hash); + break; + case onnx::AttributeProto_AttributeType_FLOATS: + UpdateHashWithContainer(attr.floats(), hash); + break; + case onnx::AttributeProto_AttributeType_INTS: + UpdateHashWithContainer(attr.ints(), hash); + break; + case onnx::AttributeProto_AttributeType_STRINGS: + UpdateHashWithContainer(attr.strings(), hash); + break; + case onnx::AttributeProto_AttributeType_TENSOR: + case onnx::AttributeProto_AttributeType_GRAPH: + case onnx::AttributeProto_AttributeType_SPARSE_TENSOR: + case onnx::AttributeProto_AttributeType_TENSORS: + case onnx::AttributeProto_AttributeType_GRAPHS: + case onnx::AttributeProto_AttributeType_SPARSE_TENSORS: + case onnx::AttributeProto_AttributeType_UNDEFINED: + break; + } + + return hash; +} + +bool SameAttribute(const NodeAttributes::value_type& lhs, const NodeAttributes::value_type& rhs) { + return lhs.first == rhs.first && AreEqual(lhs.second, rhs.second); +} + +bool SameAttributes(const NodeAttributes* lhs, const NodeAttributes* rhs) { + if (lhs == nullptr || rhs == nullptr) { + return lhs == rhs; + } + return lhs->size() == rhs->size() && + std::equal(lhs->begin(), lhs->end(), rhs->begin(), &SameAttribute); +} + +bool EquivalenceClass::operator==(const EquivalenceClass& other) const { + if (this == &other) { + return true; + } + + // Below we compare inputs_ as pointers. This is valid due to how EquivalenceClass'es are constructed: + // we'll never have two distinct but equal inputs_ here, so their addresses are effectively their value numbers. + return hash_ == other.hash_ && output_index_ == other.output_index_ && discriminator_ == other.discriminator_ && + non_op_value_ == other.non_op_value_ && + op_type_ == other.op_type_ && domain_ == other.domain_ && + inputs_ == other.inputs_ && + SameAttributes(attributes_, other.attributes_); +} + +std::size_t EquivalenceClass::CalculateHash() const { + std::size_t hash = 0; + UpdateHash(output_index_, hash); + UpdateHash(discriminator_, hash); + UpdateHash(non_op_value_, hash); + UpdateHash(op_type_, hash); + UpdateHash(domain_, hash); + if (attributes_ != nullptr) { + for (const auto& kv : *attributes_) { + UpdateHash(kv.first, hash); + UpdateHash(kv.second, &GetAttributeHash, hash); + } + } + + for (const auto& arg : inputs_) { + for (const EquivalenceClass* input : arg) { + UpdateHash(input, DeepPointerHash{}, hash); + } + } + + return hash; +} + +// Representative of an equivalence class. +// node_index and output_index define the node that produced the node_arg. +// For inputs and constant initializers, output_index == kInvalidOutputIndex. +struct Representative { + const NodeArg* node_arg; + NodeIndex node_index; + OutputIndex output_index; +}; + +struct NodeArgPtrHash { + std::size_t operator()(const NodeArg* node_arg) const { + return std::hash{}(Normalize(node_arg)); + } +}; + +struct NodeArgPtrEquality { + bool operator()(const NodeArg* lhs, const NodeArg* rhs) const { + return Normalize(lhs) == Normalize(rhs); + } +}; + +bool IsNodeSupported(const Node& node) { + return !node.ContainsSubgraph() && optimizer_utils::IsOperationDeterministic(node.Domain(), node.OpType()); +} +} // namespace + +} // namespace onnxruntime + +namespace std { +template <> +struct hash { + std::size_t operator()(const onnxruntime::EquivalenceClass& val) const noexcept { + return val.hash_; + } +}; +} // namespace std + +namespace onnxruntime { + +Status CommonSubexpressionElimination::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { + GraphViewer graph_viewer(graph); + const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); + + // Pool of equivalence classes; unique_ptr to guarantee stable address. + std::vector> unique_equivalence_classes; + unique_equivalence_classes.reserve(graph.NumberOfNodes()); + + // Maps an equivalence class of values to a representative NodeArg that belongs to this class. + std::unordered_map< + const EquivalenceClass*, + Representative, + DeepPointerHash, DeepPointerEquality> + value_to_representative; + + // Maps every NodeArg to its equivalence class of. + // This is the inverse of the above mapping, except that different NodeArgs can belong to the same + // equivalence class. In that case these NodeArgs will be "merged" into one. + std::unordered_map equivalence_classes; + + int unique_discriminator = 1; + + for (NodeIndex node_index : node_topology_list) { + Node* node = graph.GetNode(node_index); + if (node == nullptr) + continue; + + ORT_RETURN_IF_ERROR(Recurse(*node, modified, graph_level, logger)); + + std::vector input_values; + input_values.reserve(node->InputDefs().size()); + for (const NodeArg* input_def : node->InputDefs()) { + auto it = equivalence_classes.find(input_def); + if (it == equivalence_classes.end()) { + // Because nodes are processed in topological order, this will always be + // a non-op value (graph input or constant initializer). + auto value = onnxruntime::make_unique(input_def); + const auto* raw_ptr = value.get(); + unique_equivalence_classes.push_back(std::move(value)); + value_to_representative.emplace(raw_ptr, Representative{input_def, 0, kInvalidOutputIndex}); + it = equivalence_classes.emplace_hint(it, input_def, raw_ptr); + } + + input_values.push_back(it->second); + } + + int discriminator = 0; + if (!IsNodeSupported(*node)) { + discriminator = ++unique_discriminator; + } + + for (OutputIndex output_index = 0, end = static_cast(node->OutputDefs().size()); + output_index < end; ++output_index) { + const NodeArg* output_def = node->OutputDefs()[output_index]; + auto equivalence_class = onnxruntime::make_unique(*node, input_values, output_index, discriminator); + auto* raw_ptr = equivalence_class.get(); + + auto it = value_to_representative.find(raw_ptr); + if (it == value_to_representative.end()) { + unique_equivalence_classes.push_back(std::move(equivalence_class)); + it = value_to_representative.emplace_hint(it, raw_ptr, + Representative{output_def, node_index, output_index}); + } + + equivalence_classes[output_def] = it->first; + } + } + + std::unordered_set graph_outputs; + graph_outputs.insert(graph_viewer.GetOutputs().begin(), graph_viewer.GetOutputs().end()); + + for (NodeIndex node_index : node_topology_list) { + Node* node = graph.GetNode(node_index); + if (node == nullptr) + continue; + + bool node_output_replaced = false; + for (OutputIndex output_idx = 0, end = static_cast(node->OutputDefs().size()); + output_idx < end; ++output_idx) { + const NodeArg* output_def = node->OutputDefs()[output_idx]; + + const EquivalenceClass* equivalence_class = equivalence_classes.at(output_def); + const auto& representative = value_to_representative.at(equivalence_class); + if (representative.node_arg == output_def) { + // output_def is the representative of its equivalence class. All other values in the same class + // will be replaced with output_def, but output_def itself should remain. + continue; + } + + ORT_ENFORCE(representative.output_index != kInvalidOutputIndex); + + if (graph_outputs.count(output_def) > 0) { + // Currently, we don't support eliminating the graph's outputs. + LOGS(logger, VERBOSE) << "Not eliminating output " << output_def->Name() << " of node " << node->Name() << "[" << node->OpType() << "] because it's the graph's output."; + continue; + } + + Node& replacement = *graph.GetNode(representative.node_index); + OutputIndex replacement_output_idx = representative.output_index; + graph_utils::ReplaceDownstreamNodeInput(graph, *node, output_idx, replacement, replacement_output_idx); + node_output_replaced = true; + } + + if (node_output_replaced) { + modified = true; + if (optimizer_utils::CheckOutputEdges(graph, *node, 0)) { + graph.RemoveNode(node_index); + } + } + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/common_subexpression_elimination.h b/onnxruntime/core/optimizer/common_subexpression_elimination.h new file mode 100644 index 0000000000..29d87bf7b9 --- /dev/null +++ b/onnxruntime/core/optimizer/common_subexpression_elimination.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +/** +@Class CommonSubexpressionElimination +Merge nodes that always evaluate to the same result. +*/ +class CommonSubexpressionElimination : public GraphTransformer { + public: + CommonSubexpressionElimination(const std::unordered_set& compatible_execution_providers = {}) noexcept + : GraphTransformer("CommonSubexpressionElimination", compatible_execution_providers) { + } + + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 8ba911e889..dd463cea60 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/optimizer/constant_folding.h" +#include "core/optimizer/utils.h" #include "core/graph/graph_utils.h" #include "core/optimizer/optimizer_execution_frame.h" #include "core/framework/op_kernel.h" @@ -86,7 +87,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, // Check if constant folding can be applied on this node. if (!graph_utils::IsSupportedProvider(*node, GetCompatibleExecutionProviders()) || - excluded_op_types_.find(node->OpType()) != excluded_op_types_.end() || + !optimizer_utils::IsOperationDeterministic(node->Domain(), node->OpType()) || // constant folding does not support executing a node that includes subgraphs (control flow operators, // such as If/Loop/Scan, fall into this category). individual nodes in the subgraph will be processed // by the Recurse call above diff --git a/onnxruntime/core/optimizer/constant_folding.h b/onnxruntime/core/optimizer/constant_folding.h index c0b4118e69..7306cabdc4 100644 --- a/onnxruntime/core/optimizer/constant_folding.h +++ b/onnxruntime/core/optimizer/constant_folding.h @@ -23,11 +23,6 @@ class ConstantFolding : public GraphTransformer { : GraphTransformer("ConstantFolding", compatible_execution_providers), excluded_initializers_(excluded_initializers) {} private: - /** Constant folding will not be applied to nodes whose op_type is included in this set. - All non-deterministic operators should be included in this set. */ - const std::unordered_set excluded_op_types_ = - {"RandomUniform", "RandomNormal", "RandomUniformLike", "RandomNormalLike", "Multinomial"}; - Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; const std::unordered_set excluded_initializers_; diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index e9ad6b6512..c6ab466338 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -7,6 +7,7 @@ #include "core/optimizer/attention_fusion.h" #include "core/optimizer/bias_gelu_fusion.h" #include "core/optimizer/cast_elimination.h" +#include "core/optimizer/common_subexpression_elimination.h" #include "core/optimizer/constant_folding.h" #include "core/optimizer/conv_activation_fusion.h" #include "core/optimizer/conv_add_fusion.h" @@ -113,6 +114,7 @@ std::vector> GenerateTransformers(TransformerL case TransformerLevel::Level1: { std::unordered_set l1_execution_providers = {}; + transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); transformers.emplace_back(onnxruntime::make_unique(l1_execution_providers)); diff --git a/onnxruntime/core/optimizer/utils.cc b/onnxruntime/core/optimizer/utils.cc index b300cc7a82..af2d7c9e95 100644 --- a/onnxruntime/core/optimizer/utils.cc +++ b/onnxruntime/core/optimizer/utils.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. #include "core/common/make_unique.h" +#include "core/graph/constants.h" #include "core/graph/onnx_protobuf.h" #include "core/graph/graph_utils.h" #include "core/framework/tensorprotoutils.h" @@ -10,6 +11,10 @@ #include "float.h" //#include +#include +#include +#include + using namespace onnxruntime; namespace onnxruntime { @@ -248,5 +253,24 @@ bool CheckOutputEdges(const Graph& graph, const Node& node, size_t expected_outp return node.GetOutputEdgesCount() == expected_output_edges; } +// Allow certain domains/ops. We don't know anything about unknown domains/ops (e.g. custom ops), +// so we have to assume that they are not deterministic, to be on the safe side. +// We could also allow other known domains (kMSDomain, kMSNchwcDomain, kMSFeaturizersDomain), +// as long as we verify which of their operations are non-deterministic and add them in the map below. +static const std::unordered_map> kNonDeterministicOps = +{ + {kOnnxDomain, {"RandomUniform", "RandomNormal", "RandomUniformLike", "RandomNormalLike", "Multinomial"}}, +}; + +bool IsOperationDeterministic(const std::string& domain, const std::string& op) { + auto itDomain = kNonDeterministicOps.find(domain); + if (itDomain == kNonDeterministicOps.end()) { + // Unknown domain. Assume the op is not deterministic. + return false; + } + + return itDomain->second.count(op) == 0; +} + } // namespace optimizer_utils } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/utils.h b/onnxruntime/core/optimizer/utils.h index a8bb198bf5..17aa0b4b53 100644 --- a/onnxruntime/core/optimizer/utils.h +++ b/onnxruntime/core/optimizer/utils.h @@ -78,5 +78,7 @@ bool IsSupportedDataType(const Node& node, const std::vector& suppo */ bool CheckOutputEdges(const Graph& graph, const Node& node, size_t expected_output_edges); +bool IsOperationDeterministic(const std::string& domain, const std::string& op); + } // namespace optimizer_utils } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/cse_test.cc b/onnxruntime/test/optimizer/cse_test.cc new file mode 100644 index 0000000000..f0c7b24de7 --- /dev/null +++ b/onnxruntime/test/optimizer/cse_test.cc @@ -0,0 +1,248 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test/framework/test_utils.h" +#include "test/test_environment.h" +#include "core/graph/model.h" +#include "core/optimizer/common_subexpression_elimination.h" +#include "core/optimizer/graph_transformer_mgr.h" + +#include "gtest/gtest.h" + +#include +#include +#include + +namespace onnxruntime { +namespace test { + +namespace { +void ApplyCse(Model& model, unsigned num_steps = 1) { + GraphTransformerManager graph_transformation_mgr(num_steps); + ASSERT_TRUE( + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1).IsOK()); + ASSERT_TRUE( + graph_transformation_mgr.ApplyTransformers(model.MainGraph(), TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger()).IsOK()); +} + +std::vector GetSortedNames(const std::vector& node_args) { + std::vector node_arg_names; + for (const auto* node_arg : node_args) { + node_arg_names.push_back(node_arg->Name()); + } + + std::sort(node_arg_names.begin(), node_arg_names.end()); + return node_arg_names; +} + +std::vector GetNodeNames(const Graph& graph) { + std::vector res; + for (auto& node : graph.Nodes()) { + res.push_back(node.Name()); + } + return res; +} + +void GetAllNodeNamesImpl(const Graph& graph, std::vector& res) { + for (auto& node : graph.Nodes()) { + res.push_back(node.Name()); + for (const auto& subgraph : node.GetSubgraphs()) { + GetAllNodeNamesImpl(*subgraph, res); + } + } +} + +std::vector GetAllNodeNamesSorted(const Graph& graph) { + std::vector res; + GetAllNodeNamesImpl(graph, res); + std::sort(res.begin(), res.end()); + return res; +} +} // namespace + +TEST(CseTests, SimpleTest) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse1.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + ApplyCse(*model); + + Graph& graph = model->MainGraph(); + + const auto& graph_inputs = GetSortedNames(graph.GetInputs()); + ASSERT_EQ(graph_inputs, (std::vector{"x"})); + + const auto& graph_outputs = GetSortedNames(graph.GetOutputs()); + ASSERT_EQ(graph_outputs, (std::vector{"Result"})); + + auto op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count.at("MatMul"), 1); + ASSERT_EQ(op_count.at("Add"), 2); + ASSERT_EQ(op_count.at("Relu"), 1); +} + +TEST(CseTests, GraphOutput) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse_graph_output.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + ApplyCse(*model); + + Graph& graph = model->MainGraph(); + + const auto& input_names = GetSortedNames(graph.GetInputs()); + ASSERT_EQ(input_names, (std::vector{"x"})); + + const auto& output_names = GetSortedNames(graph.GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"res1", "res2"})); + + auto op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count.at("Add"), 2); +} + +TEST(CseTests, OptionalArgs) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse_optional_args.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + Graph& graph = model->MainGraph(); + auto op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count.at("Clip"), 5); + + ApplyCse(*model); + + const auto& input_names = GetSortedNames(graph.GetInputs()); + ASSERT_EQ(input_names, (std::vector{"x"})); + + const auto& output_names = GetSortedNames(graph.GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"Result"})); + + op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count.at("Clip"), 3); + auto node_names = GetNodeNames(graph); + ASSERT_EQ(std::count(node_names.begin(), node_names.end(), "clip_3"), 1); + ASSERT_EQ(std::count(node_names.begin(), node_names.end(), "clip_4"), 1); +} + +TEST(CseTests, Random) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse_random.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + Graph& graph = model->MainGraph(); + auto op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count["RandomUniform"], 4); + + ApplyCse(*model); + + const auto& input_names = GetSortedNames(graph.GetInputs()); + ASSERT_EQ(input_names, (std::vector{"x"})); + + const auto& output_names = GetSortedNames(graph.GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"Result"})); + + op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count["RandomUniform"], 4); + auto node_names = GetNodeNames(graph); + ASSERT_EQ(std::count(node_names.begin(), node_names.end(), "random_uniform_1"), 1); + ASSERT_EQ(std::count(node_names.begin(), node_names.end(), "random_uniform_2"), 1); + ASSERT_EQ(std::count(node_names.begin(), node_names.end(), "random_uniform_3"), 1); + ASSERT_EQ(std::count(node_names.begin(), node_names.end(), "random_uniform_4"), 1); +} + +TEST(CseTests, Subgraph) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse_subgraph.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + Graph& graph = model->MainGraph(); + auto node_names = GetAllNodeNamesSorted(graph); + ASSERT_EQ(node_names, (std::vector{"if_0", + "iffalse_intermediate_1", "iffalse_intermediate_2", "iffalse_res_1", "iffalse_res_2", "iffalse_res_3", + "iftrue_intermediate_1", "iftrue_intermediate_2", "iftrue_res_1", "iftrue_res_2", "iftrue_res_3"})); + + ApplyCse(*model); + + const auto& input_names = GetSortedNames(graph.GetInputs()); + ASSERT_EQ(input_names, (std::vector{"b", "x"})); + + auto output_names = GetSortedNames(graph.GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"Result1", "Result2", "Result3"})); + + const Graph* if_true_graph = nullptr; + const Graph* if_false_graph = nullptr; + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "If") { + if_true_graph = node.GetGraphAttribute("then_branch"); + if_false_graph = node.GetGraphAttribute("else_branch"); + break; + } + } + + output_names = GetSortedNames(if_true_graph->GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"Result1", "Result2", "Result3"})); + + auto op_count = CountOpsInGraph(*if_true_graph); + ASSERT_EQ(op_count["Mul"], 1); + ASSERT_EQ(op_count["Sum"], 3); + + output_names = GetSortedNames(if_false_graph->GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"Result1", "Result2", "Result3"})); + + op_count = CountOpsInGraph(*if_false_graph); + ASSERT_EQ(op_count["Mul"], 3); + ASSERT_EQ(op_count["Sum"], 1); +} + +TEST(CseTests, MergedValueAndGraphOutputAreOutputsOfSameNode) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse_only_one_graph_output.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + Graph& graph = model->MainGraph(); + ApplyCse(*model); + + const auto& input_names = GetSortedNames(graph.GetInputs()); + ASSERT_EQ(input_names, (std::vector{"x"})); + + const auto& output_names = GetSortedNames(graph.GetOutputs()); + ASSERT_EQ(output_names, (std::vector{"Add", "Split1Output2", "Split2Output2"})); + + auto op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count["ReduceSum"], 1); + ASSERT_EQ(op_count["Add"], 1); + // In current implementation we don't collapse nodes that produce graph outputs. + ASSERT_EQ(op_count["Split"], 2); +} + +TEST(CseTests, MergeConstants) { + auto model_uri = ORT_TSTR("testdata/transform/cse/cse_merge_constants.onnx"); + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()) + .IsOK()); + Graph& graph = model->MainGraph(); + GraphTransformerManager graph_transformation_mgr(1); + // In current implementation, equal constants are not merged. So CSE must precede constant folding, otherwise we end up + // with multiple copies of the same constant. + ASSERT_TRUE( + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1).IsOK()); + ASSERT_TRUE( + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1).IsOK()); + ASSERT_TRUE( + graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger()).IsOK()); + + ASSERT_EQ(graph.GetAllInitializedTensors().size(), 1U); + auto op_count = CountOpsInGraph(graph); + ASSERT_EQ(op_count.size(), 1U); + ASSERT_EQ(op_count["Add"], 2); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/testdata/transform/cse/cse1.onnx b/onnxruntime/test/testdata/transform/cse/cse1.onnx new file mode 100644 index 0000000000..28ca3dc0f0 Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse1.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/cse_graph_output.onnx b/onnxruntime/test/testdata/transform/cse/cse_graph_output.onnx new file mode 100644 index 0000000000..6a84fad9c1 Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse_graph_output.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/cse_merge_constants.onnx b/onnxruntime/test/testdata/transform/cse/cse_merge_constants.onnx new file mode 100644 index 0000000000..b831343200 Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse_merge_constants.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/cse_only_one_graph_output.onnx b/onnxruntime/test/testdata/transform/cse/cse_only_one_graph_output.onnx new file mode 100644 index 0000000000..5d8eb1deb2 Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse_only_one_graph_output.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/cse_optional_args.onnx b/onnxruntime/test/testdata/transform/cse/cse_optional_args.onnx new file mode 100644 index 0000000000..d647de4112 Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse_optional_args.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/cse_random.onnx b/onnxruntime/test/testdata/transform/cse/cse_random.onnx new file mode 100644 index 0000000000..d8397a3f00 Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse_random.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/cse_subgraph.onnx b/onnxruntime/test/testdata/transform/cse/cse_subgraph.onnx new file mode 100644 index 0000000000..aba3cf1cbb Binary files /dev/null and b/onnxruntime/test/testdata/transform/cse/cse_subgraph.onnx differ diff --git a/onnxruntime/test/testdata/transform/cse/generate.py b/onnxruntime/test/testdata/transform/cse/generate.py new file mode 100644 index 0000000000..7a676f13df --- /dev/null +++ b/onnxruntime/test/testdata/transform/cse/generate.py @@ -0,0 +1,233 @@ +import onnx +from onnx import helper, shape_inference +from onnx import AttributeProto, TensorProto, GraphProto +import os + +_this_dir = os.path.abspath(os.path.dirname(__file__)) + +def _onnx_export(graph_def, relative_path, verbose=False): + model = helper.make_model(graph_def, producer_name='makalini', opset_imports=[helper.make_operatorsetid("", 11)]) + onnx.checker.check_model(model) + inferred_model = shape_inference.infer_shapes(model) + onnx.checker.check_model(inferred_model) + model_path = os.path.join(_this_dir, relative_path) + os.makedirs(os.path.dirname(model_path), exist_ok=True) + onnx.save_model(model, model_path) + if verbose: + print() + print(inferred_model) + import onnxruntime as rt + rt.InferenceSession(model_path) + +def cse1(): + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "MatMul", inputs = ['w', 'x'], outputs = ['MatMul1'], name = 'matmul_1'), + helper.make_node(op_type = "MatMul", inputs = ['w', 'x'], outputs = ['MatMul2'], name = 'matmul_2'), + helper.make_node(op_type = "Add", inputs = ['MatMul1', 'b'], outputs = ['Add1'], name = 'add_1'), + helper.make_node(op_type = "Add", inputs = ['MatMul2', 'b'], outputs = ['Add2'], name = 'add_2'), + helper.make_node(op_type = "Relu", inputs = ['Add1'], outputs = ['Relu1'], name = 'relu_1'), + helper.make_node(op_type = "Relu", inputs = ['Add2'], outputs = ['Relu2'], name = 'relu_2'), + helper.make_node(op_type = "Add", inputs = ['Relu1', 'Relu2'], outputs = ['Result'], name = 'result') + ], + name = 'cse1', + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [5]) + ], + outputs = [ + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [2]) + ], + initializer = [ + helper.make_tensor('w', TensorProto.FLOAT, [2, 5], list(range(2*5))), + helper.make_tensor('b', TensorProto.FLOAT, [2], list(range(2))), + ] + ) + _onnx_export(graph_def, 'cse1.onnx') + +def cse_graph_output(): + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "Add", inputs = ['x', 'b'], outputs = ['res1'], name = 'add_1'), + helper.make_node(op_type = "Add", inputs = ['x', 'b'], outputs = ['res2'], name = 'add_2'), + ], + name = 'cse_graph_output', + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [5]) + ], + outputs = [ + helper.make_tensor_value_info('res1', TensorProto.FLOAT, [5]), + helper.make_tensor_value_info('res2', TensorProto.FLOAT, [5]) + ], + initializer = [ + helper.make_tensor('b', TensorProto.FLOAT, [5], list(range(5))), + ] + ) + + _onnx_export(graph_def, 'cse_graph_output.onnx') + +def cse_optional_args(): + n = 5 + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "Clip", inputs = ['x'], outputs = ['Clipped0'], name = 'clip_0'), + helper.make_node(op_type = "Clip", inputs = ['x', ''], outputs = ['Clipped1'], name = 'clip_1'), + helper.make_node(op_type = "Clip", inputs = ['x', '', ''], outputs = ['Clipped2'], name = 'clip_2'), + helper.make_node(op_type = "Clip", inputs = ['x', '', 'c'], outputs = ['Clipped3'], name = 'clip_3'), + helper.make_node(op_type = "Clip", inputs = ['x', 'c', ''], outputs = ['Clipped4'], name = 'clip_4'), + helper.make_node(op_type = "Sum", inputs = ['Clipped0', 'Clipped1', 'Clipped2', 'Clipped3', 'Clipped4'], outputs = ['Result'], name = 'sum_1') + ], + name = 'cse_optional_args', + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [n]) + ], + outputs = [ + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [n]) + ], + initializer = [ + helper.make_tensor('c', TensorProto.FLOAT, [], [0.5]) + ] + ) + _onnx_export(graph_def, 'cse_optional_args.onnx') + +def cse_subgraph(): + if_true_graph = helper.make_graph( + nodes = [ + helper.make_node(op_type = "Sum", inputs = ['x', 'x'], outputs = ['Result1'], name = 'iftrue_res_1'), + helper.make_node(op_type = "Sum", inputs = ['x', 'x'], outputs = ['Result2'], name = 'iftrue_res_2'), + helper.make_node(op_type = "Mul", inputs = ['x', 'x'], outputs = ['Intermediate1'], name = 'iftrue_intermediate_1'), + helper.make_node(op_type = "Mul", inputs = ['x', 'x'], outputs = ['Intermediate2'], name = 'iftrue_intermediate_2'), + helper.make_node(op_type = "Sum", inputs = ['Intermediate1', 'Intermediate2'], outputs = ['Result3'], name = 'iftrue_res_3'), + ], + name = 'if_true_graph', + inputs = [ + ], + outputs = [ + helper.make_tensor_value_info('Result1', TensorProto.FLOAT, [2]), + helper.make_tensor_value_info('Result2', TensorProto.FLOAT, [2]), + helper.make_tensor_value_info('Result3', TensorProto.FLOAT, [2]), + ], + initializer = [ + ] + ) + + if_false_graph = helper.make_graph( + nodes = [ + helper.make_node(op_type = "Mul", inputs = ['x', 'x'], outputs = ['Result1'], name = 'iffalse_res_1'), + helper.make_node(op_type = "Mul", inputs = ['x', 'x'], outputs = ['Result2'], name = 'iffalse_res_2'), + helper.make_node(op_type = "Sum", inputs = ['x', 'x'], outputs = ['Intermediate1'], name = 'iffalse_intermediate_1'), + helper.make_node(op_type = "Sum", inputs = ['x', 'x'], outputs = ['Intermediate2'], name = 'iffalse_intermediate_2'), + helper.make_node(op_type = "Mul", inputs = ['Intermediate1', 'Intermediate2'], outputs = ['Result3'], name = 'iffalse_res_3'), + ], + name = 'if_false_graph', + inputs = [ + ], + outputs = [ + helper.make_tensor_value_info('Result1', TensorProto.FLOAT, [2]), + helper.make_tensor_value_info('Result2', TensorProto.FLOAT, [2]), + helper.make_tensor_value_info('Result3', TensorProto.FLOAT, [2]), + ], + initializer = [ + ] + ) + + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "If", inputs = ['b'], outputs = ['Result1', 'Result2', 'Result3'], name = 'if_0', then_branch=if_true_graph, else_branch=if_false_graph), + ], + name = 'cse_subgraph', + inputs = [ + helper.make_tensor_value_info("b", TensorProto.BOOL, [1]), + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2]) + ], + outputs = [ + helper.make_tensor_value_info('Result1', TensorProto.FLOAT, [2]), + helper.make_tensor_value_info('Result2', TensorProto.FLOAT, [2]), + helper.make_tensor_value_info('Result3', TensorProto.FLOAT, [2]), + ], + initializer = [ + ] + ) + _onnx_export(graph_def, 'cse_subgraph.onnx') + +def cse_random(): + n = 5 + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "RandomUniform", inputs = [], outputs = ['Random1'], name = 'random_uniform_1', shape=[n]), + helper.make_node(op_type = "RandomUniform", inputs = [], outputs = ['Random2'], name = 'random_uniform_2', shape=[n]), + helper.make_node(op_type = "RandomUniform", inputs = [], outputs = ['Random3'], name = 'random_uniform_3', shape=[n], seed=1.0), + helper.make_node(op_type = "RandomUniform", inputs = [], outputs = ['Random4'], name = 'random_uniform_4', shape=[n], seed=1.0), + helper.make_node(op_type = "Sum", inputs = ['x', 'Random1', 'Random2', 'Random3', 'Random4'], outputs = ['Result'], name = 'sum_1') + ], + name = 'cse_random', + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [n]) + ], + outputs = [ + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [n]) + ], + initializer = [ + ] + ) + _onnx_export(graph_def, 'cse_random.onnx') + +def cse_merge_constants(): + n = 3 + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "Add", inputs = ['c', 'c'], outputs = ['Add1'], name = 'add_1'), + helper.make_node(op_type = "Add", inputs = ['c', 'c'], outputs = ['Add2'], name = 'add_2'), + helper.make_node(op_type = "Add", inputs = ['Add1', 'x'], outputs = ['Add3'], name = 'add_3'), + helper.make_node(op_type = "Add", inputs = ['Add2', 'x'], outputs = ['Add4'], name = 'add_4'), + helper.make_node(op_type = "Add", inputs = ['Add3', 'Add4'], outputs = ['Result'], name = 'add_5'), + ], + name = 'cse_merge_constants', + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [n]) + ], + outputs = [ + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [n]) + ], + initializer = [ + helper.make_tensor('c', TensorProto.FLOAT, [n], list(range(n))) + ] + ) + _onnx_export(graph_def, 'cse_merge_constants.onnx') + +def cse_only_one_graph_output(): + graph_def = helper.make_graph( + nodes = [ + helper.make_node(op_type = "Split", inputs = ['x'], outputs = ['Split1Output1', 'Split1Output2'], name = 'split_1'), + helper.make_node(op_type = "Split", inputs = ['x'], outputs = ['Split2Output1', 'Split2Output2'], name = 'split_2'), + helper.make_node(op_type = "ReduceSum", inputs = ['Split1Output1'], outputs = ['ReduceSum1'], name = 'reducesum_1'), + helper.make_node(op_type = "ReduceSum", inputs = ['Split2Output1'], outputs = ['ReduceSum2'], name = 'reducesum_2'), + helper.make_node(op_type = "Add", inputs = ['ReduceSum1', 'ReduceSum2'], outputs = ['Add'], name = 'add_1'), + ], + name = 'cse_only_one_graph_output', + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 4]) + ], + outputs = [ + helper.make_tensor_value_info('Split1Output2', TensorProto.FLOAT, [2, 4]), + helper.make_tensor_value_info('Split2Output2', TensorProto.FLOAT, [2, 4]), + helper.make_tensor_value_info('Add', TensorProto.FLOAT, [1, 1]), + ], + initializer = [ + ] + ) + _onnx_export(graph_def, 'cse_only_one_graph_output.onnx') + + +def generate_all(): + cse1() + cse_graph_output() + cse_optional_args() + cse_subgraph() + cse_random() + cse_merge_constants() + cse_only_one_graph_output() + +if __name__ == '__main__': + generate_all() + +