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?

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.

Any suggestion is appreciated, thanks very much!!


bool trigger_free_memory_callbacks(AllocParams& p) {
    bool freed_memory = false;
    for (const auto& name : FreeCudaMemoryCallbacksRegistry()->Keys()) {
      freed_memory |=
    return freed_memory;

I never did it before. But I think it is doable.

The FreeCudaMemoryCallbacksRegistry is defined here and used to register FreeMemoryCallbacks

FreeCudaMemoryCallbacksRegistry is actually a c10::registry that defined here:

FreeMemoryCallback is defined here:

So to add your own free memory callback you can do the following:

  1. Write your own FreeMemoryCallback by inheriting FreeMemoryCallback.
    Implement your own execute()
  2. Use FreeCudaMemoryCallbacksRegistry()->register(…) to register your own
    see c10::Registry to learn what parameter should be passed in, it basically the
    type of your customized memory callback class and a creator func ptr.

Your customized free memory callback will be called in trigger_free_memory_callbacks then.

Thanks @glaringlee ! 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?

I think this depends on where you call the malloc, either in main thread or other threads.
The memory free callback is triggered whenever you call malloc().

1 Like