diff options
| author |  Artem Belevich <tra@google.com> | 2017-11-07 14:00:29 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-11-07 14:06:32 -0800 |
| commit | c070ab2a58ab5dc66da54db2daac4fa8e9f191e9 (patch) | |
| tree | b8c8c4eeb0e076f46ba8b6c0a855883d920aa07b /tensorflow/stream_executor/stream_executor_pimpl.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/stream_executor_pimpl.cc')
| -rw-r--r-- | tensorflow/stream_executor/stream_executor_pimpl.cc | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/stream_executor/stream_executor_pimpl.cc b/tensorflow/stream_executor/stream_executor_pimpl.cc index 9dc1749327..76afb85068 100644 --- a/tensorflow/stream_executor/stream_executor_pimpl.cc +++ b/tensorflow/stream_executor/stream_executor_pimpl.cc @@ -217,6 +217,10 @@ bool StreamExecutor::GetKernel(const MultiKernelLoaderSpec &spec, return implementation_->GetKernel(spec, kernel); } +void StreamExecutor::UnloadKernel(const KernelBase *kernel) { + implementation_->UnloadKernel(kernel); +} + void StreamExecutor::Deallocate(DeviceMemoryBase *mem) { VLOG(1) << "Called StreamExecutor::Deallocate(mem=" << mem->opaque() << ") mem->size()=" << mem->size() << StackTraceIfVLOG10(); |