数据并行内核的微型编译器：从标量循环到显式并行

背景

在现代硬件架构中，SIMD（单指令多数据流）和 SIMT（单指令多线程）等技术允许硬件同时对多个值执行相同的操作。然而，开发者通常编写的是看似普通的标量代码，依赖编译器在后台将其重写，以便多个循环迭代能够并行执行。

为了深入理解这种转换背后的机制，作者构建了一个极简的编译器（仅约 180 行 Python 代码）。该编译器并非旨在从源代码直接生成机器指令，而是作为大型编译器中的一个中间步骤。它的核心任务是将“内核”（kernels，即并行计算的核心逻辑）降低（lower）为一种更简单、更明确的形式，使得数据并行性变得可见。

输入端是一个小型的手写抽象语法树（AST），输出端则是被降低后的中间表示（IR），以类似 Python 的代码形式打印。通过这种简化，作者旨在揭示将常规循环转换为并行执行模型的具体过程。

核心内容

1. 从标量循环到向量循环

以音频缩放为例，这是一个易于并行化的任务。开发者通常写出如下非显式并行的代码：

kernel scale_audio(samples, out, n, volume):
    for i in range(n):
        out[i] = samples[i] * volume

作者的编译器将其转换为显式的并行结构：

kernel scale_audio(samples, out, n, volume):
    vector_for base in range(0, n, LANES):
        let i = (base + lane_id)
        let active = (i < n)
        masked_store(out, i, (masked_load(samples, i, active) * volume), active)

这里的关键变化是用 vector_for 替换了普通的 for 循环。vector_for 允许循环的多个迭代并行执行。在分组执行中，每个位置被称为一个“车道”（lane）。

2. 车道（Lanes）与掩码（Masks）的概念

车道（Lane） 车道是分组执行中的一个独立元素位置。例如，如果一次分组操作同时处理四个值，则它有四个车道： [ 10 | 20 | 30 | 40 ] 对应的车道索引为 lane0 到 lane3。当执行乘法等操作时，每个车道处理不同的逻辑元素，但并行运行： [10 | 20 | 15 | 30] * [ 3 | 3 | 3 | 3] = [30 | 60 | 45 | 90]

掩码（Mask） 当数据量不能被车道数量整除时（例如数组长度不是 4 的倍数），需要处理边界情况。此时会应用每个车道的布尔掩码，以跳过越界的加载和存储操作。 例如，掩码 [true, true, false, false] 表示前两个车道允许运行，后两个被忽略。这主要用于数组的最后一段，其中剩余元素可能不足以填满所有车道。

后续步骤会将 vector_for 的主体转换为针对特定架构的指令。理想情况下，这能将性能提升车道数量的倍数，尽管内存访问和其他开销也会影响最终结果。

3. 确定可降低（Lower）的条件

降低过程需要回答两个核心问题：

1. 该循环的不同迭代是否可以独立运行？
2. 对于循环中的每个值，它是**均匀（uniform）的还是变化（varying）**的？

• Uniform（均匀值）：对所有车道来说，该值都是相同的。
• Varying（变化值）：每个车道可能拥有不同的值。

这一区分至关重要，因为均匀值在所有车道间共享，而变化值需要每个车道单独计算。

案例分析： 在音频缩放示例中：

• volume 是 uniform，因为每个车道都乘以相同的值。
• i 是 varying，因为每个车道处理不同的索引。
• 因此，samples[i] 也是 varying，因为每个车道读取不同的样本。

| 值种类 | 原因 | | :--- | :--- | | volume | uniform | 每个车道中的值相同 | | i | varying | 每个车道获得不同的索引 | | samples[i] | varying | 每个车道读取不同的样本 |

4. 编译器核心：AST 遍历分类器

编译器核心是一个小型的 AST 遍历分类器，用于判断表达式的性质：

# 伪代码，接近现实实现
def kind(expr, env):
    if expr is a literal:
        return UNIFORM
    if expr is a variable:
        return env[expr.name]
    if expr is a load:
        # 如果索引是 varying，则加载是 varying；否则是 uniform
        return VARYING if kind(expr.index, env) == VARYING else UNIFORM
    if expr is a binary expression:
        left = kind(expr.left, env)
        right = kind(expr.right, env)
        # 只要有一个操作数是 varying，结果就是 varying
        return VARYING if VARYING in (left, right) else UNIFORM

在循环之前，内核参数默认假设为 UNIFORM。进入降低后的循环后，循环索引被标记为 VARYING：

env = {param: UNIFORM for param in kernel.params}
env[i] = VARYING

随后，varying 属性通过表达式传播。由于 i 是 varying，samples[i] 也是 varying。进而 samples[i] * volume 也是 varying（尽管 volume 本身是 uniform）。

5. 内存访问模式：Contiguous Load vs. Gather

分类结果决定了编译器生成的代码类型：

• 连续加载（Contiguous Load / Masked Load）： 当 i 是循环的连续索引时，例如四车道组且 base = 8，车道持有的索引为 [8, 9, 10, 11]。这意味着 samples[i] 请求的是 samples[8] 到 samples[11]，这是一段连续的内存区域。 编译器将此类加载记录为 masked_load(samples, i, active)。掩码仅在最后一段数据中起作用，用于处理可能超出数组边界的索引。

• 收集加载（Gather）： 如果从每个车道各自的任意索引进行加载，则变为 gather。gather 是“每个车道从其自己的地址读取”的分组加载版本，地址可能不同且不相邻。

  例如，在未优化的 color_by_number 内核中：

  kernel color_by_number(color_number, colors, out, n):
      for i in range(n):
          number = color_number[i]
          out[i] = colors[number]
  

  每个车道加载不同的 number，因此 number 是 varying。这意味着 colors[number] 不是连续加载，而是： colors[number] -> gather(colors, number, active)

  gather 允许每个车道从自己的地址读取，同时仍作为单个分组操作执行。虽然它仍然是并行工作，但内存访问模式不如 samples[i] 那样规则，通常会导致性能较低（具体取决于架构）。

  该示例的完整编译输出为：

  kernel color_by_number(color_number, colors, out, n):
      vector_for base in range(0, n, LANES):
          let i = (base + lane_id)
          let active = (i < n)
          let number = masked_load(color_number, i, active)
          masked_store(out, i, gather(colors, number, active), active)
  

6. 为何这一步骤至关重要

经过降低阶段后，程序中的分组执行变得显式。它记录了哪些循环迭代一起运行、哪些车道是活动的，以及每个加载操作是使用相邻地址还是需要每个车道