[Relay][Runtime] Add memory manager for NDArray

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
@@ -32,6 +32,7 @@
 #include "src/runtime/threading_backend.cc"
 #include "src/runtime/thread_pool.cc"
 #include "src/runtime/ndarray.cc"
+#include "src/runtime/memory_manager.cc"
 
 // NOTE: all the files after this are optional modules
 // that you can include remove, depending on how much feature you use.

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,10 @@
 
 namespace tvm {
 namespace runtime {
+
+class Allocator;
+struct Buffer;
+
 /*!
  * \brief Managed NDArray.
  *  The array is backed by reference counted blocks.
@@ -163,11 +167,13 @@ class NDArray {
    * \param shape The shape of the new array.
    * \param dtype The data type of the new array.
    * \param ctx The context of the Array.
+   * \param allocator The memory allocator.
    * \return The created Array
    */
   TVM_DLL static NDArray Empty(std::vector<int64_t> shape,
                                DLDataType dtype,
-                               DLContext ctx);
+                               DLContext ctx,
+                               Allocator* allocator = nullptr);
   /*!
    * \brief Create a NDArray backed by a dlpack tensor.
    *
@@ -248,6 +254,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 +288,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.

src/runtime/memory_manager.cc

@@ -0,0 +1,51 @@
+/*
+ * 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 {
+
+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();
+}
+
+}  // namespace runtime
+}  // namespace tvm

src/runtime/memory_manager.h

@@ -0,0 +1,94 @@
+/*
+ * 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_MEMORY_MANAGER_H_
+#define TVM_RUNTIME_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 {
+
+struct Buffer {
+  // data pointer
+  void* data{nullptr};
+  // Buffer size in bytes
+  size_t size{0};
+  // TVM Context
+  TVMContext ctx;
+};
+
+class Allocator {
+ public:
+  explicit Allocator(TVMContext ctx) : ctx_(ctx) {}
+
+  virtual Buffer Alloc(size_t nbytes, size_t alignment, TVMType type_hint) = 0;
+  virtual void Free(const Buffer& buffer) = 0;
+  virtual size_t UsedMemory() = 0;
+  virtual ~Allocator() = default;
+
+ protected:
+  TVMContext ctx_;
+};
+
+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 runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_MEMORY_MANAGER_H_

src/runtime/naive_allocator.h

@@ -0,0 +1,64 @@
+/*
+ * 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_NAIVE_ALLOCATOR_H_
+#define TVM_RUNTIME_NAIVE_ALLOCATOR_H_
+
+#include <tvm/runtime/device_api.h>
+#include <atomic>
+
+#include "memory_manager.h"
+
+namespace tvm {
+namespace runtime {
+
+class NaiveAllocator final : public Allocator {
+ public:
+  explicit NaiveAllocator(TVMContext ctx) : Allocator(ctx), 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() override { return used_memory_.load(std::memory_order_relaxed); }
+
+ private:
+  std::atomic<size_t> used_memory_;
+};
+
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_NAIVE_ALLOCATOR_H_

src/runtime/ndarray.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
@@ -26,6 +26,8 @@
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/device_api.h>
+
+#include "memory_manager.h"
 #include "runtime_base.h"
 
 // deleter for arrays used by DLPack exporter
@@ -76,15 +78,28 @@ struct NDArray::Internal {
     }
     delete ptr;
   }
+
+  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;
+  }
   // Local create function which allocates tensor metadata
   // but does not allocate space for the data.
   static NDArray Create(std::vector<int64_t> shape,
                         DLDataType dtype,
-                        DLContext ctx) {
+                        DLContext ctx, bool with_allocator = false) {
     VerifyDataType(dtype);
     // critical zone
     NDArray::Container* data = new NDArray::Container();
-    data->deleter = DefaultDeleter;
+    if (with_allocator) {
+      data->deleter = BufferDeleter;
+    } else {
+      data->deleter = DefaultDeleter;
+    }
     NDArray ret(data);
     ret.data_ = data;
     // RAII now in effect
@@ -142,14 +157,23 @@ DLManagedTensor* NDArray::ToDLPack() const {
 
 NDArray NDArray::Empty(std::vector<int64_t> shape,
                        DLDataType dtype,
-                       DLContext ctx) {
-  NDArray ret = Internal::Create(shape, dtype, ctx);
+                       DLContext ctx,
+                       Allocator* allocator) {
+  NDArray ret = Internal::Create(shape, dtype, ctx, (allocator != nullptr));
   // setup memory content
   size_t size = GetDataSize(ret.data_->dl_tensor);
   size_t alignment = GetDataAlignment(ret.data_->dl_tensor);
-  ret.data_->dl_tensor.data =
-      DeviceAPI::Get(ret->ctx)->AllocDataSpace(
-          ret->ctx, size, alignment, ret->dtype);
+  if (allocator == nullptr) {
+    ret.data_->dl_tensor.data =
+        DeviceAPI::Get(ret->ctx)->AllocDataSpace(
+            ret->ctx, size, alignment, ret->dtype);
+  } else {
+    Buffer *buffer = new Buffer;
+    ret.data_->manager_ctx = new Buffer;
+    *buffer = allocator->Alloc(size, alignment, ret->dtype);
+    ret.data_->manager_ctx = reinterpret_cast<void*>(buffer);
+    ret.data_->dl_tensor.data = buffer->data;
+  }
   return ret;
 }