experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Simplify and optimize fastpath for initializable lookuptable.

Browse Source 
- 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
This commit is contained in:
Alkis Evlogimenos 2019-12-13 01:05:21 -08:00 committed by TensorFlower Gardener
parent b6fa31a56a
commit 27894d659f
3 changed files with 15 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/core/kernels/initializable_lookup_table.cc
View File

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

10
tensorflow/core/kernels/initializable_lookup_table.h
View File

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

10
tensorflow/core/kernels/lookup_table_op.h
View File

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