!90 Enable args reorder in split_host_device pass when run AKG operators...

!90 Enable args reorder in split_host_device pass when run AKG operators directly or run MindSpore operators. 
Merge pull request !90 from looop5/split_host_device

CMakeLists.txt

@@ -99,10 +99,6 @@ add_definitions(-DDMLC_LOG_CUSTOMIZE=1)
 if(USE_AKG_LOG)
   add_definitions(-DUSE_AKG_LOG=1)
 endif()
-if(NOT USE_CUDA
-   OR ENABLE_AKG)
-  add_definitions("-DFIX_INPUT_ORDER_TVM")
-endif()
 
 # Generic compilation options
 include(CheckCXXCompilerFlag)

third_party/incubator-tvm/src/pass/split_host_device.cc

@@ -234,32 +234,25 @@ class HostDeviceSplitter : public IRMutator {
       }
     }
 
-#ifdef FIX_INPUT_ORDER_TVM
-    std::shared_ptr<LoweredFuncNode> na = std::make_shared<LoweredFuncNode>();
-    for (unsigned i = 0; i < (unsigned)args_real.size(); i++) {
-      bool match = false;
-      for (unsigned j = 0; j < (unsigned)n->args.size(); j++) {
-        if (strcmp(args_real[i].get()->name_hint.c_str(), n->args[j].get()->name_hint.c_str()) == 0) {
-          na->args.push_back(n->args[j]);
-          match = true;
-          break;
-        } else {
-          continue;
-        }
+    // Reorder args to match args_real
+    Array<Var> ordered_args;
+    std::unordered_set<Var, NodeHash, NodeEqual> args_set;
+    std::unordered_set<Var, NodeHash, NodeEqual> args_real_set;
+    for (size_t i = 0; i < n->args.size(); ++i) {
+      args_set.insert(n->args[i]);
+    }
+    for (size_t i = 0; i < args_real.size(); ++i) {
+      args_real_set.insert(args_real[i]);
+      if (args_set.find(args_real[i]) != args_set.end()) {
+        ordered_args.push_back(args_real[i]);
       }
-
-      if (!match) {
-        na->args.push_back(args_real[i]);
-        // mark handle data type.
-        for (auto kv : handle_data_type_) {
-          if (strcmp(args_real[i].get()->name_hint.c_str(), kv.first->name_hint.c_str()) == 0) {
-            n->handle_data_type.Set(args_real[i], kv.second);
-          }
-        }
+    }
+    for (size_t i = 0; i < n->args.size(); ++i) {
+      if (args_real_set.find(n->args[i]) == args_real_set.end()) {
+        ordered_args.push_back(n->args[i]);
       }
     }
-    n->args = na->args;
-#endif
+    n->args = ordered_args;
     
     LoweredFunc f_device(n);
     Array<Expr> call_args;