分配器#

流排序内存池分配器#

简介#

Warp 0.14.0 添加了对 CUDA 数组的流排序内存池分配器的支持。从 Warp 0.15.0 开始，这些分配器默认在所有支持它们的 CUDA 设备上启用。“流排序内存池分配器”很拗口，所以让我们一点一点地分解它。

无论何时创建数组，都需要在设备上分配内存

a = wp.empty(n, dtype=float, device="cuda:0")
b = wp.zeros(n, dtype=float, device="cuda:0")
c = wp.ones(n, dtype=float, device="cuda:0")
d = wp.full(n, 42.0, dtype=float, device="cuda:0")

上面的每个调用都会分配一个足够大的设备内存块来容纳数组，并且可以选择使用指定的值初始化内容。 wp.empty() 是唯一不以任何方式初始化内容的函数，它只是分配内存。

内存池分配器从更大的预留内存池中获取内存块，这通常比向操作系统请求全新的存储块更快。这是这些池化分配器的一个重要优势——它们更快。

流排序意味着每个分配都安排在 CUDA 流上，它表示在 GPU 上按顺序执行的指令序列。主要好处是它允许在 CUDA 图中分配内存，这以前是不可能的

with wp.ScopedCapture() as capture:
    a = wp.zeros(n, dtype=float)
    wp.launch(kernel, dim=a.size, inputs=[a])

wp.capture_launch(capture.graph)

从现在开始，我们将这些分配器简称为内存池分配器。

配置#

内存池分配器是 CUDA 的一项功能，在大多数现代设备和操作系统上都受支持。但是，在某些系统（例如某些虚拟机设置）中，它们可能不受支持。Warp 的设计考虑了弹性，因此在引入这些新分配器之前编写的现有代码应该继续运行，而不管底层系统是否支持它们。

Warp 的启动消息给出了这些分配器的状态，例如

Warp 0.15.1 initialized:
CUDA Toolkit 11.5, Driver 12.2
Devices:
    "cpu"      : "x86_64"
    "cuda:0"   : "NVIDIA GeForce RTX 4090" (24 GiB, sm_89, mempool enabled)
    "cuda:1"   : "NVIDIA GeForce RTX 3090" (24 GiB, sm_86, mempool enabled)

请注意每个 CUDA 设备旁边的 mempool enabled 文本。这意味着内存池在该设备上已启用。无论何时在该设备上创建数组，都将使用内存池分配器来分配它。如果看到 mempool supported，则表示支持内存池，但启动时未启用。如果看到 mempool not supported，则表示无法在此设备上使用内存池。

有一个配置标志控制是否应在 wp.init() 期间自动启用内存池

import warp as wp

wp.config.enable_mempools_at_init = False

wp.init()

该标志默认为 True，但如果需要，可以设置为 False。在调用 wp.init() 后更改此配置标志无效。

在 wp.init() 之后，您可以检查每个设备上是否启用了内存池，如下所示

if wp.is_mempool_enabled("cuda:0"):
    ...

您还可以独立控制每个设备上的启用

if wp.is_mempool_supported("cuda:0"):
    wp.set_mempool_enabled("cuda:0", True)

可以使用作用域管理器临时启用或禁用内存池

with wp.ScopedMempool("cuda:0", True):
    a = wp.zeros(n, dtype=float, device="cuda:0")

with wp.ScopedMempool("cuda:0", False):
    b = wp.zeros(n, dtype=float, device="cuda:0")

在上面的代码段中，数组 a 将使用内存池分配器分配，数组 b 将使用默认分配器分配。

在大多数情况下，没有必要摆弄这些启用功能，但如果您需要它们，它们就在那里。默认情况下，如果支持内存池，Warp 将在启动时启用内存池，这将自动带来提高分配速度的好处。大多数 Warp 代码应该在有或没有内存池分配器的情况下继续运行，但图捕获期间的内存分配除外，如果未启用内存池，则会引发异常。

