mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Convert constant folding to a transformer (#866)
This commit is contained in:
parent
5bcd77e488
commit
1b7d1f2645
7 changed files with 102 additions and 75 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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>();
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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: {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue