cupy.cuda.MemoryHook#

class cupy.cuda.MemoryHook[source]#

内存分配钩子的基类。

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

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

具体来说，所有注册的内存钩子的 alloc_preprocess() (或 alloc_postprocess()) 在从 GPU 设备分配内存之前 (或之后) 被调用。

同样，所有注册的内存钩子的 malloc_preprocess() (或 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_preprocess() (或 free_postprocess()) 在内存被释放回内存池之前 (或之后) 被调用。

下面是一个伪代码，描述了 free 和钩子的工作原理

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。