cupy.cuda.MemoryHook#

class cupy.cuda.MemoryHook[源代码]#

用于内存分配的钩子基类。

MemoryHook 是一个回调对象。注册的内存钩子在从GPU设备分配内存之前和之后，以及从内存池检索内存和将内存释放到内存池时被调用。

派生自 MemoryHook 的内存钩子需要实现六个方法：alloc_preprocess()、alloc_postprocess()、malloc_preprocess()、malloc_postprocess()、free_preprocess() 和 free_postprocess()。默认情况下，这些方法不执行任何操作。

具体来说，在从GPU设备分配内存之前（相应地，之后），会调用所有注册的内存钩子的 alloc_preprocess() 方法（相应地， alloc_postprocess() 方法）。

同样地，所有注册的内存钩子中的 malloc_postprocess()）在从内存池获取内存之前（分别地，之后）被调用。

下面是一个伪代码，用于描述 malloc 和钩子如何工作。请注意，如果在缓存的空闲块中找到，则不会调用 alloc_preprocess() 和 alloc_postprocess():

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

此外，所有注册的内存钩子中的 free_postprocess()）在内存释放到内存池之前（分别对应之后）被调用。

下面是一个伪代码，用于描述自由和钩子如何工作:

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

要注册一个内存钩子，请使用 with 语句。内存钩子会在 with 语句内的所有方法调用中注册，并在 with 语句结束时注销。

备注

CuPy 将已注册的函数钩子字典存储为线程本地对象。因此，注册的内存钩子可能会因线程而异。

方法

__enter__(self)#

__exit__(self, *_)#

alloc_postprocess(self, **kwargs)#

从GPU设备分配内存后调用的回调函数。

关键字参数:

device_id (int) – CUDA 设备 ID
mem_size (int) – 圆整后的内存分配字节大小
mem_ptr (int) – 获取的内存指针。如果分配过程中发生错误，则为0。

alloc_preprocess(self, **kwargs)#

从GPU设备分配内存之前调用的回调函数。

关键字参数:

device_id (int) – CUDA 设备 ID
mem_size (int) – 要分配的圆整内存字节大小

free_postprocess(self, **kwargs)#

回调函数在将内存释放回内存池后调用。

关键字参数:

device_id (int) – CUDA 设备 ID
mem_size (int) – 内存字节大小
mem_ptr (int) – 要释放的内存指针
pmem_id (int) – 池化内存对象ID。

free_preprocess(self, **kwargs)#

在将内存释放回内存池之前调用的回调函数。

关键字参数:

device_id (int) – CUDA 设备 ID
mem_size (int) – 内存字节大小
mem_ptr (int) – 要释放的内存指针
pmem_id (int) – 池化内存对象ID。

malloc_postprocess(self, **kwargs)#

从内存池中检索内存后调用的回调函数。

关键字参数:

device_id (int) – CUDA 设备 ID
size (int) – 请求分配的内存字节大小
mem_size (int) – 圆整后的内存分配字节大小
mem_ptr (int) – 获取的内存指针。如果 malloc 中发生错误，则为 0。
pmem_id (int) – 池化内存对象ID。如果在 malloc 中发生错误，则为0。