From a4489e31a4b9bc8371198537a0a15af6011ef8ae Mon Sep 17 00:00:00 2001 From: Chen Bangduo Date: Fri, 12 Apr 2024 14:22:32 +0800 Subject: [PATCH] [EVAllocator] Fix the bug in configuring ARENA_ARRAY_SIZE. (#986) Signed-off-by: chenbangduo.cbd --- tensorflow/core/framework/ev_allocator.h | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/framework/ev_allocator.h b/tensorflow/core/framework/ev_allocator.h index d3251b14782..5082ee04b72 100644 --- a/tensorflow/core/framework/ev_allocator.h +++ b/tensorflow/core/framework/ev_allocator.h @@ -546,15 +546,15 @@ class EVAllocatorImpl { page_map_ = new PageMap(); page_map_->Init(); - int64 arena_array_size = ARENA_ARRAY_SIZE; + arena_array_size_ = ARENA_ARRAY_SIZE; Status s = ReadInt64FromEnvVar("ARENA_ARRAY_SIZE", - ARENA_ARRAY_SIZE, &arena_array_size); + ARENA_ARRAY_SIZE, &arena_array_size_); if (!s.ok()) { LOG(ERROR) << "Read ARENA_ARRAY_SIZE env error: " << s.error_message(); } - LOG(INFO) << "EVAllocator set arena array size: " << arena_array_size; + LOG(INFO) << "EVAllocator set arena array size: " << arena_array_size_; - arenas_ = new std::vector>(arena_array_size, page_map_); + arenas_ = new std::vector>(arena_array_size_, page_map_); arena_cur_index = 0; } @@ -602,7 +602,7 @@ class EVAllocatorImpl { { mutex_lock l(mu_arena_index_); ret = &((*arenas_)[arena_cur_index]); - arena_cur_index = (arena_cur_index + 1) % ARENA_ARRAY_SIZE; + arena_cur_index = (arena_cur_index + 1) % arena_array_size_; } return ret; @@ -619,6 +619,7 @@ class EVAllocatorImpl { PageMap* page_map_ = nullptr; std::vector> *arenas_ = nullptr; int arena_cur_index GUARDED_BY(mu_arena_index_); + int64 arena_array_size_; }; template