You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/05/21 08:44:42 UTC
[GitHub] [tvm] manupa-arm commented on a change in pull request #8096: Decoupling AOT from graph memory planner
manupa-arm commented on a change in pull request #8096:
URL: https://github.com/apache/tvm/pull/8096#discussion_r636732925
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -46,13 +47,171 @@ namespace backend {
using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap = std::unordered_map<Expr, std::vector<std::vector<int>>, runtime::ObjectPtrHash,
+ runtime::ObjectPtrEqual>;
+/**
+ * This is an on demand allocator for AOT. A new temporary
+ * (storage allocator identifier) is allocated for each operation.
+ */
+class AOTOnDemandAllocator : public ExprVisitor {
+ public:
+ // run the visitor on a function.
+ StorageMap Run(const Function& func) {
+ node_device_map_ = CollectDeviceInfo(func);
+
+ for (Expr param : func->params) {
+ CreateSid(param.operator->());
+ }
+
+ GetSid(func->body);
+ return storage_device_map_;
+ }
+
+ void VisitExpr_(const ConstantNode* op) final { CreateSid(op); }
+
+ void VisitExpr_(const CallNode* op) final {
+ // create token for the call node.
+ CreateSid(op);
+ for (Expr arg : op->args) {
+ GetSid(arg);
+ }
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const FunctionNode* op) final {
+ // do not recurse into sub function.
+ }
+
+ void VisitExpr_(const GlobalVarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const OpNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const TupleNode* op) final {
+ std::vector<int> field_ids;
+ std::vector<int> field_sizes;
+ std::vector<int> field_types;
+ Expr expr = GetRef<Expr>(op);
+ for (Expr field : op->fields) {
+ auto sids = GetSid(field);
+ field_ids.insert(field_ids.end(), sids[0].begin(), sids[0].end());
+ field_types.insert(field_types.end(), sids[1].begin(), sids[1].end());
+ field_sizes.insert(field_sizes.end(), sids[2].begin(), sids[2].end());
+ }
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = field_ids;
+ storage_device_map_[expr][1] = field_sizes;
+ storage_device_map_[expr][2] = field_types;
+ }
+
+ void VisitExpr_(const TupleGetItemNode* op) final {
+ Expr expr = GetRef<Expr>(op);
+ const auto& sids = GetSid(op->tuple);
+ ICHECK_LT(static_cast<size_t>(op->index), sids.size());
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = {sids[0][op->index]};
+ storage_device_map_[expr][1] = {sids[1][op->index]};
+ storage_device_map_[expr][2] = {sids[2][op->index]};
+ }
+
+ 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:
+ /*!
+ * \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.
+ */
+ std::vector<std::vector<int>> 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) {
+ std::vector<int> sids;
+ std::vector<int> sizes;
+ std::vector<int> types;
+ 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);
+ sids.push_back(sid_++);
+ types.push_back(device_type);
+ sizes.push_back(GetMemorySize(ttype));
+ }
+ } else {
+ const auto* ttype = op->checked_type().as<TensorTypeNode>();
+ ICHECK(ttype);
+ sids.push_back(sid_++);
+ types.push_back(device_type);
+ sizes.push_back(GetMemorySize(ttype));
+ }
+ InitStorage(expr);
+ storage_device_map_[expr][0] = sids;
+ storage_device_map_[expr][1] = types;
+ storage_device_map_[expr][2] = sizes;
Review comment:
This seems inconsistent with L110.
Maybe its better to have a struct describing the attributes rather than a vector. WDYT?
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -46,13 +47,171 @@ namespace backend {
using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap = std::unordered_map<Expr, std::vector<std::vector<int>>, runtime::ObjectPtrHash,
+ runtime::ObjectPtrEqual>;
+/**
+ * This is an on demand allocator for AOT. A new temporary
+ * (storage allocator identifier) is allocated for each operation.
+ */
+class AOTOnDemandAllocator : public ExprVisitor {
+ public:
+ // run the visitor on a function.
+ StorageMap Run(const Function& func) {
+ node_device_map_ = CollectDeviceInfo(func);
+
+ for (Expr param : func->params) {
+ CreateSid(param.operator->());
+ }
+
+ GetSid(func->body);
+ return storage_device_map_;
+ }
+
+ void VisitExpr_(const ConstantNode* op) final { CreateSid(op); }
+
+ void VisitExpr_(const CallNode* op) final {
+ // create token for the call node.
+ CreateSid(op);
+ for (Expr arg : op->args) {
+ GetSid(arg);
+ }
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const FunctionNode* op) final {
+ // do not recurse into sub function.
+ }
+
+ void VisitExpr_(const GlobalVarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const OpNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const TupleNode* op) final {
+ std::vector<int> field_ids;
+ std::vector<int> field_sizes;
+ std::vector<int> field_types;
+ Expr expr = GetRef<Expr>(op);
+ for (Expr field : op->fields) {
+ auto sids = GetSid(field);
+ field_ids.insert(field_ids.end(), sids[0].begin(), sids[0].end());
+ field_types.insert(field_types.end(), sids[1].begin(), sids[1].end());
+ field_sizes.insert(field_sizes.end(), sids[2].begin(), sids[2].end());
+ }
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = field_ids;
+ storage_device_map_[expr][1] = field_sizes;
+ storage_device_map_[expr][2] = field_types;
Review comment:
Is this device type ? (I am wondering why storage_device_map_ does not have any device)
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -46,13 +47,171 @@ namespace backend {
using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap = std::unordered_map<Expr, std::vector<std::vector<int>>, runtime::ObjectPtrHash,
Review comment:
what does std::vector<std::vector<int>> represent? May be better to describe what they are?
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -46,13 +47,171 @@ namespace backend {
using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap = std::unordered_map<Expr, std::vector<std::vector<int>>, runtime::ObjectPtrHash,
+ runtime::ObjectPtrEqual>;
+/**
+ * This is an on demand allocator for AOT. A new temporary
+ * (storage allocator identifier) is allocated for each operation.
+ */
+class AOTOnDemandAllocator : public ExprVisitor {
+ public:
+ // run the visitor on a function.
+ StorageMap Run(const Function& func) {
+ node_device_map_ = CollectDeviceInfo(func);
+
+ for (Expr param : func->params) {
+ CreateSid(param.operator->());
+ }
+
+ GetSid(func->body);
+ return storage_device_map_;
+ }
+
+ void VisitExpr_(const ConstantNode* op) final { CreateSid(op); }
+
+ void VisitExpr_(const CallNode* op) final {
+ // create token for the call node.
+ CreateSid(op);
+ for (Expr arg : op->args) {
+ GetSid(arg);
+ }
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const FunctionNode* op) final {
+ // do not recurse into sub function.
+ }
+
+ void VisitExpr_(const GlobalVarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const OpNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const TupleNode* op) final {
+ std::vector<int> field_ids;
+ std::vector<int> field_sizes;
+ std::vector<int> field_types;
+ Expr expr = GetRef<Expr>(op);
+ for (Expr field : op->fields) {
+ auto sids = GetSid(field);
+ field_ids.insert(field_ids.end(), sids[0].begin(), sids[0].end());
+ field_types.insert(field_types.end(), sids[1].begin(), sids[1].end());
+ field_sizes.insert(field_sizes.end(), sids[2].begin(), sids[2].end());
+ }
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = field_ids;
+ storage_device_map_[expr][1] = field_sizes;
+ storage_device_map_[expr][2] = field_types;
+ }
+
+ void VisitExpr_(const TupleGetItemNode* op) final {
+ Expr expr = GetRef<Expr>(op);
+ const auto& sids = GetSid(op->tuple);
+ ICHECK_LT(static_cast<size_t>(op->index), sids.size());
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = {sids[0][op->index]};
+ storage_device_map_[expr][1] = {sids[1][op->index]};
+ storage_device_map_[expr][2] = {sids[2][op->index]};
+ }
+
+ 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:
+ /*!
+ * \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.
+ */
+ std::vector<std::vector<int>> 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) {
+ std::vector<int> sids;
+ std::vector<int> sizes;
+ std::vector<int> types;
+ 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);
+ sids.push_back(sid_++);
+ types.push_back(device_type);
+ sizes.push_back(GetMemorySize(ttype));
+ }
+ } else {
+ const auto* ttype = op->checked_type().as<TensorTypeNode>();
+ ICHECK(ttype);
+ sids.push_back(sid_++);
+ types.push_back(device_type);
+ sizes.push_back(GetMemorySize(ttype));
+ }
+ InitStorage(expr);
+ storage_device_map_[expr][0] = sids;
+ storage_device_map_[expr][1] = types;
+ storage_device_map_[expr][2] = sizes;
+ }
+
+ void InitStorage(Expr expr) {
+ if (storage_device_map_[expr].empty()) {
+ storage_device_map_[expr].push_back(std::vector<int>());
+ storage_device_map_[expr].push_back(std::vector<int>());
+ storage_device_map_[expr].push_back(std::vector<int>());
+ }
+ }
+
+ StorageMap storage_device_map_;
+ Map<Expr, Integer> node_device_map_;
+ int sid_{0};
+};
class AotReturnSidVisitor : public ExprVisitor {
public:
- explicit AotReturnSidVisitor(Map<Expr, Array<IntegerArray>> storage_device_map)
+ explicit AotReturnSidVisitor(StorageMap storage_device_map)
Review comment:
Its a bit unclear to me purpose of this pass. It might be better to add a brief describing the purpose and how it is used in the code generation.
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -46,13 +47,171 @@ namespace backend {
using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;
+using StorageMap = std::unordered_map<Expr, std::vector<std::vector<int>>, runtime::ObjectPtrHash,
+ runtime::ObjectPtrEqual>;
+/**
+ * This is an on demand allocator for AOT. A new temporary
+ * (storage allocator identifier) is allocated for each operation.
+ */
+class AOTOnDemandAllocator : public ExprVisitor {
+ public:
+ // run the visitor on a function.
+ StorageMap Run(const Function& func) {
+ node_device_map_ = CollectDeviceInfo(func);
+
+ for (Expr param : func->params) {
+ CreateSid(param.operator->());
+ }
+
+ GetSid(func->body);
+ return storage_device_map_;
+ }
+
+ void VisitExpr_(const ConstantNode* op) final { CreateSid(op); }
+
+ void VisitExpr_(const CallNode* op) final {
+ // create token for the call node.
+ CreateSid(op);
+ for (Expr arg : op->args) {
+ GetSid(arg);
+ }
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const FunctionNode* op) final {
+ // do not recurse into sub function.
+ }
+
+ void VisitExpr_(const GlobalVarNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const OpNode* op) final {
+ // Do nothing.
+ }
+
+ void VisitExpr_(const TupleNode* op) final {
+ std::vector<int> field_ids;
+ std::vector<int> field_sizes;
+ std::vector<int> field_types;
+ Expr expr = GetRef<Expr>(op);
+ for (Expr field : op->fields) {
+ auto sids = GetSid(field);
+ field_ids.insert(field_ids.end(), sids[0].begin(), sids[0].end());
+ field_types.insert(field_types.end(), sids[1].begin(), sids[1].end());
+ field_sizes.insert(field_sizes.end(), sids[2].begin(), sids[2].end());
+ }
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = field_ids;
+ storage_device_map_[expr][1] = field_sizes;
+ storage_device_map_[expr][2] = field_types;
+ }
+
+ void VisitExpr_(const TupleGetItemNode* op) final {
+ Expr expr = GetRef<Expr>(op);
+ const auto& sids = GetSid(op->tuple);
+ ICHECK_LT(static_cast<size_t>(op->index), sids.size());
+ if (storage_device_map_[expr].empty()) {
+ InitStorage(expr);
+ }
+ storage_device_map_[expr][0] = {sids[0][op->index]};
+ storage_device_map_[expr][1] = {sids[1][op->index]};
+ storage_device_map_[expr][2] = {sids[2][op->index]};
+ }
+
+ 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:
+ /*!
+ * \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.
+ */
+ std::vector<std::vector<int>> 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) {
+ std::vector<int> sids;
+ std::vector<int> sizes;
+ std::vector<int> types;
+ 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);
+ sids.push_back(sid_++);
+ types.push_back(device_type);
+ sizes.push_back(GetMemorySize(ttype));
+ }
+ } else {
+ const auto* ttype = op->checked_type().as<TensorTypeNode>();
+ ICHECK(ttype);
+ sids.push_back(sid_++);
+ types.push_back(device_type);
+ sizes.push_back(GetMemorySize(ttype));
+ }
+ InitStorage(expr);
+ storage_device_map_[expr][0] = sids;
+ storage_device_map_[expr][1] = types;
Review comment:
Where is this used in this code (types.push_back(device_type);) ?
--
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