remove const cast

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<NodeArg*>& MutableOutputDefs() noexcept {
+    return definitions_.output_defs;
+  }
+
   /** Gets the count of arguments for each of the Node's explicit inputs. */
   const std::vector<int>& InputArgCount() const noexcept { return definitions_.input_arg_count; }
 

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<const onnxruntime::NodeArg*, onnxruntime::NodeArg*> replacement_defs;
     bool casted = false;
     for (auto input : inputs) {
       if (NeedInsertCast(node, input)) {
-        auto src_arg = const_cast<onnxruntime::NodeArg*>(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<MLFloat16>() &&
           casted) {
         //insert cast op to cast output back to float16
-        auto dst_arg = const_cast<onnxruntime::NodeArg*>(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<onnxruntime::NodeArg*>(output_node.OutputDefs()[0])] =
-                const_cast<onnxruntime::NodeArg*>(input);
+            replacement_defs[const_cast<onnxruntime::NodeArg*>(output_node.OutputDefs()[0])] = input;
             for (auto next_it = output_node.OutputNodesBegin(); next_it != output_node.OutputNodesEnd(); ++next_it) {
               const_cast<onnxruntime::Node*>(&(*next_it))->ReplaceDefs(replacement_defs);
             }

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<onnxruntime::NodeArg*>(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 <typename NodeArgSetType>

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<onnxruntime::Node*, NodeCompare> provider_nodes_;
-  std::set<const onnxruntime::NodeArg*, NodeArgCompare> non_provider_input_defs_;   // all input defs of non-provider nodes
-  std::set<const onnxruntime::NodeArg*, NodeArgCompare> non_provider_output_defs_;  // all output defs of non-provider nodes
-  std::set<const onnxruntime::NodeArg*, NodeArgCompare> provider_input_defs_;       // all input defs of provider nodes that should be in provider allocator
-  std::set<const onnxruntime::NodeArg*, NodeArgCompare> provider_output_defs_;      // all output defs of provider nodes that should be in provider allocator
+  std::set<const onnxruntime::NodeArg*, NodeArgCompare> non_provider_input_defs_;  // all input defs of non-provider nodes
+  std::set<onnxruntime::NodeArg*, NodeArgCompare> non_provider_output_defs_;       // all output defs of non-provider nodes
+  std::set<const onnxruntime::NodeArg*, NodeArgCompare> provider_input_defs_;      // all input defs of provider nodes that should be in provider allocator
+  std::set<onnxruntime::NodeArg*, NodeArgCompare> provider_output_defs_;           // all output defs of provider nodes that should be in provider allocator
   std::map<const onnxruntime::NodeArg*, onnxruntime::NodeArg*> replacements_;
   onnxruntime::Graph& graph_;
   std::string provider_;

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<NodeArg*>(conv_output_def);
+          def = conv_output_def;
         }
       }
     }

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<NodeArg*>(conv_output_def);
+          def = conv_output_def;
         }
       }
     }

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<NodeArg*>(conv_output_def);
+          def = conv_output_def;
         }
       }
     }

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<NodeArg*>(input_def);
+          def = input_def;
         }
       }
     }