[Embedding] Refactor embedding storage code by cleaning up useless StorageManager class. (#858)

Signed-off-by: lixy9474 <[email protected]>

Author: lixy9474

tensorflow/core/framework/embedding/bloom_filter_policy.h

@@ -20,10 +20,6 @@ limitations under the License.
 #include "tensorflow/core/framework/embedding/filter_policy.h"
 
 namespace tensorflow {
-namespace embedding{
-template <class K, class V>
-class StorageManager;
-}
 
 namespace {
 const static std::vector<int64> default_seeds = {
@@ -35,9 +31,8 @@ const static std::vector<int64> default_seeds = {
 template<typename K, typename V, typename EV>
 class BloomFilterPolicy : public FilterPolicy<K, V, EV> {
  public:
-  BloomFilterPolicy(const EmbeddingConfig& config, EV* ev,
-      embedding::StorageManager<K, V>* storage_manager) :
-      config_(config), ev_(ev), storage_manager_(storage_manager) {
+  BloomFilterPolicy(const EmbeddingConfig& config, EV* ev)
+      : config_(config), ev_(ev) {
     switch (config_.counter_type){
       case DT_UINT64:
         VLOG(2) << "The type of bloom counter is uint64";
@@ -349,7 +344,6 @@ class BloomFilterPolicy : public FilterPolicy<K, V, EV> {
   EmbeddingConfig config_;
   EV* ev_;
   std::vector<int64> seeds_;
-  embedding::StorageManager<K, V>* storage_manager_;
 };
 } // tensorflow
 

tensorflow/core/framework/embedding/cache_factory.h

@@ -25,11 +25,11 @@ class CacheFactory {
   static BatchCache<K>* Create(CacheStrategy cache_strategy, std::string name) {
     switch (cache_strategy) {
       case CacheStrategy::LRU:
-        LOG(INFO) << " Use StorageManager::LRU in multi-tier EmbeddingVariable "
+        LOG(INFO) << " Use Storage::LRU in multi-tier EmbeddingVariable "
                 << name;
         return new LRUCache<K>();
       case CacheStrategy::LFU:
-        LOG(INFO) << " Use StorageManager::LFU in multi-tier EmbeddingVariable "
+        LOG(INFO) << " Use Storage::LFU in multi-tier EmbeddingVariable "
                 << name;
         return new LFUCache<K>();
       default:

tensorflow/core/framework/embedding/counter_filter_policy.h

@@ -20,17 +20,12 @@ limitations under the License.
 #include "tensorflow/core/framework/embedding/filter_policy.h"
 
 namespace tensorflow {
-namespace embedding{
-template <class K, class V>
-class StorageManager;
-}
 
 template<typename K, typename V, typename EV>
 class CounterFilterPolicy : public FilterPolicy<K, V, EV> {
  public:
-  CounterFilterPolicy(const EmbeddingConfig& config,
-      EV* ev, embedding::StorageManager<K, V>* storage_manager)
-       : config_(config), ev_(ev), storage_manager_(storage_manager) {
+  CounterFilterPolicy(const EmbeddingConfig& config, EV* ev)
+      : config_(config), ev_(ev){
   }
 
   Status Lookup(EV* ev, K key, V* val, const V* default_value_ptr,
@@ -165,7 +160,6 @@ class CounterFilterPolicy : public FilterPolicy<K, V, EV> {
 
  private:
   EmbeddingConfig config_;
-  embedding::StorageManager<K, V>* storage_manager_;
   EV* ev_;
 };
 

tensorflow/core/framework/embedding/embedding_var.cu.cc

@@ -66,7 +66,7 @@ void EmbeddingVar<K, V>::SetDefaultValueOfNewFeatures(
           reinterpret_cast<ValuePtr<V>*>(memcpy_address[*it]);
       value_address[i] =
           *((V**)((char*)(value_ptr->GetPtr()) + sizeof(FixedLengthHeader))) +
-          storage_manager_->GetOffset(emb_config_.emb_index);
+          storage_->GetOffset(emb_config_.emb_index);
       default_value_address[i] = get_default_v_fn(
           default_values, keys[*it], *it, GetDefaultValueDim(), ValueLen());
     }
@@ -86,7 +86,7 @@ void EmbeddingVar<K, V>::SetDefaultValueOfNewFeatures(
       value_ptr->SetInitialized(emb_config_.emb_index);
       memcpy_address[*it] = value_ptr->GetValue(
           emb_config_.emb_index,
-          storage_manager_->GetOffset(emb_config_.emb_index));
+          storage_->GetOffset(emb_config_.emb_index));
     }
     TypedAllocator::Deallocate(alloc_, dev_value_address, total * 2);
     TypedAllocator::Deallocate(cpu_allocator(), value_address, total * 2);
@@ -148,12 +148,12 @@ void EmbeddingVar<K, V>::CopyEmbeddingsFromCPUToGPU(
     int64* output_value_ptrs) {
   if (copyback_cursor.size() > 0) {
     int64 total = copyback_cursor.size();
-    size_t value_len = emb_config_.total_num(storage_manager_->GetAllocLen());
+    size_t value_len = emb_config_.total_num(storage_->GetAllocLen());
     V* memcpy_buffer_gpu = nullptr;
     ValuePtr<V>** gpu_value_ptrs = new ValuePtr<V>*[total];
     memcpy_buffer_gpu = (V*)alloc_->AllocateRaw(Allocator::kAllocatorAlignment,
                                                 total * value_len * sizeof(V));
-    storage_manager_->CopyEmbeddingsFromCPUToGPU(
+    storage_->CopyEmbeddingsFromCPUToGPU(
         total, keys, copyback_cursor, memcpy_address, value_len, gpu_value_ptrs,
         memcpy_buffer_gpu, compute_stream, event_mgr, worker_threads);
 
@@ -185,7 +185,7 @@ void EmbeddingVar<K, V>::CopyEmbeddingsFromCPUToGPU(
     auto do_insert = [this, copyback_keys, gpu_value_ptrs, value_len](
                          int64 start, int64 limit) {
       for (int64 i = start; i < limit; i++)
-        storage_manager_->Insert(copyback_keys[i], gpu_value_ptrs[i]);
+        storage_->Insert(copyback_keys[i], gpu_value_ptrs[i]);
     };
     Shard(worker_threads->num_threads, worker_threads->workers,
           copyback_keys.size(), 100000, do_insert);