数据并行内核的微型编译器

# 一个用于数据并行内核的小型编译器 来源：https://healeycodes.com/a-tiny-compiler-for-data-parallel-kernels 很多高效的代码最初只是一个枯燥的循环。现代硬件可以同时对多个值执行相同操作（例如，SIMD (https://en.wikipedia.org/wiki/Single_instruction,_multiple_data) 和 SIMT (https://en.wikipedia.org/wiki/Single_instruction,_multiple_threads)），有时我们直接针对这些执行模型编写代码，但其他时候，编译器会从看似普通的代码开始，并将其重写，使多个循环迭代能够一起运行。 我构建了一个小型编译器 (https://github.com/healeycodes/kernel-lowering)（约 180 行 Python 代码），旨在理解这种转换是什么样的。我的编译器对内核进行降级（将其重写为更简单、更明确的形式，使数据并行性可见）。输入是一个手写的 AST，输出是降级后的 IR，我将其打印为类似 Python 的代码。你可以把这个编译器看作更大编译器中的一个中间步骤，而不是从源代码一直降级到指令。 让我们看一个例子。缩放音频容易并行化，但编写非显式并行的代码仍然很常见，例如： ``` 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) ``` 目标是替换 `for` 循环为 `vector_for` 循环，从而允许循环的多个迭代并行执行。这种分组执行中的每个位置称为一个*lane*。 ## 通道与掩码 通道是分组执行中的一个独立元素位置。例如，如果一个分组操作同时处理四个值，那么它有四个通道： ``` [ 10 | 20 | 30 | 40 ] ^ ^ ^ ^ lane0 lane1 lane2 lane3 ``` 每个槽位都是一个通道，因此当你在分组上执行乘法等操作时，看起来像这样： ``` [10 | 20 | 15 | 30] * [ 3 | 3 | 3 | 3] = [30 | 60 | 45 | 90] ``` 当为每个通道执行代码时，它们都有一个唯一的偏移量，以知道要操作哪个数据。每个通道处理不同的逻辑元素，分组操作将它们并排运行。 为了处理数据量不能被通道槽位数整除的情况，会应用一个每通道布尔掩码来跳过越界的加载和存储。你可以将掩码想象为 `[true, true, false, false]`：前两个通道允许运行，后两个被忽略。这对于数组的最后一块很有用，此时可能没有足够的剩余元素来填充每个通道。 后续步骤会将 `vector_for` 的主体转换为针对特定架构的指令。在理想情况下，这可以将性能提升数倍于通道数，尽管内存访问和其他开销也会起作用。 ## 决定哪些可以被降级 降级阶段需要回答两个问题： 1. 这个循环的不同迭代能否独立运行？ 2. 对于循环中的每个值，它是`uniform`还是`varying`？ `uniform`值在每个通道中都是相同的。`varying`值在每个通道中可能不同。这种区分很重要，因为 uniform 值在所有通道间共享，而 varying 值需要每个通道单独计算。 在音频示例中，`volume`是 uniform 的，每个通道都乘以相同的值。但 `i` 是 varying 的，因为每个通道处理不同的索引，这意味着 `samples[i]` 也是 varying 的。 ``` 值 类型 原因 --------- ------- ---------------------------- volume uniform 每个通道的值相同 i varying 每个通道得到不同的索引 samples[i] varying 每个通道读取不同的样本 ``` 我编译器的核心是一个小型 AST 遍历分类器： ``` # （伪代码，但与实际相差不远） def kind(expr, env): if expr 是字面量: return UNIFORM if expr 是变量: return env[expr.name] if expr 是加载: return VARYING if kind(expr.index, env) == VARYING else UNIFORM if expr 是二元表达式: left = kind(expr.left, env) right = kind(expr.right, env) 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]` 是 varying 的，`samples[i] * volume` 也是 varying 的（即使 `volume` 本身是 uniform 的）。 这种分类告诉编译器要发出什么。当循环被降级时，`i` 变成 `base + lane_id`。对于一个四通道分组，`base = 8`，通道包含 `[8, 9, 10, 11]`。因此 `samples[i]` 请求 `samples[8]`、`samples[9]`、`samples[10]` 和 `samples[11]`。一次连续的内存访问！编译器记录 `i` 是循环的连续索引。使用该索引的加载变成掩码加载。掩码仅对最后一个分组重要，其中某些通道索引可能超出数组范围： ``` samples[i] -> masked_load(samples, i, active) ``` 但从任意每通道索引进行 varying 加载则变成 gather 操作。gather 是分组加载的版本，相当于“每个通道从其自己的地址读取”，因此地址可能不同且不连续。例如，在这个未优化的代码中： ``` 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]` 那样规则（通道读取相邻元素）。这通常使 gather 变慢，尽管成本取决于架构。 上述 `color_by_number` 示例的完整编译输出： ``` 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) ``` 真实的编译器需要处理比这多得多的内容（类型、别名、控制流、目标特定指令），包括确定循环迭代是否可以安全地独立执行。我的编译器主要假设坚持编写写入 `out[i]` 的简单内核，但这对于我想理解的部分已经足够了。 ## 这一步骤为何重要 经过这个降级阶段后，分组执行在程序中变为显式。它记录了哪些循环迭代一起运行、哪些通道处于活动状态，以及每个加载是使用连续地址还是需要每个通道地址。这为后续的代码生成阶段提供了所需的结构，以发出更好的指令。 `masked_load` 可以变成掩码向量加载，`gather` 可以变成 gather 指令，`vector_for` 可以变成围绕这些操作的循环结构。没有这种分析，这些事实在程序中就不存在，因此代码仍然保持为普通的标量操作，并错过了使用更快指令的机会。 ## 结束语 我很享受解决这个依赖分析以及何时以及如何降级的问题。我避免了一些在之前的项目中采用过的低效步骤，例如在不需要时编写解析器。如果你检查编译器的源代码 (https://github.com/healeycodes/kernel-lowering)，输入只是硬编码的树： ``` AST = Kernel( "color_by_number", ["color_number", "colors", "out", "n"], For( "i", Lit(0), V("n"), Let( "number", Load("color_number", V("i")), Store("out", V("i"), Load("colors", V("number"))), ), ), ) ``` 我最初走得更远，生成了 C SIMD 代码，并在编译/未编译的代码上运行基准测试，但这将焦点从我想要学习和理解的核心思想上移开了。所以我将其精简为更纯粹、更高级的抽象。作为奖励，我为这个网站添加了自定义语法高亮，以支持内核片段。