More efficient rule-based transformer (#815)

Introduce a quick pre-filtering of rules based on the node op types they are targeting.
The goal is to avoid evaluating all rules for all nodes. Instead, for each node, we will only be evaluating the rules associated with its op type.

include/onnxruntime/core/optimizer/graph_transformer.h

@@ -16,8 +16,8 @@ The interface for in-place transformation of a Graph.
 */
 class GraphTransformer {
  public:
-  GraphTransformer(const std::string& name, const std::string& desc, const std::unordered_set<std::string>& compatible_execution_providers = {})
-      : name_(name), desc_(desc), compatible_provider_types_(compatible_execution_providers) {
+  GraphTransformer(const std::string& name, const std::unordered_set<std::string>& compatible_execution_providers = {})
+      : name_(name), compatible_provider_types_(compatible_execution_providers) {
   }
 
   virtual ~GraphTransformer() = default;
@@ -27,11 +27,6 @@ class GraphTransformer {
     return name_;
   }
 
-  /** Gets the description of this graph transformer. */
-  const std::string& Description() const noexcept {
-    return desc_;
-  }
-
   const std::unordered_set<std::string>& GetCompatibleExecutionProviders() const noexcept {
     return compatible_provider_types_;
   }
@@ -67,7 +62,6 @@ class GraphTransformer {
   virtual common::Status ApplyImpl(Graph& graph, bool& modified, int graph_level = 0) const = 0;
 
   const std::string name_;
-  const std::string desc_;
   const std::unordered_set<std::string> compatible_provider_types_;
 };
 }  // namespace onnxruntime

include/onnxruntime/core/optimizer/graph_transformer_utils.h

@@ -23,8 +23,5 @@ std::vector<std::unique_ptr<RewriteRule>> GenerateRewriteRules(TransformerLevel
 std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerLevel level,
                                                                     const std::vector<std::string>& rules_and_transformers_to_enable = {});
 
-/** Given a TransformerLevel, this method generates a name for the rule-based graph transformer of that level. */
-std::string GenerateRuleBasedTransformerName(TransformerLevel level);
-
 }  // namespace transformer_utils
 }  // namespace onnxruntime

include/onnxruntime/core/optimizer/rewrite_rule.h

