跳转至

MSCCLANG DSL

The MSCCLang DSL is a chunk-oriented language for specifying chunk routing through GPUs. The language is embedded in Python as a traced DSL with a fluent API that gives users flexibility in how to express algorithms. This section explains the core components of the MSCCLang DSL for implementing a collective algorithm by chunk routing. Furthermore, Section 5 discusses scheduling extensions that further optimize programs.

MSCCLang DSL是一种面向块的语言，用于指定通过GPU的块路由。该语言嵌入在Python中，作为一个可追踪的DSL，并具有流畅的API，为用户提供了灵活的算法表达方式。本节解释了MSCCLang DSL的核心组件，用于通过块路由实现集体算法。此外，第5节讨论了进一步优化程序的调度扩展。

Buffers and Chunks

MSCCLang exposes GPU memory as named buffers, three of which are available on each rank:

  • Input is a buffer containing input data. 输入是包含输入数据的缓冲区。

  • Output is an uninitialized buffer for storing output data. 输出是用于存储输出数据的未初始化缓冲区。

  • Scratch is an uninitialized buffer used for temporary storage. 暂存 是用于临时存储的未初始化缓冲区。

Buffers divide into chunks representing contiguous spans of elements with a uniform size. The user controls the number of chunks a buffer divides into, but the size (i.e., number of bytes) is abstract. The size is known at runtime when the concrete buffers are passed into the program. Additionally, users can specify that the input and output buffers are aliased to support in-place collective algorithms. The purpose of the MSCCLang program is to ensure that the output buffer on each GPU has the correct chunks for the collective at the end. Chunks can take three forms:

  • Input chunks represent chunks initialized at runtime. The pair (rank, index) uniquely identifies the input chunk in the input buffer.

  • Reduction chunks result from combining two chunks through a point-wise reduction (e.g., addition). The list of input chunks that combine to form a reduction chunk uniquely identifies it.

  • Uninitialized chunks are a unit type that stores uninitialized data. At the start of the program, the output and scratch buffers hold uninitialized chunks.

缓冲区被分成块，代表具有统一大小的连续元素跨度。用户可以控制缓冲区分成的块数，但块的大小（即字节数）是抽象的。块的大小在运行时具体的缓冲区传入程序时确定。此外，用户可以指定输入和输出缓冲区别名，以支持就地集体算法。MSCCLang程序的目的是确保每个GPU上的输出缓冲区在结束时具有正确的集体块。块可以有以下三种形式：

  • 输入块表示在运行时初始化的块。通过（rank, index）对唯一标识输入缓冲区中的输入块。
  • 化简块是通过逐点化简（例如，加法）组合两个块产生的。用于组合成化简块的输入块列表唯一标识该化简块。
  • 未初始化块是一种存储未初始化数据的单元类型。在程序开始时，输出缓冲区和临时缓冲区包含未初始化块。

Collectives

MSCCLang programs are associated with a collective that defines a precondition and a postcondition. The precondition determines the starting state of the input buffer in terms of unique input chunks. The postcondition sets the desired state of the output buffer — for each index of the output buffer, the postcondition specifies either an input chunk or a reduction chunk to correctly implement the collective. Defining a collective’s postcondition allows MSCCLang to automatically validate that a prospective algorithm is correct.

MSCCLang 程序与定义前置条件和后置条件的集体通信操作相关联。

前置条件决定输入缓冲区的初始状态，描述独特的输入块。

后置条件设定输出缓冲区的期望状态——对于输出缓冲区的每个索引，后置条件要么指定一个输入块，要么指定一个化简块，以正确实现集体通信操作。定义集体通信操作的后置条件使得 MSCCLang 能够自动验证预期算法的正确性。

Postcondition

在 MSCCLang 程序中，后置条件（postcondition）设置了输出缓冲区的期望状态：每个输出缓冲区的索引位置，后置条件都会指定一个输入块或化简块作为“期望”。这个定义对于确保集体通信操作的正确实现至关重要。

作用

  • 设定输出缓冲区的期望状态

后置条件描述了在集体通信操作完成后，输出缓冲区应该是什么样的状态。

换句话说，它告诉系统，在程序运行结束时，每个输出位置（索引）应该包含什么数据。这个数据可以是 直接从输入缓冲区得到的输入块(input chunk) ，也可以是 通过某种计算（比如加法）得到的化简块(reduction chunk) 。

  • 自动验证算法的正确性