warp.is_mempool_supported(device)[source]#

检查设备上是否可以使用 CUDA 内存池分配器。

参数:: device (Device | str | None) – 要执行查询的 Device 或设备标识符。如果为 None，将使用默认设备。
返回类型:: bool

warp.is_mempool_enabled(device)[source]#

检查设备上是否启用了 CUDA 内存池分配器。

参数:: device (Device | str | None) – 要执行查询的 Device 或设备标识符。如果为 None，将使用默认设备。
返回类型:: bool

warp.set_mempool_enabled(device, enable)[source]#

在设备上启用或禁用 CUDA 内存池分配器。

池化分配器通常更快，并且允许在图捕获期间分配内存。

它们通常应该启用，但有一个罕见的注意事项。如果在图捕获期间，使用池化分配器分配了内存，并且未在两个 GPU 之间启用内存池访问，则在不同的 GPU 之间复制数据可能会失败。这是一个与 Warp 无关的内部 CUDA 限制。首选解决方案是使用 set_mempool_access_enabled() 启用内存池访问。如果不支持对等访问，则必须使用默认的 CUDA 分配器在图捕获之前预先分配内存。

参数:

device (Device | str | None) – 要执行操作的 Device 或设备标识符。如果为 None，将使用默认设备。
enable (bool)

返回类型:

None

查询内存使用情况#

可以使用 wp.get_mempool_used_mem_current() 查询应用程序当前从特定内存池使用的内存量。这可能与为池本身保留的内存量不同。同样，可以使用 wp.get_mempool_used_mem_high() 查询已使用内存的最高水位线。

warp.get_mempool_used_mem_current(device=None)[source]#

获取应用程序当前正在使用的设备内存池中的内存量。

参数:

device (Device | str | None) – 要执行查询的 Device 或设备标识符。如果为 None，将使用默认设备。

返回:

使用的内存量（以字节为单位）。

引发:

ValueError – 如果 device 不是 CUDA 设备。
RuntimeError – 如果 device 是 CUDA 设备，但不支持内存池。

返回类型:

int

warp.get_mempool_used_mem_high(device=None)[source]#

获取应用程序在设备 CUDA 内存池中使用的内存峰值。

参数:

device (Device | str | None) – 要执行查询的 Device 或设备标识符。如果为 None，将使用默认设备。

返回:

内存池中已使用的内存峰值，以字节为单位。

引发:

ValueError – 如果 device 不是 CUDA 设备。
RuntimeError – 如果 device 是 CUDA 设备，但不支持内存池。

返回类型:

int

分配性能#

分配和释放内存是相当昂贵的操作，会增加程序的开销。我们无法避免它们，因为我们需要为我们的数据分配存储空间，但是有一些简单的策略可以减少分配对性能的总体影响。

考虑以下示例

for i in range(100):
    a = wp.zeros(n, dtype=float, device="cuda:0")
    wp.launch(kernel, dim=a.size, inputs=[a], device="cuda:0")

在循环的每次迭代中，我们都会分配一个数组并在数据上运行一个内核。该程序有 100 个分配和 100 个释放。当我们为 a 分配一个新值时，先前的值会被 Python 回收器回收，从而触发释放。

重用内存#

如果数组的大小保持不变，请考虑在后续迭代中重用内存。我们可以只分配一次数组，然后在每次迭代时重新初始化其内容

# pre-allocate the array
a = wp.empty(n, dtype=float, device="cuda:0")
for i in range(100):
    # reset the contents
    a.zero_()
    wp.launch(kernel, dim=a.size, inputs=[a], device="cuda:0")

如果数组大小在每次迭代中不更改，则此方法效果很好。如果大小发生变化但已知上限，我们仍然可以预先分配一个足够大的缓冲区来存储任何迭代中的所有元素。

