[Allocator] Fix PMEM allocator build failure and refactor code.

tensorflow/core/framework/allocator.cc

@@ -94,6 +94,7 @@ Allocator* pmem_allocator() {
 }
 
 Allocator* experimental_pmem_allocator(const std::string& pmem_path, size_t allocator_size) {
+#ifdef TF_ENABLE_PMEM
   static Allocator* experimental_pmem_allocator =
       AllocatorFactoryRegistry::singleton()->GetExperimentalPMEMAllocator(pmem_path, allocator_size);
   if (experimental_pmem_allocator && cpu_allocator_collect_full_stats &&
@@ -102,6 +103,9 @@ Allocator* experimental_pmem_allocator(const std::string& pmem_path, size_t allo
         new TrackingAllocator(experimental_pmem_allocator, true);
   }
   return experimental_pmem_allocator;
+#else
+  return nullptr;
+#endif
 }
 
 Allocator* ev_allocator() {

tensorflow/core/framework/allocator_registry.cc

@@ -16,8 +16,10 @@ limitations under the License.
 #include <string>
 
 #include "tensorflow/core/framework/allocator_registry.h"
+#ifdef TF_ENABLE_PMEM
+#include "tensorflow/core/framework/experimental_pmem_allocator.h"
+#endif
 #include "tensorflow/core/platform/logging.h"
-#include "experimental_pmem_allocator.h"
 
 namespace tensorflow {
 
@@ -88,7 +90,9 @@ Allocator* AllocatorFactoryRegistry::GetAllocator() {
   }
 }
 
-Allocator* AllocatorFactoryRegistry::GetExperimentalPMEMAllocator(const std::string& pmem_path, size_t pmem_size) {
+#ifdef TF_ENABLE_PMEM
+Allocator* AllocatorFactoryRegistry::GetExperimentalPMEMAllocator(
+    const std::string& pmem_path, size_t pmem_size) {
   mutex_lock l(mu_);
   first_alloc_made_ = true;
   FactoryEntry* best_entry = nullptr;
@@ -101,7 +105,8 @@ Allocator* AllocatorFactoryRegistry::GetExperimentalPMEMAllocator(const std::str
 
   if (best_entry) {
     if (!best_entry->allocator) {
-      static_cast<ExperimentalPMEMAllocatorFactory*>(best_entry->factory.get())->Init(pmem_path, pmem_size);
+      static_cast<ExperimentalPMEMAllocatorFactory*>(
+          best_entry->factory.get())->Init(pmem_path, pmem_size);
       best_entry->allocator.reset(best_entry->factory->CreateAllocator());
     }
     return best_entry->allocator.get();
@@ -110,6 +115,7 @@ Allocator* AllocatorFactoryRegistry::GetExperimentalPMEMAllocator(const std::str
     return nullptr;
   }
 }
+#endif
 
 Allocator* AllocatorFactoryRegistry::GetPMEMAllocator() {
   mutex_lock l(mu_);

tensorflow/core/framework/allocator_registry.h

@@ -80,7 +80,9 @@ class AllocatorFactoryRegistry {
   //If use PMEMallocator, then factory pick this one
   Allocator* GetPMEMAllocator();
 
+#ifdef TF_ENABLE_PMEM
   Allocator* GetExperimentalPMEMAllocator(const std::string& pmem_path, size_t pmem_size);
+#endif
 
   Allocator* GetEVAllocator();
 

tensorflow/core/framework/experimental_pmem_allocator.cc

@@ -1,11 +1,7 @@
-#include "experimental_pmem_allocator.h"
-
+#include "libpmem.h"
 #include <string.h>
 #include <unistd.h>
-
-#include <mutex>
-
-#include "libpmem.h"
+#include "tensorflow/core/framework/experimental_pmem_allocator.h"
 
 namespace tensorflow {
 

tensorflow/core/framework/experimental_pmem_allocator.h

@@ -1,17 +1,14 @@
 #ifndef TENSORFLOW_CORE_FRAMEWORK_EXPERIMENTAL_PMEM_ALLOCATOR_H_
 #define TENSORFLOW_CORE_FRAMEWORK_EXPERIMENTAL_PMEM_ALLOCATOR_H_
 
-#include <assert.h>
+#include <atomic>
 #include <fcntl.h>
 #include <sys/mman.h>
-
-#include <atomic>
-#include <thread>
 #include <vector>
 
-#include "experimental_pmem_allocator_utils.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/allocator_registry.h"
+#include "tensorflow/core/framework/experimental_pmem_allocator_utils.h"
 #include "tensorflow/core/lib/core/spin_lock.h"
 #include "tensorflow/core/platform/env.h"
 
@@ -393,4 +390,4 @@ class ExperimentalPMEMAllocatorFactory : public AllocatorFactory {
 };
 }  // namespace tensorflow
 
-#endif
+#endif

tensorflow/core/framework/experimental_pmem_allocator_utils.cc

@@ -1,8 +1,6 @@
 #include <assert.h>
 
-#include <mutex>
-
-#include "experimental_pmem_allocator.h"
+#include "tensorflow/core/framework/experimental_pmem_allocator.h"
 #include "tensorflow/core/lib/core/spin_lock.h"
 namespace tensorflow {
 
@@ -45,4 +43,4 @@ void ThreadManager::Release(const AllocatorThread& t) {
     usable_id_.insert(t.id);
   }
 }
-}  // namespace tensorflow
+}  // namespace tensorflow

tensorflow/core/framework/experimental_pmem_allocator_utils.h

@@ -1,11 +1,10 @@
 #ifndef TENSORFLOW_CORE_FRAMEWORK_EXPERIMENTAL_PMEM_ALLOCATOR_UTILS_H_
 #define TENSORFLOW_CORE_FRAMEWORK_EXPERIMENTAL_PMEM_ALLOCATOR_UTILS_H_
 
-#include <string.h>
-#include <sys/stat.h>
-
 #include <atomic>
 #include <memory>
+#include <string.h>
+#include <sys/stat.h>
 #include <unordered_set>
 
 #include "tensorflow/core/lib/core/spin_lock.h"
@@ -108,4 +107,4 @@ inline int create_dir_if_missing(const std::string& name) {
 
 }  // namespace tensorflow
 
-#endif
+#endif