This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new b8abff9789 [VM] Memory Manager moved up to runtime (#15833)
b8abff9789 is described below
commit b8abff978972ac8648ee94aeafd43a1993f31a63
Author: Siva <[email protected]>
AuthorDate: Wed Oct 4 04:01:58 2023 +0530
[VM] Memory Manager moved up to runtime (#15833)
* [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.
* * update dependents with new file addition.
* * define memory_manager under new namespace
* * use ShapeTuple across vm executor and memory_manager
* * ShapeTuple across the Allocators
* * GetDataSize is moved to DeviceAPI and memory_manager uses this
interface.
* * review comments
* * Make compiler happy with unused variables
* * lint
* Update src/runtime/memory/memory_manager.cc
Co-authored-by: Egor Churaev <[email protected]>
* * allow multiple allocators to coexist for the same device.
Using available allocator instead of requested is leading to an unpexpected
crash
---------
Co-authored-by: Egor Churaev <[email protected]>
---
CMakeLists.txt | 1 +
apps/android_camera/app/src/main/jni/tvm_runtime.h | 1 +
apps/android_deploy/app/src/main/jni/tvm_runtime.h | 1 +
apps/android_rpc/app/src/main/jni/tvm_runtime.h | 1 +
apps/bundle_deploy/runtime.cc | 1 +
apps/howto_deploy/tvm_runtime_pack.cc | 1 +
golang/src/tvm_runtime_pack.cc | 1 +
include/tvm/runtime/device_api.h | 8 +++
.../tvm/runtime/{vm => memory}/memory_manager.h | 45 +++++++++--------
include/tvm/runtime/vm/vm.h | 9 +++-
src/relay/backend/vm/compiler.h | 2 +-
src/runtime/c_runtime_api.cc | 14 +++++
src/runtime/graph_executor/graph_executor.cc | 3 +-
src/runtime/graph_executor/graph_executor.h | 4 ++
src/runtime/{vm => memory}/memory_manager.cc | 59 +++++++++++++---------
src/runtime/{vm => memory}/naive_allocator.h | 29 ++++++-----
src/runtime/{vm => memory}/pooled_allocator.h | 20 ++++----
src/runtime/vm/vm.cc | 14 +++--
.../runtime/{vm => memory}/memory_manager_tests.cc | 39 +++++++-------
19 files changed, 159 insertions(+), 94 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4f989a3d90..d4d5991345 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -343,6 +343,7 @@ list(APPEND COMPILER_SRCS
"src/target/datatype/myfloat/myfloat.cc")
tvm_file_glob(GLOB RUNTIME_SRCS
src/runtime/*.cc
src/runtime/vm/*.cc
+ src/runtime/memory/*.cc
src/runtime/disco/*.cc
src/runtime/minrpc/*.cc
)
diff --git a/apps/android_camera/app/src/main/jni/tvm_runtime.h
b/apps/android_camera/app/src/main/jni/tvm_runtime.h
index 0aac7f170a..e843b56360 100644
--- a/apps/android_camera/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_camera/app/src/main/jni/tvm_runtime.h
@@ -40,6 +40,7 @@
#include "../src/runtime/graph_executor/graph_executor.cc"
#include "../src/runtime/library_module.cc"
#include "../src/runtime/logging.cc"
+#include "../src/runtime/memory/memory_manager.cc"
#include "../src/runtime/minrpc/minrpc_logger.cc"
#include "../src/runtime/module.cc"
#include "../src/runtime/ndarray.cc"
diff --git a/apps/android_deploy/app/src/main/jni/tvm_runtime.h
b/apps/android_deploy/app/src/main/jni/tvm_runtime.h
index a2f10701d6..9eda834eb4 100644
--- a/apps/android_deploy/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_deploy/app/src/main/jni/tvm_runtime.h
@@ -35,6 +35,7 @@
#include "../src/runtime/graph_executor/graph_executor.cc"
#include "../src/runtime/library_module.cc"
#include "../src/runtime/logging.cc"
+#include "../src/runtime/memory/memory_manager.cc"
#include "../src/runtime/module.cc"
#include "../src/runtime/ndarray.cc"
#include "../src/runtime/object.cc"
diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
index 260c8d0cd8..fb14d84b79 100644
--- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -42,6 +42,7 @@
#include "../src/runtime/graph_executor/graph_executor_factory.cc"
#include "../src/runtime/library_module.cc"
#include "../src/runtime/logging.cc"
+#include "../src/runtime/memory/memory_manager.cc"
#include "../src/runtime/minrpc/minrpc_logger.cc"
#include "../src/runtime/module.cc"
#include "../src/runtime/ndarray.cc"
diff --git a/apps/bundle_deploy/runtime.cc b/apps/bundle_deploy/runtime.cc
index 393fc3489a..e52a4796bb 100644
--- a/apps/bundle_deploy/runtime.cc
+++ b/apps/bundle_deploy/runtime.cc
@@ -29,6 +29,7 @@
#include "../../src/runtime/graph_executor/graph_executor.cc"
#include "../../src/runtime/library_module.cc"
#include "../../src/runtime/logging.cc"
+#include "../../src/runtime/memory/memory_manager.cc"
#include "../../src/runtime/module.cc"
#include "../../src/runtime/ndarray.cc"
#include "../../src/runtime/object.cc"
diff --git a/apps/howto_deploy/tvm_runtime_pack.cc
b/apps/howto_deploy/tvm_runtime_pack.cc
index 0ee61a7e50..25e768302c 100644
--- a/apps/howto_deploy/tvm_runtime_pack.cc
+++ b/apps/howto_deploy/tvm_runtime_pack.cc
@@ -64,6 +64,7 @@
// Graph executor
#include "../../src/runtime/graph_executor/graph_executor.cc"
#include "../../src/runtime/graph_executor/graph_executor_factory.cc"
+#include "../../src/runtime/memory/memory_manager.cc"
// Uncomment the following lines to enable RPC
// #include "../../src/runtime/rpc/rpc_session.cc"
diff --git a/golang/src/tvm_runtime_pack.cc b/golang/src/tvm_runtime_pack.cc
index c2add6a367..e4056742ee 100644
--- a/golang/src/tvm_runtime_pack.cc
+++ b/golang/src/tvm_runtime_pack.cc
@@ -46,6 +46,7 @@
// Graph executor
#include "src/runtime/graph_executor/graph_executor.cc"
+#include "src/runtime/memory/memory_manager.cc"
// Uncomment the following lines to enable RPC
// #include "../../src/runtime/rpc/rpc_session.cc"
diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index cb0eb7c21f..e33539dadd 100644
--- a/include/tvm/runtime/device_api.h
+++ b/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.
*
diff --git a/include/tvm/runtime/vm/memory_manager.h
b/include/tvm/runtime/memory/memory_manager.h
similarity index 82%
rename from include/tvm/runtime/vm/memory_manager.h
rename to include/tvm/runtime/memory/memory_manager.h
index feafc01f63..8b38fbf6f0 100644
--- a/include/tvm/runtime/vm/memory_manager.h
+++ b/include/tvm/runtime/memory/memory_manager.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/runtime/vm/memory_manager.h
+ * \file tvm/runtime/memory/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_MEMORY_MANAGER_H_
+#define TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/ndarray.h>
@@ -37,22 +37,22 @@
namespace tvm {
namespace runtime {
-namespace vm {
+namespace memory {
+
+enum AllocatorType {
+ kNaive = 1,
+ kPooled,
+};
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 shape of the tensor. */
- std::vector<int64_t> shape;
/*! \brief The context of the allocated buffers. */
Device device;
-};
-
-enum AllocatorType {
- kNaive = 1,
- kPooled,
+ /*! \brief The allocator that created this buffer. */
+ AllocatorType alloc_type;
};
class Allocator {
@@ -63,9 +63,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 The device memory scope hint.
* \return The empty NDArray.
*/
- NDArray Empty(std::vector<int64_t> shape, DLDataType dtype, Device dev);
+ NDArray Empty(ShapeTuple 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.
@@ -76,13 +78,12 @@ class Allocator {
*/
virtual Buffer 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,
+ virtual Buffer Alloc(ShapeTuple shape, DLDataType type_hint,
const std::string& mem_scope = "") = 0;
/*! \brief Free a buffer allocated by the allocator.
* \param buffer The buffer to free.
@@ -94,7 +95,7 @@ class Allocator {
virtual size_t UsedMemory() const = 0;
protected:
- virtual Buffer Alloc(Device dev, int ndims, int64_t* shape, DLDataType
type_hint,
+ virtual Buffer Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
const std::string& mem_scope);
private:
@@ -114,16 +115,18 @@ class MemoryManager {
/*!
* \brief Get an allocator given the context.
* \param dev The TVM device
+ * \param type The allocator type
* \return The memory allocator.
*/
- static Allocator* GetAllocator(Device dev);
+ static Allocator* GetAllocator(Device dev, AllocatorType type);
private:
MemoryManager() {}
protected:
std::mutex mu_;
- std::unordered_map<Device, std::unique_ptr<Allocator>> allocators_;
+ std::unordered_map<Device, std::unordered_map<AllocatorType,
std::unique_ptr<Allocator>>>
+ allocators_;
};
/*! \brief An object representing a storage allocation. */
@@ -133,13 +136,13 @@ class StorageObj : public Object {
Buffer buffer;
/*! \brief Allocate an NDArray from a given piece of storage. */
- NDArray AllocNDArray(size_t offset, std::vector<int64_t> shape, DLDataType
dtype);
+ NDArray AllocNDArray(size_t offset, ShapeTuple shape, DLDataType dtype);
/*! \brief The deleter for an NDArray when allocated from underlying
storage. */
static void Deleter(Object* ptr);
~StorageObj() {
- auto alloc = MemoryManager::Global()->GetAllocator(buffer.device);
+ auto alloc = MemoryManager::Global()->GetAllocator(buffer.device,
buffer.alloc_type);
alloc->Free(buffer);
}
@@ -156,8 +159,8 @@ class Storage : public ObjectRef {
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Storage, ObjectRef, StorageObj);
};
-} // namespace vm
+} // namespace memory
} // namespace runtime
} // namespace tvm
-#endif // TVM_RUNTIME_VM_MEMORY_MANAGER_H_
+#endif // TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
diff --git a/include/tvm/runtime/vm/vm.h b/include/tvm/runtime/vm/vm.h
index c2adc3b2a0..a5fe91186d 100644
--- a/include/tvm/runtime/vm/vm.h
+++ b/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/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>
@@ -41,6 +41,13 @@
namespace tvm {
namespace runtime {
+
+using memory::Allocator;
+using memory::AllocatorType;
+using memory::MemoryManager;
+using memory::Storage;
+using memory::StorageObj;
+
namespace vm {
/*!
diff --git a/src/relay/backend/vm/compiler.h b/src/relay/backend/vm/compiler.h
index 5009d90849..acb4d2d1d2 100644
--- a/src/relay/backend/vm/compiler.h
+++ b/src/relay/backend/vm/compiler.h
@@ -41,7 +41,7 @@
#include <utility>
#include <vector>
-#include "../../../runtime/vm/naive_allocator.h"
+#include "../../../runtime/memory/naive_allocator.h"
#include "../../../runtime/vm/profiler/vm.h"
#include "../../transforms/pass_utils.h"
#include "../te_compiler.h"
diff --git a/src/runtime/c_runtime_api.cc b/src/runtime/c_runtime_api.cc
index 980447214a..3471adefeb 100644
--- a/src/runtime/c_runtime_api.cc
+++ b/src/runtime/c_runtime_api.cc
@@ -152,6 +152,20 @@ static size_t GetDataAlignment(const DLDataType dtype) {
return align;
}
+size_t DeviceAPI::GetDataSize(const DLTensor& arr, Optional<String> mem_scope)
{
+ if (!mem_scope.defined() || mem_scope.value().empty() || 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::AllocDataSpace(Device dev, int ndim, const int64_t* shape,
DLDataType dtype,
Optional<String> mem_scope) {
if (!mem_scope.defined() || mem_scope.value() == "" || mem_scope.value() ==
"global") {
diff --git a/src/runtime/graph_executor/graph_executor.cc
b/src/runtime/graph_executor/graph_executor.cc
index 777a5a442a..5bd7967cab 100644
--- a/src/runtime/graph_executor/graph_executor.cc
+++ b/src/runtime/graph_executor/graph_executor.cc
@@ -466,7 +466,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));
}
}
diff --git a/src/runtime/graph_executor/graph_executor.h
b/src/runtime/graph_executor/graph_executor.h
index 2f6b8b8147..08e06f4e6b 100644
--- a/src/runtime/graph_executor/graph_executor.h
+++ b/src/runtime/graph_executor/graph_executor.h
@@ -28,6 +28,7 @@
#include <dlpack/dlpack.h>
#include <dmlc/json.h>
#include <dmlc/memory_io.h>
+#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/packed_func.h>
@@ -42,6 +43,9 @@
namespace tvm {
namespace runtime {
+using memory::AllocatorType;
+using memory::MemoryManager;
+
/*! \brief macro to do C API call */
#define TVM_CCALL(func) \
{ \
diff --git a/src/runtime/vm/memory_manager.cc
b/src/runtime/memory/memory_manager.cc
similarity index 76%
rename from src/runtime/vm/memory_manager.cc
rename to src/runtime/memory/memory_manager.cc
index cb52a4a443..e72934ed2e 100644
--- a/src/runtime/vm/memory_manager.cc
+++ b/src/runtime/memory/memory_manager.cc
@@ -18,10 +18,10 @@
*/
/*!
- * \file tvm/runtime/vm/memory_manager.cc
+ * \file tvm/runtime/memory/memory_manager.cc
* \brief Allocate and manage memory for the runtime.
*/
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory/memory_manager.h>
#include <memory>
#include <utility>
@@ -31,17 +31,23 @@
namespace tvm {
namespace runtime {
-namespace vm {
+namespace memory {
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);
- MemoryManager::GetAllocator(buffer->device)->Free(*(buffer));
+ MemoryManager::GetAllocator(buffer->device,
buffer->alloc_type)->Free(*(buffer));
delete buffer;
delete ptr;
}
+Storage::Storage(Buffer buffer) {
+ auto n = make_object<StorageObj>();
+ n->buffer = std::move(buffer);
+ data_ = std::move(n);
+}
+
void StorageObj::Deleter(Object* obj) {
auto* ptr = static_cast<NDArray::Container*>(obj);
// When invoking AllocNDArray we don't own the underlying allocation
@@ -76,7 +82,7 @@ inline size_t GetDataAlignment(const DLTensor& arr) {
return align;
}
-NDArray StorageObj::AllocNDArray(size_t offset, std::vector<int64_t> shape,
DLDataType dtype) {
+NDArray StorageObj::AllocNDArray(size_t offset, ShapeTuple shape, DLDataType
dtype) {
VerifyDataType(dtype);
// crtical zone: allocate header, cannot throw
@@ -85,7 +91,7 @@ NDArray StorageObj::AllocNDArray(size_t offset,
std::vector<int64_t> shape, DLDa
container->dl_tensor.byte_offset = offset;
container->SetDeleter(StorageObj::Deleter);
- size_t needed_size = GetDataSize(container->dl_tensor);
+ size_t needed_size =
DeviceAPI::Get(this->buffer.device)->GetDataSize(container->dl_tensor);
this->IncRef();
// The manager context pointer must continue to point to the storage object
// which owns the backing memory, and keeps track of the reference count.
@@ -116,6 +122,9 @@ Allocator* MemoryManager::GetOrCreateAllocator(Device dev,
AllocatorType type) {
MemoryManager* m = MemoryManager::Global();
std::lock_guard<std::mutex> lock(m->mu_);
if (m->allocators_.find(dev) == m->allocators_.end()) {
+ m->allocators_.emplace(dev, std::unordered_map<AllocatorType,
std::unique_ptr<Allocator>>());
+ }
+ if (m->allocators_.at(dev).find(type) == m->allocators_.at(dev).end()) {
std::unique_ptr<Allocator> alloc;
switch (type) {
case kNaive: {
@@ -132,51 +141,55 @@ Allocator* MemoryManager::GetOrCreateAllocator(Device
dev, AllocatorType type) {
LOG(FATAL) << "Unknown allocator type: " << type;
}
auto ret = alloc.get();
- m->allocators_.emplace(dev, std::move(alloc));
+ m->allocators_.at(dev).emplace(type, std::move(alloc));
return ret;
}
- auto alloc = m->allocators_.at(dev).get();
- if (alloc->type() != type) {
+ auto alloc = m->allocators_.at(dev).at(type).get();
+ /*if (alloc->type() != type) {
LOG(WARNING) << "The type of existing allocator for " << dev
<< " is different from the request type (" << alloc->type()
<< " vs " << type
<< ")";
- }
+ }*/
return alloc;
}
-Allocator* MemoryManager::GetAllocator(Device dev) {
+Allocator* MemoryManager::GetAllocator(Device dev, AllocatorType type) {
MemoryManager* m = MemoryManager::Global();
std::lock_guard<std::mutex> lock(m->mu_);
auto it = m->allocators_.find(dev);
if (it == m->allocators_.end()) {
LOG(FATAL) << "Allocator for " << dev << " has not been created yet.";
}
- return it->second.get();
+ if (it->second.find(type) == it->second.end()) {
+ LOG(FATAL) << "Allocator for " << dev << " of type " << type << " has not
been created yet.";
+ }
+ return it->second.at(type).get();
}
-NDArray Allocator::Empty(std::vector<int64_t> shape, DLDataType dtype,
DLDevice dev) {
+NDArray Allocator::Empty(ShapeTuple 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);
size_t alignment = GetDataAlignment(container->dl_tensor);
Buffer* buffer = new Buffer;
- *buffer = this->Alloc(size, alignment, dtype);
+ if (!mem_scope.defined() || mem_scope.value().empty() || mem_scope.value()
== "global") {
+ *buffer = this->Alloc(size, alignment, dtype);
+ } else {
+ *buffer = this->Alloc(shape, 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,
+Buffer Allocator::Alloc(Device dev, ShapeTuple 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;
- for (int i = 0; i < ndims; ++i) {
- s.push_back(shape[i]);
- }
- NDArray::Container container(nullptr, s, type_hint, dev);
- size_t size = GetDataSize(container.dl_tensor);
+ NDArray::Container container(nullptr, shape, type_hint, dev);
+ size_t size = DeviceAPI::Get(dev)->GetDataSize(container.dl_tensor);
size_t alignment = GetDataAlignment(container.dl_tensor);
return Alloc(size, alignment, type_hint);
}
@@ -185,6 +198,6 @@ Buffer Allocator::Alloc(Device dev, int ndims, int64_t*
shape, DLDataType type_h
return {};
}
-} // namespace vm
+} // namespace memory
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/vm/naive_allocator.h
b/src/runtime/memory/naive_allocator.h
similarity index 78%
rename from src/runtime/vm/naive_allocator.h
rename to src/runtime/memory/naive_allocator.h
index 799f16ad60..4ab96bdfd5 100644
--- a/src/runtime/vm/naive_allocator.h
+++ b/src/runtime/memory/naive_allocator.h
@@ -18,20 +18,20 @@
*/
/*!
- * \file src/runtime/naive_allocator.h
+ * \file src/runtime/memory/naive_allocator.h
*/
-#ifndef TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
-#define TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
+#ifndef TVM_RUNTIME_MEMORY_NAIVE_ALLOCATOR_H_
+#define TVM_RUNTIME_MEMORY_NAIVE_ALLOCATOR_H_
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory/memory_manager.h>
#include <atomic>
#include <string>
namespace tvm {
namespace runtime {
-namespace vm {
+namespace memory {
class NaiveAllocator final : public Allocator {
public:
@@ -41,34 +41,35 @@ class NaiveAllocator final : public Allocator {
Buffer buf;
buf.device = device_;
buf.size = nbytes;
+ buf.alloc_type = kNaive;
buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, 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;
}
- Buffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
- const std::string& mem_scope) override {
+ Buffer Alloc(ShapeTuple shape, DLDataType type_hint, const std::string&
mem_scope) override {
Buffer buf;
size_t nbytes = 1;
- for (int i = 0; i < ndims; ++i) {
- buf.shape.push_back(shape[i]);
+ for (int i = 0; i < static_cast<int>(shape.size()); ++i) {
nbytes *= static_cast<size_t>(shape[i]);
}
nbytes *= (type_hint.bits * type_hint.lanes + 7) / 8;
buf.device = device_;
if (mem_scope.empty() || mem_scope == "global") {
- auto tmp_buf = Allocator::Alloc(device_, ndims, shape, type_hint,
mem_scope);
+ auto tmp_buf = Allocator::Alloc(device_, shape, type_hint, mem_scope);
buf.size = tmp_buf.size;
buf.data = tmp_buf.data;
+ buf.alloc_type = kNaive;
return buf;
}
buf.size = nbytes;
- buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, ndims, shape,
type_hint,
- String(mem_scope));
+ buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, shape.size(),
shape.data(),
+ type_hint,
String(mem_scope));
used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_
<< " B";
+ buf.alloc_type = kNaive;
return buf;
}
@@ -85,8 +86,8 @@ class NaiveAllocator final : public Allocator {
Device device_;
};
-} // namespace vm
+} // namespace memory
} // namespace runtime
} // namespace tvm
-#endif // TVM_RUNTIME_VM_NAIVE_ALLOCATOR_H_
+#endif // TVM_RUNTIME_MEMORY_NAIVE_ALLOCATOR_H_
diff --git a/src/runtime/vm/pooled_allocator.h
b/src/runtime/memory/pooled_allocator.h
similarity index 88%
rename from src/runtime/vm/pooled_allocator.h
rename to src/runtime/memory/pooled_allocator.h
index ea6059e0c6..01dded966b 100644
--- a/src/runtime/vm/pooled_allocator.h
+++ b/src/runtime/memory/pooled_allocator.h
@@ -18,13 +18,13 @@
*/
/*!
- * \file runtime/pooled_allocator.h
+ * \file src/runtime/memory/pooled_allocator.h
*/
-#ifndef TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
-#define TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
+#ifndef TVM_RUNTIME_MEMORY_POOLED_ALLOCATOR_H_
+#define TVM_RUNTIME_MEMORY_POOLED_ALLOCATOR_H_
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory/memory_manager.h>
#include <atomic>
#include <mutex>
@@ -34,7 +34,7 @@
namespace tvm {
namespace runtime {
-namespace vm {
+namespace memory {
class PooledAllocator final : public Allocator {
public:
@@ -58,6 +58,7 @@ class PooledAllocator final : public Allocator {
Buffer buf;
buf.device = device_;
buf.size = size;
+ buf.alloc_type = kPooled;
try {
buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, size,
alignment, type_hint);
} catch (InternalError& err) {
@@ -72,10 +73,9 @@ class PooledAllocator final : public Allocator {
return buf;
}
- Buffer Alloc(int ndims, int64_t* shape, DLDataType type_hint,
- const std::string& mem_scope) override {
+ Buffer Alloc(ShapeTuple 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);
+ return Allocator::Alloc(device_, shape, type_hint, mem_scope);
}
LOG(FATAL) << "This alloc should be implemented";
return {};
@@ -114,8 +114,8 @@ class PooledAllocator final : public Allocator {
Device device_;
};
-} // namespace vm
+} // namespace memory
} // namespace runtime
} // namespace tvm
-#endif // TVM_RUNTIME_VM_POOLED_ALLOCATOR_H_
+#endif // TVM_RUNTIME_MEMORY_POOLED_ALLOCATOR_H_
diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc
index 188a4153e1..66857ca734 100644
--- a/src/runtime/vm/vm.cc
+++ b/src/runtime/vm/vm.cc
@@ -94,8 +94,8 @@ inline ObjectRef CopyTo(ObjectRef src, const DLDevice& dev,
Optional<String> mem
}
}
-std::vector<int64_t> ToShape(NDArray shape_tensor) {
- std::vector<int64_t> shape;
+ShapeTuple ToShape(NDArray shape_tensor) {
+ std::vector<ShapeTuple::index_type> shape;
auto rank = shape_tensor.Shape().size();
auto dtype = shape_tensor.DataType();
@@ -121,7 +121,7 @@ std::vector<int64_t> ToShape(NDArray shape_tensor) {
LOG(FATAL) << "invalid shape tensor datatype: " << dtype;
}
- return shape;
+ return ShapeTuple(shape);
}
void VirtualMachine::OpStartHook(Instruction instr) {}
@@ -839,9 +839,13 @@ void VirtualMachine::RunLoop(const std::vector<Index>&
output_tensor_reg_indices
<< ", dtype_hint=" <<
DLDataType2String(instr.alloc_storage.dtype_hint)
<< ", device_index=" << instr.alloc_storage.device_index
<< ", memory_scope=" << mem_scope;
+
+ std::vector<ShapeTuple::index_type> shape_;
+ shape_.resize(instr.alloc_storage.ndim);
+ shape_.assign(instr.alloc_storage.shape,
+ instr.alloc_storage.shape + instr.alloc_storage.ndim);
storage_obj->buffer =
- allocator->Alloc(instr.alloc_storage.ndim,
instr.alloc_storage.shape,
- instr.alloc_storage.dtype_hint, mem_scope);
+ allocator->Alloc(ShapeTuple(shape_),
instr.alloc_storage.dtype_hint, mem_scope);
} else {
auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
auto alignment = instr.alloc_storage.alignment;
diff --git a/tests/cpp/runtime/vm/memory_manager_tests.cc
b/tests/cpp/runtime/memory/memory_manager_tests.cc
similarity index 86%
rename from tests/cpp/runtime/vm/memory_manager_tests.cc
rename to tests/cpp/runtime/memory/memory_manager_tests.cc
index ac1ff201cf..b51be91d74 100644
--- a/tests/cpp/runtime/vm/memory_manager_tests.cc
+++ b/tests/cpp/runtime/memory/memory_manager_tests.cc
@@ -19,15 +19,15 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-#include <tvm/runtime/vm/memory_manager.h>
+#include <tvm/runtime/memory/memory_manager.h>
#include <exception>
-#include "../../../../src/runtime/vm/pooled_allocator.h"
+#include "../../../../src/runtime/memory/pooled_allocator.h"
namespace tvm {
namespace runtime {
-namespace vm {
+namespace memory {
// MemoryManangerWrapper is necessary because in class MemoryManager we don't
have access to its
// protected members. In this class we add a new method which allow us to
clear internal state of
@@ -77,7 +77,7 @@ TEST_F(TvmVMMemoryManagerTest, NaiveEmptyBasic) {
EXPECT_EQ(allocator->UsedMemory(), 0);
auto dt = DataType::Float(32);
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
- std::vector<int64_t> shape = {1, 3, 6, 6};
+ ShapeTuple shape = {1, 3, 6, 6};
{
auto ndarray = allocator->Empty(shape, dt, dev);
EXPECT_EQ(allocator->UsedMemory(), nbytes);
@@ -93,7 +93,7 @@ TEST_F(TvmVMMemoryManagerTest, PooledEmptyBasic) {
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
size_t page_size = PooledAllocator::kDefaultPageSize;
size_t size = ((nbytes + page_size - 1) / page_size) * page_size;
- std::vector<int64_t> shape = {1, 3, 6, 6};
+ ShapeTuple shape = {1, 3, 6, 6};
{
auto ndarray = allocator->Empty(shape, dt, dev);
EXPECT_EQ(allocator->UsedMemory(), size);
@@ -107,14 +107,15 @@ TEST_F(TvmVMMemoryManagerTest, NaiveAllocWithShape) {
EXPECT_EQ(allocator->UsedMemory(), 0);
auto dt = DataType::Float(32);
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
- std::vector<int64_t> shape = {1, 3, 6, 6};
- auto buff = allocator->Alloc(shape.size(), shape.data(), dt);
+ ShapeTuple shape = {1, 3, 6, 6};
+ auto buff = allocator->Alloc(shape, dt);
EXPECT_EQ(allocator->UsedMemory(), nbytes);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), 0);
try {
- auto texture = allocator->Alloc(shape.size(), shape.data(), dt,
"global.texture");
+ auto texture = allocator->Alloc(shape, dt, "global.texture");
+ (void)texture;
FAIL();
} catch (std::exception& e) {
std::string pattern =
@@ -132,14 +133,15 @@ TEST_F(TvmVMMemoryManagerTest, PooledAllocWithShape) {
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
size_t page_size = PooledAllocator::kDefaultPageSize;
size_t size = ((nbytes + page_size - 1) / page_size) * page_size;
- std::vector<int64_t> shape = {1, 3, 6, 6};
- auto buff = allocator->Alloc(shape.size(), shape.data(), dt);
+ ShapeTuple shape = {1, 3, 6, 6};
+ auto buff = allocator->Alloc(shape, dt);
EXPECT_EQ(allocator->UsedMemory(), size);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), size);
try {
- auto texture = allocator->Alloc(shape.size(), shape.data(), dt,
"global.texture");
+ auto texture = allocator->Alloc(shape, dt, "global.texture");
+ (void)texture;
FAIL();
} catch (std::exception& e) {
std::string pattern = "This alloc should be implemented";
@@ -159,13 +161,13 @@ TEST_F(TvmVMMemoryManagerTest, NaiveAllocOpenCLTexture) {
EXPECT_EQ(allocator->UsedMemory(), 0);
auto dt = DataType::Float(32);
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
- std::vector<int64_t> shape = {1, 3, 6, 6};
- auto buff = allocator->Alloc(shape.size(), shape.data(), dt);
+ ShapeTuple shape = {1, 3, 6, 6};
+ auto buff = allocator->Alloc(shape, dt);
EXPECT_EQ(allocator->UsedMemory(), nbytes);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), 0);
- auto texture = allocator->Alloc(shape.size(), shape.data(), dt,
"global.texture");
+ auto texture = allocator->Alloc(shape, dt, "global.texture");
EXPECT_EQ(allocator->UsedMemory(), nbytes);
allocator->Free(texture);
EXPECT_EQ(allocator->UsedMemory(), 0);
@@ -184,14 +186,15 @@ TEST_F(TvmVMMemoryManagerTest, PooledAllocOpenCLTexture) {
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
size_t page_size = PooledAllocator::kDefaultPageSize;
size_t size = ((nbytes + page_size - 1) / page_size) * page_size;
- std::vector<int64_t> shape = {1, 3, 6, 6};
- auto buff = allocator->Alloc(shape.size(), shape.data(), dt);
+ ShapeTuple shape = {1, 3, 6, 6};
+ auto buff = allocator->Alloc(shape, dt);
EXPECT_EQ(allocator->UsedMemory(), size);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), size);
try {
- auto texture = allocator->Alloc(shape.size(), shape.data(), dt,
"global.texture");
+ auto texture = allocator->Alloc(shape, dt, "global.texture");
+ (void)texture;
FAIL();
} catch (std::exception& e) {
std::string pattern = "This alloc should be implemented";
@@ -199,6 +202,6 @@ TEST_F(TvmVMMemoryManagerTest, PooledAllocOpenCLTexture) {
EXPECT_NE(what.find(pattern), std::string::npos) << what;
}
}
-} // namespace vm
+} // namespace memory
} // namespace runtime
} // namespace tvm
