use IsGPUTask function

apache · Dec 5, 2020 · b2532d3 · b2532d3
1 parent c89fcd2
commit b2532d3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/auto_scheduler/feature.cc b/src/auto_scheduler/feature.cc
@@ -41,6 +41,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "search_policy/utils.h"
 #include "utils.h"
 
 namespace tvm {
@@ -1296,8 +1297,7 @@ void GetPerStoreFeaturesWorkerFunc(const SearchTask& task, const State& state, i
     }
     auto mod = IRModule(Map<GlobalVar, BaseFunc>({{global_var, f}}));
 
-    auto device_type = task->target->kind->device_type;
-    if (device_type == kDLGPU || device_type == kDLROCM) {
+    if (IsGPUTask(task)) {
       auto pass_list = Array<tvm::transform::Pass>();
       // Phase 0
       pass_list.push_back(tir::transform::InjectPrefetch());