Why does the `tvm_call_packed_lowered` intrinsic assume an `int32` return type?

Hi all!

I’m trying to generate code which calls a PackedFunc registered from Python. This means the final generated code will use the tvm_call_packed_lowered intrinsic (which is generated when a tvm_call_packed intrinsic is lowered, unsurprisingly!).

The issue I was having is that I’d like to get something back that isn’t int32, however this line lowers all tvm_call_packed usages to tvm_call_packed_lowered calls with a return type of Int(32).

As an experiment, I made the following change:

modified   src/pass/lower_tvm_builtin.cc
@@ -274,9 +274,9 @@ class BuiltinLower : public IRMutator {
       ConstInt32(arg_stack_begin),
       ConstInt32(arg_stack_begin + op->args.size() - 1)
     };
-    return Call::make(
-        Int(32), intrinsic::tvm_call_packed_lowered,
-        packed_args, Call::Intrinsic);
+    CHECK(op->type.code() == kDLInt || op->type.code() == kDLUInt || op->type.code() == kDLFloat);
+    return Call::make(op->type, intrinsic::tvm_call_packed_lowered, packed_args, Call::Intrinsic);
   }
 
   Expr MakeCallTracePacked(const Call *op, const Expr &e) {

Instead of just Int(32), I’m allowing it to return Ints, UInts, and Floats of any width. After making this change, I’m able to get meaningful values out, and it doesn’t seem to break any of the Python unittests.

My question is: why did the code originally assumed an Int(32) return? And can it safely be changed?

The tvm return value convention allows you to pass int64 or double. So we should be able to relax the support to int and float scalar types. It does not support vector types though