Skip to content

Commit a4489e3

Browse files
authored
[EVAllocator] Fix the bug in configuring ARENA_ARRAY_SIZE. (#986)
Signed-off-by: chenbangduo.cbd <[email protected]>
1 parent d5f7f6a commit a4489e3

File tree

1 file changed

+6
-5
lines changed

1 file changed

+6
-5
lines changed

tensorflow/core/framework/ev_allocator.h

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -546,15 +546,15 @@ class EVAllocatorImpl {
546546
page_map_ = new PageMap<ChunkType>();
547547
page_map_->Init();
548548

549-
int64 arena_array_size = ARENA_ARRAY_SIZE;
549+
arena_array_size_ = ARENA_ARRAY_SIZE;
550550
Status s = ReadInt64FromEnvVar("ARENA_ARRAY_SIZE",
551-
ARENA_ARRAY_SIZE, &arena_array_size);
551+
ARENA_ARRAY_SIZE, &arena_array_size_);
552552
if (!s.ok()) {
553553
LOG(ERROR) << "Read ARENA_ARRAY_SIZE env error: " << s.error_message();
554554
}
555-
LOG(INFO) << "EVAllocator set arena array size: " << arena_array_size;
555+
LOG(INFO) << "EVAllocator set arena array size: " << arena_array_size_;
556556

557-
arenas_ = new std::vector<Arena<ChunkType>>(arena_array_size, page_map_);
557+
arenas_ = new std::vector<Arena<ChunkType>>(arena_array_size_, page_map_);
558558
arena_cur_index = 0;
559559
}
560560

@@ -602,7 +602,7 @@ class EVAllocatorImpl {
602602
{
603603
mutex_lock l(mu_arena_index_);
604604
ret = &((*arenas_)[arena_cur_index]);
605-
arena_cur_index = (arena_cur_index + 1) % ARENA_ARRAY_SIZE;
605+
arena_cur_index = (arena_cur_index + 1) % arena_array_size_;
606606
}
607607

608608
return ret;
@@ -619,6 +619,7 @@ class EVAllocatorImpl {
619619
PageMap<ChunkType>* page_map_ = nullptr;
620620
std::vector<Arena<ChunkType>> *arenas_ = nullptr;
621621
int arena_cur_index GUARDED_BY(mu_arena_index_);
622+
int64 arena_array_size_;
622623
};
623624

624625
template<typename ChunkType>

0 commit comments

Comments
 (0)