1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_KERNELS_DATA_CACHE_OPS_H_ 17 #define TENSORFLOW_CORE_KERNELS_DATA_CACHE_OPS_H_ 18 19 #include "tensorflow/core/framework/resource_mgr.h" 20 #include "tensorflow/core/kernels/data/dataset_utils.h" 21 22 namespace tensorflow { 23 namespace data { 24 25 // A thread-safe data structure for caching dataset elements. 26 // 27 // The expected use is that a single `MemoryWriterIterator` populates the 28 // cache with dataset elements. Once all elements are cached, the cache can 29 // be used by one or more `MemoryReaderIterator`s. 30 class MemoryCache { 31 public: 32 MemoryCache() = default; 33 34 // Marks the cache as completed. 35 void Complete(std::vector<std::vector<Tensor>>&& cache); 36 37 // Returns whether the cache is completed. 38 bool IsCompleted(); 39 40 // Resets the cache. 41 void Reset(); 42 43 // Returns the element at the given index. 44 const std::vector<Tensor>& at(int64 index); 45 46 // Returns the size of the cache. 47 size_t size(); 48 49 // Returns a reference to the cache's data. The returned reference will be 50 // invalidated by any call to Reset(). 51 const std::vector<std::vector<Tensor>>& data(); 52 53 private: 54 mutex mu_; 55 // Determines whether all elements of the dataset have been cached. 56 bool completed_ TF_GUARDED_BY(mu_) = false; 57 std::vector<std::vector<Tensor>> cache_ TF_GUARDED_BY(mu_); 58 }; 59 60 // A resource wrapping a shared instance of a memory cache. 61 class MemoryCacheManager : public ResourceBase { 62 public: MemoryCacheManager()63 MemoryCacheManager() : cache_(std::make_shared<MemoryCache>()) {} 64 65 string DebugString() const override; 66 get()67 std::shared_ptr<MemoryCache> get() { return cache_; } 68 69 private: 70 std::shared_ptr<MemoryCache> cache_; 71 }; 72 73 // Creates an instance of cache resource and transfers ownership to the caller. 74 class AnonymousMemoryCacheHandleOp 75 : public AnonymousResourceOp<MemoryCacheManager> { 76 public: 77 explicit AnonymousMemoryCacheHandleOp(OpKernelConstruction* ctx); 78 79 private: 80 string name() override; 81 Status CreateResource(OpKernelContext* ctx, 82 std::unique_ptr<FunctionLibraryDefinition> flib_def, 83 std::unique_ptr<ProcessFunctionLibraryRuntime> pflr, 84 FunctionLibraryRuntime* lib, 85 MemoryCacheManager** manager) override; 86 }; 87 88 // Deletes an instance of cache resource. 89 class DeleteMemoryCacheOp : public OpKernel { 90 public: DeleteMemoryCacheOp(OpKernelConstruction * ctx)91 explicit DeleteMemoryCacheOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} 92 93 void Compute(OpKernelContext* ctx) override; 94 }; 95 96 } // namespace data 97 } // namespace tensorflow 98 99 #endif // TENSORFLOW_CORE_KERNELS_DATA_CACHE_OPS_H_ 100