设备#
Warp 为系统中所有支持的计算设备分配唯一的字符串别名。当前暴露了一个 CPU 设备，别名为 "cpu"。每个支持 CUDA 的 GPU 都会获得一个形式为 "cuda:i" 的别名，其中 i 是 CUDA 设备的序号。对于 PyTorch 等其他流行框架的用户来说，这种约定应该很熟悉。
可以通过 device 参数在每次 Warp API 调用中显式指定目标设备
a = wp.zeros(n, device="cpu")
wp.launch(kernel, dim=a.size, inputs=[a], device="cpu")
b = wp.zeros(n, device="cuda:0")
wp.launch(kernel, dim=b.size, inputs=[b], device="cuda:0")
c = wp.zeros(n, device="cuda:1")
wp.launch(kernel, dim=c.size, inputs=[c], device="cuda:1")
注意
一个 Warp CUDA 设备（"cuda:i"）对应于设备 i 的主 CUDA 上下文。这与 PyTorch 等使用 CUDA Runtime API 的框架及其他软件兼容。这使得互操作性变得容易，因为 GPU 资源（如内存）可以与 Warp 共享。
- class warp.context.Device(runtime, alias, ordinal=-1, is_primary=False, context=None)[source]#
一个用于分配 Warp 数组和启动内核的设备。
- is_mempool_supported#
指示设备是否支持使用
cuMemAllocAsync和cuMemPool系列 API 进行流有序内存分配。CPU 设备为False。- 类型：
- is_ipc_supported#
指示设备是否支持 IPC。
如果支持，则为
True。如果不支持，则为
False。如果无法确定 IPC 支持（例如 CUDA 11），则为
None。
- 类型：
Optional[bool]
- pci_bus_id#
CUDA 设备的标识符，格式为
[domain]:[bus]:[device]，其中domain、bus和device均为十六进制值。CPU 设备为None。- 类型：
- get_allocator(pinned=False)[source]#
获取此设备的内存分配器。
- 参数：
pinned (bool) – 如果为
True，将返回一个用于固定内存的分配器。仅当此设备是 CPU 设备时适用。
- property stream: Stream[source]#
与 CUDA 设备关联的流。
- 引发：
RuntimeError – 该设备不是 CUDA 设备。
- set_stream(stream, sync=True)[source]#
设置此 CUDA 设备的当前流。
当前流将默认用于此设备上的所有内核启动和内存操作。
如果这是外部流，调用者负责保证流的生命周期。
考虑改用
warp.ScopedStream。
Warp 还提供了可用于查询系统中可用设备的函数
默认设备#
为了简化代码编写，Warp 引入了**默认设备**的概念。当 Warp API 调用中省略 device 参数时，将使用默认设备。
不带参数调用 wp.get_device() 将返回默认设备的 warp.context.Device 实例。
在 Warp 初始化期间，如果 CUDA 可用，默认设备将设置为 "cuda:0"。否则，默认设备为 "cpu"。如果默认设备被更改，可以使用 wp.get_preferred_device() 来获取原始默认设备。
wp.set_device() 可用于更改默认设备
wp.set_device("cpu")
a = wp.zeros(n)
wp.launch(kernel, dim=a.size, inputs=[a])
wp.set_device("cuda:0")
b = wp.zeros(n)
wp.launch(kernel, dim=b.size, inputs=[b])
wp.set_device("cuda:1")
c = wp.zeros(n)
wp.launch(kernel, dim=c.size, inputs=[c])
注意
对于 CUDA 设备，wp.set_device() 执行两项操作：设置 Warp 默认设备，并使设备的 CUDA 上下文成为当前上下文。这有助于最小化针对单个设备的代码块中 CUDA 上下文切换的次数。
对于 PyTorch 用户，此函数类似于 torch.cuda.set_device()。仍然可以在单个 API 调用中指定不同的设备，如下面的代码片段所示
# set default device
wp.set_device("cuda:0")
# use default device
a = wp.zeros(n)
# use explicit devices
b = wp.empty(n, device="cpu")
c = wp.empty(n, device="cuda:1")
# use default device
wp.launch(kernel, dim=a.size, inputs=[a])
wp.copy(b, a)
wp.copy(c, a)
作用域设备#
管理默认设备的另一种方法是使用 wp.ScopedDevice 对象。它们可以任意嵌套，并在退出时恢复先前的默认设备
with wp.ScopedDevice("cpu"):
# alloc and launch on "cpu"
a = wp.zeros(n)
wp.launch(kernel, dim=a.size, inputs=[a])
with wp.ScopedDevice("cuda:0"):
# alloc on "cuda:0"
b = wp.zeros(n)
with wp.ScopedDevice("cuda:1"):
# alloc and launch on "cuda:1"
c = wp.zeros(n)
wp.launch(kernel, dim=c.size, inputs=[c])
# launch on "cuda:0"
wp.launch(kernel, dim=b.size, inputs=[b])
- class warp.ScopedDevice(device)[source]#
一个上下文管理器，用于临时更改当前默认设备。
对于 CUDA 设备，此上下文管理器会使设备的 CUDA 上下文成为当前上下文，并在退出时恢复先前的 CUDA 上下文。这在将 Warp 脚本作为更大管道的一部分运行时很方便，因为它避免了在封闭代码中更改 CUDA 上下文的任何副作用。
- 参数：
device (Devicelike)
示例：将 wp.ScopedDevice 与多个 GPU 结合使用#
以下示例展示了如何在所有可用 CUDA 设备上分配数组和启动内核。
import warp as wp
@wp.kernel
def inc(a: wp.array(dtype=float)):
tid = wp.tid()
a[tid] = a[tid] + 1.0
# get all CUDA devices
devices = wp.get_cuda_devices()
device_count = len(devices)
# number of launches
iters = 1000
# list of arrays, one per device
arrs = []
# loop over all devices
for device in devices:
# use a ScopedDevice to set the target device
with wp.ScopedDevice(device):
# allocate array
a = wp.zeros(250 * 1024 * 1024, dtype=float)
arrs.append(a)
# launch kernels
for _ in range(iters):
wp.launch(inc, dim=a.size, inputs=[a])
# synchronize all devices
wp.synchronize()
# print results
for i in range(device_count):
print(f"{arrs[i].device} -> {arrs[i].numpy()}")
当前 CUDA 设备#
Warp 使用设备别名 "cuda" 来指定当前 CUDA 设备。这允许外部代码管理执行 Warp 脚本的 CUDA 设备。这类似于 PyTorch 的 "cuda" 设备，Torch 用户应该对此熟悉，并且可以简化互操作。
在此代码片段中，我们使用 PyTorch 管理当前 CUDA 设备并在该设备上调用 Warp 内核
def example_function():
# create a Torch tensor on the current CUDA device
t = torch.arange(10, dtype=torch.float32, device="cuda")
a = wp.from_torch(t)
# launch a Warp kernel on the current CUDA device
wp.launch(kernel, dim=a.size, inputs=[a], device="cuda")
# use Torch to set the current CUDA device and run example_function() on that device
torch.cuda.set_device(0)
example_function()
# use Torch to change the current CUDA device and re-run example_function() on that device
torch.cuda.set_device(1)
example_function()
注意
如果在代码运行环境中，另一部分代码可能不可预测地更改 CUDA 上下文，使用设备别名 "cuda" 可能会有问题。建议使用显式 CUDA 设备（如 "cuda:i"）来避免此类问题。
设备同步#
CUDA 内核启动和内存操作可以异步执行。这允许在不同设备上重叠计算和内存操作。Warp 允许主机与特定设备上未完成的异步操作进行同步
wp.synchronize_device("cuda:1")
wp.synchronize_device() 提供了比 wp.synchronize() 更细粒度的同步，因为后者等待所有设备完成其工作。
自定义 CUDA 上下文#
Warp 设计用于处理任意 CUDA 上下文，因此可以轻松集成到不同的工作流程中。
基于 CUDA Runtime API 构建的应用程序针对每个设备的主上下文。Runtime API 在底层隐藏了 CUDA 上下文管理。在 Warp 中，设备 "cuda:i" 代表设备 i 的主上下文，这与 CUDA Runtime API 一致。
基于 CUDA Driver API 构建的应用程序直接使用 CUDA 上下文，并可以在任何设备上创建自定义 CUDA 上下文。可以创建具有特定亲和性或互操作性功能的自定义 CUDA 上下文，这些功能对应用程序有利。Warp 也可以使用这些 CUDA 上下文。
特殊的设备别名 "cuda" 可用于指定当前的 CUDA 上下文，无论这是主上下文还是自定义上下文。
此外，Warp 允许使用 wp.map_cuda_device() 为自定义 CUDA 上下文注册新的设备别名，以便可以通过名称显式指定它们。如果 CUcontext 指针可用，可以使用它来创建新的设备别名，如下所示
wp.map_cuda_device("foo", ctypes.c_void_p(context_ptr))
或者，如果自定义 CUDA 上下文已被应用程序设为当前上下文，则可以省略该指针
wp.map_cuda_device("foo")
无论哪种情况，映射自定义 CUDA 上下文都允许我们使用分配的别名直接指定该上下文
with wp.ScopedDevice("foo"):
a = wp.zeros(n)
wp.launch(kernel, dim=a.size, inputs=[a])
CUDA 对等访问#
如果系统硬件配置支持，CUDA 允许不同 GPU 之间直接访问内存。通常，GPU 应为同类型，并且可能需要特殊的互连（例如 NVLINK 或 PCIe 拓扑）。
在初始化期间，Warp 会报告多 GPU 系统上是否支持对等访问
Warp 0.15.1 initialized:
CUDA Toolkit 11.5, Driver 12.2
Devices:
"cpu" : "x86_64"
"cuda:0" : "NVIDIA L40" (48 GiB, sm_89, mempool enabled)
"cuda:1" : "NVIDIA L40" (48 GiB, sm_89, mempool enabled)
"cuda:2" : "NVIDIA L40" (48 GiB, sm_89, mempool enabled)
"cuda:3" : "NVIDIA L40" (48 GiB, sm_89, mempool enabled)
CUDA peer access:
Supported fully (all-directional)
如果消息报告 CUDA 对等访问为 Supported fully，则表示系统中的每个 CUDA 设备都可以访问其他所有 CUDA 设备。如果报告为 Supported partially，则其后将显示访问矩阵，显示哪些设备可以相互访问。如果报告为 Not supported，则表示任意设备之间都不支持访问。
在代码中，我们可以像这样检查支持情况并启用对等访问
if wp.is_peer_access_supported("cuda:0", "cuda:1"):
wp.set_peer_access_enabled("cuda:0", "cuda:1", True):
这将允许设备 cuda:0 的内存可以在设备 cuda:1 上直接访问。对等访问是单向的，这意味着从 cuda:1 启用对设备 cuda:0 的访问并不会自动启用从 cuda:0 对设备 cuda:1 的访问。
启用对等访问的好处是允许设备之间进行直接内存传输 (DMA)。这通常是复制数据的更快方法，否则需要使用 CPU 暂存缓冲区进行传输。
缺点是启用对等访问可能会降低分配和释放的性能。不依赖对等内存传输的程序应保持此设置禁用。
可以使用作用域管理器临时启用或禁用对等访问
with wp.ScopedPeerAccess("cuda:0", "cuda:1", True):
...
- warp.is_peer_access_supported(target_device, peer_device)[source]#
检查此系统上 peer_device 是否可以直接访问 target_device 的内存。
这适用于使用默认 CUDA 分配器分配的内存。对于使用 CUDA 池化分配器分配的内存，请使用
is_mempool_access_supported()。
- warp.is_peer_access_enabled(target_device, peer_device)[source]#
检查 peer_device 当前是否可以访问 target_device 的内存。
这适用于使用默认 CUDA 分配器分配的内存。对于使用 CUDA 池化分配器分配的内存，请使用
is_mempool_access_enabled()。
- warp.set_peer_access_enabled(target_device, peer_device, enable)[源]#
启用或禁用从 peer_device 到 target_device 内存的直接访问。
启用对等访问可以提高对等内存传输的速度，但可能对内存消耗和分配性能产生负面影响。
这适用于使用默认 CUDA 分配器分配的内存。对于使用 CUDA 池化分配器分配的内存，请使用
set_mempool_access_enabled()。