Eliminate redundant subexpressions (#3047)

* Eliminate redundant subexpressions

Apply local value numbering to merge graph nodes that will always
evaluate to the same value.

* Rename cpp->cc

* Handle optional arguments

* Add test models

* Add more tests with optional arguments

* Fix processing of subgraphs

Also, be resilient to possible mixture of optional and variadic
parameters

* Fix random operators

* Address PR comments

* Minor changes and a test

* Move CSE before constant folding

* Random* operators are always non-deterministic

Even when seed is provided.

* Fix a CSE test

* Reuse the list of non-deterministic operators with constant folding pass

* Address PR comments

* Fix formatting

* Address PR comment

* Minor cleanup / comments

* Fix build failure in Linux

* Reuse existing optimizer/utils file.

Also, check for graph outputs when removing a node.

* Add a test

* Fix compiler warnings

* Fix build in older compilers

* More compatibility with old STL versions

include/onnxruntime/core/graph/constants.h

@@ -3,8 +3,6 @@
 
 #pragma once
 
-#include <string>
-#include <unordered_set>
 #include <vector>
 
 #include "core/common/common.h"

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 <memory>
+#include <type_traits>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+// 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 <typename Ptr>
+  bool operator()(const Ptr& lhs, const Ptr& rhs) const {
+    if (lhs == nullptr || rhs == nullptr) {
+      return lhs == rhs;
+    }
+    return *lhs == *rhs;
+  }
+};
+
+struct DeepPointerHash {
+  template <typename Ptr>
+  std::size_t operator()(const Ptr& value) const {
+    if (value == nullptr) {
+      return 0;
+    }
+    return std::hash<typename std::decay<decltype(*value)>::type>{}(*value);
+  }
+};
+
+template <typename T, typename Hasher>
+void UpdateHash(const T& x, Hasher hasher, std::size_t& hash) {
+  constexpr std::size_t kPrime = 31013;
+  hash = hash * kPrime + hasher(x);
+}
+
+template <typename T>
+void UpdateHash(const T& x, std::size_t& hash) {
+  UpdateHash(x, std::hash<T>{}, hash);
+}
+
+template <typename Container>
+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<EquivalenceClass>;
+  friend std::vector<std::vector<const EquivalenceClass*>> Normalize(const Node& node, const std::vector<const EquivalenceClass*>& 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<const EquivalenceClass*>& 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<std::vector<const EquivalenceClass*>> 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<std::vector<const EquivalenceClass*>> Normalize(const Node& node, const std::vector<const EquivalenceClass*>& inputs) {
+  const auto& arg_count = node.InputArgCount();
+  auto input_iter = inputs.begin();
+  std::vector<std::vector<const EquivalenceClass*>> 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 <typename Range>
+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<std::underlying_type<ONNX_NAMESPACE::AttributeProto_AttributeType>::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<const NodeArg*>{}(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<onnxruntime::EquivalenceClass> {
+  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<std::unique_ptr<EquivalenceClass>> 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<const NodeArg*, const EquivalenceClass*, NodeArgPtrHash, NodeArgPtrEquality> 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<const EquivalenceClass*> 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<EquivalenceClass>(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<int>(node->OutputDefs().size());
+         output_index < end; ++output_index) {
+      const NodeArg* output_def = node->OutputDefs()[output_index];
+      auto equivalence_class = onnxruntime::make_unique<EquivalenceClass>(*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<const NodeArg*> 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<int>(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

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<std::string>& 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

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

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<std::string> 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<std::string> excluded_initializers_;

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<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
     case TransformerLevel::Level1: {
       std::unordered_set<std::string> l1_execution_providers = {};
 
+      transformers.emplace_back(onnxruntime::make_unique<CommonSubexpressionElimination>(l1_execution_providers));
       transformers.emplace_back(onnxruntime::make_unique<ConstantFolding>(l1_execution_providers));
       transformers.emplace_back(onnxruntime::make_unique<MatMulAddFusion>(l1_execution_providers));
       transformers.emplace_back(onnxruntime::make_unique<ReshapeFusion>(l1_execution_providers));

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 <deque>
 
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
 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<std::string, std::unordered_set<std::string>> 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

onnxruntime/core/optimizer/utils.h

@@ -78,5 +78,7 @@ bool IsSupportedDataType(const Node& node, const std::vector<std::string>& 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

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 <algorithm>
+#include <string>
+#include <vector>
+
+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<CommonSubexpressionElimination>(), TransformerLevel::Level1).IsOK());
+  ASSERT_TRUE(
+      graph_transformation_mgr.ApplyTransformers(model.MainGraph(), TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger()).IsOK());
+}
+
+std::vector<std::string> GetSortedNames(const std::vector<const NodeArg*>& node_args) {
+  std::vector<std::string> 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<std::string> GetNodeNames(const Graph& graph) {
+  std::vector<std::string> res;
+  for (auto& node : graph.Nodes()) {
+    res.push_back(node.Name());
+  }
+  return res;
+}
+
+void GetAllNodeNamesImpl(const Graph& graph, std::vector<std::string>& res) {
+  for (auto& node : graph.Nodes()) {
+    res.push_back(node.Name());
+    for (const auto& subgraph : node.GetSubgraphs()) {
+      GetAllNodeNamesImpl(*subgraph, res);
+    }
+  }
+}
+
+std::vector<std::string> GetAllNodeNamesSorted(const Graph& graph) {
+  std::vector<std::string> 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> 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<std::string>{"x"}));
+
+  const auto& graph_outputs = GetSortedNames(graph.GetOutputs());
+  ASSERT_EQ(graph_outputs, (std::vector<std::string>{"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> 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<std::string>{"x"}));
+
+  const auto& output_names = GetSortedNames(graph.GetOutputs());
+  ASSERT_EQ(output_names, (std::vector<std::string>{"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> 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<std::string>{"x"}));
+
+  const auto& output_names = GetSortedNames(graph.GetOutputs());
+  ASSERT_EQ(output_names, (std::vector<std::string>{"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> 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<std::string>{"x"}));
+
+  const auto& output_names = GetSortedNames(graph.GetOutputs());
+  ASSERT_EQ(output_names, (std::vector<std::string>{"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> 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<std::string>{"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<std::string>{"b", "x"}));
+
+  auto output_names = GetSortedNames(graph.GetOutputs());
+  ASSERT_EQ(output_names, (std::vector<std::string>{"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<std::string>{"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<std::string>{"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> 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<std::string>{"x"}));
+
+  const auto& output_names = GetSortedNames(graph.GetOutputs());
+  ASSERT_EQ(output_names, (std::vector<std::string>{"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> 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<CommonSubexpressionElimination>(), TransformerLevel::Level1).IsOK());
+  ASSERT_TRUE(
+      graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(), 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

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()
+
+