# pre-allocate a big enough buffer
buffer = wp.empty(MAX_N, dtype=float, device="cuda:0")
for i in range(100):
    # get a buffer slice of size n <= MAX_N
    n = get_size(i)
    a = buffer[:n]
    # reset the contents
    a.zero_()
    wp.launch(kernel, dim=a.size, inputs=[a], device="cuda:0")

以这种方式重用内存可以提高性能，但也可能为我们的代码增加不必要的复杂性。内存池分配器有一个有用的特性，可以在不以任何方式修改原始代码的情况下提高分配性能。

释放阈值#

内存池释放阈值确定分配器应保留多少预留内存，然后再将其释放回操作系统。对于频繁分配和释放内存的程序，设置更高的释放阈值可以提高分配性能。

默认情况下，释放阈值设置为 0。如果先前已获取内存并将其返回到池中，则将其设置为更高的数字将降低分配成本。

# set the release threshold to reduce re-allocation overhead
wp.set_mempool_release_threshold("cuda:0", 1024**3)

for i in range(100):
    a = wp.zeros(n, dtype=float, device="cuda:0")
    wp.launch(kernel, dim=a.size, inputs=[a], device="cuda:0")

0 到 1 之间的阈值被解释为可用内存的分数。例如，0.5 表示设备物理内存的一半，1.0 表示所有内存。更大的值被解释为绝对字节数。例如，1024**3 表示 1 GiB 的内存。

这是一个简单的优化，可以在不以任何方式修改现有代码的情况下提高程序的性能。

warp.get_mempool_release_threshold(device=None)[source]#

获取设备上的 CUDA 内存池释放阈值。

参数:

device (Device | str | None) – 要执行查询的 Device 或设备标识符。如果为 None，将使用默认设备。

返回:

内存池释放阈值，以字节为单位。

引发:

ValueError – 如果 device 不是 CUDA 设备。
RuntimeError – 如果 device 是 CUDA 设备，但不支持内存池。

返回类型:

int

warp.set_mempool_release_threshold(device, threshold)[source]#

在设备上设置 CUDA 内存池释放阈值。

这是在尝试将内存释放回操作系统之前要保留的预留内存量。当内存池持有的字节数超过此数量时，分配器将尝试在下次调用流、事件或设备同步时将内存释放回操作系统。

0 到 1 之间的值被解释为可用内存的分数。例如，0.5 表示设备物理内存的一半。更大的值被解释为绝对字节数。例如，1024**3 表示 1 GiB 的内存。

参数:

device (Device | str | None) – 要执行操作的 Device 或设备标识符。如果为 None，将使用默认设备。
threshold (int | float) – 一个整数，表示字节数，或者一个介于 0 和 1 之间的 float，指定所需的释放阈值。

引发:

ValueError – 如果 device 不是 CUDA 设备。
RuntimeError – 如果 device 是 CUDA 设备，但不支持内存池。
RuntimeError – 无法设置内存池释放阈值。

返回类型:

None

图分配#

内存池分配器可以在 CUDA 图中使用，这意味着您可以捕获创建数组的 Warp 代码

with wp.ScopedCapture() as capture:
    a = wp.full(n, 42, dtype=float)

wp.capture_launch(capture.graph)

print(a)

捕获分配类似于捕获其他操作，例如内核启动或内存复制。在捕获期间，操作实际上不会执行，而是会被记录下来。要执行捕获的操作，我们必须使用 wp.capture_launch() 启动该图。如果您想使用在图捕获期间分配的数组，请记住这一点很重要。该数组在启动捕获的图之前实际上并不存在。在上面的代码片段中，如果在调用 wp.capture_launch() 之前尝试打印该数组，我们会收到错误。

更一般而言，在图捕获期间分配内存的能力大大增加了可以在图中捕获的代码范围。这包括任何创建临时分配的代码。 CUDA 图可用于以最小的 CPU 开销重新运行操作，从而可以显着提高性能。

内存池访问#

在支持对等访问的多 GPU 系统上，我们可以启用从不同设备直接访问内存池的功能

