diff options
author | 2017-11-07 14:00:29 -0800 | |
---|---|---|
committer | 2017-11-07 14:06:32 -0800 | |
commit | c070ab2a58ab5dc66da54db2daac4fa8e9f191e9 (patch) | |
tree | b8c8c4eeb0e076f46ba8b6c0a855883d920aa07b /tensorflow/stream_executor/kernel.cc | |
parent | 7183348f3270b7f9c1b333970e4f9abf6b3c4d8a (diff) |
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
Diffstat (limited to 'tensorflow/stream_executor/kernel.cc')
-rw-r--r-- | tensorflow/stream_executor/kernel.cc | 6 |
1 files changed, 5 insertions, 1 deletions
diff --git a/tensorflow/stream_executor/kernel.cc b/tensorflow/stream_executor/kernel.cc index b09fa8aca5..e1b3635d52 100644 --- a/tensorflow/stream_executor/kernel.cc +++ b/tensorflow/stream_executor/kernel.cc @@ -65,7 +65,11 @@ KernelBase::KernelBase(StreamExecutor *parent, internal::KernelInterface *implementation) : parent_(parent), implementation_(implementation) {} -KernelBase::~KernelBase() {} +KernelBase::~KernelBase() { + if (parent_) { + parent_->UnloadKernel(this); + } +} unsigned KernelBase::Arity() const { return implementation_->Arity(); } |