Base class of hooks for Memory allocations.
MemoryHookis an callback object. Registered memory hooks are invoked before and after memory is allocated from GPU device, and memory is retrieved from memory pool, and memory is released to memory pool.
Memory hooks that derive
MemoryHookare required to implement six methods:
free_postprocess(), By default, these methods do nothing.
def malloc(size): Call malloc_preprocess of all memory hooks Try to find a cached free chunk from memory pool if chunk is not found: Call alloc_preprocess for all memory hooks Invoke actual memory allocation to get a new chunk Call alloc_postprocess for all memory hooks Call malloc_postprocess for all memory hooks
Below is a pseudo code to descirbe how free and hooks work:
def free(ptr): Call free_preprocess of all memory hooks Push a memory chunk of a given pointer back to memory pool Call free_postprocess for all memory hooks
To register a memory hook, use
withstatement. Memory hooks are registered to all method calls within
withstatement and are unregistered at the end of
CuPy stores the dictionary of registered function hooks as a thread local object. So, memory hooks registered can be different depending on threads.
Callback function invoked after allocating memory from GPU device.
Callback function invoked before allocating memory from GPU device.
Callback function invoked after releasing memory to memory pool.
Callback function invoked before releasing memory to memory pool.
Callback function invoked after retrieving memory from memory pool.
- Keyword Arguments
Callback function invoked before retrieving memory from memory pool.