[Unity] Introduce call_dps_packed

[yongwww]

Introduce call_dps_packed to call packed function instead tir via destination passing convention

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[slyubomirsky]

Do we have any way to enforce that the function passed to call_tir is a PrimFunc? I believe making call_tir more specific was discussed before. (Should we still permit passing a PrimFunc to call_dps_packed?)

[psrivas2]

Do we have any way to enforce that the function passed to call_tir is a PrimFunc? I believe making call_tir more specific was discussed before. (Should we still permit passing a PrimFunc to call_dps_packed?)

We can enforce it in MakeCallTIR and MakeCallDPSPacked.
For now I don't see a reason to allow call_dps_packed to call Primfunc.
cc @tqchen

[slyubomirsky]

Do we have any way to enforce that the function passed to call_tir is a PrimFunc? I believe making call_tir more specific was discussed before. (Should we still permit passing a PrimFunc to call_dps_packed?)

We can enforce it in MakeCallTIR and MakeCallDPSPacked. For now I don't see a reason to allow call_dps_packed to call Primfunc. cc @tqchen

It would have to be in the implementation as well, since you can do something like assign a TIR function to a local var and pass that local var elsewhere. In a higher-order function, we wouldn't be able to detect that at compile time unless we make TIR functions part of the type system (though that could be an option).

[tqchen]

We can have wellformess to enforce such invariant

[tqchen]

please run the following commend to update to latest change

git rebase --onto upstream/unity upstream/unity-rebase-backup-2023-03-05 

[slyubomirsky]

It would not always be possible to check statically that the operand to call_tir is a PrimFunc. Here is a case where it wouldn't work:

@R.function
def higher_order(f: R.Object) -> R.Object:
    return R.call_tir(f, ...)

@R.function
def g() -> R.Object:
    if some_condition:
        y = some_packed_func
    else:
        y = some_primfunc
    return higher_order(y)

We would not be able to determine statically that the call to higher_order in g will result in a non-PrimFunc being passed to call_tir. If we care about enforcing the distinction, we would have to check for such cases dynamically.

Edit: Otherwise, we could ask for only global vars to be the arguments to call_tir, which I don't think would be the worst thing in the world, but this limits the extent to which PrimFuncs and PackedFuncs are first class. (I don't think they really need to be first class, but in past discussions, we've said that that was desirable.)

[yongwww]

Have updated well-form check to just allow GlobalVar as the first argument of call_tir. Actually current inline PrimFunc was disallowed in Relax IR, y = some_primfunc in the example has to be y = global_var_of_some_primfunc. I guess we won't run into issues at this moment, correct me if i am wrong

[slyubomirsky]

Yes, the RHS of that assignment would have to be a global var; we do not permit inline PrimFunc definitions. Perhaps we should get wider agreement if we want to have a restriction of using only global vars for call_tir or call_dps_packed? I can update the spec to reflect whatever we decide.

[slyubomirsky]

My concerns have been addressed and everything seems to be reasonable.

[psrivas2]

Overall LGTM! Thanks!
Just some minor nits which can be handled in follow up PRs

legalize_ops would throw a warning for call_dps_packed: https://github.com/psrivas2/tvm/blob/emit-te/src/relax/transform/legalize_ops.cc#L105

Would there be follow up PRs for cleaning up transform/backend passes which earlier checked if call_tir is calling an extern func such as

• https://github.com/psrivas2/tvm/blob/emit-te/src/relax/backend/task_extraction.cc#L77
• https://github.com/psrivas2/tvm/blob/emit-te/src/relax/transform/fold_constant.cc#L181
• https://github.com/psrivas2/tvm/blob/emit-te/src/relax/transform/fuse_ops.cc#L188
• https://github.com/psrivas2/tvm/blob/emit-te/src/relax/transform/fuse_tir.cc#L305
• https://github.com/psrivas2/tvm/blob/emit-te/src/relax/transform/rewrite_dataflow_reshape.cc#L78

[psrivas2]

nit: should we rename this file and pass to reflect that it handles both call_dps_packed and call_tir?

[yongwww]

I was thinking about this, but don't have a better name, any suggestions? if we change this, then the pass name CallTIRRewrite needs to be updated accordingly. We can keep it untouched for now if we don't have an alternative names,

[psrivas2]

DPSRewrite would be my choice but we can discuss and change it later.

[yongwww]

Thanks for the suggestion! We can discuss it later, it might need to get wider agreement for pass renaming.

[yongwww]

Overall LGTM! Thanks! Just some minor nits which can be handled in follow up PRs

legalize_ops would throw a warning for call_dps_packed

Would there be follow up PRs for cleaning up transform/backend passes which earlier checked if call_tir is calling an extern func such as

Have updated for all passes, thanks a lot for pointing them out!!!