Commit 8928013d by Haibin Lin Committed by Tianqi Chen

add kDynamicStorageID to plan mem (#127)

parent ec04b943
...@@ -24,6 +24,8 @@ class GraphAllocator { ...@@ -24,6 +24,8 @@ class GraphAllocator {
static const StorageID kBadStorageID = -1; static const StorageID kBadStorageID = -1;
// external storage id // external storage id
static const StorageID kExternalStorageID = -2; static const StorageID kExternalStorageID = -2;
// dynamic storage id
static const StorageID kDynamicStorageID = -3;
// request a free storage // request a free storage
StorageID Request(int dev_id, int dtype, TShape shape, uint32_t node_id) { StorageID Request(int dev_id, int dtype, TShape shape, uint32_t node_id) {
...@@ -66,7 +68,7 @@ class GraphAllocator { ...@@ -66,7 +68,7 @@ class GraphAllocator {
// release a memory space. // release a memory space.
void Release(StorageID id, uint32_t node_id) { void Release(StorageID id, uint32_t node_id) {
CHECK_NE(id, kBadStorageID); CHECK_NE(id, kBadStorageID);
if (id == kExternalStorageID) return; if (id == kExternalStorageID || id == kDynamicStorageID) return;
StorageEntry *e = data_[id].get(); StorageEntry *e = data_[id].get();
e->released_by_node = node_id; e->released_by_node = node_id;
free_.insert({e->max_bytes, e}); free_.insert({e->max_bytes, e});
...@@ -209,6 +211,7 @@ size_t AllocMemory(const Graph& ret, const IndexedGraph& idx, StorageVector* sto ...@@ -209,6 +211,7 @@ size_t AllocMemory(const Graph& ret, const IndexedGraph& idx, StorageVector* sto
std::multimap<size_t, uint32_t> eids; std::multimap<size_t, uint32_t> eids;
for (uint32_t index = 0; index < inode.source->num_outputs(); ++index) { for (uint32_t index = 0; index < inode.source->num_outputs(); ++index) {
uint32_t eid = idx.entry_id(nid, index); uint32_t eid = idx.entry_id(nid, index);
// only request memory for kBadStorageID
if (storage[eid] == GraphAllocator::kBadStorageID) { if (storage[eid] == GraphAllocator::kBadStorageID) {
auto &eshape = shape_vec[eid]; auto &eshape = shape_vec[eid];
size_t esize = 0; size_t esize = 0;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment