[GitHub] [tvm] mbs-octoml commented on a change in pull request #9395: [AOT][DeviceAPI] Connecting devices structure to relevant operators

[GitBox <gi...@apache.org>]

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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org