[Discuss][UX] Syntactic sugar for giving a function type to packed calls

[slyubomirsky]

Currently, in R.call_packed, we give a type argument indicating the return type, but we do not indicate the argument types, which results in less informative type checking. I think there should be a mechanism to indicate the expected argument types for a call to a packed function.

I think an easy way to achieve this might be by adding some syntactic sugar for wrapping packed calls in an ordinary function call, like so:

@tvm.script.ir_module
class ExampleMod:
    my_func = R.wrap_packed("packed_func_name", [arg_type_0, arg_type_1, ..., arg_type_n], ret_type, attrs)

    @R.function
    def main():
        ...
        x : ret_type = my_func(arg_0, arg_1, ..., arg_n)
        ...

R.wrap_packed could desugar to

@R.function
def my_func(arg_0 : arg_type_0, arg_1 : arg_type_1, ..., arg_n : arg_type_n) -> ret_type:
    return R.call_packed("packed_func_name", arg_0, arg_1, ..., arg_n, type_args=(ret_type))

We could inline the calls to the wrapper func at a later stage of compilation, but an advantage would be that there would be type checking for the argument expressions. Would the proposed syntactic sugar be feasible with the ir_module decorator? I think it should be. If we encourage users to include such signatures, it would also mean that calls to packed funcs would also look like ordinary function calls and they wouldn't need to use call_packed directly as much.