[VM] memory Manager moved up to runtime

Now graph runtime also uses the same memory manager
This acommodates a common memory manager with pooled and naive support.

As a follow up we can move the WorkspacePool to use this common memory manager.

Author: srkreddy1238

include/tvm/runtime/vm/memory_manager.h renamed to include/tvm/runtime/memory_manager.h

@@ -18,11 +18,11 @@
  */
 
 /*!
- * \file tvm/runtime/vm/memory_manager.h
+ * \file tvm/runtime/memory_manager.h
  * \brief Abstract device memory management API
  */
-#ifndef TVM_RUNTIME_VM_MEMORY_MANAGER_H_
-#define TVM_RUNTIME_VM_MEMORY_MANAGER_H_
+#ifndef TVM_RUNTIME_MEMORY_MANAGER_H_
+#define TVM_RUNTIME_MEMORY_MANAGER_H_
 
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/ndarray.h>
@@ -37,9 +37,8 @@
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
-struct Buffer {
+struct MBuffer {
   /*! \brief The pointer to the allocated block of memory. */
   void* data{nullptr};
   /*! \brief The size of the block. */
@@ -63,9 +62,11 @@ class Allocator {
    *  \param shape The shape of the NDArray.
    *  \param dtype The datatype of the NDArray.
    *  \param dev The device where the array is allocated.
+   *  \param mem_scope is the device memory scope hint.
    *  \return The empty NDArray.
    */
-  NDArray Empty(std::vector<int64_t> shape, DLDataType dtype, Device dev);
+  NDArray Empty(std::vector<int64_t> shape, DLDataType dtype, Device dev,
+                Optional<String> mem_scope);
   /*! \brief Return the allocator type. */
   inline AllocatorType type() const { return type_; }
   /*! \brief Allocate a buffer given a size, alignment and type.
@@ -74,28 +75,28 @@ class Allocator {
    *  \param type_hint A type hint to the allocator.
    *  \return A sized allocation in the form of a buffer.
    */
-  virtual Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) = 0;
+  virtual MBuffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) = 0;
   /*! \brief Allocate a buffer given a shape and type.
    *  \param ndims The rank of the tensor.
    *  \param shape The shape of the tensor.
    *  \param type_hint A type hint to the allocator.
    *  \param mem_scope A memory scope of the buffer.
    *  \return A sized allocation in the form of a buffer.
    */
-  virtual Buffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
-                       const std::string& mem_scope = "") = 0;
+  virtual MBuffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
+                        const std::string& mem_scope = "") = 0;
   /*! \brief Free a buffer allocated by the allocator.
    *  \param buffer The buffer to free.
    */
-  virtual void Free(const Buffer& buffer) = 0;
+  virtual void Free(const MBuffer& buffer) = 0;
   /*! \brief The amount of memory currently allocated.
    *  \return The amount of memory currently allocated.
    */
   virtual size_t UsedMemory() const = 0;
 
  protected:
-  virtual Buffer Alloc(Device dev, int ndims, int64_t* shape, DLDataType type_hint,
-                       const std::string& mem_scope);
+  virtual MBuffer Alloc(Device dev, int ndims, int64_t* shape, DLDataType type_hint,
+                        const std::string& mem_scope);
 
  private:
   AllocatorType type_;
@@ -130,7 +131,7 @@ class MemoryManager {
 class StorageObj : public Object {
  public:
   /*! \brief The index into the VM function table. */
-  Buffer buffer;
+  MBuffer buffer;
 
   /*! \brief Allocate an NDArray from a given piece of storage. */
   NDArray AllocNDArray(size_t offset, std::vector<int64_t> shape, DLDataType dtype);
@@ -151,13 +152,12 @@ class StorageObj : public Object {
 /*! \brief reference to storage. */
 class Storage : public ObjectRef {
  public:
-  explicit Storage(Buffer buffer);
+  explicit Storage(MBuffer buffer);
 
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Storage, ObjectRef, StorageObj);
 };
 
-}  // namespace vm
 }  // namespace runtime
 }  // namespace tvm
 
-#endif  // TVM_RUNTIME_VM_MEMORY_MANAGER_H_
+#endif  // TVM_RUNTIME_MEMORY_MANAGER_H_

include/tvm/runtime/vm/vm.h

@@ -25,13 +25,13 @@
 #define TVM_RUNTIME_VM_VM_H_
 
 #include <tvm/runtime/container/closure.h>
+#include <tvm/runtime/memory_manager.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/object.h>
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/runtime/vm/bytecode.h>
 #include <tvm/runtime/vm/executable.h>
