[GitHub] zheng-da commented on a change in pull request #11658: Optimize cached op static memory allocation

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

zheng-da commented on a change in pull request #11658: Optimize cached op static memory allocation
URL: https://github.com/apache/incubator-mxnet/pull/11658#discussion_r212437803
 
 

 ##########
 File path: src/imperative/imperative_utils.h
 ##########
 @@ -789,22 +789,36 @@ inline MemoryPlanVector PlanMemory(
 }
 
 
-inline std::multimap<size_t, NDArray> AllocateMemory(
+inline NDArray AllocateMemory(
     const nnvm::Graph& g,
     const nnvm::IndexedGraph& idx,
     const Context& default_ctx,
     const uint32_t entry_start, const uint32_t entry_end,
     const MemoryPlanVector& mem_plan,
     const std::vector<NDArray*>& arrays,
     std::vector<OpReqType> *array_reqs,
-    std::multimap<size_t, NDArray>&& pool = std::multimap<size_t, NDArray>()) {
+    const bool use_pool = false,
+    NDArray pool = NDArray()) {
   using namespace nnvm;
   const auto& dtypes = g.GetAttr<DTypeVector>("dtype");
   const auto& shapes = g.GetAttr<ShapeVector>("shape");
   const auto& stypes = g.GetAttr<StorageTypeVector>("storage_type");
+  const size_t page_size = dmlc::GetEnv("MXNET_GPU_MEM_POOL_PAGE_SIZE", 4096);
 
-  std::multimap<size_t, NDArray> new_pool;
+  size_t total_size = 0;
+  for (uint32_t i = entry_start; i < entry_end; ++i) {
+    if (mem_plan[i].storage_id == exec::kExternalStorageID ||
+        mem_plan[i].storage_id == exec::kDynamicStorageID ||
+        mem_plan[i].root != i) continue;
+    total_size += (mem_plan[i].size + page_size - 1) / page_size * page_size;
 
 Review comment:
   every NDArray is aligned to the page size? I thought you want to use the alignment for the default GPU memory allocation.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on 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


With regards,
Apache Git Services