Backport: Change Partition API's options_map to std::unordered_map #1…

…8929 (#18964) Signed-off-by: Serge Panev <[email protected]>
apache · Aug 20, 2020 · 9445a2d · 9445a2d
1 parent 8dcc653
commit 9445a2d
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 7 deletions.
diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc
@@ -28,6 +28,7 @@
 #include <mutex>
 #include <memory>
 #include <functional>
+#include <unordered_map>
 #include <utility>
 #include "dmlc/base.h"
 #include "dmlc/logging.h"
@@ -1262,8 +1263,8 @@ void registerPasses(void *lib, int verbose, mxnet::ext::msgSize_t msgSize,
       // get pass name
       const char* pass_name = g.GetAttr<const char*>("pass_name");
       // get options
-      const std::vector<std::pair<std::string, std::string>>& options_map =
-            g.GetAttr<const std::vector<std::pair<std::string, std::string>>>("options_map");
+      const std::unordered_map<std::string, std::string>& options_map =
+            g.GetAttr<const std::unordered_map<std::string, std::string>>("options_map");
       // convert options_map_ to char* to pass to backend library
       std::vector<const char*> opt_keys, opt_vals;
       for (auto& kv : options_map) {

diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc
@@ -1480,9 +1480,9 @@ int MXOptimizeForBackend(SymbolHandle sym_handle,
     g.attrs["in_aux_names"] = std::make_shared<nnvm::any>(aux_names);
   }
   // create a data structure from pointer array
-  std::vector<std::pair<std::string, std::string>> options_map;
+  std::unordered_map<std::string, std::string> options_map;
   for (mx_uint i = 0; i < num_options; ++i)
-    options_map.emplace_back(keys[i], vals[i]);
+     options_map.emplace(keys[i], vals[i]);
 
   if (mxnet::op::SubgraphBackendRegistry::Get()->backend_map_.count(backend_name) > 0) {
     // use subgraph backend

diff --git a/src/operator/subgraph/partitioner/custom_subgraph_property.h b/src/operator/subgraph/partitioner/custom_subgraph_property.h
@@ -189,7 +189,7 @@ class  CustomSubgraphProperty: public SubgraphProperty {
   }
 
   void PrePartition(const nnvm::Graph& g,
-    const std::vector<std::pair<std::string, std::string>>& options_map) {
+    const std::unordered_map<std::string, std::string>& options_map) {
     // clear supported_nodes to remove state from previous calls
     supported_nodes.clear();
     // get input args and arg names

diff --git a/src/operator/subgraph/subgraph_property.h b/src/operator/subgraph/subgraph_property.h
@@ -268,7 +268,7 @@ class SubgraphProperty {
   }
 
   virtual void PrePartition(const nnvm::Graph& g,
-    const std::vector<std::pair<std::string, std::string>>& options_map) {}
+    const std::unordered_map<std::string, std::string>& options_map) {}
 
   virtual void PostPartition(const nnvm::Graph& g) {}
 

diff --git a/src/operator/subgraph/tensorrt/tensorrt-inl.h b/src/operator/subgraph/tensorrt/tensorrt-inl.h
@@ -269,7 +269,7 @@ class TensorrtProperty : public SubgraphProperty {
   }
 
   void PrePartition(const nnvm::Graph& g,
-    const std::vector<std::pair<std::string, std::string>>& options_map) override {
+    const std::unordered_map<std::string, std::string>& options_map) override {
     auto& in_arg_names = g.GetAttr<std::vector<std::string>>("in_arg_names");
     auto& in_aux_names = g.GetAttr<std::vector<std::string>>("in_aux_names");
     NDArray **in_args_ptr = g.GetAttr<NDArray**>("in_args");