NVSHMEM 和 CUDA 模型

本节讨论 CUDA 抽象机模型和 NVSHMEM 之间的交互。

CUDA 执行模型

CUDA 中的工作提交

在 CUDA 模型中,所有工作(即 CUDA 任务)都通过 CUDA 流提交到 GPU,这些流按照先进先出 (FIFO) 的顺序执行各自的任务。为了实现并发,应用程序可以创建多个 CUDA 流,并在不同的流上排队工作,以表明任务可以并行处理。CUDA 事件可用于通过在一个流上记录事件并在另一个流中等待该事件来识别跨流的依赖关系。

基于 NVSHMEM 流的操作(例如 nvshmemx_putmem_on_stream)在 CUDA 流上将相应的操作排队并立即返回。NVSHMEM 操作只有到达流的头部并由 CUDA 运行时执行后才会执行。所有 NVSHMEM 流上操作都遵循 NVSHMEM 内存模型。例如,nvshmemx_quiet_on_stream 操作可分别用于排序或完成操作。

流上操作也具有类似线程的语义,因为 CUDA 层可以并行执行来自不同流的操作,以及主机 CPU 执行的工作。但是,使用基于流的操作不需要使用额外的线程支持来初始化 NVSHMEM。用户必须小心避免流上操作可能违反共享对象的线程要求的情况。例如,使用同一团队的集合通信操作不得并行执行。当在 CPU 上并行执行集合通信操作(例如 nvshmem_malloc)与在流上使用同一团队的操作(例如 nvshmem_barrier_all_on_stream)时,可能会发生这种情况。在这种情况下,两个操作都使用 NVSHMEM_TEAM_WORLD。当在不同的流上提交使用同一团队的集合通信操作,而没有任何同步来阻止它们并行执行时,也可能发生这种情况。

NVSHMEM 的设备端 API 可以由在 GPU 上执行的内核调用。在 CUDA 模型中,所有内核启动都在用户指定的流上排队,如果未指定,则在默认流上排队。因此,NVSHMEM 设备发起操作的使用也必须考虑与 CUDA 流语义和 CUDA 线程执行模型的交互。

CUDA 抽象机

除了用户显式指定的任何依赖关系(通过流和事件)之外,CUDA 抽象机模型允许 CUDA 在 CUDA 任务之间插入虚假依赖关系。但是,CUDA 添加的这些虚假依赖关系不得向 CUDA 任务图添加循环。这允许 CUDA 层管理任务的执行并在共享资源之间调度任务。例如,在 CUDA 流上排队的任务通过分配给 CUDA 上下文的有限数量的硬件流提交到 GPU。硬件流也按 FIFO 顺序处理。因此,当 CUDA 层将来自多个 CUDA 流的任务插入到同一硬件流中时,它会在任务之间引入虚假依赖关系。

NVSHMEM 操作可能会在 CUDA 任务之间引入依赖关系。例如,内核可以执行点对点同步,而流上集合通信需要多个 PE 的参与。由于在引入虚假依赖关系时,CUDA 层看不到这些 NVSHMEM 依赖关系,因此它们可能会导致执行图中的循环并导致死锁。在以下部分中,我们将重点介绍可能发生这种情况的几种情况,并讨论避免死锁的解决方案。

非本地操作和 CUDA 执行模型

NVSHMEM 提供了可能会阻塞的操作,直到执行一个或多个附加操作。我们将此类操作称为非本地操作。

示例包括 NVSHMEM 点对点同步操作,这些操作可能会阻塞,直到本地 PE 或远程 PE 执行的一个或多个 NVSHMEM 操作更新同步变量以满足等待条件。NVSHMEM 集合通信操作也可能会阻塞,直到团队中的所有 PE 都执行对集合通信操作的匹配调用。

NVSHMEM 提供内核发起和基于流的非本地操作。当执行基于流的非本地 NVSHMEM 操作时,它具有阻止同一 CUDA 流或依赖 CUDA 流中的后续任务被执行(即阻塞 CUDA 流)的效果,直到它完成。当执行内核发起的非本地 NVSHMEM 操作时,它具有保持 CUDA 执行资源的效果,同时也会阻塞在其上排队的 CUDA 流。

CUDA 流和循环依赖

用户必须确保在流上排队的 NVSHMEM 非本地操作不会形成循环依赖,这可能会导致死锁。例如,考虑 PE 0 在流上排队以下操作的情况(最左边的操作位于流的头部)

PE 0: [ nvshmemx_barrier_all_on_stream, nvshmemx_putmem_signal_on_stream ]

PE 1 在流上排队以下操作

PE 1: [ nvshmemx_signal_wait_until_on_stream, nvshmemx_barrier_all_on_stream ]

PE 1 上的信号等待操作将由 PE 0 排队的 putmem-with-signal 操作满足。但是,PE 0 被阻塞在 barrier 操作中,阻止它执行 putmem-with-signal 操作。两个 PE 都无法向前推进,从而导致死锁。

CUDA 流顺序和执行资源

同样,考虑消息交换,其中名为 notify_kernel 的内核调用 nvshmem_putmem_signal 向对等 PE 发送消息,而 wait_kernel 执行相应的 nvshmem_signal_wait_until 操作以等待来自对等 PE 的消息。这些操作可以按如下方式在 PE 0 和 PE 1 中排队到单独的流中

