Remove LoweredModule #8886
Remove LoweredModule #8886
Conversation
electriclilies
force-pushed
the
remove-lowered-module-final-draft
branch
from
3fb6beb
to
a6aa3b8
Compare
electriclilies
force-pushed
the
remove-lowered-module-final-draft
branch
from
a6aa3b8
to
6ed813b
Compare
electriclilies
requested review from
anijain2305,
areusch,
jroesch,
junrushao,
jwfromm,
kazum,
kparzysz-quic,
liangfu,
MarisaKirisame,
masahi,
mbrookhart,
slyubomirsky,
tmoreau89,
tqchen,
vinx13,
wweic,
yzhliu,
zhiics and
ZihengJiang
as code owners
|
@manupa-arm @mbaret @Mousius could you guys take a review pass on this if you have cycles? should be a quick review, just consolidates everything into an IRModule. |
There was a problem hiding this comment.
Overall I'd say that this is a good change! I've put in a few suggestions for improvements but as this is evolving in the future I'm happy to see this merged and see if the code even exists in the next patch 😸
P.S. Thanks for addressing that nit @electriclilies 🥇 !
| function_metadata_.Set(runtime::symbol::tvm_module_main, main_func_info.value()); | ||
|
|
||
| // Get only the Relay functions out of the lowered module so we can run type inference on them | ||
| IRModule main_module = tec::GetMainModule(lowered_mod); |
There was a problem hiding this comment.
Ok in a follow up if you like, but I think folding type inf into LowerTEPass makes sense to account for the rewritten calls. Someday we'll figure out how to gracefully do that incrementally since it's really overkill for just xfering known types from old to new CallNodes!
There was a problem hiding this comment.
I will add it as a TODO
There was a problem hiding this comment.
Actually there's a slight problem with this -- you can't run InferType on the whole lowered_mod becauselowered_mod has functions with GlobalVars that are in the IRModule's global_var_map, but you can't find them using Lookup because they are not the same object (it fails on line 210)..
https://github.com/apache/tvm/blob/main/src/relay/transforms/type_infer.cc#L209:L215
My workaround was to only apply type inferencing to the main_module (since that was what was done before this PR).
I'm not sure if this is a bug in how the type inferencer is dealing with GlobalVars (maybe it should be doing lookup by name hint, not pointer equality) or if it's a bug in how those GlobalVars are being made / propagated
There was a problem hiding this comment.
Then the other problem is even if you do successfully manage to lookup the function, there are PrimFuncs in the module which the type inferencer doesn't know how to deal with. We could just skip PrimFuncs we find during type inference..
There was a problem hiding this comment.
OK actually I think I fixed this in #8399
|
LGTM. So much clearer, thank you Lily. |
|
@Mousius @mbs-octoml Thanks for the review! |
|
This is merged now, thanks @electriclilies @Mousius and @mbs-octoml! |
|
thanks @leandron @Mousius @mbs-octoml! |
* Remove LoweredModule * Clean up some comments * QEMU flaky tests * Don't add external functions to the LoweredFunctions module * QEMU flaky test * Respond to feedback * flaky test
* Remove LoweredModule * Clean up some comments * QEMU flaky tests * Don't add external functions to the LoweredFunctions module * QEMU flaky test * Respond to feedback * flaky test
In this PR, I remove LoweredModule. I do this by annotating all functions with their targets and annotating external modules and function info on the module.
There is still some cleanup work to be done, specifically with removing per_target_modules_ in the interpreter. I'm going to remove those in a follow-up PR.
@mbs-octoml @mikepapadim @jroesch