Clean up kernels cached by CUDAExecutor.

Notify CUDA executor on deallocation of previously loaded GPUExecutable and
unload corresponding module when all executables that were using it are gone.

PiperOrigin-RevId: 174908193

1 files changed, 2 insertions, 0 deletions

diff --git a/tensorflow/stream_executor/stream_executor_internal.h b/tensorflow/stream_executor/stream_executor_internal.h
index 12593e31d4..14445a7657 100644
--- a/tensorflow/stream_executor/stream_executor_internal.h
+++ b/tensorflow/stream_executor/stream_executor_internal.h
@@ -169,6 +169,8 @@ class StreamExecutorInterface {
                       const KernelArgsArrayBase &args) {
     return false;
   }
+  // Releases any state associated with the kernel.
+  virtual void UnloadKernel(const KernelBase *kernel) {}
   virtual void *Allocate(uint64 size) = 0;
   virtual void *AllocateSubBuffer(DeviceMemoryBase *parent, uint64 offset,
                                   uint64 size) = 0;