[Relax][Bugfix] Annotate ComputePrimValue output as host function (#17032)

The `ComputePrimValue` transform is used to compute the value of
symbolic expressions that may appear within a Relax function.  For
example, to compute a boolean condition used for a `relax::If` node.
These functions are used for small host-side computations, prior to
launching a device kernel.

This commit updates `ComputePrimValue` to annotate the generated
`PrimFunc` with `tir::attr::kIsHostFunc`.  This annotation is required
for correct behavior in `tvm.dlight.ApplyDefaultSchedule`, to avoid
erroneous scheduling of this function for the GPU, and for
`tir::transform::BindTarget`, to ensure that the function is compiled
for execution on the host.

Co-authored-by: Chris Sullivan <[email protected]>

Author: Lunderberg

src/relax/transform/compute_prim_value.cc

@@ -45,7 +45,8 @@ class PrimValueComputeInjector : public ExprMutator {
     auto param_vars = tir::UndefinedVars(node->value);
     tir::Stmt body = tir::Evaluate(tir::Call(ret_dtype, tir::builtin::ret(), {node->value}));
 
-    tir::PrimFunc func(param_vars, body, PrimType(ret_dtype));
+    tir::PrimFunc func(param_vars, body, PrimType(ret_dtype), {},
+                       DictAttrs({{tir::attr::kIsHostFunc, Bool(true)}}));
     func = tir::RenewDefs(func);
 
     auto callee = builder_->AddFunction(func, "compute_symbolic_expr");

tests/python/relax/test_transform_compute_prim_value.py

@@ -44,6 +44,7 @@ def main(A: R.Tensor(["N"])):
 
         @T.prim_func(private=True)
         def compute_symbolic_expr(N: T.int64) -> T.bool:
+            T.func_attr({"tir.is_host_func": True})
             T.ret(N % 16 == 0)
 
 
@@ -73,6 +74,7 @@ def main(A: R.Tensor(["N"])):
 
         @T.prim_func(private=True)
         def compute_symbolic_expr(N: T.int64) -> T.bool:
+            T.func_attr({"tir.is_host_func": True})
             T.ret(N % 16 == 0)
 
 
@@ -97,6 +99,7 @@ def main(_N: R.Prim(value="N"), _M: R.Prim(value="M")) -> R.Prim(value="N*M"):
 
         @T.prim_func(private=True)
         def compute_symbolic_expr(N: T.int64, M: T.int64) -> T.int64:
+            T.func_attr({"tir.is_host_func": True})
             T.ret(N * M)