mbs-octoml commented on a change in pull request #9395:
URL: https://github.com/apache/tvm/pull/9395#discussion_r746793358
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -637,6 +681,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
main_signature_.push_back(tir::Var("output", DataType::Handle()));
}
+ CollectDeviceVariables(lowered_mod->GetAttr<Map<GlobalVar,
String>>("device_contexts").value());
Review comment:
I think the blocker for pulling the device (or, after #9326 the SEScope)
for the callee is just that we're not carrying those annotations over into the
PrimFuncs, right? Would you consider trying that after #9326? I've
unfortunately hit my limit for any more hackery within the te_compiler.cc.
Then, afaikt you're doing what I've also had to do in the VM compiller.cc,
which is allocate globally unique tags (just ints in the VM) for every unique
SEScope, and build a global map from those tags to their actual devices (just
an array of Devices in the VM). Everything you've done to achieve that part
looks perfectly sensible to me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