-#include <tvm/runtime/vm/memory_manager.h>
 
 #include <memory>
 #include <string>

src/relay/backend/vm/compiler.h

@@ -41,7 +41,7 @@
 #include <utility>
 #include <vector>
 
-#include "../../../runtime/vm/naive_allocator.h"
+#include "../../../runtime/naive_allocator.h"
 #include "../../../runtime/vm/profiler/vm.h"
 #include "../../transforms/pass_utils.h"
 #include "../te_compiler.h"

src/runtime/graph_executor/graph_executor.cc

@@ -26,6 +26,7 @@
 #include <tvm/runtime/container/string.h>
 #include <tvm/runtime/data_type.h>
 #include <tvm/runtime/device_api.h>
+#include <tvm/runtime/memory_manager.h>
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/profiling.h>
@@ -466,7 +467,8 @@ void GraphExecutor::SetupStorage() {
       if (!pit.scope.empty()) {
         mem_scope = String(pit.scope);
       }
-      storage_pool_.push_back(NDArray::Empty(shape, pit.dtype, dev, mem_scope));
+      storage_pool_.push_back(MemoryManager::GetOrCreateAllocator(dev, AllocatorType::kNaive)
+                                  ->Empty(shape, pit.dtype, dev, mem_scope));
     }
   }
 

src/runtime/vm/memory_manager.cc renamed to src/runtime/memory_manager.cc

@@ -18,10 +18,10 @@
  */
 
 /*!
- * \file tvm/runtime/vm/memory_manager.cc
+ * \file tvm/runtime/memory_manager.cc
  * \brief Allocate and manage memory for the runtime.
  */
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory_manager.h>
 
 #include <memory>
 #include <utility>
