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

jroesch · 2019-04-30T13:37:22Z

This factors out the memory management code from the previous runtime into a separate PR.

See #2889 and #3120.

@wweic

Credit to @icemelon9 and @wweic

src/runtime/naive_allocator.h

src/runtime/pooled_allocator.h

wweic · 2019-04-30T14:27:44Z

src/runtime/pooled_allocator.h

+    buf.size = size;
+    buf.data = DeviceAPI::Get(ctx_)->AllocDataSpace(ctx_, size, alignment, type_hint);
+    used_memory_.fetch_add(size, std::memory_order_relaxed);
+    LOG(INFO) << "allocate " << size << " B, used memory " << used_memory_ << " B";


Suggested change

LOG(INFO) << "allocate " << size << " B, used memory " << used_memory_ << " B";

LOG(DEBUG) << "allocate " << size << " B, used memory " << used_memory_ << " B";

Co-Authored-By: jroesch <[email protected]>

weberlo · 2019-04-30T20:11:53Z

lgtm

src/runtime/pooled_allocator.h

include/tvm/runtime/ndarray.h

src/runtime/ndarray.cc

tqchen · 2019-04-30T23:35:13Z

src/runtime/memory_manager.h

+
+ private:
+  std::mutex mu_;
+  std::unordered_map<TVMContext, std::unique_ptr<Allocator>> allocators_;


(Add space, so that it is compatible with old runtime)

tqchen · 2019-04-30T23:35:51Z

src/runtime/memory_manager.h

+
+  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;


UsedMemory(Const)?

src/runtime/pooled_allocator.h

tqchen · 2019-04-30T23:37:47Z

include/tvm/runtime/ndarray.h