@@ -11,31 +11,31 @@ namespace onnxruntime {
 /**
 @class RewriteRule 
 
-The base class for a rewrite rule. A rewrite rule represents a semantics-preserving
-transformation of a computation graph. It can be used to represent, for example,
-the elimination of operators that serve as no-ops (e.g., dropout during
-inference), as well as inlining of "function" definitions or the dual (replacing
-a complex expression by an equivalent function-call). Unlike the more general
-IGraphTransformer, a rewrite rule is applied at a single node, representing the
-root of an expression that is rewritten.
+The base class for a rewrite rule. A rewrite rule represents a semantics-preserving transformation of a 
+computation graph. It can be used to represent, for example, the elimination of operators that serve as 
+no-ops (e.g., dropout during inference), as well as inlining of "function" definitions or the dual operation 
+of replacing a complex expression by an equivalent function-call). Unlike the more general GraphTransformer, 
+a rewrite rule is a more local transformation that is triggered on a particular node of the graph. 
 
-When creating a new rewrite rule, two main function have to be implemented: SatisfyCondition and Apply.
-- SatisfyCondition determines whether the rule will be triggered, and can include multiple condition checks.
-It is advisable to add the more selective checks first, because those will lead to discarding fast rules that 
-cannot be applied on a node.
-- Apply is the actual body of the rule that will be executed if the checks in SatisfyCondition are passed
-successfully. Note that additional, more complex checks can be included in the Apply if putting them in the
-SatisfyCondition would lead to duplicate work (e.g., when we make a check on a Node attribute but we need
-that attribute to execute the rule too).
+Each rule has a set of conditions and a body. The conditions have to be satisfied for the body of the rule 
+to be triggered. Therefore, when creating a new rewrite rule, two main functions have to be implemented: 
+- SatisfyCondition defines the condition checks. It is advisable to add the more selective checks first, 
+  because those will lead to discarding fast rules that cannot be applied on a node.
+- Apply is the actual body of the rule that will be executed if SatisfyCondition returns true for a particular
+  node. Note that additional, more complex checks can be included in the Apply if putting them in the
+  SatisfyCondition would lead to duplicate work (e.g., when we make a check on a Node attribute but we need
+  that attribute to execute the rule too).
+In general, simple fast checks are a better fit for SatisfyCondition, whereas more complex ones can be added 
+in the Apply.
 
-In general, simple fast checks are a better fit for SatisfyCondition, whereas more complex ones can be 
-added in the Apply.
+In order to avoid evaluating the SatisfyCondition for each rule and each node of the graph, each rewrite rule
+should specify the target op types for which a rule will be evaluated, by overriding the TargetOpTypes() function.
+If the op type of a node is not included in the target op types of a rule, that rule would not be considered at all.
+If the list of op types is left empty, that rule will be triggered for every op type.
 */
 class RewriteRule {
  public:
-  RewriteRule(const std::string& name, const std::string& desc)
-      : name_(name), desc_(desc) {
-  }
+  RewriteRule(const std::string& name) : name_(name) {}
 
   virtual ~RewriteRule() = default;
 
@@ -44,17 +44,17 @@ class RewriteRule {
     return name_;
   }
 
-  /** Gets the description of this rewrite rule. */
-  const std::string& Description() const noexcept {
-    return desc_;
-  }
+  /** Returns the node op types for which this rule will be triggered. If the op type of a node is not included in the
+      target op types of a rule, that rule would not be considered at all. Returning an empty list indicates that we
+      will attempt to trigger the rule for every op type. */
+  virtual std::vector<std::string> TargetOpTypes() const noexcept = 0;
 
-  /** Checks if the condition of the rule is satisfied, and if so applies the rule.
-  @param[in] graph The Graph.
-  @param[in] node The Node to apply the rewrite to.
-  @param[out] modified Set to indicate whether the node was modified or not.
-  @param[out] deleted Set to indicate if the node was deleted. 
-  @returns Status indicating success or providing error information */
+  /** Checks if the condition of the rule is satisfied, and if so applies the body of the rule.
+      @param[in] graph The Graph.
+      @param[in] node The Node to apply the rewrite to.
+      @param[out] modified Set to indicate whether the node was modified or not.
+      @param[out] deleted Set to indicate if the node was deleted.
+      @returns Status indicating success or providing error information */
   common::Status CheckConditionAndApply(Graph& graph, Node& node, bool& modified, bool& deleted) {
     return SatisfyCondition(graph, node) ? Apply(graph, node, modified, deleted) : Status::OK();
   }
@@ -63,20 +63,17 @@ class RewriteRule {
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RewriteRule);
 
   const std::string name_;
-  const std::string desc_;
 
-  /** Check if the Node of the given Graph satisfies a condition.
-  The rewrite rule is applied if the condition function returns true. This can include
-  a more complex pattern matching (conditions on the ascending or descending nodes of the
-  node for which this rule was triggered) or some other properties of the nodes. */
+  /** Checks if the Node of the given Graph satisfies the conditions of this rule. The body of the rule will be 
+      evaluated if this condition function returns true. This can include a more complex pattern matching (conditions 
+      on the ascending or descending nodes of the node for which this rule was triggered) or some other properties 
+      of the nodes. */
   virtual bool SatisfyCondition(const Graph& graph, const Node& node) = 0;
 
-  /**
-  Apply the rewrite rule to a specific node.
-  The transformation happens in-place. The return-value of node may be different
-  from the input-value due to rewriting.
-  The value of "modified" indicates if the graph was modified or not. 
-  The value of "deleted" indicates if the node was deleted or not. */
+  /** This is the actual body of the rule that performs the graph transformation. The transformation happens in-place. 
+      The return-value of node may be different from the input-value due to rewriting.
+      The value of "modified" indicates if the graph was modified or not.
+      The value of "deleted" indicates if the node was deleted or not. */
   virtual common::Status Apply(Graph& graph, Node& node, bool& modified, bool& deleted) = 0;
 };
 }  // namespace onnxruntime

include/onnxruntime/core/optimizer/rule_based_graph_transformer.h

