Simplify and optimize fastpath for initializable lookuptable.

- is_initialized_ is made private to avoid wrong uses
- is_initialized_ is set once with release semantics
- is_initialized_ is read with acquire semantics
This makes readers that check is_initialized() guarantee to see an initialized table. On x86 this generates the same assembly. On arm this is better because it does not require a full memory barrier.
Also check is_initialized in MemoryUsed().

PiperOrigin-RevId: 285354674
Change-Id: Ic91054b93c8da345a03196c5726ca1a484b6d389

tensorflow/core/kernels/initializable_lookup_table.cc

@@ -68,10 +68,7 @@ Status InitializableLookupTable::Initialize(InitTableIterator& iter) {
     return iter.status();
   }
 
-  // Prevent compiler/memory reordering of is_initialized and
-  // the initialization itself.
-  std::atomic_thread_fence(std::memory_order_release);
-  is_initialized_ = true;
+  is_initialized_.store(true, std::memory_order_release);
   return Status::OK();
 }
 

tensorflow/core/kernels/initializable_lookup_table.h

@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_KERNELS_INITIALIZABLE_LOOKUP_TABLE_H_
 #define TENSORFLOW_CORE_KERNELS_INITIALIZABLE_LOOKUP_TABLE_H_
 
+#include <atomic>
+
 #include "tensorflow/core/framework/lookup_interface.h"
 #include "tensorflow/core/platform/macros.h"
 
@@ -71,7 +73,9 @@ class InitializableLookupTable : public LookupInterface {
   TensorShape value_shape() const final { return TensorShape(); }
 
   // Returns whether the table was initialized and is ready to serve lookups.
-  bool is_initialized() const { return is_initialized_; }
+  bool is_initialized() const {
+    return is_initialized_.load(std::memory_order_acquire);
+  }
 
   // Initializes the table from the given init table iterator.
   //
@@ -156,7 +160,9 @@ class InitializableLookupTable : public LookupInterface {
   virtual Status AreEntriesSame(const InitTableIterator& iter, bool* result);
 
   mutex mu_;
-  bool is_initialized_ = false;
+
+ private:
+  std::atomic<bool> is_initialized_{false};
 };
 
 // Iterator to initialize tables given 'keys' and 'values' tensors.

tensorflow/core/kernels/lookup_table_op.h

@@ -180,15 +180,14 @@ class HashTable : public InitializableLookupTable {
 
   size_t size() const override {
     // return the size of the table only if it's initialized, otherwise 0.
-    if (!is_initialized_) {
+    if (!is_initialized()) {
       return 0;
     }
-    std::atomic_thread_fence(std::memory_order_acquire);
     return table_ ? table_->size() : 0;
   }
 
   Status ExportValues(OpKernelContext* context) override {
-    if (!is_initialized_) {
+    if (!is_initialized()) {
       return errors::Aborted("HashTable is not initialized.");
     }
 
@@ -217,7 +216,7 @@ class HashTable : public InitializableLookupTable {
 
  protected:
   Status DoPrepare(size_t unused) override {
-    if (is_initialized_) {
+    if (is_initialized()) {
       return errors::Aborted("HashTable already initialized.");
     }
     if (!table_) {
@@ -266,6 +265,9 @@ class HashTable : public InitializableLookupTable {
   }
 
   int64 MemoryUsed() const override {
+    if (!is_initialized()) {
+      return 0;
+    }
     if (table_) {
       const int64 num_elements = table_->size();
       return num_elements * (sizeof(K) + sizeof(V));