jroesch commented on a change in pull request #7518:
URL: https://github.com/apache/tvm/pull/7518#discussion_r664941440
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -176,112 +179,88 @@ class GraphOpNode : public GraphNode {
const std::string op_type_name_{"tvm_op"};
};
-/*! \brief Code generator for graph executor */
+/*! \brief Code generator for the graph executor, produces a module containing
the graph JSON,
+ * module, and parameters.
+ */
class GraphExecutorCodegen : public
backend::MemoizedExprTranslator<std::vector<GraphNodeRef>> {
public:
- GraphExecutorCodegen(runtime::Module* mod, const TargetsMap& targets) :
mod_(mod) {
- compile_engine_ = CompileEngine::Global();
+ GraphExecutorCodegen(runtime::Module* mod, const TargetMap& targets) :
mod_(mod) {
targets_ = targets;
}
- /*!
- * \brief Update the "main" control function's metadata
- *
- * \param func The main function that contains calls to relay primitive
functions
- */
- void UpdateMainWorkspaceSize(const Function& func) {
- // This is a Map<device,Map<storage_id, size>>
- std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
- // This is a Map<device, size_of_inputs_and_outputs>
- std::unordered_map<int, int> device_io;
- // This is a Map<device, size_of_constants>
- std::unordered_map<int, int> device_consts;
-
- // Initialize the maps to zero
- for (const auto& kv : storage_device_map_) {
- auto sids = kv.second[0];
- auto devices = kv.second[1];
- CHECK_EQ(sids.size(), devices.size());
- for (uint32_t i = 0; i < sids.size(); i++) {
- sid_workspace[devices[i]][sids[i]] = 0;
- device_io[devices[i]] = 0;
- device_consts[devices[i]] = 0;
- }
- }
-
- // Collect sizes of tensors
- for (const auto& kv : storage_device_map_) {
- auto size_bytes = CalculateRelayExprSizeBytes(kv.first->checked_type());
- auto sids = kv.second[0];
- auto devices = kv.second[1];
- if (kv.first->IsInstance<ConstantNode>()) {
- for (const auto& dev : devices) {
- device_consts[dev] += size_bytes;
- }
- continue;
- } else if (kv.first->IsInstance<VarNode>() || kv.first == func->body) {
- for (const auto& dev : devices) {
- device_io[dev] += size_bytes;
- }
- continue;
- }
- for (uint32_t i = 0; i < sids.size(); i++) {
- // Here we record the largest size of the tensor
- // that share the same storage id, because storage_id will
- // be shared between multiple tensors that are not live simultaneously.
- if (size_bytes > sid_workspace[devices[i]][sids[i]]) {
- sid_workspace[devices[i]][sids[i]] = size_bytes;
- }
- }
- }
-
- // This is a Map<device, workspace_size>
- std::unordered_map<int, int> device_workspace;
- // Once we know the sizes of sids, we need to accumulate per device
- for (const auto& dev_sid_size : sid_workspace) {
- auto dev = dev_sid_size.first;
- device_workspace[dev] = 0;
- for (const auto& sid_size : dev_sid_size.second) {
- device_workspace[dev] += sid_size.second;
- }
- }
-
- // Populate FunctionInfo
- auto fi_node = make_object<FunctionInfoNode>();
- // Initialize all target workspaces to zero
- for (const auto& kv : targets_) {
- auto tgt = kv.second;
- fi_node->workspace_sizes.Set(tgt, 0);
- }
- for (const auto& dev_and_size : device_workspace) {
- auto tgt = GetTargetFromInteger(dev_and_size.first);
- fi_node->workspace_sizes.Set(tgt, dev_and_size.second);
- fi_node->relay_primfuncs.Set(tgt, func);
- }
- for (const auto& dev_and_size : device_io) {
- auto tgt = GetTargetFromInteger(dev_and_size.first);
- fi_node->io_sizes.Set(tgt, dev_and_size.second);
- }
- for (const auto& dev_and_size : device_consts) {
- auto tgt = GetTargetFromInteger(dev_and_size.first);
- fi_node->constant_sizes.Set(tgt, dev_and_size.second);
- }
-
- function_metadata_.Set(String(runtime::symbol::tvm_module_main),
FunctionInfo(fi_node));
+ StorageInfo GetStorageInfo(const Expr& e) {
+ size_t count = memory_plan_->expr_to_storage_info.count(e);
+ ICHECK_GT(count, 0) << "Expr is not existing in storage plan";
+ auto storage_info = memory_plan_->expr_to_storage_info[e];
+ return storage_info;
}
LoweredOutput Codegen(relay::Function func, String mod_name) {
- auto pf = GetPackedFunc("relay.backend.GraphPlanMemory");
- storage_device_map_ = (*pf)(func);
mod_name_ = mod_name;
- UpdateMainWorkspaceSize(func);
+
+ std::cout << "MODULE_NAME: " << mod_name_ << std::endl;
Review comment:
Remove
--
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]