@@ -29,36 +29,47 @@ each with different trade offs. At the moment, we define one that performs top-d
 */
 class RuleBasedGraphTransformer : public GraphTransformer {
  public:
-  RuleBasedGraphTransformer(const std::string& name, const std::string& desc, const std::unordered_set<std::string>& compatible_execution_providers = {})
-      : GraphTransformer(name, desc, compatible_execution_providers) {}
+  RuleBasedGraphTransformer(const std::string& name,
+                            const std::unordered_set<std::string>& compatible_execution_providers = {})
+      : GraphTransformer(name, compatible_execution_providers) {}
 
-  /**
-  Register a rewrite rule in this transformer.
-  */
+  /** Registers a rewrite rule in this transformer. */
   Status Register(std::unique_ptr<RewriteRule> rule);
 
-  /**
-  Gets the list of registered rewrite rules in this rule-based transformer.
-  @returns a reference to the vector containing all the registered rewrite rules.
-  */
-  const std::vector<std::unique_ptr<RewriteRule>>& GetRewriteRules() const {
-    return rules_;
+  /** Gets the list of registered rewrite rules that will be triggered on nodes with the given op type 
+      by this rule-based transformer.
+      @returns a pointer to the vector containing all the registered rewrite rules. */
+  const std::vector<std::unique_ptr<RewriteRule>>* GetRewriteRulesForOpType(const std::string& op_type) const {
+    auto rules = op_type_to_rules_.find(op_type);
+    return (rules != op_type_to_rules_.cend()) ? &rules->second : nullptr;
   }
 
+  /** Gets the rewrite rules that are evaluated on all nodes irrespective of their op type.
+      @returns a pointer to the vector containing all such rewrite rules or nullptr if no such rule. */
+  const std::vector<std::unique_ptr<RewriteRule>>* GetAnyOpRewriteRules() const {
+    return &any_op_type_rules_;
+  }
+
+  /** Returns the total number of rules that are registered in this transformer. */
+  size_t RulesCount() const;
+
  protected:
-  /** Apply the given set of rewrite rules on the Node of this Graph.
-  @param[in] graph The Graph.
-  @param[in] node The Node to apply the rules to.
-  @param[in] rules The vector of RewriteRules that will be applied to the Node.
-  @param[out] modified Set to indicate whether the node was modified or not.
-  @param[out] deleted Set to indicate if the node was deleted.
-  @returns Status indicating success or providing error information */
+  /** Applies the given set of rewrite rules on the Node of this Graph.
+      @param[in] graph The Graph.
+      @param[in] node The Node to apply the rules to.
+      @param[in] rules The vector of RewriteRules that will be applied to the Node.
+      @param[out] modified Set to indicate whether the node was modified or not.
+      @param[out] deleted Set to indicate if the node was deleted.
+      @returns Status indicating success or providing error information. */
   common::Status ApplyRulesOnNode(Graph& graph, Node& node,
                                   const std::vector<std::unique_ptr<RewriteRule>>& rules,
                                   bool& modified, bool& deleted) const;
 
  private:
-  std::vector<std::unique_ptr<RewriteRule>> rules_;
+  // Map that associates a node's op type with the vector of rules that are registered to be triggered for that node.
+  std::unordered_map<std::string, std::vector<std::unique_ptr<RewriteRule>>> op_type_to_rules_;
+  // Rules that will be evaluated regardless of the op type of the node.
+  std::vector<std::unique_ptr<RewriteRule>> any_op_type_rules_;
 
   // Performs a single top-down traversal of the graph and applies all registered rules.
   common::Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override;

onnxruntime/core/optimizer/constant_folding.cc

@@ -10,8 +10,6 @@
 namespace onnxruntime {
 
 Status ConstantFolding::Apply(Graph& graph, Node& node, bool& modified, bool& deleted) {
-  // TODO Check if we need default transformers any more. I dont think we do...
-
   // Create execution frame for executing constant nodes.
   OptimizerExecutionFrame::Info info({&node}, graph.GetAllInitializedTensors());
 

onnxruntime/core/optimizer/constant_folding.h

@@ -14,10 +14,16 @@ namespace onnxruntime {
 Rewrite rule that performs constant folding to the graph.
 The rule gets applied to nodes that have only initializers as inputs. It statically computes these 
 nodes and replaces their output with an initializer that corresponds to the result of the computation.
+
+It is attempted to be triggered on all nodes irrespective of their op type.
 */
 class ConstantFolding : public RewriteRule {
  public:
-  ConstantFolding() noexcept : RewriteRule("ConstantFolding", "Constant folding") {}
+  ConstantFolding() noexcept : RewriteRule("ConstantFolding") {}
+
+  std::vector<std::string> TargetOpTypes() const noexcept override {
+    return std::vector<std::string>();
+  }
 
  private:
   /** Constant folding will not be applied to nodes whose op_type is included in this set.

onnxruntime/core/optimizer/conv_activation_fusion.h

@@ -9,8 +9,8 @@ namespace onnxruntime {
 
 class ConvActivationFusion : public GraphTransformer {
  public:
-  ConvActivationFusion(const std::unordered_set<std::string>& compatible_execution_providers = {}) noexcept 
-      : GraphTransformer("ConvActivationFusion", "Fusing Activation into Conv", compatible_execution_providers) {}
+  ConvActivationFusion(const std::unordered_set<std::string>& compatible_execution_providers = {}) noexcept
+      : GraphTransformer("ConvActivationFusion", compatible_execution_providers) {}
 
  private:
   Status ApplyImpl(onnxruntime::Graph& graph, bool& modified, int graph_level) const override;

onnxruntime/core/optimizer/conv_add_fusion.h

@@ -9,7 +9,7 @@ namespace onnxruntime {
 
 class ConvAddFusion : public onnxruntime::GraphTransformer {
  public:
-  ConvAddFusion() noexcept : onnxruntime::GraphTransformer("ConvAddFusion", "Fusing Add into Conv") {}
+  ConvAddFusion() noexcept : onnxruntime::GraphTransformer("ConvAddFusion") {}
 
  private:
   Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override;

onnxruntime/core/optimizer/conv_bn_fusion.h

@@ -9,7 +9,7 @@ namespace onnxruntime {
 
 class ConvBNFusion : public onnxruntime::GraphTransformer {
  public:
-  ConvBNFusion() noexcept : onnxruntime::GraphTransformer("ConvBNFusion", "Fusing BN into Conv") {}
+  ConvBNFusion() noexcept : onnxruntime::GraphTransformer("ConvBNFusion") {}
 
  private:
   Status ApplyImpl(onnxruntime::Graph& graph, bool& modified, int graph_level) const override;

onnxruntime/core/optimizer/conv_mul_fusion.h

@@ -8,7 +8,7 @@ namespace onnxruntime {
 
 class ConvMulFusion : public onnxruntime::GraphTransformer {
  public:
-  ConvMulFusion() noexcept : onnxruntime::GraphTransformer("ConvMulFusion", "Fusing Mul into Conv") {}
+  ConvMulFusion() noexcept : onnxruntime::GraphTransformer("ConvMulFusion") {}
 
  private:
   Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override;

onnxruntime/core/optimizer/gemm_activation_fusion.h

@@ -10,7 +10,7 @@ namespace onnxruntime {
 class GemmActivationFusion : public GraphTransformer {
  public:
   GemmActivationFusion(const std::unordered_set<std::string>& compatible_execution_providers = {}) noexcept
-      : GraphTransformer("GemmActivationFusion", "Fusing Activation into Gemm", compatible_execution_providers) {}
+      : GraphTransformer("GemmActivationFusion", compatible_execution_providers) {}
 
   Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override;
 };

onnxruntime/core/optimizer/graph_transformer_utils.cc

@@ -16,6 +16,11 @@ namespace onnxruntime {
 
 namespace transformer_utils {
 
+/** Given a TransformerLevel, this method generates a name for the rule-based graph transformer of that level. */
+static std::string GenerateRuleBasedTransformerName(TransformerLevel level) {
+  return "Level" + std::to_string(static_cast<uint32_t>(level)) + "_RuleBasedTransformer";
+}
+
 std::vector<std::unique_ptr<RewriteRule>> GenerateRewriteRules(TransformerLevel level,
                                                                const std::vector<std::string>& rules_to_enable) {
   std::vector<std::unique_ptr<RewriteRule>> rules;
@@ -57,8 +62,6 @@ std::unique_ptr<RuleBasedGraphTransformer> GenerateRuleBasedGraphTransformer(Tra
 
   std::unique_ptr<RuleBasedGraphTransformer> rule_transformer =
       std::make_unique<RuleBasedGraphTransformer>(transformer_utils::GenerateRuleBasedTransformerName(level),
-                                                  "Apply rewrite rules for Level" +
-                                                      std::to_string(static_cast<uint32_t>(level)),
                                                   compatible_execution_providers);
   for (auto& entry : rewrite_rules_to_register) {
     rule_transformer->Register(std::move(entry));
@@ -125,9 +128,5 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
   }
 }
 
-std::string GenerateRuleBasedTransformerName(TransformerLevel level) {
-  return "Level" + std::to_string(static_cast<uint32_t>(level)) + "_RuleBasedTransformer";
-}
-
 }  // namespace transformer_utils
 }  // namespace onnxruntime

onnxruntime/core/optimizer/identity_elimination.cc

@@ -19,8 +19,7 @@ Status EliminateIdentity::Apply(Graph& graph, Node& node, bool& modified, bool&
 }
 
 bool EliminateIdentity::SatisfyCondition(const Graph& graph, const Node& node) {
-  return node.OpType() == included_op_type_ &&
-         graph_utils::IsSingleInSingleOutNode(node) &&
+  return graph_utils::IsSingleInSingleOutNode(node) &&
          !graph.IsNodeOutputsInGraphOutputs(node);
 }
 

onnxruntime/core/optimizer/identity_elimination.h

@@ -7,18 +7,25 @@
 
 namespace onnxruntime {
 
-// Rewrite rule that eliminates the identity node.
+/**
+@Class EliminateIdentity
+
+Rewrite rule that eliminates the identity node.
+
+It is attempted to be triggered only on nodes with op type "Identity".
+*/
 class EliminateIdentity : public RewriteRule {
  public:
-  EliminateIdentity() noexcept : RewriteRule("EliminateIdentity", "Eliminate identity node") {}
+  EliminateIdentity() noexcept : RewriteRule("EliminateIdentity") {}
+
+  std::vector<std::string> TargetOpTypes() const noexcept override {
+    return {"Identity"};
+  }
 
  private:
-  /** Apply rule when op type is the following. */
-  const std::string included_op_type_ = "Identity";
-
   bool SatisfyCondition(const Graph& graph, const Node& node) override;
 
   Status Apply(Graph& graph, Node& node, bool& modified, bool& deleted) override;
-};
+};  // namespace onnxruntime
 
 }  // namespace onnxruntime

onnxruntime/core/optimizer/insert_cast_transformer.cc

@@ -91,10 +91,10 @@ Status ForceSingleNodeCPUFloat16ToFloat32(onnxruntime::Graph& graph) {
   return Status::OK();
 }
 
+/** Transformer to remove duplicate Cast nodes. */
 class RemoveDuplicateCastTransformer : public GraphTransformer {
  public:
-  RemoveDuplicateCastTransformer() : GraphTransformer("RemoveDuplicateCastTransformer",
-                                                      "Transformer to remove duplicate Cast nodes.") {
+  RemoveDuplicateCastTransformer() : GraphTransformer("RemoveDuplicateCastTransformer") {
   }
 
  private:

onnxruntime/core/optimizer/insert_cast_transformer.h

@@ -8,10 +8,16 @@
 #include "core/optimizer/graph_transformer.h"
 
 namespace onnxruntime {
+
+/**
+@Class InsertCastTransformer
+
+Transformer to insert cast node that casts float16 to float for cpu nodes
+*/
 class InsertCastTransformer : public onnxruntime::GraphTransformer {
  public:
   InsertCastTransformer(const std::string& name)
-      : onnxruntime::GraphTransformer(name, "Transformer to insert cast node that casts float16 to float for cpu nodes"),
+      : onnxruntime::GraphTransformer(name),
         force_cpu_fp32_(true) {
   }
 

onnxruntime/core/optimizer/matmul_add_fusion.h

@@ -10,7 +10,7 @@ namespace onnxruntime {
 class MatMulAddFusion : public GraphTransformer {
  public:
   MatMulAddFusion(const std::unordered_set<std::string>& compatible_execution_providers = {}) noexcept 
-      : GraphTransformer("MatMulAddFusion", "Fusing MatMul and Add into Gemm", compatible_execution_providers) {}
+      : GraphTransformer("MatMulAddFusion", compatible_execution_providers) {}
 
   Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override;
 };

onnxruntime/core/optimizer/rule_based_graph_transformer.cc

@@ -9,7 +9,14 @@ using namespace ::onnxruntime::common;
 namespace onnxruntime {
 
 Status RuleBasedGraphTransformer::Register(std::unique_ptr<RewriteRule> rule) {
-  rules_.push_back(std::move(rule));
+  auto op_types = rule->TargetOpTypes();
+  // If the target op types are empty, this rule will be evaluated for all op types.
+  if (op_types.empty()) {
+    any_op_type_rules_.push_back(std::move(rule));
+  } else {
+    std::for_each(op_types.cbegin(), op_types.cend(),
+                  [&](const auto& op_type) { op_type_to_rules_[op_type].push_back(std::move(rule)); });
+  }
   return Status::OK();
 }
 
@@ -40,10 +47,23 @@ Status RuleBasedGraphTransformer::ApplyImpl(Graph& graph, bool& modified, int gr
       continue;
     }
 
-    // Apply rewrite rules on current node, then recursively apply rules to subgraphs (if any).
+    // First apply rewrite rules that are registered for the op type of the current node; then apply rules that are
+    // registered to be applied regardless of the op type; then recursively apply rules to subgraphs (if any).
     // Stop further rule application for the current node, if the node gets removed by a rule.
     bool deleted = false;
-    ORT_RETURN_IF_ERROR(ApplyRulesOnNode(graph, *node, GetRewriteRules(), modified, deleted));
+    const std::vector<std::unique_ptr<RewriteRule>>* rules = nullptr;
+
+    rules = GetRewriteRulesForOpType(node->OpType());
+    if (rules) {
+      ORT_RETURN_IF_ERROR(ApplyRulesOnNode(graph, *node, *rules, modified, deleted));
+    }
+
+    if (!deleted) {
+      rules = GetAnyOpRewriteRules();
+      if (rules) {
+        ORT_RETURN_IF_ERROR(ApplyRulesOnNode(graph, *node, *rules, modified, deleted));
+      }
+    }
 
     if (!deleted) {
       ORT_RETURN_IF_ERROR(Recurse(*node, modified, graph_level));
@@ -52,4 +72,11 @@ Status RuleBasedGraphTransformer::ApplyImpl(Graph& graph, bool& modified, int gr
 
   return Status::OK();
 }
+
+size_t RuleBasedGraphTransformer::RulesCount() const {
+  return any_op_type_rules_.size() +
+         std::accumulate(op_type_to_rules_.cbegin(), op_type_to_rules_.cend(), size_t(0),
+                         [](size_t sum, const auto& rules) { return sum + rules.second.size(); });
+}
+
 }  // namespace onnxruntime

onnxruntime/core/optimizer/slice_elimination.cc

@@ -17,8 +17,7 @@ Status EliminateSlice::Apply(Graph& graph, Node& node, bool& modified, bool& rem
 }
 
 bool EliminateSlice::SatisfyCondition(const Graph& graph, const Node& node) {
-  if (node.OpType() != included_op_type_ ||
-      !graph_utils::IsSingleInSingleOutNode(node) ||
+  if (!graph_utils::IsSingleInSingleOutNode(node) ||
       graph.IsNodeOutputsInGraphOutputs(node)) {
     return false;
   }

onnxruntime/core/optimizer/slice_elimination.h

@@ -7,15 +7,22 @@
 
 namespace onnxruntime {
 
-// Rewrite rule that eliminates a slice operator if it is redundant (does not lead to data reduction).
+/**
+@Class EliminateSlice
+
+Rewrite rule that eliminates a slice operator if it is redundant (does not lead to data reduction).
+
+It is attempted to be triggered only on nodes with op type "Slice".
+*/
 class EliminateSlice : public RewriteRule {
  public:
-  EliminateSlice() noexcept : RewriteRule("EliminateSlice", "Eliminate slice node") {}
+  EliminateSlice() noexcept : RewriteRule("EliminateSlice") {}
+
+  std::vector<std::string> TargetOpTypes() const noexcept override {
+    return {"Slice"};
+  }
 
  private:
-  /** Apply rule when op type is the following. */
-  const std::string included_op_type_ = "Slice";
-
   bool SatisfyCondition(const Graph& graph, const Node& node) override;
 
   Status Apply(Graph& graph, Node& node, bool& modified, bool& removed) override;

onnxruntime/core/optimizer/transformer_memcpy.h

@@ -10,13 +10,15 @@
 
 namespace onnxruntime {
 
-// implements MemCpy node insertion in graph transform
+/**
+@Class MemcpyTransformer
+
+Transformer that inserts nodes to copy memory between devices when needed.
+*/
 class MemcpyTransformer : public GraphTransformer {
  public:
   MemcpyTransformer(const std::vector<std::string>& provider_types, const KernelRegistryManager& registry_manager)
-      : GraphTransformer("MemcpyTransformer", "Insert nodes to copy memory between devices when needed"),
-        provider_types_{provider_types},
-        registry_manager_{registry_manager} {}
+      : GraphTransformer("MemcpyTransformer"), provider_types_{provider_types}, registry_manager_{registry_manager} {}
 
  private:
   common::Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override;

onnxruntime/core/optimizer/unsqueeze_elimination.cc

@@ -74,8 +74,8 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, bool& modified, boo
 }  // namespace onnxruntime
 
 bool UnsqueezeElimination::SatisfyCondition(const Graph& graph, const Node& node) {
-  return node.OpType() == included_op_type_ &&
-         node.GetInputEdgesCount() == 0 &&
+  // Attempt to remove an Unsqueeze operator only if it gets an initializer as input.
+  return node.GetInputEdgesCount() == 0 &&
          !graph.IsNodeOutputsInGraphOutputs(node);
 }
 

onnxruntime/core/optimizer/unsqueeze_elimination.h

@@ -7,14 +7,22 @@
 
 namespace onnxruntime {
 
+/**
+@Class UnsqueezeElimination
+
+Rewrite rule that eliminates an unsqueeze operator that takes as input an initializer.
+
+It is attempted to be triggered only on nodes with op type "Unsqueeze".
+*/
 class UnsqueezeElimination : public RewriteRule {
  public:
-  UnsqueezeElimination() noexcept : RewriteRule("UnsqueezeElimination", "Eliminate unsqueeze node") {}
+  UnsqueezeElimination() noexcept : RewriteRule("UnsqueezeElimination") {}
+
+  std::vector<std::string> TargetOpTypes() const noexcept override {
+    return {"Unsqueeze"};
+  }
 
  private:
-  /** Apply rule when op type is the following. */
-  const std::string included_op_type_ = "Unsqueeze";
-
   bool SatisfyCondition(const Graph& graph, const Node& node) override;
 
   Status Apply(Graph& graph, Node& node, bool& modified, bool& deleted) override;

onnxruntime/test/optimizer/dummy_graph_transformer.h

@@ -10,7 +10,7 @@ namespace test {
 // Dummy graph transformer that does nothing, but just sets the modified value
 class DummyGraphTransformer : public GraphTransformer {
  public:
-  DummyGraphTransformer(const std::string& name) noexcept : GraphTransformer(name, "Dummy transformer for testing"),
+  DummyGraphTransformer(const std::string& name) noexcept : GraphTransformer(name),
                                                             transformer_invoked_(false) {}
 
   bool IsTransformerInvoked() const {
@@ -23,32 +23,34 @@ class DummyGraphTransformer : public GraphTransformer {
   Status ApplyImpl(Graph& /*graph*/, bool& /*modified*/, int /*graph_level*/) const override {
     transformer_invoked_ = true;
     return Status::OK();
-  }  
+  }
 };
 
 // Dummy graph transformer that does nothing, but just sets the modified value
 // This is currently used to test custom transformer selection feature
 class DummyRewriteRule : public RewriteRule {
  public:
-  DummyRewriteRule(const std::string& name) noexcept : RewriteRule(name, "Dummy transformer for testing"),
-                                                       rewrite_rule_invoked_(false) {}
+  DummyRewriteRule(const std::string& name) noexcept : RewriteRule(name), rewrite_rule_invoked_(false) {}
 
   bool IsRewriteRuleInvoked() const {
     return rewrite_rule_invoked_;
   }
 
+  std::vector<std::string> TargetOpTypes() const noexcept override {
+    return std::vector<std::string>();
+  }
+
  private:
-  const std::string included_op_type_ = "Dummy";
   bool rewrite_rule_invoked_;
 
   bool SatisfyCondition(const Graph& /*graph*/, const Node& /*node*/) override {
-      return true;
-  }  
+    return true;
+  }
 
   Status Apply(Graph& /*graph*/, Node& /*node*/, bool& /*modified*/, bool& /*deleted*/) override {
     rewrite_rule_invoked_ = true;
     return Status::OK();
-  } 
+  }
 };
 
 }  // namespace test

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -52,7 +52,7 @@ TEST(GraphTransformationTests, IdentityElimination) {
   std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
   ASSERT_TRUE(op_to_count["Identity"] == 1);
 
-  auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1", "First rule transformer");
+  auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
   rule_transformer->Register(std::make_unique<EliminateIdentity>());
   onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
   graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1);
@@ -70,8 +70,7 @@ TEST(GraphTransformationTests, SliceElimination) {
   std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
   ASSERT_TRUE(op_to_count["Slice"] == 5);
 
-  std::unique_ptr<RuleBasedGraphTransformer> rule_transformer =
-      std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1", "First rule transformer");
+  auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
   rule_transformer->Register(std::make_unique<EliminateSlice>());
   onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
   graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1);
@@ -89,12 +88,9 @@ TEST(GraphTransformationTests, ConstantFolding1) {
   std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
   ASSERT_TRUE(op_to_count["Unsqueeze"] == 2);
 
-  std::unique_ptr<RuleBasedGraphTransformer> rule_transformer =
-      std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1", "First rule transformer");
-
+  auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
   rule_transformer->Register(std::make_unique<ConstantFolding>());
   onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
-
   graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1);
   ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1).IsOK());
 
@@ -129,8 +125,7 @@ TEST(GraphTransformationTests, FuseConvBNMulAddUnsqueeze) {
     Graph& graph = p_model->MainGraph();
 
     onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
-    std::unique_ptr<RuleBasedGraphTransformer> rule_transformer =
-        std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1", "First rule transformer");
+    auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
     rule_transformer->Register(std::make_unique<UnsqueezeElimination>());
     graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1);
     graph_transformation_mgr.Register(std::make_unique<ConvBNFusion>(), TransformerLevel::Level2);
@@ -183,8 +178,7 @@ TEST(GraphTransformationTests, FuseConvMulNoBias) {
   Graph& graph = p_model->MainGraph();
 
   onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
-  std::unique_ptr<RuleBasedGraphTransformer> rule_transformer =
-      std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1", "First rule transformer");
+  auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
   rule_transformer->Register(std::make_unique<UnsqueezeElimination>());
   graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1);
   graph_transformation_mgr.Register(std::make_unique<ConvMulFusion>(), TransformerLevel::Level2);
@@ -205,8 +199,7 @@ TEST(GraphTransformationTests, FuseConvAddNoBias) {
   Graph& graph = p_model->MainGraph();
 
   onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
-  std::unique_ptr<RuleBasedGraphTransformer> rule_transformer =
-      std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1", "First rule transformer");
+  auto rule_transformer = std::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
   rule_transformer->Register(std::make_unique<UnsqueezeElimination>());
   graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1);
   graph_transformation_mgr.Register(std::make_unique<ConvAddFusion>(), TransformerLevel::Level2);

onnxruntime/test/optimizer/graph_transform_utils_test.cc

@@ -54,16 +54,16 @@ TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers) {
 
 TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers_CustomList) {
   // custom list of rules and transformers
-  std::string l1_rule = "EliminateIdentity";
+  std::string l1_rule1 = "EliminateIdentity";
+  std::string l1_rule2 = "ConstantFolding";
   std::string l2_transformer = "ConvAddFusion";
-  std::vector<std::string> custom_list = {l1_rule, l2_transformer};
+  std::vector<std::string> custom_list = {l1_rule1, l1_rule2, l2_transformer};
 
   auto transformers = transformer_utils::GenerateTransformers(TransformerLevel::Level1, custom_list);
   ASSERT_TRUE(transformers.size() == 1);
   auto rule_transformer = dynamic_cast<RuleBasedGraphTransformer*>(transformers[0].get());
-  ASSERT_TRUE(rule_transformer->GetRewriteRules().size() == 1);
-  ASSERT_TRUE(rule_transformer->GetRewriteRules()[0]->Name() == l1_rule);
-
+  ASSERT_TRUE(rule_transformer->RulesCount() == 2);
+  
   transformers = transformer_utils::GenerateTransformers(TransformerLevel::Level2, custom_list);
   ASSERT_TRUE(transformers.size() == 1);
   ASSERT_TRUE(transformers[0]->Name() == l2_transformer);

onnxruntime/test/optimizer/rule_based_graph_transformer_test.cc

@@ -29,15 +29,14 @@ TEST(RuleBasedGraphTransformerTest, TestCompatibleProviders) {
   auto dummy_rule = std::make_unique<DummyRewriteRule>("DummyRule");
   const auto* dummy_rule_ptr = dummy_rule.get();
 
-  auto graph_transformer = std::make_unique<RuleBasedGraphTransformer>("CUDATopDownTransformer", "Registered for CUDA",
-                                                                       compatible_provider);
+  auto graph_transformer = std::make_unique<RuleBasedGraphTransformer>("CUDATopDownTransformer", compatible_provider);
   graph_transformer->Register(std::move(dummy_rule));  
 
   // Create rule based transformer with a dummy rewrite rule and register it with CPU as compatible provider
   auto dummy_rule1 = std::make_unique<DummyRewriteRule>("DummyRule1");
   const auto* dummy_rule1_ptr = dummy_rule1.get();
 
-  auto graph_transformer1 = std::make_unique<RuleBasedGraphTransformer>("CPUTopDownTransformer", "Registered for CPU");
+  auto graph_transformer1 = std::make_unique<RuleBasedGraphTransformer>("CPUTopDownTransformer");
 
   graph_transformer1->Register(std::move(dummy_rule1));