Code review fixes

docs/api/python/contrib/tensorrt.md

@@ -87,7 +87,7 @@ sym, arg_params, aux_params = mx.model.load_checkpoint(model_prefix, epoch)
 
 all_params = merge_dicts(arg_params, aux_params)
 ```
-This `all_params` dictionary cn be seem in use in the `simple_bind` call in `#2`.
+This `all_params` dictionary can be seen in use in the `simple_bind` call in `#2`.
 4. Once the symbol is bound, we need to feed the data and run the `forward()` method. Let's say we're using a test set data iterator called `test_iter`. We can run inference as follows:
 ```python
 for idx, dbatch in enumerate(test_iter):

example/image-classification/tensorrt/test_tensorrt_resnet50.py

@@ -18,13 +18,10 @@
 from __future__ import print_function
 
 import os.path
-import subprocess
 import mxnet as mx
 import numpy as np
 from time import time
 import sys
-import urllib
-
 
 def get_use_tensorrt():
     return int(os.environ.get("MXNET_USE_TENSORRT", 0))

include/mxnet/executor.h

@@ -164,7 +164,7 @@ class Executor {
                               std::vector<NDArray>* arg_grads,
                               std::vector<NDArray>* aux_states,
                               std::unordered_map<std::string, NDArray>*
-                              shared_data_arrays = nullptr,
+                                shared_data_arrays = nullptr,
                               Executor* shared_exec = nullptr);
   /*!
    * \brief the prototype of user-defined monitor callback

python/mxnet/base.py

@@ -680,3 +680,19 @@ def write_all_str(module_file, module_all_list):
     module_op_file.close()
     write_all_str(module_internal_file, module_internal_all)
     module_internal_file.close()
+
+def cint(init_val=0):
+    """create a C int with an optional initial value"""
+    return C.c_int(init_val)
+
+def int_addr(x):
+    """given a c_int, return it's address as an int ptr"""
+    x_addr = C.addressof(x)
+    int_p = C.POINTER(C.c_int)
+    x_int_addr = C.cast(x_addr, int_p)
+    return x_int_addr
+
+def checked_call(f, *args):
+    """call a cuda function and check for success"""
+    error_t = f(*args)
+    assert error_t == 0, "Failing cuda call %s returns %s." % (f.__name__, error_t)

python/mxnet/cuda_utils.py

@@ -25,22 +25,7 @@
 # As a stand-alone program, it prints a list of unique cuda SM architectures
 import ctypes as C
 from ctypes.util import find_library
-
-def cint(init_val=0):
-    """create a C int with an optional initial value"""
-    return C.c_int(init_val)
-
-def int_addr(x):
-    """given a c_int, return it's address as an int ptr"""
-    x_addr = C.addressof(x)
-    INTP = C.POINTER(C.c_int)
-    x_int_addr = C.cast(x_addr, INTP)
-    return x_int_addr
-
-def checked_call(f, *args):
-    """call a cuda function and check for success"""
-    error_t = f(*args)
-    assert error_t == 0, "Failing cuda call %s returns %s." % (f.__name__, error_t)
+from .base import cint, int_addr, checked_call
 
 def find_cuda_lib(candidates):
     for candidate in candidates:
@@ -84,7 +69,3 @@ def unique_sm_arches():
     for device_id in range(device_count):
         archs.add(get_sm_arch(device_id))
     return sorted(archs)
-
-# print a list of unique cuda SM architectures on the system
-if __name__ == '__main__':
-    print(' '.join(str(x) for x in unique_sm_arches()))

src/common/serialization.h

@@ -138,7 +138,7 @@ inline void deserialize(std::tuple<Args...>* obj, const std::string& buffer, siz
 
 
 template<typename T>
-struct is_cont {
+struct is_container {
   static const bool value = !std::is_pod<T>::value;
 };
 
@@ -149,7 +149,7 @@ inline size_t serialized_size(const T& obj) {
 
 template<typename T>
 inline size_t serialized_size(const nnvm::Tuple<T>& obj) {
-  if (is_cont<T>::value) {
+  if (is_container<T>::value) {
     size_t sum_val = 4;
     for (auto& el : obj) {
       sum_val += serialized_size(el);
@@ -162,7 +162,7 @@ inline size_t serialized_size(const nnvm::Tuple<T>& obj) {
 
 template<typename T>
 inline size_t serialized_size(const std::vector<T>& obj) {
-  if (is_cont<T>::value) {
+  if (is_container<T>::value) {
     size_t sum_val = 4;
     for (T i : obj) {
       sum_val += serialized_size(i);
@@ -181,16 +181,16 @@ inline size_t serialized_size(const std::pair<K, V>& obj) {
 template<typename K, typename V>
 inline size_t serialized_size(const std::map<K, V>& obj) {
   size_t sum_val = 4;
-  if (is_cont<K>::value && is_cont<V>::value) {
+  if (is_container<K>::value && is_container<V>::value) {
     for (auto p : obj) {
       sum_val += serialized_size(p.first) + serialized_size(p.second);
     }
-  } else if (is_cont<K>::value) {
+  } else if (is_container<K>::value) {
     for (auto p : obj) {
       sum_val += serialized_size(p.first);
     }
     sum_val += sizeof(V) * obj.size();
-  } else if (is_cont<K>::value) {
+  } else if (is_container<K>::value) {
     for (auto p : obj) {
       sum_val += serialized_size(p.second);
     }
@@ -204,16 +204,16 @@ inline size_t serialized_size(const std::map<K, V>& obj) {
 template<typename K, typename V>
 inline size_t serialized_size(const std::unordered_map<K, V>& obj) {
   size_t sum_val = 4;
-  if (is_cont<K>::value && is_cont<V>::value) {
+  if (is_container<K>::value && is_container<V>::value) {
     for (auto p : obj) {
       sum_val += serialized_size(p.first) + serialized_size(p.second);
     }
-  } else if (is_cont<K>::value) {
+  } else if (is_container<K>::value) {
     for (auto p : obj) {
       sum_val += serialized_size(p.first);
     }
     sum_val += sizeof(V) * obj.size();
-  } else if (is_cont<K>::value) {
+  } else if (is_container<K>::value) {
     for (auto p : obj) {
       sum_val += serialized_size(p.second);
     }
@@ -226,7 +226,7 @@ inline size_t serialized_size(const std::unordered_map<K, V>& obj) {
 
 template<typename K>
 inline size_t serialized_size(const std::set<K>& obj) {
-  if (is_cont<K>::value) {
+  if (is_container<K>::value) {
     size_t sum_val = 4;
     for (auto& el : obj) {
       sum_val += serialized_size(el);
@@ -239,7 +239,7 @@ inline size_t serialized_size(const std::set<K>& obj) {
 
 template<typename K>
 inline size_t serialized_size(const std::unordered_set<K>& obj) {
-  if (is_cont<K>::value) {
+  if (is_container<K>::value) {
     size_t sum_val = 4;
     for (auto& el : obj) {
       sum_val += serialized_size(el);
@@ -279,7 +279,7 @@ inline size_t serialized_size(const std::tuple<Args...>& obj) {
 //  SERIALIZE
 
 template<typename T>
-inline size_t serialize_cont_size(const T& obj, char** buffer) {
+inline size_t serialize_container_size(const T& obj, char** buffer) {
   uint32_t size = obj.size();
   std::memcpy(*buffer, &size, 4);
   *buffer += 4;
@@ -304,8 +304,8 @@ inline void serialize(const nnvm::Tuple<T>& obj, char** buffer) {
 
 template<typename T>
 inline void serialize(const std::vector<T>& obj, char** buffer) {
-  auto size = serialize_cont_size(obj, buffer);
-  if (is_cont<T>::value) {
+  auto size = serialize_container_size(obj, buffer);
+  if (is_container<T>::value) {
     for (auto& el : obj) {
       serialize(el, buffer);
     }
@@ -323,7 +323,7 @@ inline void serialize(const std::pair<K, V>& obj, char** buffer) {
 
 template<typename K, typename V>
 inline void serialize(const std::map<K, V>& obj, char** buffer) {
-  serialize_cont_size(obj, buffer);
+  serialize_container_size(obj, buffer);
   for (auto& p : obj) {
     serialize(p.first, buffer);
     serialize(p.second, buffer);
@@ -332,7 +332,7 @@ inline void serialize(const std::map<K, V>& obj, char** buffer) {
 
 template<typename K, typename V>
 inline void serialize(const std::unordered_map<K, V>& obj, char** buffer) {
-  serialize_cont_size(obj, buffer);
+  serialize_container_size(obj, buffer);
   for (auto& p : obj) {
     serialize(p.first, buffer);
     serialize(p.second, buffer);
@@ -341,23 +341,23 @@ inline void serialize(const std::unordered_map<K, V>& obj, char** buffer) {
 
 template<typename K>
 inline void serialize(const std::set<K>& obj, char** buffer) {
-  serialize_cont_size(obj, buffer);
+  serialize_container_size(obj, buffer);
   for (auto& el : obj) {
     serialize(el, buffer);
   }
 }
 
 template<typename K>
 inline void serialize(const std::unordered_set<K>& obj, char** buffer) {
-  serialize_cont_size(obj, buffer);
+  serialize_container_size(obj, buffer);
   for (auto& el : obj) {
     serialize(el, buffer);
   }
 }
 
 template<>
 inline void serialize(const std::string& obj, char** buffer) {
-  auto size = serialize_cont_size(obj, buffer);
+  auto size = serialize_container_size(obj, buffer);
   std::memcpy(*buffer, &obj[0], size);
   *buffer += size;
 }
@@ -387,7 +387,7 @@ inline void serialize(const std::tuple<Args...>& obj, char** buffer) {
 
 // Deserializer
 template<typename T>
-inline size_t deserialize_cont_size(T* obj, const std::string& buffer, size_t* curr_pos) {
+inline size_t deserialize_container_size(T* obj, const std::string& buffer, size_t* curr_pos) {
   uint32_t size = obj->size();
   std::memcpy(&size, &buffer[*curr_pos], 4);
   *curr_pos += 4;
@@ -413,9 +413,9 @@ inline void deserialize(nnvm::Tuple<T>* obj, const std::string& buffer, size_t*
 
 template<typename T>
 inline void deserialize(std::vector<T>* obj, const std::string& buffer, size_t* curr_pos) {
-  auto size = deserialize_cont_size(obj, buffer, curr_pos);
+  auto size = deserialize_container_size(obj, buffer, curr_pos);
   obj->resize(size);
-  if (is_cont<T>::value) {
+  if (is_container<T>::value) {
     for (size_t i = 0; i < size; ++i) {
       deserialize((*obj)[i], buffer, curr_pos);
     }
@@ -433,7 +433,7 @@ inline void deserialize(std::pair<K, V>* obj, const std::string& buffer, size_t*
 
 template<typename K, typename V>
 inline void deserialize(std::map<K, V>* obj, const std::string& buffer, size_t* curr_pos) {
-  auto size = deserialize_cont_size(obj, buffer, curr_pos);
+  auto size = deserialize_container_size(obj, buffer, curr_pos);
   K first;
   for (size_t i = 0; i < size; ++i) {
     deserialize(&first, buffer, curr_pos);
@@ -444,7 +444,7 @@ inline void deserialize(std::map<K, V>* obj, const std::string& buffer, size_t*
 template<typename K, typename V>
 inline void deserialize(std::unordered_map<K, V>* obj,
                         const std::string& buffer, size_t* curr_pos) {
-  auto size = deserialize_cont_size(obj, buffer, curr_pos);
+  auto size = deserialize_container_size(obj, buffer, curr_pos);
   K first;
   for (size_t i = 0; i < size; ++i) {
     deserialize(first, buffer, curr_pos);
@@ -454,7 +454,7 @@ inline void deserialize(std::unordered_map<K, V>* obj,
 
 template<typename K>
 inline void deserialize(std::set<K>* obj, const std::string& buffer, size_t* curr_pos) {
-  auto size = deserialize_cont_size(obj, buffer, curr_pos);
+  auto size = deserialize_container_size(obj, buffer, curr_pos);
   K first;
   for (size_t i = 0; i < size; ++i) {
     deserialize(first, buffer, curr_pos);
@@ -464,7 +464,7 @@ inline void deserialize(std::set<K>* obj, const std::string& buffer, size_t* cur
 
 template<typename K>
 inline void deserialize(std::unordered_set<K>* obj, const std::string& buffer, size_t* curr_pos) {
-  auto size = deserialize_cont_size(obj, buffer, curr_pos);
+  auto size = deserialize_container_size(obj, buffer, curr_pos);
   K first;
   for (size_t i = 0; i < size; ++i) {
     deserialize(first, buffer, curr_pos);
@@ -474,7 +474,7 @@ inline void deserialize(std::unordered_set<K>* obj, const std::string& buffer, s
 
 template<>
 inline void deserialize(std::string* obj, const std::string& buffer, size_t* curr_pos) {
-  auto size = deserialize_cont_size(obj, buffer, curr_pos);
+  auto size = deserialize_container_size(obj, buffer, curr_pos);
   obj->resize(size);
   std::memcpy(&(obj->front()), &buffer[*curr_pos], size);
   *curr_pos += size;

src/executor/exec_pass.h

@@ -213,7 +213,7 @@ bool DefaultStorageType(const nnvm::NodeAttrs& attrs,
  * \brief Replace subgraphs by TRT (forward only)
  */
 Graph ReplaceSubgraph(Graph&& g,
-                      std::unordered_set<nnvm::Node*> set_subgraph,
+                      const std::unordered_set<nnvm::Node*>& set_subgraph,
                       std::unordered_map<std::string, NDArray>* const params_map);
 
 std::vector<std::unordered_set<nnvm::Node*>> GetTrtCompatibleSubsets(const Graph& g,

src/executor/tensorrt_pass.cc

@@ -88,18 +88,18 @@ class BidirectionalGraph {
   template <typename FVisit>
   void DFS(const std::vector<Node*>& heads, bool reverse, FVisit fvisit) {
     std::unordered_set<Node*> visited;
-    std::deque<Node*> stack(heads.begin(), heads.end());
+    std::vector<Node*> vec(heads.begin(), heads.end());
     visited.reserve(heads.size());
-    while (!stack.empty()) {
-      Node* vertex = stack.back();
-      stack.pop_back();
+    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) {
-            stack.emplace_back(node);
+            vec.emplace_back(node);
           }
         }
       }
@@ -337,12 +337,12 @@ Graph UpdateSubgraphAttrs(Graph&& subgraph, const Graph& g,
   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");
+  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());
@@ -453,7 +453,7 @@ void dispNodesSet(Graph g, std::unordered_set<nnvm::Node*> s) {
  * \brief Replace a set of nodes by a TensorRT node
  */
 Graph ReplaceSubgraph(Graph&& g,
-                      std::unordered_set<nnvm::Node*> set_subgraph,
+                      const std::unordered_set<nnvm::Node*>& set_subgraph,
                       std::unordered_map<std::string, NDArray>* const params_map) {
   // Create MXNet subgraph
   Graph subgraph;

src/operator/contrib/tensorrt-inl.h

@@ -105,33 +105,6 @@ struct TRTEngineParam {
   std::vector<std::pair<uint32_t, tensorrt::TypeIO> > binding_map;
 };
 
-OpStatePtr TRTCreateState(const nnvm::NodeAttrs& attrs, Context ctx,
-                          const std::vector<TShape>& ishape,
-                          const std::vector<int>& itype);
-
-template<typename xpu>
-void TRTCompute(const OpStatePtr& state, const OpContext& ctx,
-                const std::vector<TBlob>& inputs, const std::vector<OpReqType>& req,
-                const std::vector<TBlob>& outputs);
-
-inline bool TRTInferShape(const NodeAttrs& attrs,
-                          std::vector<TShape> *in_shape,
-                          std::vector<TShape> *out_shape);
-
-inline bool TRTInferStorageType(const NodeAttrs& attrs,
-                                const int dev_mask,
-                                DispatchMode* dispatch_mode,
-                                std::vector<int> *in_storage_type,
-                                std::vector<int> *out_storage_type);
-
-inline bool TRTInferType(const NodeAttrs& attrs,
-                         std::vector<int> *in_dtype,
-                         std::vector<int> *out_dtype);
-
-inline std::vector<std::string> TRTListInputNames(const NodeAttrs& attrs);
-
-inline std::vector<std::string> TRTListOutputNames(const NodeAttrs& attrs);
-
 }  // namespace op
 }  // namespace mxnet