jroesch commented on a change in pull request #8697:
URL: https://github.com/apache/tvm/pull/8697#discussion_r686181338
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -684,35 +582,69 @@ class AOTExecutorCodegen : public ExprVisitor {
/*! \brief mapping sid -> tir::Var */
std::unordered_map<int, te::Var> sids_table_;
/*! \brief lowered funcs */
- std::unordered_map<std::string, IRModule> lowered_funcs_;
- /*! \brief lowered funcs */
Map<String, FunctionInfo> function_metadata_;
- /*! \brief compile engine */
- CompileEngine compile_engine_;
/*! \brief the set of statements that make the program */
std::vector<tir::Stmt> stmts_;
/*! \brief the list of return sids (note that the function might return more
then one output */
std::vector<int> return_sid_;
- /*! \brief the module name we use to mangle the function names */
- String mod_name_;
public:
- AOTExecutorCodegen(runtime::Module* mod, const TargetsMap& targets, Target
target_host)
+ AOTExecutorCodegen(runtime::Module* mod, const tec::TargetMap& targets,
Target target_host)
: mod_(mod),
targets_(targets),
target_host_(target_host),
-
use_unpacked_api_(target_host->GetAttr<Bool>("unpacked-api").value_or(Bool(false))),
- compile_engine_(CompileEngine::Global()) {}
+
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);
- // Retrieve the storage map
- storage_device_map_ = aot_allocator.GetStorageMap();
- mod_name_ = mod_name;
+ // 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;
Review comment:
@mbs-octoml I don't think @Mousius needs to do this in this patch, but
this is where we should split device planning and storage planning. I think we
can remove the need to pre-storage plan at all if we can obtain the device
information pre-lowering, then storage plan after the lowering.
--
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]
