From 1b7d1f264562a99dee16d3419aebacb70d02db61 Mon Sep 17 00:00:00 2001 From: Konstantinos Karanasos Date: Mon, 29 Apr 2019 18:12:49 -0700 Subject: [PATCH] Convert constant folding to a transformer (#866) --- include/onnxruntime/core/graph/graph.h | 7 ++ .../core/optimizer/graph_transformer_utils.h | 3 + .../core/optimizer/constant_folding.cc | 114 ++++++++++-------- onnxruntime/core/optimizer/constant_folding.h | 24 ++-- .../core/optimizer/graph_transformer_utils.cc | 8 +- .../test/optimizer/graph_transform_test.cc | 5 +- .../optimizer/graph_transform_utils_test.cc | 16 ++- 7 files changed, 102 insertions(+), 75 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 4a0153b86d..b0cc98f47f 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -270,6 +270,13 @@ class Node { */ Graph* GetMutableGraphAttribute(const std::string& attr_name); + /** Checks if the Node contains at least one subgraph (this is the case for control flow operators, such as If, Scan, Loop). + @returns true if the Node contains a subgraph. + */ + bool ContainsSubgraph() const { + return !attr_to_subgraph_map_.empty(); + } + /** Gets a map of attribute name to the mutable Graph instances for all subgraphs of the Node. @returns Map of the attribute name that defines the subgraph to the subgraph's Graph instance. nullptr if the Node has no subgraphs. diff --git a/include/onnxruntime/core/optimizer/graph_transformer_utils.h b/include/onnxruntime/core/optimizer/graph_transformer_utils.h index 6d658fee39..e592b1f098 100644 --- a/include/onnxruntime/core/optimizer/graph_transformer_utils.h +++ b/include/onnxruntime/core/optimizer/graph_transformer_utils.h @@ -23,5 +23,8 @@ 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/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 99513fc624..ac277dfff9 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -7,63 +7,83 @@ #include "core/framework/op_kernel.h" #include "core/framework/ml_value.h" +using namespace onnxruntime::common; + namespace onnxruntime { -Status ConstantFolding::Apply(Graph& graph, Node& node, bool& modified, bool& deleted) { - // Create execution frame for executing constant nodes. - OptimizerExecutionFrame::Info info({&node}, graph.GetAllInitializedTensors()); +Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level) const { + GraphViewer graph_viewer(graph); + auto& order = graph_viewer.GetNodesInTopologicalOrder(); - std::vector fetch_mlvalue_idxs; - for (const auto* node_out : node.OutputDefs()) { - fetch_mlvalue_idxs.push_back(info.GetMLValueIndex(node_out->Name())); + for (NodeIndex i : order) { + auto* node = graph.GetNode(i); + if (!node) { + continue; + } + + ORT_RETURN_IF_ERROR(Recurse(*node, modified, 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() || + // constant folding is not currently supported for nodes that include subgraphs (control flow operators, + // such as If/Loop/Scan, fall into this category). + node->ContainsSubgraph() || + // if the node output is in the graph output, we will get a graph with no nodes. + // TODO check if this is allowed in ONNX and ORT. + graph.IsNodeOutputsInGraphOutputs(*node) || + !graph_utils::AllNodeInputsAreConstant(graph, *node)) { + continue; + } + + // Create execution frame for executing constant nodes. + OptimizerExecutionFrame::Info info({node}, graph.GetAllInitializedTensors()); + + std::vector fetch_mlvalue_idxs; + for (const auto* node_out : node->OutputDefs()) { + fetch_mlvalue_idxs.push_back(info.GetMLValueIndex(node_out->Name())); + } + + OptimizerExecutionFrame frame(info, fetch_mlvalue_idxs); + + auto* kernel = info.GetKernel(node->Index()); + OpKernelContext op_kernel_context(&frame, kernel, ::onnxruntime::logging::LoggingManager::DefaultLogger()); + + kernel->Compute(&op_kernel_context); + + std::vector fetches; + frame.GetOutputs(fetches); + + // Go over all output node args and substitute them with the newly computed tensors, which will be + // added to the graph as initializers. + ORT_ENFORCE(fetches.size() == node->OutputDefs().size()); + for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { + MLValue& mlvalue = fetches[fetch_idx]; + + // Build the TensorProto that corresponds to the computed MLValue and add it as initializer to the graph. + ONNX_NAMESPACE::TensorProto out_tensorproto; + const auto* constant_arg_out = node->OutputDefs()[fetch_idx]; + BuildTensorProtoForInitializer(mlvalue, *constant_arg_out, out_tensorproto); + + graph.AddInitializedTensor(out_tensorproto); + } + + // Remove the output edges of the constant node and then remove the node itself. + graph_utils::RemoveNodeOutputEdges(graph, *node); + graph.RemoveNode(node->Index()); + + // The output nodes already have the right input arg, since we used the same name in the initializer. + // We could remove unused graph initializers here, but Graph::Resolve() will take care of it. + + modified = true; } - OptimizerExecutionFrame frame(info, fetch_mlvalue_idxs); - - auto* kernel = info.GetKernel(node.Index()); - OpKernelContext op_kernel_context(&frame, kernel, ::onnxruntime::logging::LoggingManager::DefaultLogger()); - - kernel->Compute(&op_kernel_context); - - std::vector fetches; - frame.GetOutputs(fetches); - - // Go over all output node args and substitute them with the newly computed tensors, which will be - // added to the graph as initializers. - ORT_ENFORCE(fetches.size() == node.OutputDefs().size()); - for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { - MLValue& mlvalue = fetches[fetch_idx]; - - // Build the TensorProto that corresponds to the computed MLValue and add it as initializer to the graph. - ONNX_NAMESPACE::TensorProto out_tensorproto; - const auto* constant_arg_out = node.OutputDefs()[fetch_idx]; - BuildTensorProtoForInitializer(mlvalue, *constant_arg_out, out_tensorproto); - - graph.AddInitializedTensor(out_tensorproto); - } - - // Remove the output edges of the constant node and then remove the node itself. - graph_utils::RemoveNodeOutputEdges(graph, node); - graph.RemoveNode(node.Index()); - - // The output nodes already have the right input arg, since we used the same name in the initializer. - // We could remove unused graph initializers here, but Graph::Resolve() will take care of it. - - modified = deleted = true; - return Status::OK(); } // namespace onnxruntime -bool ConstantFolding::SatisfyCondition(const Graph& graph, const Node& node) { - return excluded_op_types_.find(node.OpType()) == excluded_op_types_.end() && - // constant folding is not currently supported in subgraphs. - node.ImplicitInputDefs().size() == 0 && - graph_utils::AllNodeInputsAreConstant(graph, node); -} - void ConstantFolding::BuildTensorProtoForInitializer(const MLValue& mlvalue, const NodeArg& constant_node_arg, - ONNX_NAMESPACE::TensorProto& tensorproto) { + ONNX_NAMESPACE::TensorProto& tensorproto) const { ORT_ENFORCE(mlvalue.IsTensor()); const Tensor& out_tensor = mlvalue.Get(); diff --git a/onnxruntime/core/optimizer/constant_folding.h b/onnxruntime/core/optimizer/constant_folding.h index 9082952662..29371e4c2c 100644 --- a/onnxruntime/core/optimizer/constant_folding.h +++ b/onnxruntime/core/optimizer/constant_folding.h @@ -3,7 +3,7 @@ #pragma once -#include "core/optimizer/rewrite_rule.h" +#include "core/optimizer/graph_transformer.h" #include "core/framework/ml_value.h" namespace onnxruntime { @@ -11,19 +11,13 @@ namespace onnxruntime { /** @class ConstantFolding -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. +Transformer that traverses the graph top-down and performs constant folding, i.e., +it statically computes parts of the graph that rely only on constant initializers. */ -class ConstantFolding : public RewriteRule { +class ConstantFolding : public GraphTransformer { public: - ConstantFolding() noexcept : RewriteRule("ConstantFolding") {} - - std::vector TargetOpTypes() const noexcept override { - return std::vector(); - } + ConstantFolding(const std::unordered_set& compatible_execution_providers = {}) noexcept : + GraphTransformer("ConstantFolding", compatible_execution_providers) {} private: /** Constant folding will not be applied to nodes whose op_type is included in this set. @@ -31,15 +25,13 @@ class ConstantFolding : public RewriteRule { const std::unordered_set excluded_op_types_ = {"RandomUniform", "RandomNormal", "RandomUniformLike", "RandomNormalLike", "Multinomial"}; - bool SatisfyCondition(const Graph& graph, const Node& node) override; - - Status Apply(Graph& graph, Node& node, bool& modified, bool& deleted) override; + Status ApplyImpl(Graph& graph, bool& modified, int graph_level) const override; /** Create a TensorProto that has the same value as the given MLValue and the same type and dimensions as the given NodeArg. */ void BuildTensorProtoForInitializer(const MLValue& mlvalue, const NodeArg& constant_node_arg, - ONNX_NAMESPACE::TensorProto& tensorproto); + ONNX_NAMESPACE::TensorProto& tensorproto) const; }; } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 1b46aea1da..23f9718004 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -16,8 +16,7 @@ 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) { +std::string GenerateRuleBasedTransformerName(TransformerLevel level) { return "Level" + std::to_string(static_cast(level)) + "_RuleBasedTransformer"; } @@ -28,7 +27,6 @@ std::vector> GenerateRewriteRules(TransformerLevel case TransformerLevel::Level1: rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); - rules.push_back(std::make_unique()); break; case TransformerLevel::Level2: @@ -77,8 +75,10 @@ std::vector> GenerateTransformers(TransformerL switch (level) { case TransformerLevel::Level1: { std::unordered_set l1_execution_providers = {}; + + transformers.emplace_back(std::make_unique(l1_execution_providers)); + rule_transformer = GenerateRuleBasedGraphTransformer(level, transformers_and_rules_to_enable, l1_execution_providers); - // At the moment, we have only a rule-based transformer for Level1. } break; case TransformerLevel::Level2: { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 039a4bac94..a2d96d6730 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -88,10 +88,9 @@ TEST(GraphTransformationTests, ConstantFolding1) { std::map op_to_count = CountOpsInGraph(graph); ASSERT_TRUE(op_to_count["Unsqueeze"] == 2); - 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); + graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level1); + ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1).IsOK()); op_to_count = CountOpsInGraph(graph); diff --git a/onnxruntime/test/optimizer/graph_transform_utils_test.cc b/onnxruntime/test/optimizer/graph_transform_utils_test.cc index 336175d011..8d3e2a0d99 100644 --- a/onnxruntime/test/optimizer/graph_transform_utils_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_utils_test.cc @@ -55,14 +55,20 @@ TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers) { TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers_CustomList) { // custom list of rules and transformers std::string l1_rule1 = "EliminateIdentity"; - std::string l1_rule2 = "ConstantFolding"; + std::string l1_transformer = "ConstantFolding"; std::string l2_transformer = "ConvAddFusion"; - std::vector custom_list = {l1_rule1, l1_rule2, l2_transformer}; + std::vector custom_list = {l1_rule1, l1_transformer, 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->RulesCount() == 2); + ASSERT_TRUE(transformers.size() == 2); + auto l1_rule_transformer_name = transformer_utils::GenerateRuleBasedTransformerName(TransformerLevel::Level1); + RuleBasedGraphTransformer* rule_transformer = nullptr; + for (const auto& transformer : transformers) { + if (transformer->Name() == l1_rule_transformer_name) { + rule_transformer = dynamic_cast(transformers[0].get()); + } + } + ASSERT_TRUE(rule_transformer && rule_transformer->RulesCount() == 1); transformers = transformer_utils::GenerateTransformers(TransformerLevel::Level2, custom_list); ASSERT_TRUE(transformers.size() == 1);