Merge branch 'apache:main' into compute_at_floordivmod

Author: abhikran-quic

src/tir/transforms/split_host_device.cc

@@ -108,12 +108,12 @@ PrimFunc SplitHostDevice(PrimFunc func, IRModule* device_mod, const GlobalVar& g
 
   HostDeviceSplitter splitter(device_mod, name_prefix);
 
-  auto body = splitter(func->body);
-
-  if (!body.same_as(func->body)) {
+  if (auto body = splitter(func->body); !body.same_as(func->body)) {
     func.CopyOnWrite()->body = body;
-    auto target_host = target->GetHost().value_or(Target("llvm"));
-    func = WithAttr(std::move(func), tvm::attr::kTarget, target_host);
+  }
+
+  if (auto target_host = target->GetHost()) {
+    func = WithAttr(std::move(func), tvm::attr::kTarget, target_host.value());
   }
 
   return func;

tests/python/unittest/test_tir_transform_split_host_device.py

@@ -168,5 +168,21 @@ def main_kernel(n: T.int32):
         return mod
 
 
+class TestSplitHostDevice(BaseCompare):
+    """Like TestSplitHostDevice, but no device regions to extract
+
+    Even if there are no device regions, the host-side function should
+    still have its "target" attribute updated.
+    """
+
+    def before():
+        T.func_attr({"target": T.target("ext_dev", host="llvm")})
+        T.evaluate(0)
+
+    def expected():
+        T.func_attr({"target": T.target("llvm")})
+        T.evaluate(0)
+
+
 if __name__ == "__main__":
     tvm.testing.main()