@@ -31,12 +31,11 @@
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
 static void BufferDeleter(Object* obj) {
   auto* ptr = static_cast<NDArray::Container*>(obj);
   ICHECK(ptr->manager_ctx != nullptr);
-  Buffer* buffer = reinterpret_cast<Buffer*>(ptr->manager_ctx);
+  MBuffer* buffer = reinterpret_cast<MBuffer*>(ptr->manager_ctx);
   MemoryManager::GetAllocator(buffer->device)->Free(*(buffer));
   delete buffer;
   delete ptr;
@@ -154,21 +153,26 @@ Allocator* MemoryManager::GetAllocator(Device dev) {
   return it->second.get();
 }
 
-NDArray Allocator::Empty(std::vector<int64_t> shape, DLDataType dtype, DLDevice dev) {
+NDArray Allocator::Empty(std::vector<int64_t> shape, DLDataType dtype, DLDevice dev,
+                         Optional<String> mem_scope) {
   VerifyDataType(dtype);
   NDArray::Container* container = new NDArray::Container(nullptr, shape, dtype, dev);
   container->SetDeleter(BufferDeleter);
   size_t size = GetDataSize(container->dl_tensor);
   size_t alignment = GetDataAlignment(container->dl_tensor);
-  Buffer* buffer = new Buffer;
-  *buffer = this->Alloc(size, alignment, dtype);
+  MBuffer* buffer = new MBuffer;
+  if (!mem_scope.defined() || mem_scope == "global") {
+    *buffer = this->Alloc(size, alignment, dtype);
+  } else {
+    *buffer = this->Alloc(shape.size(), shape.data(), dtype, mem_scope.value());
+  }
   container->manager_ctx = reinterpret_cast<void*>(buffer);
   container->dl_tensor.data = buffer->data;
   return NDArray(GetObjectPtr<Object>(container));
 }
 
-Buffer Allocator::Alloc(Device dev, int ndims, int64_t* shape, DLDataType type_hint,
-                        const std::string& mem_scope) {
+MBuffer Allocator::Alloc(Device dev, int ndims, int64_t* shape, DLDataType type_hint,
+                         const std::string& mem_scope) {
   if (mem_scope.empty() || mem_scope == "global") {
     // by default, we can always redirect to the flat memory allocations
     std::vector<int64_t> s;
@@ -185,6 +189,5 @@ Buffer Allocator::Alloc(Device dev, int ndims, int64_t* shape, DLDataType type_h
   return {};
 }
 
-}  // namespace vm
 }  // namespace runtime
 }  // namespace tvm

src/runtime/vm/naive_allocator.h renamed to src/runtime/naive_allocator.h

@@ -20,25 +20,24 @@
 /*!
  * \file src/runtime/naive_allocator.h
  */
-#ifndef TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
-#define TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
+#ifndef TVM_RUNTIME_NAIVE_ALLOCATOR_H_
+#define TVM_RUNTIME_NAIVE_ALLOCATOR_H_
 
 #include <tvm/runtime/device_api.h>
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory_manager.h>
 
 #include <atomic>
 #include <string>
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
 class NaiveAllocator final : public Allocator {
  public:
   explicit NaiveAllocator(Device dev) : Allocator(kNaive), used_memory_(0), device_(dev) {}
 
-  Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
-    Buffer buf;
+  MBuffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
+    MBuffer buf;
     buf.device = device_;
     buf.size = nbytes;
     buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, nbytes, alignment, type_hint);
@@ -47,9 +46,9 @@ class NaiveAllocator final : public Allocator {
     return buf;
   }
 
-  Buffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
-               const std::string& mem_scope) override {
-    Buffer buf;
+  MBuffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
+                const std::string& mem_scope) override {
+    MBuffer buf;
     size_t nbytes = 1;
     for (int i = 0; i < ndims; ++i) {
       buf.shape.push_back(shape[i]);
@@ -72,7 +71,7 @@ class NaiveAllocator final : public Allocator {
     return buf;
   }
 
-  void Free(const Buffer& buffer) override {
+  void Free(const MBuffer& buffer) override {
     DeviceAPI::Get(device_)->FreeDataSpace(buffer.device, buffer.data);
     used_memory_.fetch_sub(buffer.size, std::memory_order_relaxed);
     DLOG(INFO) << "free " << buffer.size << " B, used memory " << used_memory_ << " B";
@@ -85,8 +84,7 @@ class NaiveAllocator final : public Allocator {
   Device device_;
 };
 
-}  // namespace vm
 }  // namespace runtime
 }  // namespace tvm
 
-#endif  // TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
+#endif  // TVM_RUNTIME_NAIVE_ALLOCATOR_H_

src/runtime/vm/pooled_allocator.h renamed to src/runtime/pooled_allocator.h

@@ -20,11 +20,11 @@
 /*!
  * \file runtime/pooled_allocator.h
  */
-#ifndef TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
-#define TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
+#ifndef TVM_RUNTIME_POOLED_ALLOCATOR_H_
+#define TVM_RUNTIME_POOLED_ALLOCATOR_H_
 
 #include <tvm/runtime/device_api.h>
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory_manager.h>
 
 #include <atomic>
 #include <mutex>
@@ -34,7 +34,6 @@
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
 class PooledAllocator final : public Allocator {
  public:
@@ -45,7 +44,7 @@ class PooledAllocator final : public Allocator {
 
   ~PooledAllocator() { ReleaseAll(); }
 
-  Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
+  MBuffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
     std::lock_guard<std::recursive_mutex> lock(mu_);
     size_t size = ((nbytes + page_size_ - 1) / page_size_) * page_size_;
     auto&& it = memory_pool_.find(size);
@@ -55,7 +54,7 @@ class PooledAllocator final : public Allocator {
       pool.pop_back();
       return ret;
     }
-    Buffer buf;
+    MBuffer buf;
     buf.device = device_;
     buf.size = size;
     try {
@@ -72,19 +71,19 @@ class PooledAllocator final : public Allocator {
     return buf;
   }
 
-  Buffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
-               const std::string& mem_scope) override {
+  MBuffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
+                const std::string& mem_scope) override {
     if (mem_scope.empty() || mem_scope == "global") {
       return Allocator::Alloc(device_, ndims, shape, type_hint, mem_scope);
     }
     LOG(FATAL) << "This alloc should be implemented";
     return {};
   }
 
-  void Free(const Buffer& buffer) override {
+  void Free(const MBuffer& buffer) override {
     std::lock_guard<std::recursive_mutex> lock(mu_);
     if (memory_pool_.find(buffer.size) == memory_pool_.end()) {
-      memory_pool_.emplace(buffer.size, std::vector<Buffer>{});
+      memory_pool_.emplace(buffer.size, std::vector<MBuffer>{});
     }
     memory_pool_.at(buffer.size).push_back(buffer);
     VLOG(1) << "reclaim buffer " << buffer.size;
@@ -109,13 +108,12 @@ class PooledAllocator final : public Allocator {
  private:
   size_t page_size_;
   std::atomic<size_t> used_memory_;
-  std::unordered_map<size_t, std::vector<Buffer>> memory_pool_;
+  std::unordered_map<size_t, std::vector<MBuffer>> memory_pool_;
   std::recursive_mutex mu_;
   Device device_;
 };
 
-}  // namespace vm
 }  // namespace runtime
 }  // namespace tvm
 
-#endif  // TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
+#endif  // TVM_RUNTIME_POOLED_ALLOCATOR_H_