如何清除CuPy所有缓存内核以强制每次调用时重新编译？

CuPy文档提到：“CuPy会在进程内缓存发送至GPU设备的内核代码，从而减少后续调用时的内核编译时间。”这意味着首次调用CuPy函数后，后续调用速度会大幅提升。

以下是验证该机制的示例代码：

import cupy as cp
from timeit import default_timer as timer
import time

mempool = cp.get_default_memory_pool()
pinned_mempool = cp.get_default_pinned_memory_pool()


def multiply():
    rand = cp.random.default_rng()                             # 使用CuPy快速创建大型数组
    arr = rand.integers(0, 100_000, (10000, 1000))        # 创建数组
    y = cp.multiply(arr, 42) ## 乘以42（随机选择的数值）
    return y

if __name__ == '__main__':
    times = []
    start = timer()
    for i in range(21):
        mempool.free_all_blocks()
        pinned_mempool.free_all_blocks()
        start = timer()
        multiply()
        times.append(timer()-start) 

    print(times)

运行后得到的时间结果：

[0.17462146899993058, 0.0006819850000283623, 0.0006159440001738403, 0.0006145069999092811, 0.000610309999956371, 0.0006169410000893549, 0.0006062159998236893, 0.0006096620002153941, 0.0006096250001519365, 0.0006106630000886071, 0.0006063629998607212, 0.0006168999998408253, 0.0006058349999875645, 0.0006090080000831222, 0.0005964219999441411, 0.0006113049998930364, 0.0005968339999071759, 0.0005951619998540991, 0.0005980400001135422, 0.0005941219999385794, 0.0006568090000200755]

从结果可以看到，仅第一次调用包含内核编译时间，后续调用速度极快。问题是：有没有办法清空所有缓存，使得每次调用multiply()函数时都强制重新编译内核？

解决方法 #

1. 使用cp.clear_memo()清除内核缓存 #

CuPy提供了cp.clear_memo()方法，专门用于清空进程内的所有内核编译缓存。只需在每次调用目标函数前加入这个调用，就能强制内核重新编译。

修改后的主函数代码如下：

if __name__ == '__main__':
    times = []
    for i in range(21):
        mempool.free_all_blocks()
        pinned_mempool.free_all_blocks()
        # 清空内核缓存，强制重新编译
        cp.clear_memo()
        start = timer()
        multiply()
        times.append(timer()-start) 

    print(times)

2. 效果验证 #

修改后运行代码，所有调用的耗时都会接近第一次的数值，说明每次调用都触发了内核重新编译。

注意事项 #

• cp.clear_memo()仅清除当前进程内的内核缓存，不会影响其他进程的CuPy运行。
• 频繁清除缓存会显著降低性能，仅建议在测试内核编译耗时等特定场景下使用。

内容的提问来源于stack exchange，提问作者JOKKINATOR