由于后置条件明确了输出缓冲区的期望状态，MSCCLang系统可以自动检查一个给定的算法是否正确实现了集体通信操作。系统会对比算法执行后的实际输出状态和后置条件设定的期望状态。如果两者匹配，则说明算法正确；否则，说明算法有误。

AllReduce操作的前置条件和后置条件

为了更好地理解precondition和postcondition在MSCCLang中的作用，我们可以用AllReduce操作来具体说明

前置条件

前置条件描述了在操作开始时，各个GPU的输入缓冲区的状态。对于一个具有 R 个排名（rank）的AllReduce操作，前置条件如下：

  • 每个排名的输入缓冲区包含 C 个唯一的输入块。这些输入块可以表示为： \((c_{0}^r, c_{1}^r, ..., c_{C-1}^r)\) 这里，r 表示排名，\(c_{i}^r\) 表示在排名 r 的GPU上，输入缓冲区的第 i 个块。

后置条件

后置条件描述了在操作完成后，各个GPU的输出缓冲区的期望状态。对于AllReduce操作，后置条件如下：

  • 所有排名的输出缓冲区应当相同，并且包含 \( C \) 个化简块。这些化简块表示为： \(\left(\sum_{r=0}^{R-1} c_{0}^r, \sum_{r=0}^{R-1} c_{1}^r, ..., \sum_{r=0}^{R-1} c_{C-1}^r\right)\) 这里，\(\sum_{r=0}^{R-1} c_{i}^r\) 表示对所有排名的第 i 个输入块进行求和得到的化简块。

示例说明

假设我们有3个GPU（即 R = 3），每个GPU的输入缓冲区有2个块（即 C = 2），如下：

  1. 前置条件：
  2. GPU 0: 输入缓冲区包含块 \(c_{0}^0\) 和 \(c_{1}^0\)
  3. GPU 1: 输入缓冲区包含块 \(c_{0}^1\) 和 \(c_{1}^1\)

  4. GPU 2: 输入缓冲区包含块 \(c_{0}^2\) 和 \(c_{1}^2\)

  5. 后置条件：

  6. 所有GPU的输出缓冲区应该相同，包含化简块：
    • 块 \(\sum_{r=0}^{2} c_{0}^r = c_{0}^0 + c_{0}^1 + c_{0}^2\)
    • 块 \(\sum_{r=0}^{2} c_{1}^r = c_{1}^0 + c_{1}^1 + c_{1}^2\)

具体过程

在AllReduce操作中，系统将根据定义好的前置条件和后置条件进行数据传输和计算：

  1. 初始化：
  2. 每个GPU的输入缓冲区按前置条件初始化。
  3. 数据传输和计算：
  4. 通过网络和GPU之间的通信，将输入块在GPU之间传输，并执行化简操作（例如加法）。
  5. 结果验证：
  6. 在操作完成后，系统检查每个GPU的输出缓冲区，确保其状态符合后置条件。
  7. 具体来说，系统会检查每个GPU的输出缓冲区是否包含 \(c_{0}^0 + c_{0}^1 + c_{0}^2\) 和 \(c_{1}^0 + c_{1}^1 + c_{1}^2\)

通过这种方式，MSCCLang可以确保所定义的集体通信算法正确实现了预期的操作，并能自动验证算法的正确性，避免手工检查的复杂性和潜在错误。

The algorithm, not the collective, determines the number of chunks and whether the algorithm the input and output buffer alias each other (i.e. the algorithm is in-place). For example, the hierarchical AllReduce algorithm (Figure 3) is an in-place algorithm that uses 𝑁 ×𝐺 chunks. For programmability, MSCCLang automatically deduces the number of chunks in the scratch buffer based on the highest scratch indices accessed in the program.

算法（而不是集体操作）决定了块的数量以及输入和输出缓冲区是否彼此别名（即，算法是否是就地算法）。例如，分层 AllReduce 算法（图 3）是一个就地算法，使用 N×G 个块。

为了便于编程，MSCCLang 会根据程序中访问的 最高临时缓冲区 索引 自动推断 出 临时缓冲区中块的数量。

The algorithm, not the collective, determines the number of chunks and whether the algorithm the input and output buffer alias each other (i.e. the algorithm is in-place)
  • 算法决定块的数量：

在进行集体通信操作时，数据通常会被划分成多个较小的块（chunks）以便于并行处理。 不同的算法可能会对数据块进行不同的划分，因此 数据被划分成多少块是由具体算法决定的，而不是集体操作本身。

  • 算法决定缓冲区是否别名：

就地算法（in-place algorithm）指的是输入缓冲区和输出缓冲区实际上是同一个缓冲区，即 输入数据会被直接覆盖成输出数据。 这种方式可以节省内存和提高效率。

