Add enum AllocatorStrategy

test=develop

paddle/fluid/memory/allocation/CMakeLists.txt

@@ -43,6 +43,7 @@ cc_library(aligned_allocator SRCS aligned_allocator.cc DEPS allocator)
 cc_library(auto_increment_allocator SRCS auto_increment_allocator.cc DEPS allocator)
 cc_library(zero_size_allocator SRCS zero_size_allocator.cc DEPS allocator)
 cc_library(conditional_allocator SRCS conditional_allocator.cc DEPS allocator)
+cc_library(allocator_strategy SRCS allocator_strategy.cc DEPS gflags)
 cc_library(allocator_facade SRCS allocator_facade.cc DEPS
         ${AllocatorFacadeDeps}
         cpu_allocator
@@ -54,7 +55,9 @@ cc_library(allocator_facade SRCS allocator_facade.cc DEPS
         zero_size_allocator
         conditional_allocator
         retry_allocator
-        buffered_allocator)
+        buffered_allocator
+        allocator_strategy
+        )
 
 nv_test(allocation_and_eigen_test SRCS allocation_and_eigen_test.cu DEPS allocator_facade)
 

paddle/fluid/memory/allocation/allocator_strategy.cc

@@ -0,0 +1,39 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/memory/allocation/allocator_strategy.h"
+#include "gflags/gflags.h"
+
+DEFINE_string(
+    allocator_strategy, "legacy",
+    "The allocation strategy. Legacy means the original allocator of Fluid."
+    "New means the experimental allocators of Fluid. in [legacy, new]");
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+static AllocatorStrategy GetStrategyFromFlag() {
+  return FLAGS_allocator_strategy == "legacy"
+             ? AllocatorStrategy::kLegacy
+             : AllocatorStrategy::kNaiveBestFit;
+}
+
+AllocatorStrategy GetAllocatorStrategy() {
+  static AllocatorStrategy strategy = GetStrategyFromFlag();
+  return strategy;
+}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_strategy.h

@@ -0,0 +1,27 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+enum class AllocatorStrategy { kLegacy, kNaiveBestFit };
+
+extern AllocatorStrategy GetAllocatorStrategy();
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/malloc.cc

@@ -16,10 +16,10 @@ limitations under the License. */
 
 #include "glog/logging.h"
 #include "paddle/fluid/memory/allocation/allocator_facade.h"
-#include "paddle/fluid/memory/malloc.h"
-
+#include "paddle/fluid/memory/allocation/allocator_strategy.h"
 #include "paddle/fluid/memory/detail/buddy_allocator.h"
 #include "paddle/fluid/memory/detail/system_allocator.h"
+#include "paddle/fluid/memory/malloc.h"
 #include "paddle/fluid/platform/gpu_info.h"
 
 DEFINE_bool(init_allocated_mem, false,
@@ -30,11 +30,6 @@ DEFINE_bool(init_allocated_mem, false,
             "during unit testing.");
 DECLARE_double(fraction_of_gpu_memory_to_use);
 
-DEFINE_string(
-    allocator_strategy, "legacy",
-    "The allocation strategy. Legacy means the original allocator of Fluid."
-    "New means the experimental allocators of Fluid. in [legacy, new]");
-
 namespace paddle {
 namespace memory {
 
@@ -288,7 +283,8 @@ class LegacyAllocation : public Allocation {
 
 std::shared_ptr<Allocation> AllocShared(const platform::Place& place,
                                         size_t size, Allocator::Attr attr) {
-  if (FLAGS_allocator_strategy == "legacy") {
+  if (allocation::GetAllocatorStrategy() ==
+      allocation::AllocatorStrategy::kLegacy) {
     void* p = boost::apply_visitor(legacy::AllocVisitor(size), place);
     return std::shared_ptr<Allocation>(
         new legacy::LegacyAllocation(p, size, place));
@@ -300,7 +296,8 @@ std::shared_ptr<Allocation> AllocShared(const platform::Place& place,
 
 std::unique_ptr<Allocation> Alloc(const platform::Place& place, size_t size,
                                   Allocator::Attr attr) {
-  if (FLAGS_allocator_strategy == "legacy") {
+  if (allocation::GetAllocatorStrategy() ==
+      allocation::AllocatorStrategy::kLegacy) {
     void* p = boost::apply_visitor(legacy::AllocVisitor(size), place);
     return std::unique_ptr<Allocation>(
         new legacy::LegacyAllocation(p, size, place));