mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-30 23:18:20 +00:00
243 lines
10 KiB
C++
243 lines
10 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/graph/function_impl.h"
|
|
#include "core/graph/graph_viewer.h"
|
|
#include "core/graph/model.h"
|
|
#include "onnx/shape_inference/implementation.h"
|
|
|
|
namespace onnxruntime {
|
|
void TypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
|
|
std::unique_ptr<ONNX_NAMESPACE::OpSchema>& op_schema_,
|
|
/*out*/
|
|
std::unordered_map<std::string, int>& input_name_idx_map,
|
|
std::unordered_map<std::string, int>& output_name_idx_map) {
|
|
std::vector<std::pair<std::string, std::string>> input_types_list(onnx_func_proto_->input_size());
|
|
std::vector<std::pair<std::string, std::string>> output_types_list(onnx_func_proto_->output_size());
|
|
std::unordered_map<std::string, std::vector<std::string>> type_constraint_map;
|
|
for (int i = 0; i < onnx_func_proto_->input_size(); ++i) {
|
|
input_name_idx_map[onnx_func_proto_->input().Get(i)] = i;
|
|
}
|
|
for (int i = 0; i < onnx_func_proto_->output_size(); ++i) {
|
|
output_name_idx_map[onnx_func_proto_->output().Get(i)] = i;
|
|
}
|
|
auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
|
|
for (auto& node : onnx_func_proto_->node()) {
|
|
const auto node_op_schema = schema_registry->GetSchema(node.op_type(), (int)onnx_func_proto_->since_version(), node.domain());
|
|
for (int i = 0; i < node.input_size(); ++i) {
|
|
auto& in_name = node.input().Get(i);
|
|
auto iter = input_name_idx_map.find(in_name);
|
|
if (iter != input_name_idx_map.end()) {
|
|
int idx = iter->second;
|
|
const auto& p = node_op_schema->inputs().at(i);
|
|
std::string type_str = p.GetTypeStr() + "in" + std::to_string(i);
|
|
input_types_list[idx] = std::make_pair(in_name, type_str);
|
|
if (!type_constraint_map.count(type_str)) {
|
|
for (auto s : p.GetTypes()) {
|
|
type_constraint_map[type_str].emplace_back(*s);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
for (int i = 0; i < node.output_size(); ++i) {
|
|
auto& out_name = node.output().Get(i);
|
|
auto iter = output_name_idx_map.find(out_name);
|
|
if (iter != output_name_idx_map.end()) {
|
|
int idx = iter->second;
|
|
const auto& p = node_op_schema->outputs().at(i);
|
|
std::string type_str = p.GetTypeStr() + "out" + std::to_string(i);
|
|
output_types_list[idx] = std::make_pair(out_name, type_str);
|
|
if (!type_constraint_map.count(type_str)) {
|
|
for (auto s : p.GetTypes()) {
|
|
type_constraint_map[type_str].emplace_back(*s);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
int i = 0;
|
|
for (auto& input : input_types_list) {
|
|
op_schema_->Input(i, input.first, "", input.second);
|
|
++i;
|
|
}
|
|
i = 0;
|
|
for (auto& output : output_types_list) {
|
|
op_schema_->Output(i, output.first, "", output.second);
|
|
++i;
|
|
}
|
|
|
|
for (auto& tc : type_constraint_map) {
|
|
op_schema_->TypeConstraint(tc.first, tc.second, "");
|
|
}
|
|
}
|
|
|
|
FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
|
|
std::unique_ptr<IndexedSubGraph> customized_func)
|
|
: parent_graph_(&graph), onnx_func_proto_{nullptr} {
|
|
customized_func_body_ = std::move(customized_func);
|
|
auto meta_def = customized_func_body_->GetMetaDef();
|
|
op_schema_ = std::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;
|
|
for (auto& input : meta_def->inputs) {
|
|
auto input_type = parent_graph_->GetNodeArg(input)->Type();
|
|
op_schema_->Input(i, input, "", *input_type);
|
|
++i;
|
|
}
|
|
i = 0;
|
|
for (auto& output : meta_def->outputs) {
|
|
auto output_type = parent_graph_->GetNodeArg(output)->Type();
|
|
op_schema_->Output(i, output, "", *output_type);
|
|
++i;
|
|
}
|
|
op_schema_->Finalize();
|
|
//construct body
|
|
body_ = std::make_unique<onnxruntime::Model>("fused_function_subgraph", false, onnxruntime::ModelMetaData(),
|
|
IOnnxRuntimeOpSchemaRegistryList({graph.GetSchemaRegistry()}),
|
|
graph.DomainToVersionMap());
|
|
|
|
auto& sub_graph = body_->MainGraph();
|
|
//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) {
|
|
auto node = parent_graph_->GetNode(node_index);
|
|
std::vector<onnxruntime::NodeArg*> inputs, outputs;
|
|
for (auto input : node->InputDefs()) {
|
|
auto& n_input = sub_graph.GetOrCreateNodeArg(input->Name(), input->TypeAsProto());
|
|
inputs.push_back(&n_input);
|
|
}
|
|
for (auto output : node->OutputDefs()) {
|
|
auto& n_output = sub_graph.GetOrCreateNodeArg(output->Name(), output->TypeAsProto());
|
|
outputs.push_back(&n_output);
|
|
}
|
|
sub_graph.AddNode(node->Name(), node->OpType(), node->Description(), inputs, outputs, &node->GetAttributes(), node->Domain());
|
|
}
|
|
|
|
for (auto input : meta_def->inputs) {
|
|
const onnx::TensorProto* initializer = nullptr;
|
|
if (graph.GetInitializedTensor(input, initializer)) {
|
|
sub_graph.AddInitializedTensor(*initializer);
|
|
}
|
|
}
|
|
|
|
//TODO: if we reuse the nodes in parent graph, maybe we don't need to resolve it.
|
|
ORT_ENFORCE(sub_graph.Resolve().IsOK());
|
|
}
|
|
|
|
FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
|
|
const onnxruntime::NodeIndex& node_index,
|
|
const ONNX_NAMESPACE::FunctionProto* onnx_func_proto)
|
|
: parent_graph_(&graph) {
|
|
onnx_func_proto_ = onnx_func_proto;
|
|
auto node_in_parent_graph = parent_graph_->GetNode(node_index);
|
|
op_schema_ = std::make_unique<onnx::OpSchema>();
|
|
op_schema_->SetName(onnx_func_proto_->name());
|
|
op_schema_->SetDomain(onnx_func_proto_->node().Get(0).domain());
|
|
op_schema_->SetDoc(onnx_func_proto_->doc_string());
|
|
op_schema_->SinceVersion((ONNX_NAMESPACE::OperatorSetVersion)onnx_func_proto_->since_version());
|
|
std::unordered_map<std::string, int> input_name_idx_map;
|
|
std::unordered_map<std::string, int> output_name_idx_map;
|
|
TypeConstraintHelper(onnx_func_proto_, this->op_schema_, input_name_idx_map, output_name_idx_map);
|
|
|
|
op_schema_->TypeAndShapeInferenceFunction(
|
|
[this](ONNX_NAMESPACE::InferenceContext& ctx) {
|
|
auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
|
|
const ONNX_NAMESPACE::FunctionProto* func_ptr = this->GetFuncProto();
|
|
if (nullptr != func_ptr) {
|
|
ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(*func_ptr, schema_registry, ctx);
|
|
}
|
|
});
|
|
|
|
op_schema_->Finalize();
|
|
//construct body
|
|
std::unordered_map<std::string, int> domain_to_version;
|
|
//TODO: set correct domain and version
|
|
domain_to_version[onnxruntime::kOnnxDomain] = (int)onnx_func_proto_->since_version();
|
|
body_ = std::make_unique<onnxruntime::Model>(onnx_func_proto_->name(), false, onnxruntime::ModelMetaData(),
|
|
IOnnxRuntimeOpSchemaRegistryList(), domain_to_version);
|
|
auto& sub_graph = body_->MainGraph();
|
|
// Add node and node args into subgraph
|
|
// The subgraph preserved the input/output tensor names
|
|
// in the parent graph for later inlining purpose
|
|
auto attr_map = node_in_parent_graph->GetAttributes();
|
|
for (auto& node : onnx_func_proto_->node()) {
|
|
std::vector<onnxruntime::NodeArg*> inputs, outputs;
|
|
for (int idx = 0; idx < node.input_size(); ++idx) {
|
|
std::string tensor_name = node.input().Get(idx);
|
|
auto iter = input_name_idx_map.find(tensor_name);
|
|
if (iter != input_name_idx_map.end()) {
|
|
// Preserving NodeArg and input/output names
|
|
ONNX_NAMESPACE::NodeProto temp_node_proto;
|
|
node_in_parent_graph->ToProto(temp_node_proto);
|
|
const onnxruntime::NodeArg* node_arg = parent_graph_->GetNodeArg(temp_node_proto.input().Get(input_name_idx_map[tensor_name]));
|
|
auto& n_input = sub_graph.GetOrCreateNodeArg(
|
|
temp_node_proto.input().Get(iter->second), node_arg->TypeAsProto());
|
|
inputs.push_back(&n_input);
|
|
} else {
|
|
auto& n_input = sub_graph.GetOrCreateNodeArg(
|
|
tensor_name + "_" + std::to_string(node_index), nullptr);
|
|
inputs.push_back(&n_input);
|
|
}
|
|
}
|
|
for (int idx = 0; idx < node.output_size(); ++idx) {
|
|
std::string tensor_name = node.output().Get(idx);
|
|
auto iter = output_name_idx_map.find(tensor_name);
|
|
if (iter != output_name_idx_map.end()) {
|
|
// Preserving NodeArg and input/output names
|
|
ONNX_NAMESPACE::NodeProto temp_node_proto;
|
|
node_in_parent_graph->ToProto(temp_node_proto);
|
|
const onnxruntime::NodeArg* node_arg = parent_graph_->GetNodeArg(temp_node_proto.output().Get(output_name_idx_map[tensor_name]));
|
|
auto& n_output = sub_graph.GetOrCreateNodeArg(
|
|
temp_node_proto.output().Get(iter->second), node_arg->TypeAsProto());
|
|
outputs.push_back(&n_output);
|
|
} else {
|
|
auto& n_output = sub_graph.GetOrCreateNodeArg(
|
|
tensor_name + "_" + std::to_string(node_index), nullptr);
|
|
outputs.push_back(&n_output);
|
|
}
|
|
}
|
|
|
|
onnxruntime::NodeAttributes new_attr_map;
|
|
for (auto& attr : node.attribute()) {
|
|
if (attr.has_ref_attr_name()) {
|
|
if (attr_map.count(attr.ref_attr_name())) {
|
|
new_attr_map[attr.name()] = attr_map[attr.ref_attr_name()];
|
|
}
|
|
} else {
|
|
new_attr_map[attr.name()] = attr;
|
|
}
|
|
}
|
|
sub_graph.AddNode(node.name() + "_" + std::to_string(node_index), node.op_type(), node.doc_string(), inputs, outputs, &new_attr_map, node.domain());
|
|
}
|
|
auto status = sub_graph.Resolve();
|
|
ORT_ENFORCE(status.IsOK());
|
|
}
|
|
|
|
FunctionImpl::~FunctionImpl() = default;
|
|
|
|
const ONNX_NAMESPACE::OpSchema& FunctionImpl::OpSchema() const {
|
|
return *op_schema_;
|
|
}
|
|
|
|
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) {
|
|
return std::make_unique<FunctionImpl>(graph, std::move(customized_func));
|
|
}
|
|
} // namespace onnxruntime
|