@@ -32,6 +32,10 @@

 namespace tvm {
 namespace runtime {
+
+class Allocator;


remove Allocator(see my other comment about reverse the dependency)

allocator depend on ndarry.h

tqchen · 2019-04-30T23:38:17Z

golang/src/tvm_runtime_pack.cc

@@ -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"


for now, do not include memory_manager.cc, consider add back once we have all the vm changes

It doesn't link correctly without this line

if you break the dependency of the NDArrary from memory_manager(move allocator aware alloca to allocator->Empty), you should be able to break the link dp

tqchen · 2019-04-30T23:38:22Z

src/runtime/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


consider moving all the memory manager into src/runtime/vm, because it is used for dynamic memory allocation, and basic version of runtime does not need it

src/runtime/memory_manager.h

tqchen · 2019-04-30T23:46:28Z

Summary of my main review comments:

Currently, NDArray needs to be aware of Allocator, it does not have to be the case. Instead, we can implement allocator->Empty(args), which creates an empty NDArray using the allocator. This means that we need to create array creation logic into memory_manager.cc.
- The advantage is that the memory manager will become a plugin
Consider, move memory manager related logic into src/vm, given that they are being used by vm and not necessarily static graph runtime.

tqchen · 2019-05-01T19:34:33Z

include/tvm/runtime/ndarray.h

@@ -296,6 +299,19 @@ class NDArray::Container {
    dl_tensor.strides = nullptr;
    dl_tensor.byte_offset = 0;
  }
+
+  Container(void* data,


can we just use new and set it up in the impl? To minimize the constructor logic in here

Why minimize? it creates even more chances to do it wrong code for construction should be constructors

There are certain cases where the shape vector is not needed, when converting from DLPack. Anyway if we want to keep it in this way , try to use copy initializers to initialize shape, this will remove the need for move

I think we can keep the constructor this way because it's non-trivial to correctly initialize the shape in container and DLTensor in my mind. If there's use case to create Container without shape in the future, we can revisit this or add a new constructor.

tqchen · 2019-05-01T19:36:41Z

src/runtime/vm/memory_manager.cc

+  delete ptr;
+}
+
+NDArray Allocator::EmptyNDArray(std::vector<int64_t> shape, DLDataType dtype, DLContext ctx) {


debating whether to use Allocator->EmptyNDArray or Allocator->Empty, if we assume Empty always corresponds to NDArray, perhaps we can remove the NDArray

tqchen · 2019-05-01T19:37:01Z

some nits, looks good in general, please see if you can fix

src/runtime/vm/memory_manager.cc

icemelon · 2019-05-01T22:35:10Z

include/tvm/runtime/ndarray.h

@@ -296,6 +299,19 @@ class NDArray::Container {
    dl_tensor.strides = nullptr;
    dl_tensor.byte_offset = 0;
  }
+
+  Container(void* data,


I think we can keep the constructor this way because it's non-trivial to correctly initialize the shape in container and DLTensor in my mind. If there's use case to create Container without shape in the future, we can revisit this or add a new constructor.

tqchen · 2019-05-01T23:00:21Z

Please fix @icemelon9 's comment and retrigger the CI. this is almost ready to be merged

tqchen · 2019-05-02T13:56:51Z

@icemelon9 please double check and merge it in if you think it is ready

icemelon · 2019-05-02T16:11:20Z

Thanks everyone. This is now merged.

@wweic

* 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 <[email protected]> * 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

@wweic

* 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 <[email protected]> * 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

Add support for custom NDArray memory management

5d5ae6d

Credit to @icemelon9 and @wweic

jroesch force-pushed the runtime-mm branch from c8ea99b to 5d5ae6d Compare April 30, 2019 13:37

jroesch added 2 commits April 30, 2019 09:58

Fix copy-paste issue

6decb1d

Fix naive allocator.h

e4d9e2d

wweic requested changes Apr 30, 2019

View reviewed changes

jroesch and others added 4 commits April 30, 2019 10:44

Remove buffer field

41af59d

Apply Wei's suggestions.

d71d875

Co-Authored-By: jroesch <[email protected]>

Fix Wei's suggestion

0591b9a

Fix go rts

d00950d

wweic approved these changes Apr 30, 2019

View reviewed changes

tqchen requested changes Apr 30, 2019

View reviewed changes

tqchen added status: need review status: need update need update based on feedbacks labels Apr 30, 2019

tqchen self-assigned this May 1, 2019

jroesch added 6 commits May 1, 2019 10:18

Break MM dependency

f26cb52

Add docs and clean up diff

048b003

Add more docs

15763c8

Move to VM folder

76c201c

Fix lint

52cf9fa

Remove Go dep.

ce21917

tqchen reviewed May 1, 2019

View reviewed changes

tqchen requested changes May 1, 2019

View reviewed changes

Rename to Empty

d317631

icemelon reviewed May 1, 2019

View reviewed changes

Address Haichen's comments

c07f6e6

jroesch force-pushed the runtime-mm branch from d822a5a to c07f6e6 Compare May 2, 2019 05:24

tqchen approved these changes May 2, 2019

View reviewed changes

icemelon approved these changes May 2, 2019

View reviewed changes

icemelon merged commit 83cb872 into apache:master May 2, 2019

icemelon added status: accepted and removed status: need review status: need update need update based on feedbacks labels May 2, 2019

tqchen mentioned this pull request Nov 8, 2019

[RELEASE][DRAFT] TVM v0.6 Release candidate #4259

Closed

jroesch deleted the runtime-mm branch February 4, 2021 04:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

jroesch commented Apr 30, 2019

wweic Apr 30, 2019

weberlo commented Apr 30, 2019

tqchen Apr 30, 2019

tqchen Apr 30, 2019

tqchen Apr 30, 2019

tqchen Apr 30, 2019

jroesch May 1, 2019

tqchen May 1, 2019

tqchen Apr 30, 2019

tqchen commented Apr 30, 2019

tqchen May 1, 2019

jroesch May 1, 2019

tqchen May 1, 2019

icemelon May 1, 2019

tqchen May 1, 2019

tqchen commented May 1, 2019

icemelon May 1, 2019

tqchen commented May 1, 2019

tqchen commented May 2, 2019

icemelon commented May 2, 2019

	LOG(INFO) << "allocate " << size << " B, used memory " << used_memory_ << " B";
	LOG(DEBUG) << "allocate " << size << " B, used memory " << used_memory_ << " B";

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

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

Conversation

jroesch commented Apr 30, 2019

Choose a reason for hiding this comment

weberlo commented Apr 30, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Apr 30, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented May 1, 2019

Choose a reason for hiding this comment

tqchen commented May 1, 2019

tqchen commented May 2, 2019

icemelon commented May 2, 2019