From 805b7ccc2ec86c9dd59fa3550c57109a4a71c0d3 Mon Sep 17 00:00:00 2001 From: Skye Wanderman-Milne Date: Fri, 31 May 2019 11:49:17 -0700 Subject: [PATCH] [XLA:Python] Allow growth in BFCAllocator. With this change, the BFCAllocator no longer allocates 90% of available GPU memory by default. Instead it grows the total allocation dynamically. PiperOrigin-RevId: 250927329 --- tensorflow/compiler/xla/python/local_client.cc | 7 +++---- tensorflow/compiler/xla/python/local_client.h | 5 +++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorflow/compiler/xla/python/local_client.cc b/tensorflow/compiler/xla/python/local_client.cc index 7ab95b4050330c..0d0784db76f2e1 100644 --- a/tensorflow/compiler/xla/python/local_client.cc +++ b/tensorflow/compiler/xla/python/local_client.cc @@ -208,12 +208,11 @@ static StatusOr> CreateBFCAllocator( device_ordinal); } size_t allocator_memory = free_memory * memory_fraction; - LOG(INFO) << "XLA backend reserving " << allocator_memory << " out of " - << total_memory << " bytes on device " << device_ordinal - << " for BFCAllocator."; + LOG(INFO) << "XLA backend will use up to " << allocator_memory + << " bytes on device " << device_ordinal << " for BFCAllocator."; auto gpu_bfc_allocator = absl::make_unique( - sub_allocator.release(), allocator_memory, /*allow_growth=*/false, + sub_allocator.release(), allocator_memory, /*allow_growth=*/true, absl::StrCat("GPU_", device_ordinal, "_bfc")); allocators.emplace_back(std::move(gpu_bfc_allocator)); } diff --git a/tensorflow/compiler/xla/python/local_client.h b/tensorflow/compiler/xla/python/local_client.h index e3a5e1880cb34e..a7e5c29346b96c 100644 --- a/tensorflow/compiler/xla/python/local_client.h +++ b/tensorflow/compiler/xla/python/local_client.h @@ -197,8 +197,9 @@ struct AllocatorConfig { }; Kind kind = Kind::kDefault; - // Only used if kind == kBFC. Fraction of available memory to allocate. - double memory_fraction = .9; + // Only used if kind == kBFC. The maximum fraction of available memory to + // allocate. + double memory_fraction = 1.0; }; // Encapsulates the state of Python session with XLA.