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 Int
s, UInt
s, and Float
s 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?