[Relax] Support callback as argument #16542
Conversation
Prior to this commit, calls from Relax to external PackedFuncs could only be done through the TVM global registry. While Relax functions accepting a callback could be written as `callback_arg: R.Callable(arg_struct_info, ret_struct_info)`, attempting to compile these functions would raise an error during the `CodeGenVM` step of `relax.build`. In addition, the global registry is only queried when initializing the `relax.VirtualMachine`, and so later changes requires restarting the VM. This commit updates both the `CodeGenVM` lowering pass and the relax VM to support callbacks. The is primarily intended for use with the `LazyTransformParams` pass, to improve flexibility by avoiding use of the global registry.
Lunderberg
force-pushed
the
relax_vm_callback_argument
branch
from
a264e81 to
f34323b
Compare
|
Merged from main to PR branch, to resolve CI breakage that was fixed in #16546. |
| /*! \brief The index into the packed function table. */ | ||
| /*! \brief The index of the function. | ||
| * | ||
| * For `OpCode::Call`, this is an index into the table of static |
There was a problem hiding this comment.
we should not use CallFromRegister, instead we should lower to https://github.com/apache/tvm/blob/main/src/runtime/relax_vm/builtin.cc#L373
|
Sorry for being late on this one. In this particular case, it is better to use
@Lunderberg would be great to followup on this soon given this would change some binary format of the VM https://github.com/apache/tvm/blob/main/src/runtime/relax_vm/builtin.cc#L373 |
|
cc @yongwww as well |
|
@tqchen Thank you, and I'll make a follow-up PR. I had thought that |
A follow-up PR from apache#16542, after some post-merge discussion regarding the implementation.
|
I've submitted #16573, which updates the handling of callback functions, reverts the VM changes, and passes the unit tests that were added in this PR. |
A follow-up PR from #16542, after some post-merge discussion regarding the implementation.
|
Good to know re: invoke_closure |
Prior to this commit, calls from Relax to external PackedFuncs could only be done through the TVM global registry. While Relax functions accepting a callback could be written as
callback_arg: R.Callable(arg_struct_info, ret_struct_info), attempting to compile these functions would raise an error during theCodeGenVMstep ofrelax.build. In addition, the global registry is only queried when initializing therelax.VirtualMachine, and so later changes requires restarting the VM.This commit updates both the
CodeGenVMlowering pass and the relax VM to support callbacks. The motivating use case is with theLazyTransformParamspass, to improve flexibility by avoiding use of the global registry.