diff --git a/include/onnxruntime/core/graph/graph_base.h b/include/onnxruntime/core/graph/graph_base.h index d7fecefb79..3e31fc5c8d 100644 --- a/include/onnxruntime/core/graph/graph_base.h +++ b/include/onnxruntime/core/graph/graph_base.h @@ -130,6 +130,11 @@ class Node { return definitions_.input_defs; } + /** Gets a modifiable collection of the Node's output definitions. */ + std::vector& MutableOutputDefs() noexcept { + return definitions_.output_defs; + } + /** Gets the count of arguments for each of the Node's explicit inputs. */ const std::vector& InputArgCount() const noexcept { return definitions_.input_arg_count; } diff --git a/onnxruntime/core/framework/insert_cast_transformer.cc b/onnxruntime/core/framework/insert_cast_transformer.cc index cf46bb5705..eb875f4218 100644 --- a/onnxruntime/core/framework/insert_cast_transformer.cc +++ b/onnxruntime/core/framework/insert_cast_transformer.cc @@ -106,12 +106,12 @@ Status InsertCastTransformer::Apply(onnxruntime::Graph& graph, bool& modified) c if (!node) return Status(ONNXRUNTIME, INVALID_ARGUMENT); - auto& inputs = node->InputDefs(); + auto& inputs = node->MutableInputDefs(); std::map replacement_defs; bool casted = false; for (auto input : inputs) { if (NeedInsertCast(node, input)) { - auto src_arg = const_cast(input); + auto src_arg = input; if (input_def_updates.count(src_arg)) { replacement_defs[src_arg] = input_def_updates[src_arg]; } else { @@ -136,7 +136,7 @@ Status InsertCastTransformer::Apply(onnxruntime::Graph& graph, bool& modified) c node->SetExecutionProviderType(kCpuExecutionProvider); } - auto& outputs = node->OutputDefs(); + auto& outputs = node->MutableOutputDefs(); for (auto output : outputs) { // todo: check is the kernel available // here is based on the assumption that if we cast a cpu op's input from float16 to float @@ -146,7 +146,7 @@ Status InsertCastTransformer::Apply(onnxruntime::Graph& graph, bool& modified) c DataTypeImpl::TypeFromProto(*output->TypeAsProto()) == DataTypeImpl::GetTensorType() && casted) { //insert cast op to cast output back to float16 - auto dst_arg = const_cast(output); + auto dst_arg = output; auto src_arg = AddCastNode(graph, id_generator, dst_arg, @@ -172,7 +172,7 @@ Status InsertCastTransformer::Apply(onnxruntime::Graph& graph, bool& modified) c // remove those two cast. auto src_type = node.InputDefs()[0]->Type(); auto dst_type = node.OutputDefs()[0]->Type(); - auto input = node.InputDefs()[0]; + auto input = node.MutableInputDefs()[0]; int child_removed = 0; int num_child = 0; for (auto it = node.OutputNodesBegin(); it != node.OutputNodesEnd(); ++it) { @@ -183,8 +183,7 @@ Status InsertCastTransformer::Apply(onnxruntime::Graph& graph, bool& modified) c if (src_type == dst_type1 && src_type1 == dst_type) { //node *it's output's follower could be linked with node's input. replacement_defs.clear(); - replacement_defs[const_cast(output_node.OutputDefs()[0])] = - const_cast(input); + replacement_defs[const_cast(output_node.OutputDefs()[0])] = input; for (auto next_it = output_node.OutputNodesBegin(); next_it != output_node.OutputNodesEnd(); ++next_it) { const_cast(&(*next_it))->ReplaceDefs(replacement_defs); } diff --git a/onnxruntime/core/framework/transformer_memcpy.cc b/onnxruntime/core/framework/transformer_memcpy.cc index f0d7586edd..ebff869a0e 100644 --- a/onnxruntime/core/framework/transformer_memcpy.cc +++ b/onnxruntime/core/framework/transformer_memcpy.cc @@ -80,42 +80,47 @@ void TransformerMemcpyImpl::ProcessDefs(onnxruntime::Node& node, const KernelReg return Status::OK(); }) .IsOK()); - ONNXRUNTIME_ENFORCE(onnxruntime::Node::ForEachWithIndex( - node.OutputDefs(), - [this, &output_mem_types](const onnxruntime::NodeArg& arg, size_t index) { - if (output_mem_types && MemTypeOnCpuExplicitly(*output_mem_types, index)) - non_provider_output_defs_.insert(&arg); - else - provider_output_defs_.insert(&arg); - return Status::OK(); - }) - .IsOK()); + auto& output_defs = node.MutableOutputDefs(); + for (size_t i = 0; i < output_defs.size(); ++i) { + auto arg = output_defs[i]; + if (!arg->Exists()) + continue; + + if (output_mem_types && MemTypeOnCpuExplicitly(*output_mem_types, i)) + non_provider_output_defs_.insert(arg); + else + provider_output_defs_.insert(arg); + } } else { // TODO: copy between devices? i.e. multiple GPUs if (node.GetExecutionProviderType() != onnxruntime::kCpuExecutionProvider && !node.GetExecutionProviderType().empty()) { ONNXRUNTIME_THROW("Execution type '", node.GetExecutionProviderType(), "' doesn't support memcpy "); } - node.ForEachDef([this](const onnxruntime::NodeArg& arg, bool is_input) { - if (is_input) - non_provider_input_defs_.insert(&arg); - else - non_provider_output_defs_.insert(&arg); - }); + + for (const auto* arg : node.InputDefs()) { + if (arg->Exists()) + non_provider_input_defs_.insert(arg); + } + + for (const auto* arg : node.ImplicitInputDefs()) { + if (arg->Exists()) + non_provider_input_defs_.insert(arg); + } + + for (auto* arg : node.MutableOutputDefs()) { + if (arg->Exists()) + non_provider_output_defs_.insert(arg); + } } } -void TransformerMemcpyImpl::AddCopyNode(const onnxruntime::NodeArg* arg, bool is_input) { - // TODO: eliminate the const-cast. - // The current graph API seems to only allow us to get a "const NodeArg*", but - // then we need to pass in a "NodeArg*" when we create a new node. - auto original_arg = const_cast(arg); - +void TransformerMemcpyImpl::AddCopyNode(onnxruntime::NodeArg* arg, bool is_input) { // create unique name for new def - std::string new_def_name = graph_.GenerateNodeArgName(original_arg->Name() + "_" + provider_); + std::string new_def_name = graph_.GenerateNodeArgName(arg->Name() + "_" + provider_); - auto* new_arg = &graph_.GetOrCreateNodeArg(new_def_name, original_arg->TypeAsProto()); - auto* src_arg = is_input ? original_arg : new_arg; - auto* dst_arg = is_input ? new_arg : original_arg; + auto* new_arg = &graph_.GetOrCreateNodeArg(new_def_name, arg->TypeAsProto()); + auto* src_arg = is_input ? arg : new_arg; + auto* dst_arg = is_input ? new_arg : arg; // create unique name for copy node std::string new_node_name = graph_.GenerateNodeName("Memcpy"); @@ -127,7 +132,7 @@ void TransformerMemcpyImpl::AddCopyNode(const onnxruntime::NodeArg* arg, bool is new_node.SetExecutionProviderType(provider_); // only add copy-node here; renaming references happens later - replacements_.insert(std::make_pair(original_arg, new_arg)); + replacements_.insert(std::make_pair(arg, new_arg)); } template diff --git a/onnxruntime/core/framework/transformer_memcpy.h b/onnxruntime/core/framework/transformer_memcpy.h index c2fbb7ba26..97eacb1091 100644 --- a/onnxruntime/core/framework/transformer_memcpy.h +++ b/onnxruntime/core/framework/transformer_memcpy.h @@ -20,7 +20,7 @@ class TransformerMemcpyImpl { private: void ProcessDefs(onnxruntime::Node& node, const KernelRegistryManager& kernel_registries); - void AddCopyNode(const onnxruntime::NodeArg* arg, bool is_input); + void AddCopyNode(onnxruntime::NodeArg* arg, bool is_input); void ProcessInitializers(); private: @@ -41,10 +41,10 @@ class TransformerMemcpyImpl { }; std::set provider_nodes_; - std::set non_provider_input_defs_; // all input defs of non-provider nodes - std::set non_provider_output_defs_; // all output defs of non-provider nodes - std::set provider_input_defs_; // all input defs of provider nodes that should be in provider allocator - std::set provider_output_defs_; // all output defs of provider nodes that should be in provider allocator + std::set non_provider_input_defs_; // all input defs of non-provider nodes + std::set non_provider_output_defs_; // all output defs of non-provider nodes + std::set provider_input_defs_; // all input defs of provider nodes that should be in provider allocator + std::set provider_output_defs_; // all output defs of provider nodes that should be in provider allocator std::map replacements_; onnxruntime::Graph& graph_; std::string provider_; diff --git a/onnxruntime/core/graph/conv_add_fusion.cc b/onnxruntime/core/graph/conv_add_fusion.cc index 1218115b5d..89333e275b 100644 --- a/onnxruntime/core/graph/conv_add_fusion.cc +++ b/onnxruntime/core/graph/conv_add_fusion.cc @@ -93,17 +93,16 @@ Status ConvAddFusion::Apply(onnxruntime::Graph& graph, bool& modified) const { // Replace the input of the node following add node const NodeArg* add_output_def = add_node.OutputDefs()[0]; - const NodeArg* conv_output_def = conv_node.OutputDefs()[0]; + NodeArg* conv_output_def = conv_node.MutableOutputDefs()[0]; for (auto it = add_node.OutputNodesBegin(); it != add_node.OutputNodesEnd(); ++it) { auto output_node = graph.GetNode((*it).Index()); if (!output_node) { return Status(ONNXRUNTIME, INVALID_ARGUMENT); } - auto& input_defs = output_node->MutableInputDefs(); for (auto& def : input_defs) { if (def == add_output_def) { - def = const_cast(conv_output_def); + def = conv_output_def; } } } diff --git a/onnxruntime/core/graph/conv_bn_fusion.cc b/onnxruntime/core/graph/conv_bn_fusion.cc index a416133d58..a6b9b4ce9a 100644 --- a/onnxruntime/core/graph/conv_bn_fusion.cc +++ b/onnxruntime/core/graph/conv_bn_fusion.cc @@ -155,7 +155,7 @@ Status ConvBNFusion::Apply(onnxruntime::Graph& graph, bool& modified) const { // Replace the input of the nodes following batch normalization node const NodeArg* bn_output_def = bn_node.OutputDefs()[0]; - const NodeArg* conv_output_def = conv_node.OutputDefs()[0]; + NodeArg* conv_output_def = conv_node.MutableOutputDefs()[0]; for (auto it = bn_node.OutputNodesBegin(); it != bn_node.OutputNodesEnd(); ++it) { auto output_node = graph.GetNode((*it).Index()); if (!output_node) { @@ -165,7 +165,7 @@ Status ConvBNFusion::Apply(onnxruntime::Graph& graph, bool& modified) const { auto& input_defs = output_node->MutableInputDefs(); for (auto& def : input_defs) { if (def == bn_output_def) { - def = const_cast(conv_output_def); + def = conv_output_def; } } } diff --git a/onnxruntime/core/graph/conv_mul_fusion.cc b/onnxruntime/core/graph/conv_mul_fusion.cc index e48da544df..447a57fb80 100644 --- a/onnxruntime/core/graph/conv_mul_fusion.cc +++ b/onnxruntime/core/graph/conv_mul_fusion.cc @@ -94,7 +94,7 @@ Status ConvMulFusion::Apply(onnxruntime::Graph& graph, bool& modified) const { // Replace the input of the node following mul node const NodeArg* mul_output_def = mul_node.OutputDefs()[0]; - const NodeArg* conv_output_def = conv_node.OutputDefs()[0]; + NodeArg* conv_output_def = conv_node.MutableOutputDefs()[0]; for (auto it = mul_node.OutputNodesBegin(); it != mul_node.OutputNodesEnd(); ++it) { auto output_node = graph.GetNode((*it).Index()); if (!output_node) { @@ -104,7 +104,7 @@ Status ConvMulFusion::Apply(onnxruntime::Graph& graph, bool& modified) const { auto& input_defs = output_node->MutableInputDefs(); for (auto& def : input_defs) { if (def == mul_output_def) { - def = const_cast(conv_output_def); + def = conv_output_def; } } } diff --git a/onnxruntime/core/graph/unsqueeze_elimination.cc b/onnxruntime/core/graph/unsqueeze_elimination.cc index f7b32f9643..31b0347c24 100644 --- a/onnxruntime/core/graph/unsqueeze_elimination.cc +++ b/onnxruntime/core/graph/unsqueeze_elimination.cc @@ -78,7 +78,7 @@ Status UnsqueezeElimination::Apply(onnxruntime::Graph& graph, bool& modified) co auto& input_defs = output_node->MutableInputDefs(); for (auto& def : input_defs) { if (def == output_def) { - def = const_cast(input_def); + def = input_def; } } }