[Relay][Runtime] Add memory manager for NDArray (#3121)

* Add support for custom NDArray memory management

Credit to @icemelon9 and @wweic

* Fix copy-paste issue

* Fix naive allocator.h

* Remove buffer field

* Apply Wei's suggestions.

Co-Authored-By: jroesch <roeschinc@gmail.com>

* Fix Wei's suggestion

* Fix go rts

* Break MM dependency

* Add docs and clean up diff

* Add more docs

* Move to VM folder

* Fix lint

* Remove Go dep.

* Rename to Empty

* Address Haichen's comments

golang/src/tvm_runtime_pack.cc

@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

include/tvm/runtime/ndarray.h

@@ -19,7 +19,7 @@
 
 /*!
  * \file tvm/runtime/ndarray.h
- * \brief Abstract device memory management API
+ * \brief A device-independent managed NDArray abstraction.
  */
 #ifndef TVM_RUNTIME_NDARRAY_H_
 #define TVM_RUNTIME_NDARRAY_H_
@@ -32,6 +32,7 @@
 
 namespace tvm {
 namespace runtime {
+
 /*!
  * \brief Managed NDArray.
  *  The array is backed by reference counted blocks.
@@ -248,6 +249,7 @@ class NDArray::Container {
    *  The head ptr of this struct can be viewed as DLTensor*.
    */
   DLTensor dl_tensor;
+
   /*!
    * \brief addtional context, reserved for recycling
    * \note We can attach additional content here
@@ -281,6 +283,7 @@ class NDArray::Container {
   int32_t array_type_code_{0};
   /*! \brief The internal reference counter */
   std::atomic<int> ref_counter_{0};
+
   /*!
    * \brief The shape container,
    *  can be used used for shape data.
@@ -296,6 +299,19 @@ class NDArray::Container {
     dl_tensor.strides = nullptr;
     dl_tensor.byte_offset = 0;
   }
+
+  Container(void* data,
+            std::vector<int64_t> shape,
+            DLDataType dtype,
+            DLContext ctx) {
+    dl_tensor.data = data;
+    shape_ = std::move(shape);
+    dl_tensor.shape = dmlc::BeginPtr(shape);
+    dl_tensor.ndim = static_cast<int>(shape.size());
+    dl_tensor.dtype = dtype;
+    dl_tensor.ctx = ctx;
+  }
+
   /*! \brief developer function, increases reference counter */
   void IncRef() {
     ref_counter_.fetch_add(1, std::memory_order_relaxed);

src/runtime/vm/memory_manager.cc

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file tvm/runtime/memory_manager.cc
+ * \brief Allocate and manage memory for the runtime.
+ */
+#include <utility>
+#include <memory>
+#include "memory_manager.h"
+#include "naive_allocator.h"
+#include "pooled_allocator.h"
+
+namespace tvm {
+namespace runtime {
+namespace vm {
+
+MemoryManager* MemoryManager::Global() {
+  static MemoryManager memory_manager;
+  return &memory_manager;
+}
+
+Allocator* MemoryManager::GetAllocator(TVMContext ctx) {
+  std::lock_guard<std::mutex> lock(mu_);
+  if (allocators_.find(ctx) == allocators_.end()) {
+    // LOG(INFO) << "New allocator for " << DeviceName(ctx.device_type) << "("
+    //           << ctx.device_id << ")";
+    std::unique_ptr<Allocator> alloc(new NaiveAllocator(ctx));
+    allocators_.emplace(ctx, std::move(alloc));
+  }
+  return allocators_.at(ctx).get();
+}
+
+static void BufferDeleter(NDArray::Container* ptr) {
+  CHECK(ptr->manager_ctx != nullptr);
+  Buffer* buffer = reinterpret_cast<Buffer*>(ptr->manager_ctx);
+  MemoryManager::Global()->GetAllocator(buffer->ctx)->
+      Free(*(buffer));
+  delete buffer;
+  delete ptr;
+}
+
+NDArray Allocator::Empty(std::vector<int64_t> shape, DLDataType dtype, DLContext ctx) {
+  VerifyDataType(dtype);
+  NDArray::Container* container = new NDArray::Container(nullptr, shape, dtype, ctx);
+  container->deleter = 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);
+  container->manager_ctx = reinterpret_cast<void*>(buffer);
+  container->dl_tensor.data = buffer->data;
+  return NDArray(container);
+}
+
+}  // namespace vm
+}  // namespace runtime
+}  // namespace tvm

src/runtime/vm/memory_manager.h

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file src/runtime/memory_manager.h
+ * \brief Abstract device memory management API
+ */
+#ifndef TVM_RUNTIME_VM_MEMORY_MANAGER_H_
+#define TVM_RUNTIME_VM_MEMORY_MANAGER_H_
+
+#include <tvm/runtime/c_runtime_api.h>
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+
+namespace std {
+template <>
+struct hash<TVMContext> {
+  std::size_t operator()(const TVMContext& ctx) const {
+    return ((ctx.device_id << 8) | ctx.device_type);
+  }
+};
+
+template <>
+struct equal_to<TVMContext> {
+  bool operator()(const TVMContext& lhs, const TVMContext& rhs) const {
+    return (lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id);
+  }
+};
+
+}  // namespace std
+
+namespace tvm {
+namespace runtime {
+namespace vm {
+
+struct Buffer {
+  /*! \brief The pointer to the allocated block of memory. */
+  void* data{nullptr};
+  /*! \brief The size of the block. */
+  size_t size{0};
+  /*! \brief The context of the allocated buffers. */
+  TVMContext ctx;
+};
+
+class Allocator {
+ public:
+  Allocator() {}
+
+  /*! \brief Allocate an empty NDArray using from the allocator.
+   *  \param shape The shape of the NDArray.
+   *  \param alignment The datatype of the NDArray.
+   *  \param ctx The context where the array is allocated.
+   *  \return The empty NDArray.
+   */
+  NDArray Empty(std::vector<int64_t> shape,
+                DLDataType dtype,
+                DLContext ctx);
+  /*! \brief Allocate a buffer given a size, alignment and type.
+   *  \param nbytes The size of the buffer.
+   *  \param alignment The alignment of the buffer.
+   *  \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, TVMType type_hint) = 0;
+  /*! \brief Free a buffer allocated by the allocator.
+   *  \param buffer The buffer to free.
+   */
+  virtual void Free(const Buffer& buffer) = 0;
+  /*! \brief The amount of memory currently allocated.
+   *  \return The amount of memory currently allocated.
+   */
+  virtual size_t UsedMemory() const = 0;
+  virtual ~Allocator() = default;
+};
+
+class MemoryManager {
+ public:
+  static MemoryManager* Global();
+
+  Allocator* GetAllocator(TVMContext ctx);
+
+ private:
+  MemoryManager() {}
+
+ private:
+  std::mutex mu_;
+  std::unordered_map<TVMContext, std::unique_ptr<Allocator> > allocators_;
+};
+
+}  // namespace vm
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_VM_MEMORY_MANAGER_H_

src/runtime/vm/naive_allocator.h

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file src/runtime/naive_allocator.h
+ */
+#ifndef TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
+#define TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
+
+#include <tvm/runtime/device_api.h>
+#include <atomic>
+
+#include "memory_manager.h"
+
+namespace tvm {
+namespace runtime {
+namespace vm {
+
+class NaiveAllocator final : public Allocator {
+ public:
+  explicit NaiveAllocator(TVMContext ctx) : Allocator(), used_memory_(0) {}
+
+  Buffer Alloc(size_t nbytes, size_t alignment, TVMType type_hint) override {
+    Buffer buf;
+    buf.ctx = ctx_;
+    buf.size = nbytes;
+    buf.data = DeviceAPI::Get(ctx_)->AllocDataSpace(ctx_, nbytes, alignment, type_hint);
+    used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
+    DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_ << " B";
+    return buf;
+  }
+
+  void Free(const Buffer& buffer) override {
+    DeviceAPI::Get(ctx_)->FreeDataSpace(buffer.ctx, buffer.data);
+    used_memory_.fetch_sub(buffer.size, std::memory_order_relaxed);
+    DLOG(INFO) << "free " << buffer.size << " B, used memory " << used_memory_ << " B";
+  }
+
+  size_t UsedMemory() const override {
+    return used_memory_.load(std::memory_order_relaxed);
+  }
+
+ private:
+  std::atomic<size_t> used_memory_;
+  TVMContext ctx_;
+};
+
+}  // namespace vm
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_

src/runtime/vm/pooled_allocator.h

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file runtime/pooled_allocator.h
+ */
+#ifndef TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
+#define TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
+
+#include <tvm/runtime/device_api.h>
+#include <atomic>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+
+#include "memory_manager.h"
+
+namespace tvm {
+namespace runtime {
+namespace vm {
+
+class PooledAllocator final : public Allocator {
+ public:
+  static constexpr size_t kDefaultPageSize = 4096;
+
+  explicit PooledAllocator(TVMContext ctx, size_t page_size = kDefaultPageSize)
+      : Allocator(), page_size_(page_size), used_memory_(0) {}
+
+  ~PooledAllocator() { ReleaseAll(); }
+
+  Buffer Alloc(size_t nbytes, size_t alignment, TVMType type_hint) override {
+    std::lock_guard<std::mutex> lock(mu_);
+    size_t size = ((nbytes + page_size_ - 1) / page_size_) * page_size_;
+    auto&& it = memory_pool_.find(size);
+    if (it != memory_pool_.end() && !it->second.empty()) {
+      auto&& pool = it->second;
+      auto ret = pool.back();
+      pool.pop_back();
+      return ret;
+    }
+    Buffer buf;
+    buf.ctx = ctx_;
+    buf.size = size;
+    buf.data = DeviceAPI::Get(ctx_)->AllocDataSpace(ctx_, size, alignment, type_hint);
+    used_memory_.fetch_add(size, std::memory_order_relaxed);
+    DLOG(INFO) << "allocate " << size << " B, used memory " << used_memory_ << " B";
+    return buf;
+  }
+
+  void Free(const Buffer& buffer) override {
+    std::lock_guard<std::mutex> lock(mu_);
+    if (memory_pool_.find(buffer.size) == memory_pool_.end()) {
+      memory_pool_.emplace(buffer.size, std::vector<Buffer>{});
+    }
+    memory_pool_.at(buffer.size).push_back(buffer);
+    DLOG(INFO) << "reclaim buffer " << buffer.size;
+  }
+
+  size_t UsedMemory() const override { return used_memory_.load(std::memory_order_relaxed); }
+
+ private:
+  void ReleaseAll() {
+    std::lock_guard<std::mutex> lock(mu_);
+    for (auto const& it : memory_pool_) {
+      auto const& pool = it.second;
+      for (auto const& buf : pool) {
+        DeviceAPI::Get(buf.ctx)->FreeDataSpace(buf.ctx, buf.data);
+      }
+    }
+    memory_pool_.clear();
+    used_memory_ = 0;
+    DLOG(INFO) << "release all buffers";
+  }
+
+ private:
+  size_t page_size_;
+  std::atomic<size_t> used_memory_;
+  std::unordered_map<size_t, std::vector<Buffer> > memory_pool_;
+  std::mutex mu_;
+  TVMContext ctx_;
+};
+
+}  // namespace vm
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_