* Memory Manager

Move the VM memory manager to the runtime level.
Use this memory manager for graph runtime.

Author: srkreddy1238

include/tvm/runtime/device_api.h

@@ -95,6 +95,14 @@ class TVM_DLL DeviceAPI {
    */
   virtual void GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) = 0;
 
+  /*!
+   * \brief Get the physical memory size required.
+   * \param arr the tensor object.
+   * \param mem_scope the memory scope if any
+   * \return the memory size.
+   */
+  virtual size_t GetDataSize(const DLTensor& arr, Optional<String> mem_scope = NullOpt);
+
   /*!
    * \brief Query the device for specified properties.
    *

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>
@@ -36,7 +36,6 @@
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
 struct Buffer {
   /*! \brief The pointer to the allocated block of memory. */
@@ -60,9 +59,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 optional memory scope
    *  \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 = NullOpt);
   /*! \brief Return the allocator type. */
   inline AllocatorType type() const { return type_; }
   /*! \brief Allocate a buffer given a size, alignment and type.
@@ -140,8 +141,7 @@ class Storage : public ObjectRef {
   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/c_runtime_api.cc

@@ -176,6 +176,20 @@ void* DeviceAPI::AllocDataSpaceView(Device dev, void* data, int ndim, const int6
   return data;
 }
 
+size_t DeviceAPI::GetDataSize(const DLTensor& arr, Optional<String> mem_scope) {
+  if (!mem_scope.defined() || mem_scope.value() == "global") {
+    size_t size = 1;
+    for (tvm_index_t i = 0; i < arr.ndim; ++i) {
+      size *= static_cast<size_t>(arr.shape[i]);
+    }
+    size *= (arr.dtype.bits * arr.dtype.lanes + 7) / 8;
+    return size;
+  }
+  LOG(FATAL) << "Device does not support physical mem computation with "
+             << "specified memory scope: " << mem_scope.value();
+  return 0;
+}
+
 void DeviceAPI::CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream) {
   // by default, we can always redirect to the flat memory copy operation.
   size_t nbytes = GetDataSize(*from);

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>
@@ -472,7 +473,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::kPooled)
+                                  ->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,7 +31,6 @@
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
 static void BufferDeleter(Object* obj) {
   auto* ptr = static_cast<NDArray::Container*>(obj);
@@ -157,19 +156,20 @@ 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 size = DeviceAPI::Get(dev)->GetDataSize(container->dl_tensor, mem_scope);
   size_t alignment = GetDataAlignment(container->dl_tensor);
   Buffer* buffer = new Buffer;
+
   *buffer = this->Alloc(size, alignment, dtype);
-  container->manager_ctx = reinterpret_cast<void*>(buffer);
   container->dl_tensor.data = buffer->data;
+  container->manager_ctx = reinterpret_cast<void*>(buffer);
   return NDArray(GetObjectPtr<Object>(container));
 }
 
-}  // namespace vm
 }  // namespace runtime
 }  // namespace tvm

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

@@ -20,17 +20,16 @@
 /*!
  * \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>
 
 namespace tvm {
 namespace runtime {
-namespace vm {
 
 class NaiveAllocator final : public Allocator {
  public:
@@ -59,8 +58,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_