决定是否采用就地算法的也是具体的算法，而不是集体操作本身。

MSCCLang Operations

Table 1 lists the MSCCLang operations used for manipulating chunks. The function chunk(rank, buffer, index, count=C) returns a reference to C contiguous chunks currently assigned to the named buffer starting at index; count has a default value of one when it is not set. MSCCLang raises an error if the program accesses an uninitialized chunk.

MSCCLang 中关于缓冲区块的引用机制：返回一个引用，该引用包含当前在缓冲区中的块的详细信息。具体地，这个引用包括以下几个部分：

  • rank：块所在的 GPU 的排名（rank）
  • buffer：块所在的缓冲区
  • index：块在缓冲区中的索引
  • count：缓冲区中块的数量

这个引用机制使得用户可以方便地访问和操作缓冲区中的块

The copy and reduce operations move chunks between buffers. Specifically, c1.copy(rank2, buffer, index2) copies the chunks referenced by c1 to (rank2, buffer2, index2). c1.reduce(c2) reduces two equal count chunk references c1 and c2. This is an in-place pointwise operation that overwrites c1 with the reduced chunks. Both the copy and reduce operations return references to the newly created chunks, which allows fluently chaining copy and reduce calls.

复制（copy）和归约（reduce）操作在缓冲区之间移动块。具体而言：

  • c1.copy(rank2, buffer, index2)：将 c1 引用的块复制到 (rank2, buffer2, index2)
  • c1.reduce(c2)：对两个具有相同块数量的引用 c1 和 c2 进行归约。这个操作是一个就地逐点操作，会用归约后的块覆盖 c1

这两种操作都会返回对新创建块的引用，从而允许对复制和归约操作进行流畅的链式调用。

复制操作 (copy)

  • 语法：c1.copy(rank2, buffer, index2)
  • 功能：将由 c1 引用的块复制到新的位置 (rank2, buffer, index2)
  • 返回值：返回对新位置上块的引用

示例

假设 c1 引用了缓冲区 buf1 中的块，调用 c1.copy(1, buf2, 3) 会将 c1 引用的块复制到 rank 为 1 的 GPU 上的缓冲区 buf2 的索引为 3 的位置。返回的引用可以用于进一步的操作

归约操作 (reduce)

  • 语法：c1.reduce(c2)
  • 功能：对两个引用 c1 和 c2 进行逐点归约操作，结果覆盖 c1。例如，如果归约操作是加法，那么每个对应位置的块会相加。
  • 返回值：返回对归约后块的引用。

示例

假设 c1 和 c2 引用了相同数量的块，调用 c1.reduce(c2) 会对 c1 和 c2 进行逐点加法，并用结果覆盖 c1。返回的引用可以用于进一步的操作。

链式调用

由于 copy 和 reduce 操作都返回对新创建块的引用，因此可以进行链式调用。例如：

Python
1
new_chunk = c1.copy(1, buf2, 3).reduce(c2)

这个例子中，首先将 c1 复制到 rank 为 1 的 buf2 的索引为 3 的位置，然后对结果进行逐点加法操作并覆盖 c1。

point

通过使用 copy 和 reduce 操作，MSCCLang 提供了一种灵活的方法来在缓冲区之间移动和处理数据块，并且返回的引用可以用于链式调用，使得编程更为简洁和高效

alt text

Programs manipulate references rather than chunks to prevent operations on stale data. A program can create multiple references to the same (rank, buffer, index) location potentially referring to stale chunks that are overwritten by later operations. MSCCLang only allows the latest reference for any location to be used and will generate an error otherwise. This enforces a chunk-oriented coding style with the program always operating on the latest reference, thereby making MSCCLang programs data race free by construction.

程序通过操作引用而非块本身来防止对过时数据的操作。一个程序可以创建多个引用，指向相同的 (rank, buffer, index) 位置，这些引用可能会指向被后续操作覆盖的过时块。MSCCLang 只允许使用任何位置的最新引用，否则将生成错误。这种机制强制了一种面向块的编码风格，程序始终操作最新的引用，从而使得 MSCCLang 程序在构建时就能避免数据竞争。

MSCCLang lets users express operations between buffers uniformly with copy and reduce regardless of whether they are on the same GPU or not. The next section explores how MSCCLang enables this abstraction.

MSCCLang 允许用户通过 copy 和 reduce 操作以统一的方式表达缓冲区之间的操作，无论它们是否在同一个 GPU 上。下一节将探讨 MSCCLang 如何实现这一抽象。