FreeCudaCallbackRegistry()

dfalbel: <p>It seems that the CUDA allocator in c10 has a callback registry. Would it be possible to add other callbacks to it using the C++ frontend?</p>
<p>Context: the R garbage collector is lazy and will run only when R needs more memory, however R is not aware of CUDA memory and won’t GC even if the GPU memory is used by objects that no longer have references. I’d like then, to add a callback to the CUDA memory allocator so it calls R garbage collector whenever it needs more memory.</p>
<p>Any suggestion is appreciated, thanks very much!!</p>
<p>see: <a href="https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L617-L624" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L617-L624</a></p>
<pre><code class="lang-auto">bool trigger_free_memory_callbacks(AllocParams&amp; p) {
    bool freed_memory = false;
    for (const auto&amp; name : FreeCudaMemoryCallbacksRegistry()-&gt;Keys()) {
      freed_memory |=
        FreeCudaMemoryCallbacksRegistry()-&gt;Create(name)-&gt;Execute();
    }
    return freed_memory;
}
</code></pre>

glaringlee: <p><a class="mention" href="/u/dfalbel">@dfalbel</a><br>
I never did it before. But I think it is doable.</p>
<p>The FreeCudaMemoryCallbacksRegistry is defined here and used to register FreeMemoryCallbacks<br>
‘<a href="https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L23" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L23</a>’</p>
<p>FreeCudaMemoryCallbacksRegistry is actually a c10::registry that defined here:<br>
‘<a href="https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L268" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L268</a>’<br>
‘<a href="https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L215" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L215</a>’<br>
‘<a href="https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L268" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L268</a>’</p>
<p>FreeMemoryCallback is defined here:<br>
‘<a href="https://github.com/pytorch/pytorch/blob/6f396e18c33efbb8547b7ca9d412a316bd51429b/c10/cuda/CUDACachingAllocator.h#L20" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/6f396e18c33efbb8547b7ca9d412a316bd51429b/c10/cuda/CUDACachingAllocator.h#L20</a>’</p>
<p>So to add your own free memory callback you can do the following:</p>
<ol>
<li>Write your own FreeMemoryCallback by inheriting FreeMemoryCallback.<br>
Implement your own execute()</li>
<li>Use FreeCudaMemoryCallbacksRegistry()-&gt;register(…) to register your own<br>
see c10::Registry to learn what parameter should be passed in, it basically the<br>
type of your customized memory callback class and a creator func ptr.</li>
</ol>
<p>Your customized free memory callback will be called in trigger_free_memory_callbacks then.</p>

dfalbel: <p>Thanks <a class="mention" href="/u/glaringlee">@glaringlee</a> ! This is very helpful! Just one follow up question: would you know if the allocators/callbacks are always called from the main thread or there’s no guaranty for this?</p>

glaringlee: <p><a class="mention" href="/u/dfalbel">@dfalbel</a><br>
I think this depends on where you call the malloc, either in main thread or other threads.<br>
‘<a href="https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L214" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L214</a>’<br>
The memory free callback is triggered whenever you call malloc().</p>