diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-09-19 17:28:57 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-19 17:32:07 -0700 |
commit | 5d2047029a77545c97c0fdf74d9c03c92d1dcb88 (patch) | |
tree | b5a232646fe6c0b976104c6462709b71cb01eee7 /tensorflow/core/common_runtime | |
parent | dfa0447740b42edff5ae2d76d5957aa688ae8053 (diff) |
Allow setting a global override for the "allow_growth" GPU option via the TF_FORCE_GPU_ALLOW_GROWTH environment variable.
PiperOrigin-RevId: 213728460
Diffstat (limited to 'tensorflow/core/common_runtime')
4 files changed, 95 insertions, 3 deletions
diff --git a/tensorflow/core/common_runtime/bfc_allocator.h b/tensorflow/core/common_runtime/bfc_allocator.h index 364071e066..2d74bf2b28 100644 --- a/tensorflow/core/common_runtime/bfc_allocator.h +++ b/tensorflow/core/common_runtime/bfc_allocator.h @@ -304,7 +304,7 @@ class BFCAllocator : public Allocator { }; // Returns 'bytes' rounded up to the next highest kMinAllocationSize. - size_t RoundedBytes(size_t bytes); + static size_t RoundedBytes(size_t bytes); // Try to add a new memory region that can satisfy an allocation of // 'rounded_bytes' bytes. Returns true on success and false on diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc index 44ffce77a1..42021e51f3 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc @@ -22,6 +22,39 @@ limitations under the License. namespace tensorflow { +bool GPUBFCAllocator::GetAllowGrowthValue(const GPUOptions& gpu_options) { + const char* force_allow_growth_string = + std::getenv("TF_FORCE_GPU_ALLOW_GROWTH"); + if (force_allow_growth_string == nullptr) { + return gpu_options.allow_growth(); + } + + if (strcmp("false", force_allow_growth_string) == 0) { + if (gpu_options.allow_growth()) { + LOG(WARNING) + << "Overriding allow_growth setting because the" + << " TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original" + << " config value was " << gpu_options.allow_growth() << "."; + } + return false; + } else if (strcmp("true", force_allow_growth_string) == 0) { + if (!gpu_options.allow_growth()) { + LOG(WARNING) + << "Overriding allow_growth setting because the" + << " TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original" + << " config value was " << gpu_options.allow_growth() << "."; + } + return true; + } + + LOG(ERROR) + << "The TF_FORCE_GPU_ALLOW_GROWTH environment variable is set but could" + << " not be parsed: \"" << force_allow_growth_string << "\". Valid" + << " values are \"true\" or \"false\". Using original config value" + << " of " << gpu_options.allow_growth() << "."; + return gpu_options.allow_growth(); +} + GPUBFCAllocator::GPUBFCAllocator(GPUMemAllocator* sub_allocator, size_t total_memory, const string& name) : GPUBFCAllocator(sub_allocator, total_memory, GPUOptions(), name) {} @@ -30,7 +63,7 @@ GPUBFCAllocator::GPUBFCAllocator(GPUMemAllocator* sub_allocator, size_t total_memory, const GPUOptions& gpu_options, const string& name) - : BFCAllocator(sub_allocator, total_memory, gpu_options.allow_growth(), - name) {} + : BFCAllocator(sub_allocator, total_memory, + GPUBFCAllocator::GetAllowGrowthValue(gpu_options), name) {} } // namespace tensorflow diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h index 3470f7a9f7..d4c9cee89a 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h +++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h @@ -93,6 +93,9 @@ class GPUBFCAllocator : public BFCAllocator { ~GPUBFCAllocator() override {} TF_DISALLOW_COPY_AND_ASSIGN(GPUBFCAllocator); + + private: + static bool GetAllowGrowthValue(const GPUOptions& gpu_options); }; } // namespace tensorflow diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc index e313135d8d..60e82ed13b 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc @@ -410,6 +410,8 @@ BENCHMARK(BM_AllocationDelayed)->Arg(1)->Arg(10)->Arg(100)->Arg(1000); class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test { protected: + void SetUp() override { CHECK_EQ(unsetenv("TF_FORCE_GPU_ALLOW_GROWTH"), 0); } + // The following test methods are called from tests. The reason for this is // that this class is a friend class to BFCAllocator, but tests are not, so // only methods inside this class can access private members of BFCAllocator. @@ -510,6 +512,56 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test { EXPECT_EQ(10, a.Log2FloorNonZeroSlow(1024)); EXPECT_EQ(10, a.Log2FloorNonZeroSlow(1025)); } + + void TestForceAllowGrowth() { + PlatformGpuId platform_gpu_id(0); + GPUOptions options; + // Unset flag value uses provided option. + unsetenv("TF_FORCE_GPU_ALLOW_GROWTH"); + options.set_allow_growth(true); + GPUMemAllocator* sub_allocator = new GPUMemAllocator( + GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(), + platform_gpu_id, false /*use_unified_memory*/, {}, {}); + GPUBFCAllocator unset_flag_allocator(sub_allocator, 1LL << 31, options, + "GPU_0_bfc"); + EXPECT_EQ(GPUBFCAllocator::RoundedBytes(size_t{1048576}), + unset_flag_allocator.curr_region_allocation_bytes_); + + // Unparseable flag value uses provided option. + setenv("TF_FORCE_GPU_ALLOW_GROWTH", "unparseable", 1); + options.set_allow_growth(true); + sub_allocator = new GPUMemAllocator( + GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(), + platform_gpu_id, false /*use_unified_memory*/, {}, {}); + GPUBFCAllocator unparsable_flag_allocator(sub_allocator, 1LL << 31, options, + "GPU_1_bfc"); + EXPECT_EQ(GPUBFCAllocator::RoundedBytes(size_t{1048576}), + unparsable_flag_allocator.curr_region_allocation_bytes_); + + // Max of 2GiB total memory. Env variable set forces allow_growth, which + // does an initial allocation of 1MiB. + setenv("TF_FORCE_GPU_ALLOW_GROWTH", "true", 1); + options.set_allow_growth(false); + sub_allocator = new GPUMemAllocator( + GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(), + platform_gpu_id, false /*use_unified_memory*/, {}, {}); + GPUBFCAllocator force_allow_growth_allocator(sub_allocator, 1LL << 31, + options, "GPU_2_bfc"); + EXPECT_EQ(GPUBFCAllocator::RoundedBytes(size_t{1048576}), + force_allow_growth_allocator.curr_region_allocation_bytes_); + + // If env variable forces allow_growth disabled, all available memory is + // allocated. + setenv("TF_FORCE_GPU_ALLOW_GROWTH", "false", 1); + options.set_allow_growth(true); + sub_allocator = new GPUMemAllocator( + GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(), + platform_gpu_id, false /*use_unified_memory*/, {}, {}); + GPUBFCAllocator force_no_allow_growth_allocator(sub_allocator, 1LL << 31, + options, "GPU_3_bfc"); + EXPECT_EQ(GPUBFCAllocator::RoundedBytes(1LL << 31), + force_no_allow_growth_allocator.curr_region_allocation_bytes_); + } }; TEST_F(GPUBFCAllocatorPrivateMethodsTest, BinDebugInfo) { TestBinDebugInfo(); } @@ -518,6 +570,10 @@ TEST_F(GPUBFCAllocatorPrivateMethodsTest, Log2FloorNonZeroSlow) { TestLog2FloorNonZeroSlow(); } +TEST_F(GPUBFCAllocatorPrivateMethodsTest, ForceAllowGrowth) { + TestForceAllowGrowth(); +} + } // namespace tensorflow #endif // GOOGLE_CUDA |