experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fixing a build issue in mkl_cpu_allocator

Browse Source
This commit is contained in:
Mahmoud Abuzaina 2019-02-21 09:46:18 -08:00
parent d17ac0d930
commit d427c12c97
2 changed files with 26 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
tensorflow/core/common_runtime/mkl_cpu_allocator.h
View File

@ -82,14 +82,18 @@ class MklSmallSizeAllocator : public Allocator {
port::AlignedFree(ptr);
}
void GetStats(AllocatorStats* stats) override {
absl::optional<AllocatorStats> GetStats() override {
mutex_lock l(mutex_);
*stats = stats_;
return stats_;
}
void ClearStats() override {
mutex_lock l(mutex_);
stats_.Clear();
stats_.num_allocs = 0;
stats_.peak_bytes_in_use = 0;
stats_.largest_alloc_size = 0;
stats_.bytes_in_use = 0;
stats_.bytes_limit = 0;
}
private:
@ -98,10 +102,10 @@ class MklSmallSizeAllocator : public Allocator {
mutex_lock l(mutex_);
++stats_.num_allocs;
stats_.bytes_in_use += alloc_size;
stats_.max_bytes_in_use =
std::max(stats_.max_bytes_in_use, stats_.bytes_in_use);
stats_.max_alloc_size =
std::max(alloc_size, static_cast<size_t>(stats_.max_alloc_size));
stats_.peak_bytes_in_use =
std::max(stats_.peak_bytes_in_use, stats_.bytes_in_use);
stats_.largest_alloc_size =
std::max(alloc_size, static_cast<size_t>(stats_.largest_alloc_size));
}
// Decrement statistics for the allocator handling small allocations.
@ -244,22 +248,22 @@ class MklCPUAllocator : public Allocator {
}
}
void GetStats(AllocatorStats* stats) override {
AllocatorStats l_stats, s_stats;
small_size_allocator_->GetStats(&s_stats);
large_size_allocator_->GetStats(&l_stats);
absl::optional<AllocatorStats> GetStats() override {
auto s_stats = small_size_allocator_->GetStats();
auto l_stats = large_size_allocator_->GetStats();
// Combine statistics from small-size and large-size allocator.
stats->num_allocs = l_stats.num_allocs + s_stats.num_allocs;
stats->bytes_in_use = l_stats.bytes_in_use + s_stats.bytes_in_use;
stats->max_bytes_in_use =
l_stats.max_bytes_in_use + s_stats.max_bytes_in_use;
stats_.num_allocs = l_stats->num_allocs + s_stats->num_allocs;
stats_.bytes_in_use = l_stats->bytes_in_use + s_stats->bytes_in_use;
stats_.peak_bytes_in_use =
l_stats->peak_bytes_in_use + s_stats->peak_bytes_in_use;
// Since small-size allocations go to MklSmallSizeAllocator,
// max_alloc_size from large_size_allocator would be the maximum
// size allocated by MklCPUAllocator.
stats->max_alloc_size = l_stats.max_alloc_size;
stats->bytes_limit = std::max(s_stats.bytes_limit, l_stats.bytes_limit);
stats_.largest_alloc_size = l_stats->largest_alloc_size;
stats_.bytes_limit = std::max(s_stats->bytes_limit, l_stats->bytes_limit);
return stats_;
}
void ClearStats() override {
@ -308,6 +312,7 @@ class MklCPUAllocator : public Allocator {
SubAllocator* sub_allocator_; // not owned by this class
mutable mutex mutex_;
AllocatorStats stats_ GUARDED_BY(mutex_);
// Hash map to keep track of "BFC" allocations
// We do not use BFC allocator for small allocations.

9
tensorflow/core/common_runtime/mkl_cpu_allocator_test.cc
View File

@ -24,22 +24,21 @@ limitations under the License.
namespace tensorflow {
TEST(MKLBFCAllocatorTest, TestMaxLimit) {
AllocatorStats stats;
setenv(MklCPUAllocator::kMaxLimitStr, "1000", 1);
MklCPUAllocator a;
TF_EXPECT_OK(a.Initialize());
a.GetStats(&stats);
EXPECT_EQ(stats.bytes_limit, 1000);
auto stats = a.GetStats();
EXPECT_EQ(stats->bytes_limit, 1000);
unsetenv(MklCPUAllocator::kMaxLimitStr);
TF_EXPECT_OK(a.Initialize());
a.GetStats(&stats);
stats = a.GetStats();
uint64 max_mem_bytes = MklCPUAllocator::kDefaultMaxLimit;
#if defined(_SC_PHYS_PAGES) && defined(_SC_PAGESIZE)
max_mem_bytes =
(uint64)sysconf(_SC_PHYS_PAGES) * (uint64)sysconf(_SC_PAGESIZE);
#endif
EXPECT_EQ(stats.bytes_limit, max_mem_bytes);
EXPECT_EQ(stats->bytes_limit, max_mem_bytes);
setenv(MklCPUAllocator::kMaxLimitStr, "wrong-input", 1);
EXPECT_TRUE(errors::IsInvalidArgument(a.Initialize()));