Avoid unloading kernels that haven't been loaded and fix replay_computation to

destroy the local client instance.

Add a user defined move constructor for KernelBase to nullify the parent stream
pointer. This is needed to avoid unloading kernels that haven't been loaded
when the moved-from KernelBase objects are destructed.

Add a call to ClientLibrary::DestroyLocalInstances to destroy the local client
instance used by replay_computation.

PiperOrigin-RevId: 181995818

1 files changed, 9 insertions, 0 deletions

diff --git a/tensorflow/stream_executor/kernel.cc b/tensorflow/stream_executor/kernel.cc
index e1b3635d52..81e531efb3 100644
--- a/tensorflow/stream_executor/kernel.cc
+++ b/tensorflow/stream_executor/kernel.cc
@@ -57,6 +57,15 @@ void KernelMetadata::set_shared_memory_bytes(int shared_memory_bytes) {
   has_shared_memory_bytes_ = true;
 }
 
+KernelBase::KernelBase(KernelBase &&from)
+    : parent_(from.parent_),
+      implementation_(std::move(from.implementation_)),
+      name_(std::move(from.name_)),
+      demangled_name_(std::move(from.demangled_name_)),
+      metadata_(from.metadata_) {
+  from.parent_ = nullptr;
+}
+
 KernelBase::KernelBase(StreamExecutor *parent)
     : parent_(parent),
       implementation_(parent->implementation()->CreateKernelImplementation()) {}