PE 0, Stream A: [ wait_kernel ]
PE 0, Stream B: [ notify_kernel ]
PE 1, Stream A: [ wait_kernel ]
PE 1, Stream B: [ notify_kernel ]

用户已将这些操作排队到单独的流中,因为它们可以并行执行。由于内核在单独的流中排队,因此 CUDA 运行时可以按任何顺序执行它们。如果 CUDA 运行时首先执行 wait_kernel,它将阻塞并保持 CUDA 执行资源。这些资源可能是执行 notify_kernel 所需的,从而阻止它执行。此外,即使有足够的资源可用,CUDA 运行时也可能在内核之间插入设备同步,从而阻止它们并行运行。如果两个 PE 都首先执行 wait_kernel,则可能导致死锁。可以通过将两个内核与 notify kernel 一起排队到同一流中,或通过插入 CUDA 事件以防止首先执行 wait_kernel 来防止死锁。

CUDA 流和虚假循环依赖

CUDA 模型允许用户创建任意数量的流来描述其工作负载。CUDA 运行时将来自这些流的工作推送到由 GPU 管理的 GPU 工作队列中。GPU 工作队列中的工作按 FIFO 顺序处理。在将来自流的工作分配给 GPU 工作队列时,CUDA 运行时会考虑任何流顺序和 CUDA 事件依赖关系。NVSHMEM 的非本地依赖关系对 CUDA 运行时不可见,用户必须小心地引入 CUDA 可见的依赖关系,以防止 CUDA 层以可能导致死锁的顺序序列化任务。

例如,考虑上面给出的消息交换示例

PE 0, Stream A: [ wait_kernel ] PE 0, Stream B: [ notify_kernel ]
PE 1, Stream A: [ wait_kernel ] PE 1, Stream B: [ notify_kernel ]

CUDA 运行时可能会按如下方式序列化这些任务,这将导致死锁

PE 0, GPU Work Queue: [ wait_kernel, notify_kernel ]
PE 1, GPU Work Queue: [ wait_kernel, notify_kernel ]

用户可以通过将 notify_kernel 和 wait_kernel 排队到同一流中,或通过引入 CUDA 事件来引入依赖关系,如下所示

PE 0, Stream A: [ cudaEventSynchronize(e), wait_kernel ]
PE 0, Stream B: [ notify_kernel, cudaEventRecord(e) ]
PE 1, Stream A: [ cudaEventSynchronize(e), wait_kernel ]
PE 1, Stream B: [ notify_kernel, cudaEventRecord(e) ]

内核内同步

CUDA 吞吐量计算模型允许用户指定远大于给定 GPU 上可以并行执行的网格和线程块维度。公开如此大的工作量可以使 GPU 有效地处理工作。但是,CUDA 不保证抢占式线程调度。因此,当内核中的线程执行非本地操作时,它可能会阻止同一内核中其他线程的执行。反过来,这可能会阻止将进行匹配调用的线程执行,从而导致死锁。可以通过使用 cudaLaunchCooperativeKernel API 启动内核来避免此问题,这可确保内核中的线程可以安全地相互同步而不会导致死锁。

使用 NVSHMEM 协同内核启动确保安全的非本地操作

为了简化在 CUDA 内核中使用非本地 NVSHMEM 函数,NVSHMEM 提供了内核启动例程nvshmemx_collective_launch 函数可用于在 CUDA 内核使用 NVSHMEM 同步或集合通信 API(例如,nvshmem_waitnvshmem_barriernvshmem_barrier_all 或任何其他集合通信操作)时在 GPU 上启动 CUDA 内核。不需要通过此 API 启动不使用同步 NVSHMEM API 或根本不使用 NVSHMEM API 的 CUDA 内核。

此调用在 NVSHMEM 作业中的 PE 之间是集合式的。它确保内核适合每个 PE 上的 GPU,并且内核在所有 PE 上同时启动。

隐式异步 cudaMemcpy

cudaMemcpycudaMemset 例程可能会表现出异步行为,如此处所述。当异步执行时,这些操作可能在执行任何后续 NVSHMEM 操作之前尚未完成。当不使用 VMM 时,NVSHMEM 在对称内存上设置 CU_POINTER_ATTRIBUTE_SYNC_MEMOPS 标志,强制同步执行这些操作。但是,VMM 分配当前不支持此标志。为避免可能的数据竞争,用户可以使用显式异步 cudaMemcpyAsynccudaMemsetAsync 操作并同步相应的流。

例如,以下主机代码在 cudaMemcpy 可以异步执行时包含可能的竞争。

cudaMemcpy(out_d, &out, sizeof(unsigned int), cudaMemcpyHostToDevice);
nvshmem_uint_or_reduce(NVSHMEM_TEAM_WORLD, final_out_d, out_d, 1);

可以使用 cudaMemcpyAsync 删除竞争。

cudaMemcpyAsync(out_d, &out, sizeof(unsigned int), cudaMemcpyHostToDevice, stream);
cudaStreamSynchronize(stream);
nvshmem_uint_or_reduce(NVSHMEM_TEAM_WORLD, final_out_d, out_d, 1);