mbs-octoml commented on a change in pull request #8788:
URL: https://github.com/apache/tvm/pull/8788#discussion_r694198085
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -562,32 +558,35 @@ class AOTExecutorCodegen : public ExprVisitor {
use_unpacked_api_(target_host->GetAttr<Bool>("unpacked-api").value_or(Bool(false)))
{}
LoweredOutput Codegen(relay::Function func, String mod_name) {
- auto aot_allocator = AOTOnDemandAllocator();
- aot_allocator.Run(func);
-
- // Pre-lowering storage map and memory plan
- StorageMap initial_storage_map = aot_allocator.GetStorageMap();
- StaticMemoryPlan memory_plan(initial_storage_map);
-
- // Build a map from each operation to device.
- tec::DeviceMap device_context_map;
- for (const auto& it : memory_plan->expr_to_storage_info) {
- auto expr = it.first;
- auto storage_info = it.second;
- auto device_types = storage_info->device_types;
- // CHECK_EQ(device_types.size(), 1);
- tvm::Device dev;
- dev.device_id = 0;
- dev.device_type = device_types[0];
- device_context_map.insert({expr, dev});
- }
-
// This first phase moves from implicit use of compile engine,
// to instead explicitly lowering the incoming IRModule, and then
// performing the preexisting AOT executor code generation phase.
IRModule mod = IRModule::FromExpr(func);
+
+ // Make explicit which device each expression should be executed on.
Review comment:
Ideally all the backends will treat MakeDevicesExplicit as the VM
compiler does and it's just another phase. I think the 'choose the default
device' should also be common (hopefully can do that in this PR, I think it's
just N ways of doing the same thing over the various backends). That leaves the
memory planning which I think we should tackle next.
--
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]