if wp.is_mempool_access_supported("cuda:0", "cuda:1"):
    wp.set_mempool_access_enabled("cuda:0", "cuda:1", True)

这将允许设备 cuda:0 的内存池在设备 cuda:1 上直接访问。内存池访问是定向的，这意味着启用从 cuda:1 对 cuda:0 的访问不会自动启用从 cuda:0 对 cuda:1 的访问。

启用内存池访问的好处是它允许设备之间的直接内存传输 (DMA)。这通常是复制数据的更快方法，因为否则需要使用 CPU 暂存缓冲区来完成传输。

缺点是启用内存池访问可能会稍微降低分配和释放的性能。但是，对于依赖于在设备之间复制内存的应用程序，应该有净收益。

可以使用范围管理器暂时启用或禁用内存池访问

with wp.ScopedMempoolAccess("cuda:0", "cuda:1", True):
    a0 = wp.zeros(n, dtype=float, device="cuda:0")
    a1 = wp.empty(n, dtype=float, device="cuda:1")

    # use direct memory transfer between GPUs
    wp.copy(a1, a0)

请注意，内存池访问仅适用于使用内存池分配器分配的内存。对于使用默认 CUDA 分配器分配的内存，我们可以使用 wp.set_peer_access_enabled() 启用 CUDA 对等访问，以获得类似的好处。

由于启用内存池访问可能会带来缺点，因此即使支持，Warp 也不会自动启用它。因此，不需要在 GPU 之间复制数据的程序不会受到任何影响。

warp.is_mempool_access_supported(target_device, peer_device)[source]#

检查 peer_device 是否可以直接访问 target_device 的内存池。

如果内存池访问是可能的，可以使用 set_mempool_access_enabled() 和 is_mempool_access_enabled() 进行管理。

返回:

一个布尔值，指示系统是否支持此内存池访问。

参数:

target_device (Device | str | None)
peer_device (Device | str | None)

返回类型:

bool

warp.is_mempool_access_enabled(target_device, peer_device)[source]#

检查 peer_device 当前是否可以访问 target_device 的内存池。

这适用于使用 CUDA 池化分配器分配的内存。对于使用默认 CUDA 分配器分配的内存，请使用 is_peer_access_enabled()。

返回:

一个布尔值，指示当前是否启用了此对等访问。

参数:

target_device (Device | str | None)
peer_device (Device | str | None)

返回类型:

bool

warp.set_mempool_access_enabled(target_device, peer_device, enable)[source]#

启用或禁用从 peer_device 到 target_device 内存池的访问。

这适用于使用 CUDA 池化分配器分配的内存。对于使用默认 CUDA 分配器分配的内存，请使用 set_peer_access_enabled()。

参数:

target_device (Device | str | None)
peer_device (Device | str | None)
enable (bool)

返回类型:

None

限制#

GPU 之间在图捕获期间的内存池到内存池的复制#

如果在图捕获期间，源和目标使用内存池分配器分配并且设备之间未启用内存池访问，则在不同 GPU 之间复制数据将会失败。请注意，这仅适用于捕获图中的内存池到内存池的复制；在图捕获之外完成的复制不受影响。同一内存池（即同一设备）内的复制也不受影响。

有两种解决方法。如果支持内存池访问，您只需在图捕获之前启用设备之间的内存池访问即可，如内存池访问中所示。

如果不支持内存池访问，您需要使用默认的 CUDA 分配器预先分配复制中涉及的数组。这需要在捕获开始之前完成。

# pre-allocate the arrays with mempools disabled
with wp.ScopedMempool("cuda:0", False):
    a0 = wp.zeros(n, dtype=float, device="cuda:0")
with wp.ScopedMempool("cuda:1", False):
    a1 = wp.empty(n, dtype=float, device="cuda:1")

with wp.ScopedCapture("cuda:1") as capture:
    wp.copy(a1, a0)

wp.capture_launch(capture.graph)

这是由于 CUDA 中的一个限制，我们预计将来会修复它。