Convert constant folding to a transformer (#866)

This commit is contained in:
Konstantinos Karanasos 2019-04-29 18:12:49 -07:00 committed by GitHub
parent 5bcd77e488
commit 1b7d1f2645
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 102 additions and 75 deletions

View file

@ -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.

View file

@ -23,5 +23,8 @@ 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

View file

@ -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<int> 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<int> 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<MLValue> 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<MLValue> 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<Tensor>();

View file

@ -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<std::string> TargetOpTypes() const noexcept override {
return std::vector<std::string>();
}
ConstantFolding(const std::unordered_set<std::string>& 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<std::string> 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

View file

@ -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<uint32_t>(level)) + "_RuleBasedTransformer";
}
@ -28,7 +27,6 @@ std::vector<std::unique_ptr<RewriteRule>> GenerateRewriteRules(TransformerLevel
case TransformerLevel::Level1:
rules.push_back(std::make_unique<EliminateIdentity>());
rules.push_back(std::make_unique<EliminateSlice>());
rules.push_back(std::make_unique<ConstantFolding>());
break;
case TransformerLevel::Level2:
@ -77,8 +75,10 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
switch (level) {
case TransformerLevel::Level1: {
std::unordered_set<std::string> l1_execution_providers = {};
transformers.emplace_back(std::make_unique<ConstantFolding>(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: {

View file

@ -88,10 +88,9 @@ TEST(GraphTransformationTests, ConstantFolding1) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Unsqueeze"] == 2);
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);
graph_transformation_mgr.Register(std::make_unique<ConstantFolding>(), TransformerLevel::Level1);
ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1).IsOK());
op_to_count = CountOpsInGraph(graph);

View file

@ -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<std::string> custom_list = {l1_rule1, l1_rule2, l2_transformer};
std::vector<std::string> 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<RuleBasedGraphTransformer*>(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<RuleBasedGraphTransformer*>(transformers[0].get());
}
}
ASSERT_TRUE(rule_transformer && rule_transformer->RulesCount() == 1);
transformers = transformer_utils::GenerateTransformers(TransformerLevel::Level2, custom_list);
ASSERT_TRUE(transformers.size() == 1);