From 9d33f296d1edc2f656e253cf2a015d36daedd5c1 Mon Sep 17 00:00:00 2001
From: Henry Tan <henrytan@google.com>
Date: Thu, 18 Jun 2020 15:00:50 -0700
Subject: [PATCH] Prep change for publishing TPU Ops.

PiperOrigin-RevId: 317188030
Change-Id: I29f9236c0ade6bf586c8a52ead977b5d31aec357
---
 tensorflow/core/tpu/kernels/BUILD             | 11 +++
 .../tpu_compilation_cache_entry_unloader.h    | 69 +++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h

diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index 9d38eb71f3c..a41747ee8c5 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -450,3 +450,14 @@ cc_library(
     hdrs = ["tpu_compile_op.h"],
     deps = ["//tensorflow/core:framework"],
 )
+
+cc_library(
+    name = "tpu_compilation_cache_entry_unloader",
+    hdrs = ["tpu_compilation_cache_entry_unloader.h"],
+    deps = [
+        ":tpu_compilation_cache_interface",
+        "//tensorflow/core:framework",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h b/tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h
new file mode 100644
index 00000000000..c298d8fcc12
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h
@@ -0,0 +1,69 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TPU_KERNELS_COMPILATION_CACHE_ENTRY_UNLOADER_H_
+#define TENSORFLOW_CORE_TPU_KERNELS_COMPILATION_CACHE_ENTRY_UNLOADER_H_
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/tpu/kernels/tpu_compilation_cache_interface.h"
+
+namespace tensorflow {
+namespace tpu {
+
+class TpuCompilationCacheEntryUnloader : public ResourceBase {
+ public:
+  explicit TpuCompilationCacheEntryUnloader(TpuCompilationCacheInterface* cache)
+      : cache_(cache) {
+    // Hold a reference to the cache until the unloader is destroyed.
+    cache_->Ref();
+    VLOG(1) << "Will unload compilation cache entries when session closes.";
+  }
+
+  ~TpuCompilationCacheEntryUnloader() override {
+    absl::MutexLock lock(&mu_);
+    for (int64 uid : cache_entry_uids_) {
+      Status s = cache_->MarkEntryForEviction(uid);
+      if (!s.ok()) {
+        LOG(WARNING) << "MarkEntryForEviction in "
+                        "~CompilationCacheEntryUnloader fails with error "
+                     << s;
+      }
+    }
+    // Release our reference to the cache.
+    cache_->Unref();
+  }
+
+  // Add cache entry uid to be unloaded in destructor.
+  void AddCacheEntryUid(int64 uid) {
+    absl::MutexLock lock(&mu_);
+    cache_entry_uids_.insert(uid);
+  }
+
+  std::string DebugString() const override {
+    return "CompilationCacheEntryUnloader";
+  }
+
+ private:
+  TF_DISALLOW_COPY_AND_ASSIGN(TpuCompilationCacheEntryUnloader);
+  mutable absl::Mutex mu_;
+  TpuCompilationCacheInterface* cache_;  // Not owned.
+  absl::flat_hash_set<int64> cache_entry_uids_ ABSL_GUARDED_BY(mu_);
+};
+
+}  // namespace tpu
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TPU_KERNELS_COMPILATION_CACHE_ENTRY_UNLOADER_H_