The IndexedSubGraph is used to create the Function body, but after that is invalid as the nodes it referred to have been removed from the main Graph. As such there's no need to store it in the FunctionImpl instance. (#5669)

This commit is contained in:
Scott McKay 2020-11-05 17:21:56 +10:00 committed by GitHub
parent 941e3a69f9
commit 2127a229d7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 40 additions and 48 deletions

View file

@ -26,9 +26,6 @@ class Function {
/** Gets the Graph instance for the Function body subgraph. */
virtual const onnxruntime::Graph& Body() const = 0;
/** Gets the IndexedSubGraph for the Function. */
virtual const IndexedSubGraph& GetIndexedSubGraph() const = 0;
};
/**
@ -37,6 +34,6 @@ Create a new Function instance.
@param customized_func the IndexedSubGraph to use for the Function.
*/
std::unique_ptr<Function> MakeFunction(const onnxruntime::Graph& graph,
std::unique_ptr<IndexedSubGraph> customized_func,
const IndexedSubGraph& nodes_to_fuse,
const logging::Logger& logger);
} // namespace onnxruntime

View file

@ -906,7 +906,7 @@ class Graph {
@param fused_node_name The name for the new Node.
@returns Node with fused subgraph.
*/
Node& FuseSubGraph(std::unique_ptr<IndexedSubGraph> sub_graph, const std::string& fused_node_name);
Node& FuseSubGraph(const IndexedSubGraph& sub_graph, const std::string& fused_node_name);
/**
Directly insert the nodes in the function Node provided into this Graph.

View file

@ -60,28 +60,25 @@ KernelDefBuilder& BuildFusedKernelDef(KernelDefBuilder& builder, const onnxrunti
* \param count A counter for generating fused node names. Should be unique within this subgraph
* \return Fused node. Return nullptr if there is no fuse
*/
static Node* PlaceNode(Graph& graph, std::unique_ptr<IndexedSubGraph> capability,
const KernelRegistryManager& kernel_registry_mgr, const std::string& provider_type, int& count) {
if (nullptr == capability) {
return nullptr;
}
if (nullptr == capability->GetMetaDef()) {
static Node* PlaceNode(Graph& graph, const IndexedSubGraph& capability,
const KernelRegistryManager& kernel_registry_mgr, const std::string& provider_type,
int& count) {
if (nullptr == capability.GetMetaDef()) {
// The <provider> can run a single node in the <graph> if not using meta-defs.
// A fused kernel is not supported in this case.
ORT_ENFORCE(1 == capability->nodes.size());
ORT_ENFORCE(1 == capability.nodes.size());
auto* node = graph.GetNode(capability->nodes[0]);
auto* node = graph.GetNode(capability.nodes[0]);
if (nullptr != node && node->GetExecutionProviderType().empty()) {
// The node was not fused or assigned. Assign it to this <provider>.
node->SetExecutionProviderType(provider_type);
}
} else {
// The <provider> can run a fused <sub_graph> in the <graph>.
ORT_ENFORCE(nullptr != capability->GetMetaDef());
ORT_ENFORCE(nullptr != capability.GetMetaDef());
// Check whether any node in the <sub_graph> was already assigned.
bool sub_graph_available_for_assignment = true;
for (auto node_index : capability->nodes) {
for (auto node_index : capability.nodes) {
auto node = graph.GetNode(node_index);
if (nullptr == node || !node->GetExecutionProviderType().empty()) {
// The node was fused or assigned, so that the whole sub-graph will not be assigned to this <provider>
@ -92,9 +89,9 @@ static Node* PlaceNode(Graph& graph, std::unique_ptr<IndexedSubGraph> capability
}
if (sub_graph_available_for_assignment) {
std::ostringstream oss;
oss << provider_type << "_" << capability->GetMetaDef()->name << "_" << count++;
oss << provider_type << "_" << capability.GetMetaDef()->name << "_" << count++;
std::string node_name = oss.str();
auto& fused_node = graph.FuseSubGraph(std::move(capability), node_name);
auto& fused_node = graph.FuseSubGraph(capability, node_name);
fused_node.SetExecutionProviderType(provider_type);
// searching in kernel registries, if no kernel registered for the fused_node, use compile approach
if (!KernelRegistryManager::HasImplementationOf(kernel_registry_mgr, fused_node, provider_type)) {
@ -154,7 +151,11 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f
std::vector<std::unique_ptr<ComputeCapability>> capabilities =
provider->GetCapability(graph_viewer, kernel_registry_mgr_.GetKernelRegistriesByProviderType(provider->Type()));
for (auto& capability : capabilities) {
Node* n = PlaceNode(graph, std::move(capability->sub_graph), kernel_registry_mgr_, provider->Type(), count);
if (!capability || !capability->sub_graph) { // in theory an EP could return an empty value...
continue;
}
Node* n = PlaceNode(graph, *capability->sub_graph, kernel_registry_mgr_, provider->Type(), count);
if (n != nullptr) {
nodes_need_compile.push_back(n);
}
@ -197,9 +198,9 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f
if (nullptr == node_func) {
continue;
}
nodes_need_inline.push_back(&node);
nodes_need_inline.push_back(&node);
}
}
}
for (auto* node : nodes_need_inline) {
// If the node has a functionbody with no kernel and cannot be inlined

View file

@ -145,22 +145,22 @@ static void update_subgraphs_within_function_body(ONNX_NAMESPACE::GraphProto& su
}
FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
std::unique_ptr<IndexedSubGraph> customized_func,
const IndexedSubGraph& nodes_to_fuse,
const logging::Logger& logger)
: parent_graph_(&graph),
body_("fused_function_subgraph", false, onnxruntime::ModelMetaData(),
graph.ModelPath().ToPathString(),
IOnnxRuntimeOpSchemaRegistryList({graph.GetSchemaRegistry()}),
graph.DomainToVersionMap(), {}, logger) {
customized_func_body_ = std::move(customized_func);
auto& function_body_graph = body_.MainGraph();
auto meta_def = customized_func_body_->GetMetaDef();
auto meta_def = nodes_to_fuse.GetMetaDef();
op_schema_ = onnxruntime::make_unique<ONNX_NAMESPACE::OpSchema>();
op_schema_->SetName(meta_def->name);
op_schema_->SetDomain(meta_def->domain);
op_schema_->SetDoc(meta_def->doc_string);
op_schema_->SinceVersion(meta_def->since_version);
int i = 0;
std::vector<const NodeArg*> function_body_graph_inputs;
function_body_graph_inputs.resize(meta_def->inputs.size());
@ -172,6 +172,7 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
op_schema_->Input(i, input, "", *input_arg->Type());
++i;
}
i = 0;
std::vector<const NodeArg*> function_body_graph_outputs;
function_body_graph_outputs.resize(meta_def->outputs.size());
@ -182,14 +183,16 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
op_schema_->Output(i, output, "", *output_arg->Type());
++i;
}
op_schema_->Finalize();
function_body_graph.SetInputs(function_body_graph_inputs);
function_body_graph.SetOutputs(function_body_graph_outputs);
//Add node and node args
//TODO: for better performance, we could try to transfer the nodes in parent graph to sub-graph directly,
//instead of create new nodes.
for (auto& node_index : customized_func_body_->nodes) {
for (auto& node_index : nodes_to_fuse.nodes) {
auto node = parent_graph_->GetNode(node_index);
std::vector<onnxruntime::NodeArg*> inputs;
std::vector<onnxruntime::NodeArg*> outputs;
@ -394,17 +397,13 @@ const onnxruntime::Graph& FunctionImpl::Body() const {
return body_.MainGraph();
}
const IndexedSubGraph& FunctionImpl::GetIndexedSubGraph() const {
return *customized_func_body_;
}
const ONNX_NAMESPACE::FunctionProto* FunctionImpl::GetFuncProto() const {
return &onnx_func_proto_;
}
std::unique_ptr<Function> MakeFunction(const onnxruntime::Graph& graph,
std::unique_ptr<IndexedSubGraph> customized_func,
const IndexedSubGraph& nodes_to_fuse,
const logging::Logger& logger) {
return onnxruntime::make_unique<FunctionImpl>(graph, std::move(customized_func), logger);
return onnxruntime::make_unique<FunctionImpl>(graph, nodes_to_fuse, logger);
}
} // namespace onnxruntime

View file

@ -17,7 +17,7 @@ namespace onnxruntime {
class FunctionImpl final : public Function {
public:
FunctionImpl(const onnxruntime::Graph& graph,
std::unique_ptr<IndexedSubGraph> customized_func,
const IndexedSubGraph& nodes_to_fuse,
const logging::Logger& logger);
FunctionImpl(const onnxruntime::Graph& graph,
@ -31,13 +31,10 @@ class FunctionImpl final : public Function {
const onnxruntime::Graph& Body() const override;
const IndexedSubGraph& GetIndexedSubGraph() const override;
const ONNX_NAMESPACE::FunctionProto* GetFuncProto() const;
private:
const onnxruntime::Graph* const parent_graph_;
std::unique_ptr<IndexedSubGraph> customized_func_body_;
std::unique_ptr<ONNX_NAMESPACE::OpSchema> op_schema_;
onnxruntime::Model body_;
ONNX_NAMESPACE::FunctionProto onnx_func_proto_;

View file

@ -991,7 +991,7 @@ Graph::Graph(const Model& owning_model,
// For now we convert sparse_intializer to dense tensors
// since there are currently no supported ops that consume sparse
// initializers directly. We remove them from graph_proto. We will reconstitute them
// initializers directly. We remove them from graph_proto. We will reconstitute them
// when saving to ORT format to save space on disk.
if (graph_proto_->sparse_initializer_size() > 0) {
for (const auto& sparse_tensor : graph_proto_->sparse_initializer()) {
@ -1037,8 +1037,8 @@ Graph::Graph(const Model& owning_model,
auto p = name_to_initial_tensor_.emplace(tensor.name(), &tensor);
if (!p.second) {
LOGS(logger_, WARNING) << "Duplicate initializer (dense, sparse or ConstantNode): '" << tensor.name()
<< "' the model will use the latest encountered initializer"
<< ". Please, fix your model.";
<< "' the model will use the latest encountered initializer"
<< ". Please, fix your model.";
p.first->second = &tensor;
}
@ -3317,12 +3317,10 @@ IOnnxRuntimeOpSchemaCollectionPtr Graph::GetSchemaRegistry() const {
return schema_registry_;
}
Node& Graph::FuseSubGraph(std::unique_ptr<::onnxruntime::IndexedSubGraph> sub_graph,
const std::string& fused_node_name) {
ORT_ENFORCE(nullptr != sub_graph && nullptr != sub_graph->GetMetaDef());
auto func_meta_def = sub_graph->GetMetaDef();
Node& Graph::FuseSubGraph(const IndexedSubGraph& sub_graph, const std::string& fused_node_name) {
auto* func_meta_def = sub_graph.GetMetaDef();
ORT_ENFORCE(nullptr != func_meta_def);
std::vector<NodeArg*> input_args;
std::vector<NodeArg*> output_args;
for (auto& arg_name : func_meta_def->inputs) {
@ -3341,12 +3339,11 @@ Node& Graph::FuseSubGraph(std::unique_ptr<::onnxruntime::IndexedSubGraph> sub_gr
func_meta_def->domain);
fused_node.SetNodeType(Node::Type::Fused);
function_container_.emplace_back(MakeFunction(*this, std::move(sub_graph), logger_));
function_container_.emplace_back(MakeFunction(*this, sub_graph, logger_));
fused_node.SetFunctionBody(*function_container_.back());
// Remove nodes fused above.
auto& sub_graph_ref = function_container_.back()->GetIndexedSubGraph();
for (auto node_index : sub_graph_ref.nodes) {
for (auto node_index : sub_graph.nodes) {
auto node = GetNode(node_index);
if (nullptr == node) {
continue;
@ -3357,6 +3354,7 @@ Node& Graph::FuseSubGraph(std::unique_ptr<::onnxruntime::IndexedSubGraph> sub_gr
}
RemoveNode(node_index);
}
return fused_node;
}
@ -3578,8 +3576,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::experimental::fbs::Gr
auto p = name_to_initial_tensor_.emplace(initializer->name(), initializer);
if (!p.second) {
LOGS(logger_, WARNING) << "Duplicate initializer (dense or ConstantNode): '" << initializer->name()
<< "' the model will use the latest encountered initializer"
<< ". Please, fix your model.";
<< "' the model will use the latest encountered initializer"
<< ". Please, fix your model.";
p.first->second = initializer;
}
}