Better detection of host target

apache · Jun 13, 2023 · 03c39ef · 03c39ef
1 parent c35a8b5
commit 03c39ef
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc
@@ -473,17 +473,29 @@ runtime::Module TIRToRuntime(const Map<Target, IRModule>& inputs_arg,
   }
 
   auto host_target = [&]() -> Target {
+    // All targets that contain a kIsEntryFunc=True function
     Array<Target> targets_with_entry_func;
+
+    // All targets that can run on the CPU and contain at least one
+    // function without kIsEntryFunc=False.
     Array<Target> cpu_targets;
     for (const auto& [target, mod] : split) {
-      bool contains_entry_func = std::any_of(
-          mod->functions.begin(), mod->functions.end(),
-          [](const auto& kv) { return kv.second->HasNonzeroAttr(tvm::tir::attr::kIsEntryFunc); });
+      bool contains_entry_func = false;
+      bool may_contain_entry_func = false;
+      for (const auto& [gvar, func] : mod->functions) {
+        Optional<Bool> is_entry_func = func->attrs.GetAttr<Bool>(tvm::tir::attr::kIsEntryFunc);
+        if (is_entry_func.defined() && is_entry_func.value()->value) {
+          contains_entry_func = true;
+        } else if (!is_entry_func.defined()) {
+          may_contain_entry_func = true;
+        }
+      }
+
       if (contains_entry_func) {
         targets_with_entry_func.push_back(target);
       }
 
-      if (target->HasKey("cpu")) {
+      if (may_contain_entry_func && target->HasKey("cpu")) {
         cpu_targets.push_back(target);
       }
     }

diff --git a/src/tir/transforms/split_host_device.cc b/src/tir/transforms/split_host_device.cc
@@ -84,7 +84,8 @@ class HostDeviceSplitter : public StmtMutator {
     PrimFunc device_func(params, body);
     device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget, device_target},
                                                      {tir::attr::kNoAlias, Bool(true)},
-                                                     {tir::attr::kIsGlobalFunc, Bool(true)}});
+                                                     {tir::attr::kIsGlobalFunc, Bool(true)},
+                                                     {tir::attr::kIsEntryFunc, Bool(false)}});
 
     (*device_mod_)->Add(kernel_symbol_global, device_func);
     Array<PrimExpr> args = params.Map([](const Var& var) -> PrimExpr { return var; });

diff --git a/tests/python/unittest/test_tir_transform_split_host_device.py b/tests/python/unittest/test_tir_transform_split_host_device.py
@@ -122,6 +122,7 @@ def main_kernel(n: T.int32):
                         "target": T.target("cuda"),
                         "tir.noalias": T.bool(True),
                         "tir.is_global_func": True,
+                        "tir.is_entry_func": False,
                     }
                 )
                 T.evaluate(n)
@@ -162,6 +163,7 @@ def main_kernel(n: T.int32):
                         "target": T.target("cuda"),
                         "tir.noalias": T.bool(True),
                         "tir.is_global_func": True,
+                        "tir.is_entry_func": False,
                     }
                 )
                 T.evaluate(n)