src/executor/tensorrt_pass.cc - mxnet - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 /*!
  * Copyright (c) 2018 by Contributors
  * \file tensorrt_pass.cc
  * \brief Replace TRT compatible subgraphs by TRT engines
  * \author Clement Fuji Tsang
  */

 #if MXNET_USE_TENSORRT

 #include <NvInfer.h>
 #include <mxnet/base.h>
 #include <mxnet/op_attr_types.h>
 #include <mxnet/operator.h>
 #include <nnvm/graph_attr_types.h>
 #include <onnx/onnx.pb.h>

 #include "../operator/contrib/nnvm_to_onnx-inl.h"
 #include "./exec_pass.h"
 #include "./onnx_to_tensorrt.h"

 namespace mxnet {
 namespace exec {

 using NodePtr = nnvm::NodePtr;

 /*!
  * \brief Custom graph class, which will contain bi-directional nodes
  * we need to compute DFS and reverse DFS for graph partitioning
  */
 class BidirectionalGraph {
  public:
   struct Node {
     nnvm::Node* nnvmptr;
     std::vector<Node*> inputs;
     std::vector<Node*> outputs;
   };
   std::vector<Node> nodes;
   std::unordered_map<nnvm::Node*, uint32_t> nnvm2nid;
   std::vector<Node*> outputs;
   static const std::unordered_set<std::string> unconditionalTRTop;

   explicit BidirectionalGraph(const Graph &g) {
     auto& idx = g.indexed_graph();
     auto num_nodes = idx.num_nodes();
     nodes.reserve(num_nodes);
     nnvm2nid.reserve(num_nodes);
     outputs.reserve(idx.outputs().size());
     DFSVisit(g.outputs, [this](const nnvm::NodePtr& n) {
       BidirectionalGraph::Node new_node;
       new_node.nnvmptr = n.get();
       nnvm2nid[n.get()] = static_cast<uint32_t>(nodes.size());
       nodes.emplace_back(std::move(new_node));
     });
     for (const auto& it : nnvm2nid) {
       nnvm::Node* nnvmnode = it.first;
       uint32_t nid = it.second;
       for (auto& n : nnvmnode->inputs) {
         uint32_t input_nid = nnvm2nid[n.node.get()];
         nodes[input_nid].outputs.emplace_back(&nodes[nid]);
         nodes[nid].inputs.emplace_back(&nodes[input_nid]);
       }
     }
     for (auto& e : g.outputs) {
       uint32_t nid = nnvm2nid[e.node.get()];
       outputs.emplace_back(&nodes[nid]);
     }
   }

   template <typename FVisit>
   void DFS(const std::vector<Node*>& heads, bool reverse, FVisit fvisit) {
     std::unordered_set<Node*> visited;
     std::vector<Node*> vec(heads.begin(), heads.end());
     visited.reserve(heads.size());
     while (!vec.empty()) {
       Node* vertex = vec.back();
       vec.pop_back();
       if (visited.count(vertex) == 0) {
         visited.insert(vertex);
         fvisit(vertex);
         std::vector<Node*> nexts = reverse ? vertex->inputs : vertex->outputs;
         for (Node* node : nexts) {
           if (visited.count(node) == 0) {
             vec.emplace_back(node);
           }
         }
       }
     }
   }

   using t_pairset = std::pair<std::unordered_set<Node*>, std::unordered_set<Node*>>;
   using t_pairvec = std::pair<std::vector<Node*>, std::vector<Node*>>;
   using t_uncomp_map = std::unordered_map<Node*, std::unordered_set<Node*>>;

   std::unordered_set<Node*> naive_grow_subgraph(Node* head,
                                                 std::unordered_set<Node*>* set_unused,
                                                 t_uncomp_map* uncomp_map) {
     std::unordered_set<Node*> subgraph;
     std::unordered_set<Node*> uncomp_set;
     std::deque<Node*> stack;
     stack.emplace_back(head);
     while (!stack.empty()) {
       Node* vertex = stack.back();
       stack.pop_back();
       if (set_unused->count(vertex) && !uncomp_set.count(vertex)) {
         set_unused->erase(vertex);
         subgraph.insert(vertex);
         uncomp_set.insert((*uncomp_map)[vertex].begin(), (*uncomp_map)[vertex].end());
         for (Node* input : vertex->inputs) {
           if (set_unused->count(input) && !uncomp_set.count(input)) {
             stack.emplace_back(input);
           }
         }
         for (Node* output : vertex->outputs) {
           if (set_unused->count(output) && !uncomp_set.count(output)) {
             stack.emplace_back(output);
           }
         }
       }
     }
     return subgraph;
   }

   std::vector<std::unordered_set<Node*>> get_subsets(
     std::unordered_map<std::string, NDArray>* const params_map) {
     std::vector<std::unordered_set<Node*>> subgraphs;
     std::unordered_set<Node*> set_nonTRTnodes;
     std::unordered_set<Node*> set_allnodes(nodes.size());
     std::vector<t_pairset> separation_sets;
     for (Node& node : nodes) {
       if (!IsTRTCompatible(node.nnvmptr)) {
         set_nonTRTnodes.insert(&node);
         std::unordered_set<Node*> in_graph;
         std::unordered_set<Node*> out_graph;
         std::vector<Node*> dummy_head;
         dummy_head.emplace_back(&node);
         DFS(dummy_head, false, [&out_graph](Node* node) {
           out_graph.insert(node);
         });
         DFS(dummy_head, true, [&in_graph](Node* node) {
           in_graph.insert(node);
         });
         separation_sets.emplace_back(std::make_pair(in_graph, out_graph));
       }
       set_allnodes.emplace(&node);
     }
     t_uncomp_map uncomp_map;
     std::unordered_set<Node*> set_TRTnodes;
     set_TRTnodes.insert(set_allnodes.begin(), set_allnodes.end());
     for (Node* n : set_nonTRTnodes) {
       set_TRTnodes.erase(n);
     }
     for (Node* n : set_TRTnodes) {
       for (t_pairset p : separation_sets) {
         if (p.first.count(n)) {
           uncomp_map[n].insert(p.second.begin(), p.second.end());
         } else if (p.second.count(n)) {
           uncomp_map[n].insert(p.first.begin(), p.first.end());
         }
       }
       for (Node* nonTRTn : set_nonTRTnodes) {
         uncomp_map[n].erase(nonTRTn);
       }
     }
     std::unordered_set<Node*> set_unused;
     set_unused.reserve(set_TRTnodes.size());

     for (auto& n : set_TRTnodes) {
       if (n->nnvmptr->attrs.op != nullptr || params_map->count(n->nnvmptr->attrs.name)) {
         set_unused.insert(n);
       }
     }
     std::unordered_set<Node*> visited;
     std::deque<Node*> stack(outputs.begin(), outputs.end());
     while (!stack.empty()) {
       Node* vertex = stack.front();
       stack.pop_front();
       if (!visited.count(vertex)) {
         visited.insert(vertex);
         if (set_unused.count(vertex)) {
           subgraphs.emplace_back(naive_grow_subgraph(vertex, &set_unused, &uncomp_map));
         }
         for (Node* input : vertex->inputs) {
           stack.emplace_back(input);
         }
       }
     }

     return subgraphs;
   }


  private:
   friend class Graph;

   bool IsTRTCompatible(nnvm::Node* nodeptr) {
     if (nodeptr->op() == nullptr) {
       return true;
     }

     const std::string op_name = nodeptr->op()->name;
     if (op_name == "Pooling") {
       return (nodeptr->attrs.dict.at("pool_type") == "avg" ||
           nodeptr->attrs.dict.at("pool_type") == "max");
     }

     if (unconditionalTRTop.count(op_name)) {
       return true;
     }

     if (op_name == "Activation") {
       return nodeptr->attrs.dict.at("act_type") == "relu" ||
         nodeptr->attrs.dict.at("act_type") == "tanh" ||
         nodeptr->attrs.dict.at("act_type") == "sigmoid";
     }

     return false;
   }
 };  // class BidirectionalGraph

 /*!
  * \brief function which transform std::vector<dmlc::any> back to Attrs (dmlc::any)
  */
 const std::unordered_set<std::string> BidirectionalGraph::unconditionalTRTop = {
   "Convolution",
   "BatchNorm",
   "elemwise_add",
   "elemwise_sub",
   "elemwise_mul",
   "rsqrt",
   "pad",
   "Pad",
   "mean",
   "FullyConnected",
   "Flatten",
   "SoftmaxOutput",
 };


 using NodeEntrySet = std::unordered_set<nnvm::NodeEntry, nnvm::NodeEntryHash,
                                         nnvm::NodeEntryEqual>;

 /*!
  * \brief get the output nodes of the subgraph in the main graph
  * \return a vector of the output nodes
 */
 std::vector<nnvm::NodeEntry> GetSubgraphNodeEntries(Graph g,
     std::unordered_set<nnvm::Node*> set_subgraph) {
   std::vector<nnvm::NodeEntry> outputs;
   NodeEntrySet _outputs;
   for (auto& e : g.outputs) {
     if (set_subgraph.count(e.node.get())) {
       _outputs.insert(e);
     }
   }
   DFSVisit(g.outputs, [&set_subgraph, &_outputs](const nnvm::NodePtr &node){
     if (!set_subgraph.count(node.get())) {
       for (auto& e : node->inputs) {
         if (set_subgraph.count(e.node.get())) {
           _outputs.insert(e);
         }
       }
     }
   });
   outputs.insert(outputs.begin(), _outputs.begin(), _outputs.end());
   return outputs;
 }


 /*!
  * \brief get the nodes outside of the subgraph for which outputs are used in the subgraph
  * \return a vector the nodes
 */
 std::vector<nnvm::NodeEntry> GetSubgraphInterfaceNodes(Graph g,
     std::unordered_set<nnvm::Node*> set_subgraph) {
   std::vector<nnvm::NodeEntry> inputs;
   NodeEntrySet _inputs;
   DFSVisit(g.outputs, [&set_subgraph, &_inputs](const nnvm::NodePtr &node){
     if (set_subgraph.count(node.get())) {
       for (auto& e : node->inputs) {
         if (!set_subgraph.count(e.node.get())) {
           _inputs.insert(e);
         }
       }
     }
   });
   inputs.insert(inputs.begin(), _inputs.begin(), _inputs.end());
   return inputs;
 }

 std::unordered_map<uint32_t, uint32_t> GetGraphInputsMap(const Graph& g) {
   std::unordered_map<uint32_t, uint32_t> outputs;
   auto& idx = g.indexed_graph();
   outputs.reserve(idx.num_nodes());
   std::vector<uint32_t> input_nodes = idx.input_nodes();
   for (size_t i = 0; i < input_nodes.size(); ++i) {
     outputs[input_nodes[i]] = static_cast<uint32_t>(i);
   }
   return outputs;
 }

 /*!
  * \brief Dummy function which creates a fake TensorRT Node
  */
 nnvm::NodePtr ConvertNnvmGraphToOnnx(const nnvm::Graph &g,
                                      std::unordered_map<std::string, NDArray>* const params_map) {
   auto p = nnvm::Node::Create();
   p->attrs.op = nnvm::Op::Get("_trt_op");
   op::TRTParam trt_param = op::nnvm_to_onnx::ConvertNnvmGraphToOnnx(g, params_map);
   p->attrs.dict["serialized_output_map"] = trt_param.serialized_output_map;
   p->attrs.dict["serialized_input_map"]  = trt_param.serialized_input_map;
   p->attrs.dict["serialized_onnx_graph"] = trt_param.serialized_onnx_graph;
   if (p->op()->attr_parser != nullptr) {
     p->op()->attr_parser(&(p->attrs));
   }
   return p;
 }

 /*!
  * \brief Update attributes of the graph (such as some inputs properties)
  */
 Graph UpdateSubgraphAttrs(Graph&& subgraph, const Graph& g,
                           const std::unordered_map<nnvm::Node*, nnvm::NodePtr>& old2new,
                           const nnvm::NodeEntryMap<nnvm::NodeEntry>& main_input_entry_to_sub) {
   const auto& idx     = g.indexed_graph();
   const auto& sub_idx = subgraph.indexed_graph();

   const auto& shape               = g.GetAttr<nnvm::ShapeVector>("shape");
   const auto& dtype               = g.GetAttr<nnvm::DTypeVector>("dtype");
   const auto& storage_type        = g.GetAttr<StorageTypeVector>("storage_type");
   const auto& shape_inputs        = g.GetAttr<nnvm::ShapeVector>("shape_inputs");
   const auto& dtype_inputs        = g.GetAttr<nnvm::DTypeVector>("dtype_inputs");
   const auto& storage_type_inputs = g.GetAttr<StorageTypeVector>("storage_type_inputs");

   nnvm::ShapeVector sub_shape(sub_idx.num_node_entries());
   nnvm::DTypeVector sub_dtype(sub_idx.num_node_entries());
   StorageTypeVector sub_storage_type(sub_idx.num_node_entries());
   nnvm::ShapeVector sub_shape_inputs(sub_idx.input_nodes().size());
   nnvm::DTypeVector sub_dtype_inputs(sub_idx.input_nodes().size());
   StorageTypeVector sub_storage_type_inputs(sub_idx.input_nodes().size());

   const std::unordered_map<uint32_t, uint32_t> inputsindex2pos     = GetGraphInputsMap(g);
   const std::unordered_map<uint32_t, uint32_t> sub_inputsindex2pos = GetGraphInputsMap(subgraph);
   // map attributes from graph to subgraph
   for (auto& p : old2new) {
     const uint32_t nid     = idx.node_id(p.first);
     const uint32_t sub_nid = sub_idx.node_id(p.second.get());
     const nnvm::Op* op = sub_idx[sub_nid].source->op();
     if (op == nullptr) {  // if it's an input node, there is only one output node entry
       const uint32_t sub_i       = sub_idx.entry_id(sub_nid, 0);
       const uint32_t sub_input_i = sub_inputsindex2pos.at(sub_nid);
       const uint32_t i           = idx.entry_id(nid, 0);

       sub_shape[sub_i] = shape[i];
       sub_dtype[sub_i] = dtype[i];
       sub_storage_type[sub_i]       = storage_type[i];
       sub_shape_inputs[sub_input_i] = shape_inputs[inputsindex2pos.at(nid)];
       sub_dtype_inputs[sub_input_i] = dtype_inputs[inputsindex2pos.at(nid)];
       sub_storage_type_inputs[sub_input_i] = storage_type_inputs[inputsindex2pos.at(nid)];

     } else {
       for (size_t oi = 0; oi < op->num_outputs; ++oi) {
         const uint32_t sub_i = sub_idx.entry_id(sub_nid, oi);
         const uint32_t i = idx.entry_id(nid, oi);
           sub_shape[sub_i] = shape[i];
           sub_dtype[sub_i] = dtype[i];
           sub_storage_type[sub_i] = storage_type[i];
       }
     }
   }
   // old2new doesn't contain placeholder / interfaces
   for (auto& p : main_input_entry_to_sub) {
     nnvm::NodeEntry main_entry = p.first;
     nnvm::NodeEntry sub_entry = p.second;
     const uint32_t sub_nid = sub_idx.node_id(sub_entry.node.get());
     const uint32_t sub_i = sub_idx.entry_id(sub_entry);
     const uint32_t i = idx.entry_id(main_entry);
     const uint32_t sub_input_i = sub_inputsindex2pos.at(sub_nid);
     sub_shape[sub_i] = shape[i];
     sub_dtype[sub_i] = dtype[i];
     sub_storage_type[sub_i] = storage_type[i];
     sub_shape_inputs[sub_input_i] = sub_shape[sub_i];
     sub_dtype_inputs[sub_input_i] = sub_dtype[sub_i];
     sub_storage_type_inputs[sub_input_i] = sub_storage_type[sub_i];
   }
   subgraph.attrs["shape"] =
       std::make_shared<dmlc::any>(std::move(sub_shape));
   subgraph.attrs["dtype"] =
       std::make_shared<dmlc::any>(std::move(sub_dtype));
   subgraph.attrs["storage_type"] =
       std::make_shared<dmlc::any>(std::move(sub_storage_type));
   subgraph.attrs["shape_inputs"] =
       std::make_shared<dmlc::any>(std::move(sub_shape_inputs));
   subgraph.attrs["dtype_inputs"] =
       std::make_shared<dmlc::any>(std::move(sub_dtype_inputs));
   subgraph.attrs["storage_type_inputs"] =
       std::make_shared<dmlc::any>(std::move(sub_storage_type_inputs));

   return subgraph;
 }

 /*!
  * \brief Generate a name for a new TRT node, avoid collision if some TRT_nodes are already defined
  */
 const std::string GetNewTrtName(const Graph& g, const Graph& subgraph) {
   const std::string name_prefix("TRT_node");
   std::unordered_set<std::string> name_set;
   DFSVisit(g.outputs, [&name_set, &name_prefix](const nnvm::NodePtr& node) {
     if (node->attrs.name.compare(0, name_prefix.size(), name_prefix) == 0) {
       name_set.insert(node->attrs.name);
     }
   });
   // name inside the subgraph will be avaible as they will be removed
   DFSVisit(subgraph.outputs, [&name_set, &name_prefix](const nnvm::NodePtr& node) {
     if (node->attrs.name.compare(0, name_prefix.size(), name_prefix) == 0) {
       name_set.erase(node->attrs.name);
     }
   });
   uint32_t name_suffix = 0;
   std::string full_name = name_prefix + std::to_string(name_suffix);
   while (name_set.count(full_name)) {
     full_name = name_prefix + std::to_string(++name_suffix);
   }
   return full_name;
 }

 /*!
  * \brief helper function to display what nodes are in a specific subset
  */
 void dispNodesSet(Graph g, std::unordered_set<nnvm::Node*> s) {
   DFSVisit(g.outputs, [&s](const nnvm::NodePtr n){
     if (s.count(n.get())) {
       std::cout << "  Y " << n->attrs.name << std::endl;
     } else {
       std::cout << "  N " << n->attrs.name << std::endl;
     }
   });
 }

 /*!
  * \brief Replace a set of nodes by a TensorRT node
  */
 Graph ReplaceSubgraph(Graph&& g,
                       const std::unordered_set<nnvm::Node*>& set_subgraph,
                       std::unordered_map<std::string, NDArray>* const params_map) {
   // Create MXNet subgraph
   Graph subgraph;

   const auto sub_outputs_in_main = GetSubgraphNodeEntries(g, set_subgraph);
   subgraph.outputs = sub_outputs_in_main;
   // old2new will link raw pointer of the nodes in the graph to
   // the corresponding shared_ptr of the nodes in the generated subgraph
   std::unordered_map<nnvm::Node*, nnvm::NodePtr> old2new;
   std::deque<nnvm::Node*> stack;
   std::unordered_set<nnvm::Node*> visited;
   int32_t reservation = set_subgraph.size();
   old2new.reserve(reservation);
   visited.reserve(reservation);

   // Create the shared_ptr using the same raw pointer don't really matter
   for (auto& n : set_subgraph) {
     old2new[n] = std::make_shared<nnvm::Node>(*n);
   }

   // To generate a subgraph an input have to be replace by data node (no op)
   // and it have to be agnostic to the node from which it's an output
   // (For exemple even if two inputs are two different outputs from the same node)
   nnvm::NodeEntryMap<nnvm::NodeEntry> main_input_entry_to_sub;
   for (auto& e : GetSubgraphInterfaceNodes(g, set_subgraph)) {
     auto node = nnvm::Node::Create();
     node->attrs.name = e.node->attrs.name + "_" + std::to_string(e.index);
     auto new_e = nnvm::NodeEntry{node, 0, 0};
     main_input_entry_to_sub[e] = new_e;
   }

   for (nnvm::NodeEntry& e : subgraph.outputs) {
     e.node = old2new[e.node.get()];
     stack.emplace_back(e.node.get());
   }
   // link all nodes in the subgraph to nodes in the subgraph instead of main graph
   while (!stack.empty()) {
     auto vertex = stack.front();
     stack.pop_front();
     if (!visited.count(vertex)) {
       visited.insert(vertex);
       for (auto& e : vertex->inputs) {
         auto it = main_input_entry_to_sub.find(e);
         if (it != main_input_entry_to_sub.end()) {
           e = it->second;
         } else {
           e.node = old2new[e.node.get()];
         }
       stack.emplace_back(e.node.get());
       }
     }
   }
   // Remove the control dependencies of the subgraph to nodes that are not in the subgraph
   DFSVisit(subgraph.outputs, [&set_subgraph, &old2new](const nnvm::NodePtr& node) {
     std::remove_if(node->control_deps.begin(),
                    node->control_deps.end(),
                    [&set_subgraph](nnvm::NodePtr n_ptr) {
                     return !set_subgraph.count(n_ptr.get());
                    });
     for (nnvm::NodePtr& n_ptr : node->control_deps) {
       n_ptr = old2new[n_ptr.get()];
     }
   });

   subgraph = UpdateSubgraphAttrs(std::move(subgraph), g, old2new, main_input_entry_to_sub);
   auto& sub_idx = subgraph.indexed_graph();

   auto trtnodeptr = ConvertNnvmGraphToOnnx(subgraph, params_map);
   trtnodeptr->attrs.name = GetNewTrtName(g, subgraph);

   // Insert new trt node and unplug replaced nodes
   std::unordered_map<uint32_t, nnvm::NodeEntry> sub_input_entryid_to_main;
   for (auto& p : main_input_entry_to_sub) {
     sub_input_entryid_to_main[sub_idx.entry_id(p.second)] = p.first;
   }

   // Plug the nodes from the main graph as inputs of the trt node
   trtnodeptr->inputs.resize(main_input_entry_to_sub.size());
   {
     uint32_t counter = 0;
     for (uint32_t i : sub_idx.input_nodes()) {
       auto it = sub_input_entryid_to_main.find(sub_idx.entry_id(i, 0));
       if (it != sub_input_entryid_to_main.end()) {
         trtnodeptr->inputs[counter++] = it->second;
       }
     }
   }
   nnvm::NodeEntryMap<uint32_t> sub_outputs_in_main_to_pos;
   for (uint32_t i = 0; i < sub_outputs_in_main.size(); ++i) {
     sub_outputs_in_main_to_pos[sub_outputs_in_main[i]] = i;
   }
   // Plug the trt node as inputs to the main graph nodes
   DFSVisit(g.outputs, [&sub_outputs_in_main_to_pos, &trtnodeptr](const nnvm::NodePtr& n) {
     for (auto& e : n->inputs) {
       auto it = sub_outputs_in_main_to_pos.find(e);
       if (it != sub_outputs_in_main_to_pos.end()) {
         e.index = it->second;
         e.node = trtnodeptr;
       }
     }
   });

   for (auto& output : g.outputs) {
     auto it = sub_outputs_in_main_to_pos.find(output);
     if (it != sub_outputs_in_main_to_pos.end()) {
       output.index = it->second;
       output.node = trtnodeptr;
     }
   }

   Graph new_graph;
   new_graph.outputs = g.outputs;
   return new_graph;
 }

 std::vector<std::unordered_set<nnvm::Node*>> GetTrtCompatibleSubsets(const Graph& g,
     std::unordered_map<std::string, NDArray>* const params_map) {
   BidirectionalGraph biG = BidirectionalGraph(g);
   std::vector<std::unordered_set<BidirectionalGraph::Node*>> subsets = biG.get_subsets(params_map);
   std::vector<std::unordered_set<nnvm::Node*>> nnvm_subsets(subsets.size(),
                                                             std::unordered_set<nnvm::Node*>());
   for (size_t i = 0; i < subsets.size(); ++i) {
     nnvm_subsets[i].reserve(subsets[i].size());
     for (auto& n : subsets[i]) {
       nnvm_subsets[i].insert(n->nnvmptr);
     }
   }
   return nnvm_subsets;
 }

 }  // namespace exec
 }  // namespace mxnet

 #endif  // MXNET_USE_TENSORRT
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	/*!
	* Copyright (c) 2018 by Contributors
	* \file tensorrt_pass.cc
	* \brief Replace TRT compatible subgraphs by TRT engines
	* \author Clement Fuji Tsang
	*/

	#if MXNET_USE_TENSORRT

	#include <NvInfer.h>
	#include <mxnet/base.h>
	#include <mxnet/op_attr_types.h>
	#include <mxnet/operator.h>
	#include <nnvm/graph_attr_types.h>
	#include <onnx/onnx.pb.h>

	#include "../operator/contrib/nnvm_to_onnx-inl.h"
	#include "./exec_pass.h"
	#include "./onnx_to_tensorrt.h"

	namespace mxnet {
	namespace exec {

	using NodePtr = nnvm::NodePtr;

	/*!
	* \brief Custom graph class, which will contain bi-directional nodes
	* we need to compute DFS and reverse DFS for graph partitioning
	*/
	class BidirectionalGraph {
	public:
	struct Node {
	nnvm::Node* nnvmptr;
	std::vector<Node*> inputs;
	std::vector<Node*> outputs;
	};
	std::vector<Node> nodes;
	std::unordered_map<nnvm::Node*, uint32_t> nnvm2nid;
	std::vector<Node*> outputs;
	static const std::unordered_set<std::string> unconditionalTRTop;

	explicit BidirectionalGraph(const Graph &g) {
	auto& idx = g.indexed_graph();
	auto num_nodes = idx.num_nodes();
	nodes.reserve(num_nodes);
	nnvm2nid.reserve(num_nodes);
	outputs.reserve(idx.outputs().size());
	DFSVisit(g.outputs, [this](const nnvm::NodePtr& n) {
	BidirectionalGraph::Node new_node;
	new_node.nnvmptr = n.get();
	nnvm2nid[n.get()] = static_cast<uint32_t>(nodes.size());
	nodes.emplace_back(std::move(new_node));
	});
	for (const auto& it : nnvm2nid) {
	nnvm::Node* nnvmnode = it.first;
	uint32_t nid = it.second;
	for (auto& n : nnvmnode->inputs) {
	uint32_t input_nid = nnvm2nid[n.node.get()];
	nodes[input_nid].outputs.emplace_back(&nodes[nid]);
	nodes[nid].inputs.emplace_back(&nodes[input_nid]);
	}
	}
	for (auto& e : g.outputs) {
	uint32_t nid = nnvm2nid[e.node.get()];
	outputs.emplace_back(&nodes[nid]);
	}
	}

	template <typename FVisit>
	void DFS(const std::vector<Node*>& heads, bool reverse, FVisit fvisit) {
	std::unordered_set<Node*> visited;
	std::vector<Node*> vec(heads.begin(), heads.end());
	visited.reserve(heads.size());
	while (!vec.empty()) {
	Node* vertex = vec.back();
	vec.pop_back();
	if (visited.count(vertex) == 0) {
	visited.insert(vertex);
	fvisit(vertex);
	std::vector<Node*> nexts = reverse ? vertex->inputs : vertex->outputs;
	for (Node* node : nexts) {
	if (visited.count(node) == 0) {
	vec.emplace_back(node);
	}
	}
	}
	}
	}

	using t_pairset = std::pair<std::unordered_set<Node>, std::unordered_set<Node>>;
	using t_pairvec = std::pair<std::vector<Node>, std::vector<Node>>;
	using t_uncomp_map = std::unordered_map<Node, std::unordered_set<Node>>;

	std::unordered_set<Node> naive_grow_subgraph(Node head,
	std::unordered_set<Node> set_unused,
	t_uncomp_map* uncomp_map) {
	std::unordered_set<Node*> subgraph;
	std::unordered_set<Node*> uncomp_set;
	std::deque<Node*> stack;
	stack.emplace_back(head);
	while (!stack.empty()) {
	Node* vertex = stack.back();
	stack.pop_back();
	if (set_unused->count(vertex) && !uncomp_set.count(vertex)) {
	set_unused->erase(vertex);
	subgraph.insert(vertex);
	uncomp_set.insert((uncomp_map)[vertex].begin(), (uncomp_map)[vertex].end());
	for (Node* input : vertex->inputs) {
	if (set_unused->count(input) && !uncomp_set.count(input)) {
	stack.emplace_back(input);
	}
	}
	for (Node* output : vertex->outputs) {
	if (set_unused->count(output) && !uncomp_set.count(output)) {
	stack.emplace_back(output);
	}
	}
	}
	}
	return subgraph;
	}

	std::vector<std::unordered_set<Node*>> get_subsets(
	std::unordered_map<std::string, NDArray>* const params_map) {
	std::vector<std::unordered_set<Node*>> subgraphs;
	std::unordered_set<Node*> set_nonTRTnodes;
	std::unordered_set<Node*> set_allnodes(nodes.size());
	std::vector<t_pairset> separation_sets;
	for (Node& node : nodes) {
	if (!IsTRTCompatible(node.nnvmptr)) {
	set_nonTRTnodes.insert(&node);
	std::unordered_set<Node*> in_graph;
	std::unordered_set<Node*> out_graph;
	std::vector<Node*> dummy_head;
	dummy_head.emplace_back(&node);
	DFS(dummy_head, false, [&out_graph](Node* node) {
	out_graph.insert(node);
	});
	DFS(dummy_head, true, [&in_graph](Node* node) {
	in_graph.insert(node);
	});
	separation_sets.emplace_back(std::make_pair(in_graph, out_graph));
	}
	set_allnodes.emplace(&node);
	}
	t_uncomp_map uncomp_map;
	std::unordered_set<Node*> set_TRTnodes;
	set_TRTnodes.insert(set_allnodes.begin(), set_allnodes.end());
	for (Node* n : set_nonTRTnodes) {
	set_TRTnodes.erase(n);
	}
	for (Node* n : set_TRTnodes) {
	for (t_pairset p : separation_sets) {
	if (p.first.count(n)) {
	uncomp_map[n].insert(p.second.begin(), p.second.end());
	} else if (p.second.count(n)) {
	uncomp_map[n].insert(p.first.begin(), p.first.end());
	}
	}
	for (Node* nonTRTn : set_nonTRTnodes) {
	uncomp_map[n].erase(nonTRTn);
	}
	}
	std::unordered_set<Node*> set_unused;
	set_unused.reserve(set_TRTnodes.size());

	for (auto& n : set_TRTnodes) {
	if (n->nnvmptr->attrs.op != nullptr \|\| params_map->count(n->nnvmptr->attrs.name)) {
	set_unused.insert(n);
	}
	}
	std::unordered_set<Node*> visited;
	std::deque<Node*> stack(outputs.begin(), outputs.end());
	while (!stack.empty()) {
	Node* vertex = stack.front();
	stack.pop_front();
	if (!visited.count(vertex)) {
	visited.insert(vertex);
	if (set_unused.count(vertex)) {
	subgraphs.emplace_back(naive_grow_subgraph(vertex, &set_unused, &uncomp_map));
	}
	for (Node* input : vertex->inputs) {
	stack.emplace_back(input);
	}
	}
	}

	return subgraphs;
	}


	private:
	friend class Graph;

	bool IsTRTCompatible(nnvm::Node* nodeptr) {
	if (nodeptr->op() == nullptr) {
	return true;
	}

	const std::string op_name = nodeptr->op()->name;
	if (op_name == "Pooling") {
	return (nodeptr->attrs.dict.at("pool_type") == "avg" \|\|
	nodeptr->attrs.dict.at("pool_type") == "max");
	}

	if (unconditionalTRTop.count(op_name)) {
	return true;
	}

	if (op_name == "Activation") {
	return nodeptr->attrs.dict.at("act_type") == "relu" \|\|
	nodeptr->attrs.dict.at("act_type") == "tanh" \|\|
	nodeptr->attrs.dict.at("act_type") == "sigmoid";
	}

	return false;
	}
	}; // class BidirectionalGraph

	/*!
	* \brief function which transform std::vector<dmlc::any> back to Attrs (dmlc::any)
	*/
	const std::unordered_set<std::string> BidirectionalGraph::unconditionalTRTop = {
	"Convolution",
	"BatchNorm",
	"elemwise_add",
	"elemwise_sub",
	"elemwise_mul",
	"rsqrt",
	"pad",
	"Pad",
	"mean",
	"FullyConnected",
	"Flatten",
	"SoftmaxOutput",
	};


	using NodeEntrySet = std::unordered_set<nnvm::NodeEntry, nnvm::NodeEntryHash,
	nnvm::NodeEntryEqual>;

	/*!
	* \brief get the output nodes of the subgraph in the main graph
	* \return a vector of the output nodes
	*/
	std::vector<nnvm::NodeEntry> GetSubgraphNodeEntries(Graph g,
	std::unordered_set<nnvm::Node*> set_subgraph) {
	std::vector<nnvm::NodeEntry> outputs;
	NodeEntrySet _outputs;
	for (auto& e : g.outputs) {
	if (set_subgraph.count(e.node.get())) {
	_outputs.insert(e);
	}
	}
	DFSVisit(g.outputs, [&set_subgraph, &_outputs](const nnvm::NodePtr &node){
	if (!set_subgraph.count(node.get())) {
	for (auto& e : node->inputs) {
	if (set_subgraph.count(e.node.get())) {
	_outputs.insert(e);
	}
	}
	}
	});
	outputs.insert(outputs.begin(), _outputs.begin(), _outputs.end());
	return outputs;
	}


	/*!
	* \brief get the nodes outside of the subgraph for which outputs are used in the subgraph
	* \return a vector the nodes
	*/
	std::vector<nnvm::NodeEntry> GetSubgraphInterfaceNodes(Graph g,
	std::unordered_set<nnvm::Node*> set_subgraph) {
	std::vector<nnvm::NodeEntry> inputs;
	NodeEntrySet _inputs;
	DFSVisit(g.outputs, [&set_subgraph, &_inputs](const nnvm::NodePtr &node){
	if (set_subgraph.count(node.get())) {
	for (auto& e : node->inputs) {
	if (!set_subgraph.count(e.node.get())) {
	_inputs.insert(e);
	}
	}
	}
	});
	inputs.insert(inputs.begin(), _inputs.begin(), _inputs.end());
	return inputs;
	}

	std::unordered_map<uint32_t, uint32_t> GetGraphInputsMap(const Graph& g) {
	std::unordered_map<uint32_t, uint32_t> outputs;
	auto& idx = g.indexed_graph();
	outputs.reserve(idx.num_nodes());
	std::vector<uint32_t> input_nodes = idx.input_nodes();
	for (size_t i = 0; i < input_nodes.size(); ++i) {
	outputs[input_nodes[i]] = static_cast<uint32_t>(i);
	}
	return outputs;
	}

	/*!
	* \brief Dummy function which creates a fake TensorRT Node
	*/
	nnvm::NodePtr ConvertNnvmGraphToOnnx(const nnvm::Graph &g,
	std::unordered_map<std::string, NDArray>* const params_map) {
	auto p = nnvm::Node::Create();
	p->attrs.op = nnvm::Op::Get("_trt_op");
	op::TRTParam trt_param = op::nnvm_to_onnx::ConvertNnvmGraphToOnnx(g, params_map);
	p->attrs.dict["serialized_output_map"] = trt_param.serialized_output_map;
	p->attrs.dict["serialized_input_map"] = trt_param.serialized_input_map;
	p->attrs.dict["serialized_onnx_graph"] = trt_param.serialized_onnx_graph;
	if (p->op()->attr_parser != nullptr) {
	p->op()->attr_parser(&(p->attrs));
	}
	return p;
	}

	/*!
	* \brief Update attributes of the graph (such as some inputs properties)
	*/
	Graph UpdateSubgraphAttrs(Graph&& subgraph, const Graph& g,
	const std::unordered_map<nnvm::Node*, nnvm::NodePtr>& old2new,
	const nnvm::NodeEntryMap<nnvm::NodeEntry>& main_input_entry_to_sub) {
	const auto& idx = g.indexed_graph();
	const auto& sub_idx = subgraph.indexed_graph();

	const auto& shape = g.GetAttr<nnvm::ShapeVector>("shape");
	const auto& dtype = g.GetAttr<nnvm::DTypeVector>("dtype");
	const auto& storage_type = g.GetAttr<StorageTypeVector>("storage_type");
	const auto& shape_inputs = g.GetAttr<nnvm::ShapeVector>("shape_inputs");
	const auto& dtype_inputs = g.GetAttr<nnvm::DTypeVector>("dtype_inputs");
	const auto& storage_type_inputs = g.GetAttr<StorageTypeVector>("storage_type_inputs");

	nnvm::ShapeVector sub_shape(sub_idx.num_node_entries());
	nnvm::DTypeVector sub_dtype(sub_idx.num_node_entries());
	StorageTypeVector sub_storage_type(sub_idx.num_node_entries());
	nnvm::ShapeVector sub_shape_inputs(sub_idx.input_nodes().size());
	nnvm::DTypeVector sub_dtype_inputs(sub_idx.input_nodes().size());
	StorageTypeVector sub_storage_type_inputs(sub_idx.input_nodes().size());

	const std::unordered_map<uint32_t, uint32_t> inputsindex2pos = GetGraphInputsMap(g);
	const std::unordered_map<uint32_t, uint32_t> sub_inputsindex2pos = GetGraphInputsMap(subgraph);
	// map attributes from graph to subgraph
	for (auto& p : old2new) {
	const uint32_t nid = idx.node_id(p.first);
	const uint32_t sub_nid = sub_idx.node_id(p.second.get());
	const nnvm::Op* op = sub_idx[sub_nid].source->op();
	if (op == nullptr) { // if it's an input node, there is only one output node entry
	const uint32_t sub_i = sub_idx.entry_id(sub_nid, 0);
	const uint32_t sub_input_i = sub_inputsindex2pos.at(sub_nid);
	const uint32_t i = idx.entry_id(nid, 0);

	sub_shape[sub_i] = shape[i];
	sub_dtype[sub_i] = dtype[i];
	sub_storage_type[sub_i] = storage_type[i];
	sub_shape_inputs[sub_input_i] = shape_inputs[inputsindex2pos.at(nid)];
	sub_dtype_inputs[sub_input_i] = dtype_inputs[inputsindex2pos.at(nid)];
	sub_storage_type_inputs[sub_input_i] = storage_type_inputs[inputsindex2pos.at(nid)];

	} else {
	for (size_t oi = 0; oi < op->num_outputs; ++oi) {
	const uint32_t sub_i = sub_idx.entry_id(sub_nid, oi);
	const uint32_t i = idx.entry_id(nid, oi);
	sub_shape[sub_i] = shape[i];
	sub_dtype[sub_i] = dtype[i];
	sub_storage_type[sub_i] = storage_type[i];
	}
	}
	}
	// old2new doesn't contain placeholder / interfaces
	for (auto& p : main_input_entry_to_sub) {
	nnvm::NodeEntry main_entry = p.first;
	nnvm::NodeEntry sub_entry = p.second;
	const uint32_t sub_nid = sub_idx.node_id(sub_entry.node.get());
	const uint32_t sub_i = sub_idx.entry_id(sub_entry);
	const uint32_t i = idx.entry_id(main_entry);
	const uint32_t sub_input_i = sub_inputsindex2pos.at(sub_nid);
	sub_shape[sub_i] = shape[i];
	sub_dtype[sub_i] = dtype[i];
	sub_storage_type[sub_i] = storage_type[i];
	sub_shape_inputs[sub_input_i] = sub_shape[sub_i];
	sub_dtype_inputs[sub_input_i] = sub_dtype[sub_i];
	sub_storage_type_inputs[sub_input_i] = sub_storage_type[sub_i];
	}
	subgraph.attrs["shape"] =
	std::make_shared<dmlc::any>(std::move(sub_shape));
	subgraph.attrs["dtype"] =
	std::make_shared<dmlc::any>(std::move(sub_dtype));
	subgraph.attrs["storage_type"] =
	std::make_shared<dmlc::any>(std::move(sub_storage_type));
	subgraph.attrs["shape_inputs"] =
	std::make_shared<dmlc::any>(std::move(sub_shape_inputs));
	subgraph.attrs["dtype_inputs"] =
	std::make_shared<dmlc::any>(std::move(sub_dtype_inputs));
	subgraph.attrs["storage_type_inputs"] =
	std::make_shared<dmlc::any>(std::move(sub_storage_type_inputs));

	return subgraph;
	}

	/*!
	* \brief Generate a name for a new TRT node, avoid collision if some TRT_nodes are already defined
	*/
	const std::string GetNewTrtName(const Graph& g, const Graph& subgraph) {
	const std::string name_prefix("TRT_node");
	std::unordered_set<std::string> name_set;
	DFSVisit(g.outputs, [&name_set, &name_prefix](const nnvm::NodePtr& node) {
	if (node->attrs.name.compare(0, name_prefix.size(), name_prefix) == 0) {
	name_set.insert(node->attrs.name);
	}
	});
	// name inside the subgraph will be avaible as they will be removed
	DFSVisit(subgraph.outputs, [&name_set, &name_prefix](const nnvm::NodePtr& node) {
	if (node->attrs.name.compare(0, name_prefix.size(), name_prefix) == 0) {
	name_set.erase(node->attrs.name);
	}
	});
	uint32_t name_suffix = 0;
	std::string full_name = name_prefix + std::to_string(name_suffix);
	while (name_set.count(full_name)) {
	full_name = name_prefix + std::to_string(++name_suffix);
	}
	return full_name;
	}

	/*!
	* \brief helper function to display what nodes are in a specific subset
	*/
	void dispNodesSet(Graph g, std::unordered_set<nnvm::Node*> s) {
	DFSVisit(g.outputs, [&s](const nnvm::NodePtr n){
	if (s.count(n.get())) {
	std::cout << " Y " << n->attrs.name << std::endl;
	} else {
	std::cout << " N " << n->attrs.name << std::endl;
	}
	});
	}

	/*!
	* \brief Replace a set of nodes by a TensorRT node
	*/
	Graph ReplaceSubgraph(Graph&& g,
	const std::unordered_set<nnvm::Node*>& set_subgraph,
	std::unordered_map<std::string, NDArray>* const params_map) {
	// Create MXNet subgraph
	Graph subgraph;

	const auto sub_outputs_in_main = GetSubgraphNodeEntries(g, set_subgraph);
	subgraph.outputs = sub_outputs_in_main;
	// old2new will link raw pointer of the nodes in the graph to
	// the corresponding shared_ptr of the nodes in the generated subgraph
	std::unordered_map<nnvm::Node*, nnvm::NodePtr> old2new;
	std::deque<nnvm::Node*> stack;
	std::unordered_set<nnvm::Node*> visited;
	int32_t reservation = set_subgraph.size();
	old2new.reserve(reservation);
	visited.reserve(reservation);

	// Create the shared_ptr using the same raw pointer don't really matter
	for (auto& n : set_subgraph) {
	old2new[n] = std::make_shared<nnvm::Node>(*n);
	}

	// To generate a subgraph an input have to be replace by data node (no op)
	// and it have to be agnostic to the node from which it's an output
	// (For exemple even if two inputs are two different outputs from the same node)
	nnvm::NodeEntryMap<nnvm::NodeEntry> main_input_entry_to_sub;
	for (auto& e : GetSubgraphInterfaceNodes(g, set_subgraph)) {
	auto node = nnvm::Node::Create();
	node->attrs.name = e.node->attrs.name + "_" + std::to_string(e.index);
	auto new_e = nnvm::NodeEntry{node, 0, 0};
	main_input_entry_to_sub[e] = new_e;
	}

	for (nnvm::NodeEntry& e : subgraph.outputs) {
	e.node = old2new[e.node.get()];
	stack.emplace_back(e.node.get());
	}
	// link all nodes in the subgraph to nodes in the subgraph instead of main graph
	while (!stack.empty()) {
	auto vertex = stack.front();
	stack.pop_front();
	if (!visited.count(vertex)) {
	visited.insert(vertex);
	for (auto& e : vertex->inputs) {
	auto it = main_input_entry_to_sub.find(e);
	if (it != main_input_entry_to_sub.end()) {
	e = it->second;
	} else {
	e.node = old2new[e.node.get()];
	}
	stack.emplace_back(e.node.get());
	}
	}
	}
	// Remove the control dependencies of the subgraph to nodes that are not in the subgraph
	DFSVisit(subgraph.outputs, [&set_subgraph, &old2new](const nnvm::NodePtr& node) {
	std::remove_if(node->control_deps.begin(),
	node->control_deps.end(),
	[&set_subgraph](nnvm::NodePtr n_ptr) {
	return !set_subgraph.count(n_ptr.get());
	});
	for (nnvm::NodePtr& n_ptr : node->control_deps) {
	n_ptr = old2new[n_ptr.get()];
	}
	});

	subgraph = UpdateSubgraphAttrs(std::move(subgraph), g, old2new, main_input_entry_to_sub);
	auto& sub_idx = subgraph.indexed_graph();

	auto trtnodeptr = ConvertNnvmGraphToOnnx(subgraph, params_map);
	trtnodeptr->attrs.name = GetNewTrtName(g, subgraph);

	// Insert new trt node and unplug replaced nodes
	std::unordered_map<uint32_t, nnvm::NodeEntry> sub_input_entryid_to_main;
	for (auto& p : main_input_entry_to_sub) {
	sub_input_entryid_to_main[sub_idx.entry_id(p.second)] = p.first;
	}

	// Plug the nodes from the main graph as inputs of the trt node
	trtnodeptr->inputs.resize(main_input_entry_to_sub.size());
	{
	uint32_t counter = 0;
	for (uint32_t i : sub_idx.input_nodes()) {
	auto it = sub_input_entryid_to_main.find(sub_idx.entry_id(i, 0));
	if (it != sub_input_entryid_to_main.end()) {
	trtnodeptr->inputs[counter++] = it->second;
	}
	}
	}
	nnvm::NodeEntryMap<uint32_t> sub_outputs_in_main_to_pos;
	for (uint32_t i = 0; i < sub_outputs_in_main.size(); ++i) {
	sub_outputs_in_main_to_pos[sub_outputs_in_main[i]] = i;
	}
	// Plug the trt node as inputs to the main graph nodes
	DFSVisit(g.outputs, [&sub_outputs_in_main_to_pos, &trtnodeptr](const nnvm::NodePtr& n) {
	for (auto& e : n->inputs) {
	auto it = sub_outputs_in_main_to_pos.find(e);
	if (it != sub_outputs_in_main_to_pos.end()) {
	e.index = it->second;
	e.node = trtnodeptr;
	}
	}
	});

	for (auto& output : g.outputs) {
	auto it = sub_outputs_in_main_to_pos.find(output);
	if (it != sub_outputs_in_main_to_pos.end()) {
	output.index = it->second;
	output.node = trtnodeptr;
	}
	}

	Graph new_graph;
	new_graph.outputs = g.outputs;
	return new_graph;
	}

	std::vector<std::unordered_set<nnvm::Node*>> GetTrtCompatibleSubsets(const Graph& g,
	std::unordered_map<std::string, NDArray>* const params_map) {
	BidirectionalGraph biG = BidirectionalGraph(g);
	std::vector<std::unordered_set<BidirectionalGraph::Node*>> subsets = biG.get_subsets(params_map);
	std::vector<std::unordered_set<nnvm::Node*>> nnvm_subsets(subsets.size(),
	std::unordered_set<nnvm::Node*>());
	for (size_t i = 0; i < subsets.size(); ++i) {
	nnvm_subsets[i].reserve(subsets[i].size());
	for (auto& n : subsets[i]) {
	nnvm_subsets[i].insert(n->nnvmptr);
	}
	}
	return nnvm_subsets;
	}

	} // namespace exec
	} // namespace mxnet

	#endif // MXNET_USE_TENSORRT