Restrict the int32 return type to targets that need to propagate errors

Author: Lunderberg

src/tir/transforms/split_host_device.cc

@@ -60,7 +60,7 @@ class HostDeviceSplitter : public StmtMutator {
       VarUseDefAnalyzer use_def(/*defined_vars=*/{}, /*visit_thread_extent=*/false);
       use_def(body);
 
-      // Sort first by variable typ, then by variable name
+      // Sort first by variable type, then by variable name
       std::vector<Var> params{use_def.undefined_.begin(), use_def.undefined_.end()};
       std::sort(params.begin(), params.end(), [](const Var& a, const Var& b) {
         auto sort_key = [](const Var& var) {
@@ -77,27 +77,40 @@ class HostDeviceSplitter : public StmtMutator {
     // CodeGenCPU is used for some device-side targets, such as
     // "ext_dev", and expects to be able to return a int32_t status
     // code.
-    auto error_code_dtype = DataType::Int(32);
-    IntImm success(error_code_dtype, 0);
-    body = SeqStmt::Flatten(body, Evaluate(ret(success)));
+
+    bool can_propagate_errors = [&]() {
+      auto kind = device_target->GetTargetDeviceType();
+      return kind == kDLCPU || kind == kDLExtDev || kind == kDLHexagon;
+    }();
+    IntImm success(DataType::Int(32), 0);
+    Type kernel_ret_type;
+    if (can_propagate_errors) {
+      kernel_ret_type = PrimType(DataType::Int(32));
+      body = SeqStmt::Flatten(body, Evaluate(ret(success)));
+    } else {
+      kernel_ret_type = VoidType();
+    }
 
     GlobalVar kernel_symbol_global = var_supply_();
-    PrimFunc device_func(params, body, PrimType(error_code_dtype));
+    PrimFunc device_func(params, body, kernel_ret_type);
     device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget, device_target},
                                                      {tir::attr::kNoAlias, Bool(true)},
                                                      {tir::attr::kIsGlobalFunc, Bool(true)}});
 
     (*device_mod_)->Add(kernel_symbol_global, device_func);
     Array<PrimExpr> args = params.Map([](const Var& var) -> PrimExpr { return var; });
 
-    Var kernel_error_code("kernel_error_code", error_code_dtype);
+    if (can_propagate_errors) {
+      Var kernel_error_code("kernel_error_code", success->dtype);
+      Call kernel_call(success->dtype, kernel_symbol_global, args);
+      AssertStmt assert_success(kernel_error_code == success,
+                                StringImm("Error executing compute kernel"), Evaluate(0));
+      LetStmt let_check(kernel_error_code, kernel_call, assert_success);
 
-    AssertStmt assert_success(kernel_error_code == success,
-                              StringImm("Error executing compute kernel"), Evaluate(0));
-    LetStmt kernel_call(kernel_error_code, Call(error_code_dtype, kernel_symbol_global, args),
-                        assert_success);
-
-    return std::move(kernel_call);
+      return std::move(let_check);
+    } else {
+      return Evaluate(Call(DataType::Void(), kernel_symbol_global, args));
+    }
   }
 
   // target ir module