[GitHub] [tvm] areusch commented on a change in pull request #8096: Decoupling AOT from graph memory planner

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

areusch commented on a change in pull request #8096:
URL: https://github.com/apache/tvm/pull/8096#discussion_r638207535



##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -44,52 +45,179 @@ namespace tvm {
 namespace relay {
 namespace backend {
 
+/**
+ * Struct to contain information about intermediate variables in the
+ * runner function
+ */
+struct StorageInfo {
+  /*! \brief unique integer identifier of the particular intermediate variable */

Review comment:
       it would be great to qualify the name `ids` here, since there are storage ids, buffer ids, etc. also, would it make more sense to say "intermediate buffer" here?

##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -44,52 +45,179 @@ namespace tvm {
 namespace relay {
 namespace backend {
 
+/**
+ * Struct to contain information about intermediate variables in the
+ * runner function
+ */
+struct StorageInfo {

Review comment:
       each of these vectors are of the same length with corresponding elements in matching indicies, right? does it then make sense to just declare `std::vector<StorageInfo>` and remove the vectors from here? In GraphPlanMemory, we split them because you can't send user-defined data types through PackedFunc calls.

##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -44,52 +45,179 @@ namespace tvm {
 namespace relay {
 namespace backend {
 
+/**
+ * Struct to contain information about intermediate variables in the
+ * runner function
+ */
+struct StorageInfo {
+  /*! \brief unique integer identifier of the particular intermediate variable */
+  std::vector<int> ids;
+  /*! \brief exact size of the temporary */
+  std::vector<int> sizes_bytes;
+  /*! \brief device type of the temporary variable */
+  std::vector<int> dev_types;
+};
+
 using IntegerArray = Array<Integer>;
 using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap =
+    std::unordered_map<Expr, StorageInfo, runtime::ObjectPtrHash, runtime::ObjectPtrEqual>;
 
-class AotReturnSidVisitor : public ExprVisitor {
+/**
+ * This is an on demand allocator for AOT. A new temporary
+ * (storage allocator identifier) is allocated for each operation.
+ */
+class AOTOnDemandAllocator : public ExprVisitor {
  public:
-  explicit AotReturnSidVisitor(Map<Expr, Array<IntegerArray>> storage_device_map)
-      : storage_device_map_{storage_device_map}, return_sid_{-1} {}
+  // run the visitor on a function.
+  void Run(const Function& func) {
+    node_device_map_ = CollectDeviceInfo(func);
 
-  IntegerArray FindReturnSid(Function func) {
-    VisitExpr(func->body);
-    return return_sid_;
+    for (Expr param : func->params) {
+      CreateSid(param.operator->());
+    }
+
+    GetSid(func->body);
   }
 
- protected:
-  void AssignReturnSid(Expr e) {
-    auto iter = storage_device_map_.find(e);
-    if (iter != storage_device_map_.end()) {
-      return_sid_ = (*iter).second[0];
+  std::vector<int> GetReturnIds() const { return return_ids_; }
+
+  StorageMap GetStorageMap() const { return storage_device_map_; }
+
+  void VisitExpr_(const ConstantNode* op) final {
+    CreateSid(op);
+    AssignReturnSid(GetRef<Expr>(op));
+  }
+
+  void VisitExpr_(const CallNode* op) final {
+    // create token for the call node.
+    CreateSid(op);
+    for (Expr arg : op->args) {
+      GetSid(arg);
     }
+    AssignReturnSid(GetRef<Expr>(op));
   }
 
-  void VisitExpr_(const ConstantNode* cn) override {
-    ExprVisitor::VisitExpr_(cn);
-    AssignReturnSid(GetRef<Expr>(cn));
+  void VisitExpr_(const VarNode* op) final {
+    ExprVisitor::VisitExpr_(op);
+    AssignReturnSid(GetRef<Expr>(op));
   }
 
-  void VisitExpr_(const VarNode* vn) override {
-    ExprVisitor::VisitExpr_(vn);
-    AssignReturnSid(GetRef<Expr>(vn));
+  void VisitExpr_(const FunctionNode* op) final {
+    // do not recurse into sub function.
   }
 
-  void VisitExpr_(const CallNode* cn) override {
-    ExprVisitor::VisitExpr_(cn);
-    AssignReturnSid(GetRef<Expr>(cn));
+  void VisitExpr_(const GlobalVarNode* op) final {
+    // Do nothing.
   }
 
-  void VisitExpr_(const LetNode* op) override { VisitExpr(op->body); }
+  void VisitExpr_(const OpNode* op) final {
+    // Do nothing.
+  }
+
+  void VisitExpr_(const TupleNode* op) final {
+    StorageInfo field_sid;
+    Expr expr = GetRef<Expr>(op);
+    for (Expr field : op->fields) {
+      auto sid = GetSid(field);
+      field_sid.ids.insert(field_sid.ids.end(), sid.ids.begin(), sid.ids.end());
+      field_sid.dev_types.insert(field_sid.dev_types.end(), sid.dev_types.begin(),
+                                 sid.dev_types.end());
+      field_sid.sizes_bytes.insert(field_sid.sizes_bytes.end(), sid.sizes_bytes.begin(),
+                                   sid.sizes_bytes.end());
+    }
+
+    storage_device_map_[expr] = field_sid;
+    AssignReturnSid(expr);
+  }
 
-  void VisitExpr_(const TupleNode* tn) override {
-    ExprVisitor::VisitExpr_(tn);
-    AssignReturnSid(GetRef<Expr>(tn));
+  void VisitExpr_(const TupleGetItemNode* op) final {
+    Expr expr = GetRef<Expr>(op);
+    const auto& sid = GetSid(op->tuple);
+    ICHECK_LT(static_cast<size_t>(op->index), sid.ids.size());
+    storage_device_map_[expr].ids = {sid.ids[op->index]};
+    storage_device_map_[expr].sizes_bytes = {sid.sizes_bytes[op->index]};
+    storage_device_map_[expr].dev_types = {sid.dev_types[op->index]};
+    AssignReturnSid(expr);
   }
 
+  void VisitExpr_(const IfNode* op) final { LOG(FATAL) << "if is not supported."; }
+
+  void VisitExpr_(const LetNode* op) final { LOG(FATAL) << "if is not supported."; }
+
  private:
-  Map<Expr, Array<IntegerArray>> storage_device_map_;
-  IntegerArray return_sid_;
+  void AssignReturnSid(Expr e) {
+    auto iter = storage_device_map_.find(e);
+    if (iter != storage_device_map_.end()) {
+      return_ids_ = (*iter).second.ids;
+    }
+  }
+  /*!
+   * \brief ceil(size/word_size) to get number of words.
+   * \param size The original size.
+   * \param word_size The element size.
+   */
+  static size_t DivRoundUp(size_t size, size_t word_size) {
+    return (size + word_size - 1) / word_size;
+  }
+  /*!
+   * \brief Get the memory requirement.
+   * \param prototype The prototype token.
+   * \return The required memory size.
+   */
+  size_t GetMemorySize(const TensorTypeNode* ttype) {
+    ICHECK(ttype != nullptr);
+    size_t size = 1;
+    for (IndexExpr dim : ttype->shape) {
+      const int64_t* pval = tir::as_const_int(dim);
+      ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << ttype->shape;
+      ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+      size *= static_cast<size_t>(pval[0]);
+    }
+    size *= DivRoundUp(ttype->dtype.bits() * ttype->dtype.lanes(), 8);
+    return size;
+  }
+  /*!
+   * \brief Get the necessary token.
+   * \param expr The expression.
+   * \return The corresponding token.
+   */
+  StorageInfo GetSid(const Expr& expr) {
+    this->VisitExpr(expr);
+    auto it = storage_device_map_.find(expr);
+    ICHECK(it != storage_device_map_.end());
+    return it->second;
+  }
+
+  void CreateSid(const ExprNode* op) {
+    StorageInfo sid;
+    Expr expr = GetRef<Expr>(op);
+    int device_type = node_device_map_.count(GetRef<Expr>(op)) ? node_device_map_[expr]->value : 0;
+    if (const auto* tuple_type = op->checked_type().as<TupleTypeNode>()) {
+      for (Type t : tuple_type->fields) {
+        const auto* ttype = t.as<TensorTypeNode>();
+        ICHECK(ttype);
+        sid.ids.push_back(sid_++);
+        sid.dev_types.push_back(device_type);
+        sid.sizes_bytes.push_back(GetMemorySize(ttype));
+      }
+    } else {
+      const auto* ttype = op->checked_type().as<TensorTypeNode>();
+      ICHECK(ttype);
+      sid.ids.push_back(sid_++);
+      sid.dev_types.push_back(device_type);
+      sid.sizes_bytes.push_back(GetMemorySize(ttype));
+    }
+    storage_device_map_[expr] = sid;
+  }
+  /*! \brief mapping of expression -> storageInfo*/
+  StorageMap storage_device_map_;
+  /*! \brief mapping of expression -> device type*/
+  Map<Expr, Integer> node_device_map_;
+  /*! \brief current id of the temporary allocated*/
+  int sid_{0};
+  /*! \brief the set of identifiers that are return variables */

Review comment:
       could you explain in comment what the "identifiers" are here? there are lots of those here now.

##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -364,5 +364,27 @@ def test_byoc_utvm(use_calculated_workspaces):
     compile_and_run(mod, input_list, output_list, use_calculated_workspaces)
 
 
+def test_quant_mobilenet_tfl():

Review comment:
       could you clarify which part of the PR this is testing?




-- 
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.

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