diff --git a/include/onnxruntime/core/optimizer/graph_transformer.h b/include/onnxruntime/core/optimizer/graph_transformer.h index 1f236b9059..bbe65aadaf 100644 --- a/include/onnxruntime/core/optimizer/graph_transformer.h +++ b/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& compatible_execution_providers = {}) - : name_(name), desc_(desc), compatible_provider_types_(compatible_execution_providers) { + GraphTransformer(const std::string& name, const std::unordered_set& 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& 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 compatible_provider_types_; }; } // namespace onnxruntime diff --git a/include/onnxruntime/core/optimizer/graph_transformer_utils.h b/include/onnxruntime/core/optimizer/graph_transformer_utils.h index e592b1f098..6d658fee39 100644 --- a/include/onnxruntime/core/optimizer/graph_transformer_utils.h +++ b/include/onnxruntime/core/optimizer/graph_transformer_utils.h @@ -23,8 +23,5 @@ std::vector> GenerateRewriteRules(TransformerLevel std::vector> GenerateTransformers(TransformerLevel level, const std::vector& 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 diff --git a/include/onnxruntime/core/optimizer/rewrite_rule.h b/include/onnxruntime/core/optimizer/rewrite_rule.h index 732daac7f8..4401fb0f77 100644 --- a/include/onnxruntime/core/optimizer/rewrite_rule.h +++ b/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 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 diff --git a/include/onnxruntime/core/optimizer/rule_based_graph_transformer.h b/include/onnxruntime/core/optimizer/rule_based_graph_transformer.h index 7ab912b248..b20abadec3 100644 --- a/include/onnxruntime/core/optimizer/rule_based_graph_transformer.h +++ b/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& compatible_execution_providers = {}) - : GraphTransformer(name, desc, compatible_execution_providers) {} + RuleBasedGraphTransformer(const std::string& name, + const std::unordered_set& 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 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>& 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>* 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>* 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>& rules, bool& modified, bool& deleted) const; private: - std::vector> 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>> op_type_to_rules_; + // Rules that will be evaluated regardless of the op type of the node. + std::vector> 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; diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 6a1f4dad8d..c2db61a461 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/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()); diff --git a/onnxruntime/core/optimizer/constant_folding.h b/onnxruntime/core/optimizer/constant_folding.h index e07ac924fe..9082952662 100644 --- a/onnxruntime/core/optimizer/constant_folding.h +++ b/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 TargetOpTypes() const noexcept override { + return std::vector(); + } private: /** Constant folding will not be applied to nodes whose op_type is included in this set. diff --git a/onnxruntime/core/optimizer/conv_activation_fusion.h b/onnxruntime/core/optimizer/conv_activation_fusion.h index 3bdf564ab3..b2377c95fb 100644 --- a/onnxruntime/core/optimizer/conv_activation_fusion.h +++ b/onnxruntime/core/optimizer/conv_activation_fusion.h @@ -9,8 +9,8 @@ namespace onnxruntime { class ConvActivationFusion : public GraphTransformer { public: - ConvActivationFusion(const std::unordered_set& compatible_execution_providers = {}) noexcept - : GraphTransformer("ConvActivationFusion", "Fusing Activation into Conv", compatible_execution_providers) {} + ConvActivationFusion(const std::unordered_set& compatible_execution_providers = {}) noexcept + : GraphTransformer("ConvActivationFusion", compatible_execution_providers) {} private: Status ApplyImpl(onnxruntime::Graph& graph, bool& modified, int graph_level) const override; diff --git a/onnxruntime/core/optimizer/conv_add_fusion.h b/onnxruntime/core/optimizer/conv_add_fusion.h index 48f2514f60..717e013cac 100644 --- a/onnxruntime/core/optimizer/conv_add_fusion.h +++ b/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; diff --git a/onnxruntime/core/optimizer/conv_bn_fusion.h b/onnxruntime/core/optimizer/conv_bn_fusion.h index b331e6a080..b605956f1c 100644 --- a/onnxruntime/core/optimizer/conv_bn_fusion.h +++ b/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; diff --git a/onnxruntime/core/optimizer/conv_mul_fusion.h b/onnxruntime/core/optimizer/conv_mul_fusion.h index 64a74d510b..6452360bd1 100644 --- a/onnxruntime/core/optimizer/conv_mul_fusion.h +++ b/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; diff --git a/onnxruntime/core/optimizer/gemm_activation_fusion.h b/onnxruntime/core/optimizer/gemm_activation_fusion.h index 4645333d9e..3456c6a5ac 100644 --- a/onnxruntime/core/optimizer/gemm_activation_fusion.h +++ b/onnxruntime/core/optimizer/gemm_activation_fusion.h @@ -10,7 +10,7 @@ namespace onnxruntime { class GemmActivationFusion : public GraphTransformer { public: GemmActivationFusion(const std::unordered_set& 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; }; diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 4033a61323..6ddbdde47d 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/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(level)) + "_RuleBasedTransformer"; +} + std::vector> GenerateRewriteRules(TransformerLevel level, const std::vector& rules_to_enable) { std::vector> rules; @@ -57,8 +62,6 @@ std::unique_ptr GenerateRuleBasedGraphTransformer(Tra std::unique_ptr rule_transformer = std::make_unique(transformer_utils::GenerateRuleBasedTransformerName(level), - "Apply rewrite rules for Level" + - std::to_string(static_cast(level)), compatible_execution_providers); for (auto& entry : rewrite_rules_to_register) { rule_transformer->Register(std::move(entry)); @@ -125,9 +128,5 @@ std::vector> GenerateTransformers(TransformerL } } -std::string GenerateRuleBasedTransformerName(TransformerLevel level) { - return "Level" + std::to_string(static_cast(level)) + "_RuleBasedTransformer"; -} - } // namespace transformer_utils } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/identity_elimination.cc b/onnxruntime/core/optimizer/identity_elimination.cc index da24f53775..1ba4eb64c3 100644 --- a/onnxruntime/core/optimizer/identity_elimination.cc +++ b/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); } diff --git a/onnxruntime/core/optimizer/identity_elimination.h b/onnxruntime/core/optimizer/identity_elimination.h index a910fd6e57..5f9d56dbe2 100644 --- a/onnxruntime/core/optimizer/identity_elimination.h +++ b/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 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 diff --git a/onnxruntime/core/optimizer/insert_cast_transformer.cc b/onnxruntime/core/optimizer/insert_cast_transformer.cc index 35719a5fb5..873633555c 100644 --- a/onnxruntime/core/optimizer/insert_cast_transformer.cc +++ b/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: diff --git a/onnxruntime/core/optimizer/insert_cast_transformer.h b/onnxruntime/core/optimizer/insert_cast_transformer.h index 24668f852b..9f87c9da62 100644 --- a/onnxruntime/core/optimizer/insert_cast_transformer.h +++ b/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) { } diff --git a/onnxruntime/core/optimizer/matmul_add_fusion.h b/onnxruntime/core/optimizer/matmul_add_fusion.h index 6527fedcdc..ffa5e301e0 100644 --- a/onnxruntime/core/optimizer/matmul_add_fusion.h +++ b/onnxruntime/core/optimizer/matmul_add_fusion.h @@ -10,7 +10,7 @@ namespace onnxruntime { class MatMulAddFusion : public GraphTransformer { public: MatMulAddFusion(const std::unordered_set& 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; }; diff --git a/onnxruntime/core/optimizer/rule_based_graph_transformer.cc b/onnxruntime/core/optimizer/rule_based_graph_transformer.cc index 029cbf9a72..e8b65a0a21 100644 --- a/onnxruntime/core/optimizer/rule_based_graph_transformer.cc +++ b/onnxruntime/core/optimizer/rule_based_graph_transformer.cc @@ -9,7 +9,14 @@ using namespace ::onnxruntime::common; namespace onnxruntime { Status RuleBasedGraphTransformer::Register(std::unique_ptr 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>* 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 diff --git a/onnxruntime/core/optimizer/slice_elimination.cc b/onnxruntime/core/optimizer/slice_elimination.cc index 5f3a833ad4..b7a195b571 100644 --- a/onnxruntime/core/optimizer/slice_elimination.cc +++ b/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; } diff --git a/onnxruntime/core/optimizer/slice_elimination.h b/onnxruntime/core/optimizer/slice_elimination.h index baec796054..b43af73209 100644 --- a/onnxruntime/core/optimizer/slice_elimination.h +++ b/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 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; diff --git a/onnxruntime/core/optimizer/transformer_memcpy.h b/onnxruntime/core/optimizer/transformer_memcpy.h index 815660602e..b8f9cee65f 100644 --- a/onnxruntime/core/optimizer/transformer_memcpy.h +++ b/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& 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; diff --git a/onnxruntime/core/optimizer/unsqueeze_elimination.cc b/onnxruntime/core/optimizer/unsqueeze_elimination.cc index c625a36778..a53e52d650 100644 --- a/onnxruntime/core/optimizer/unsqueeze_elimination.cc +++ b/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); } diff --git a/onnxruntime/core/optimizer/unsqueeze_elimination.h b/onnxruntime/core/optimizer/unsqueeze_elimination.h index 9ea6aef28b..f4107fba36 100644 --- a/onnxruntime/core/optimizer/unsqueeze_elimination.h +++ b/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 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; diff --git a/onnxruntime/test/optimizer/dummy_graph_transformer.h b/onnxruntime/test/optimizer/dummy_graph_transformer.h index 2c95b43108..b89f6c860b 100644 --- a/onnxruntime/test/optimizer/dummy_graph_transformer.h +++ b/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 TargetOpTypes() const noexcept override { + return std::vector(); + } + 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 diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index dd87a9672f..3453dad3fc 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -52,7 +52,7 @@ TEST(GraphTransformationTests, IdentityElimination) { std::map op_to_count = CountOpsInGraph(graph); ASSERT_TRUE(op_to_count["Identity"] == 1); - auto rule_transformer = std::make_unique("RuleTransformer1", "First rule transformer"); + auto rule_transformer = std::make_unique("RuleTransformer1"); rule_transformer->Register(std::make_unique()); 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 op_to_count = CountOpsInGraph(graph); ASSERT_TRUE(op_to_count["Slice"] == 5); - std::unique_ptr rule_transformer = - std::make_unique("RuleTransformer1", "First rule transformer"); + auto rule_transformer = std::make_unique("RuleTransformer1"); rule_transformer->Register(std::make_unique()); 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 op_to_count = CountOpsInGraph(graph); ASSERT_TRUE(op_to_count["Unsqueeze"] == 2); - std::unique_ptr rule_transformer = - std::make_unique("RuleTransformer1", "First rule transformer"); - + auto rule_transformer = std::make_unique("RuleTransformer1"); rule_transformer->Register(std::make_unique()); 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 rule_transformer = - std::make_unique("RuleTransformer1", "First rule transformer"); + auto rule_transformer = std::make_unique("RuleTransformer1"); rule_transformer->Register(std::make_unique()); graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1); graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level2); @@ -183,8 +178,7 @@ TEST(GraphTransformationTests, FuseConvMulNoBias) { Graph& graph = p_model->MainGraph(); onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; - std::unique_ptr rule_transformer = - std::make_unique("RuleTransformer1", "First rule transformer"); + auto rule_transformer = std::make_unique("RuleTransformer1"); rule_transformer->Register(std::make_unique()); graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1); graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level2); @@ -205,8 +199,7 @@ TEST(GraphTransformationTests, FuseConvAddNoBias) { Graph& graph = p_model->MainGraph(); onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; - std::unique_ptr rule_transformer = - std::make_unique("RuleTransformer1", "First rule transformer"); + auto rule_transformer = std::make_unique("RuleTransformer1"); rule_transformer->Register(std::make_unique()); graph_transformation_mgr.Register(std::move(rule_transformer), TransformerLevel::Level1); graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level2); diff --git a/onnxruntime/test/optimizer/graph_transform_utils_test.cc b/onnxruntime/test/optimizer/graph_transform_utils_test.cc index e9cd822b36..336175d011 100644 --- a/onnxruntime/test/optimizer/graph_transform_utils_test.cc +++ b/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 custom_list = {l1_rule, l2_transformer}; + std::vector 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(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); diff --git a/onnxruntime/test/optimizer/rule_based_graph_transformer_test.cc b/onnxruntime/test/optimizer/rule_based_graph_transformer_test.cc index 19ca0e2391..069553a804 100644 --- a/onnxruntime/test/optimizer/rule_based_graph_transformer_test.cc +++ b/onnxruntime/test/optimizer/rule_based_graph_transformer_test.cc @@ -29,15 +29,14 @@ TEST(RuleBasedGraphTransformerTest, TestCompatibleProviders) { auto dummy_rule = std::make_unique("DummyRule"); const auto* dummy_rule_ptr = dummy_rule.get(); - auto graph_transformer = std::make_unique("CUDATopDownTransformer", "Registered for CUDA", - compatible_provider); + auto graph_transformer = std::make_unique("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("DummyRule1"); const auto* dummy_rule1_ptr = dummy_rule1.get(); - auto graph_transformer1 = std::make_unique("CPUTopDownTransformer", "Registered for CPU"); + auto graph_transformer1 = std::make_unique("CPUTopDownTransformer"); graph_transformer1->Register(std::move(dummy_rule1));