Taichi.js 实战 WebGPU
作为一名计算机图形学和编程语言极客，我很高兴在过去两年中参与了多个 GPU 编译器的开发。
这始于 2021 年，当时我开始为 taichi 做贡献，这是一个 Python 库，可将 Python 函数编译为 CUDA、Metal 或 Vulkan 中的 GPU 内核。后来，我加入了 Meta，开始研究 SparkSL，这是一种着色器语言，为 Instagram 和 Facebook 上的 AR 效果的跨平台 GPU 编程提供支持。
除了个人乐趣之外，我一直相信，或者至少希望，这些框架非常有用。它们使非专家更容易进行 GPU 编程，使人们无需掌握复杂的 GPU 概念即可创建迷人的图形内容。
在我最新的编译器文章中，我将目光转向了 WebGPU——下一代 Web 图形 API。WebGPU 承诺通过低 CPU 开销和显式 GPU 控制带来高性能图形，这与七年前 Vulkan 和 D3D12 开创的趋势保持一致。
就像 Vulkan 一样，WebGPU 的性能优势是以陡峭的学习曲线为代价的。虽然我相信这不会阻止世界各地的优秀程序员使用 WebGPU 构建精彩的内容，但我想为人们提供一种使用 WebGPU 的方法，而不必面对它的复杂性。这就是 taichi.js 的由来。
在 taichi.js 编程模型下，程序员不必考虑设备、命令队列、绑定组等 WebGPU 概念。相反，他们编写普通的 JavaScript 函数，编译器将这些函数转换为 WebGPU 计算或渲染管道。这意味着任何人都可以通过 taichi.js 编写 WebGPU 代码，只要他们熟悉基本的 JavaScript 语法。
本文的其余部分将通过“生命游戏”程序演示 taichi.js 的编程模型。如你所见，使用不到 100 行代码，我们将创建一个完全并行的 WebGPU 程序，其中包含三个 GPU 计算管道和一个渲染管道。该演示的完整源代码可在此处找到，如果你想使用代码而不必设置任何本地环境，请转到此页面。
1、游戏
生命游戏是元胞机或细胞自动机（cellular automaton）的一个经典例子，细胞自动机是一种根据简单规则随时间演化的细胞系统。它由数学家约翰·康威于 1970 年发明，自此成为计算机科学家和数学家的最爱。游戏在二维网格上进行，每个细胞都可以是活的或死的。游戏规则很简单：
- 如果活细胞的活细胞邻居少于两个或多于三个，它就会死亡
- 如果死细胞的活细胞邻居恰好有三个，它就会复活。
尽管生命游戏很简单，但它可能会表现出令人惊讶的行为。从任何随机的初始状态开始，游戏通常会收敛到一种状态，其中一些模式占主导地位，就好像这些模式是通过进化而幸存下来的“物种”。
2、模拟
让我们深入研究使用 taichi.js 的生命游戏实现。首先，我们在简写 ti 下导入 taichi.js 库，并定义一个包含所有逻辑的异步 main() 函数。在 main() 中，我们首先调用 ti.init()，它初始化库及其 WebGPU 上下文。
import * as ti from "path/to/taichi.js"
let main = async () => {
await ti.init();
...
};
main()在 ti.init() 之后，让我们定义“生命游戏”模拟所需的数据结构：
let N = 128;
let liveness = ti.field(ti.i32, [N, N])
let numNeighbors = ti.field(ti.i32, [N, N])
ti.addToKernelScope({ N, liveness, numNeighbors });这里，我们定义了两个变量， liveness 和 numNeighbors，它们都是 ti.fields。在 taichi.js 中，“fields”本质上是一个 n 维数组，其维数由 ti.field() 的第二个参数提供。数组的元素类型在第一个参数中定义。在本例中，我们有 ti.i32，表示 32 位整数。但是，字段元素也可能是更复杂的类型，包括向量、矩阵和结构。
下一行代码 ti.addToKernelScope({...}) 确保变量 N、 liveness 和 numNeighbors 在 taichi.js“内核”中可见，这些“内核”是以 JavaScript 函数形式定义的 GPU 计算和/或渲染管道。例如，以下初始化内核用于用初始活跃度值填充我们的网格单元，其中每个单元最初有 20% 的存活几率：
let init = ti.kernel(() => {
for (let I of ti.ndrange(N, N)) {
liveness[I] = 0
let f = ti.random()
if (f < 0.2) {
liveness[I] = 1
}
}
})
init()init() 内核是通过调用 ti.kernel() 并以 JavaScript lambda 作为参数创建的。在底层，taichi.js 将查看此 lambda 的 JavaScript 字符串表示并将其逻辑编译为 WebGPU 代码。在这里，lambda 包含一个 for 循环，其循环索引 I 通过 ti.ndrange(N, N) 进行迭代。这意味着我将取 NxN 个不同的值，范围从 [0, 0] 到 [N-1, N-1]。
神奇的部分来了 — 在 taichi.js 中，内核中的所有顶层 for 循环都将并行化。更具体地说，对于循环索引的每个可能值，taichi.js 将分配一个 WebGPU 计算着色器线程来执行它。在这种情况下，我们为“生命游戏”模拟中的每个单元专门分配一个 GPU 线程，将其初始化为随机活动状态。随机性来自 ti.random() 函数，这是 taichi.js 库中提供的供内核使用的众多函数之一。这些内置实用程序的完整列表可在 taichi.js 文档中找到。
创建游戏的初始状态后，让我们继续定义游戏如何演变。以下是定义此演变的两个 taichi.js 内核：
let countNeighbors = ti.kernel(() => {
for (let I of ti.ndrange(N, N)) {
let neighbors = 0
for (let delta of ti.ndrange(3, 3)) {
let J = (I + delta - 1) % N
if ((J.x != I.x || J.y != I.y) && liveness[J] == 1) {
neighbors = neighbors + 1;
}
}
numNeighbors[I] = neighbors
}
});
let updateLiveness = ti.kernel(() => {
for (let I of ti.ndrange(N, N)) {
let neighbors = numNeighbors[I]
if (liveness[I] == 1) {
if (neighbors < 2 || neighbors > 3) {
liveness[I] = 0;
}
}
else {
if (neighbors == 3) {
liveness[I] = 1;
}
}
}
})与我们之前看到的 init() 内核一样，这两个内核也具有顶层 for 循环，可遍历每个网格单元，这些循环由编译器并行化。在 countNeighbors() 中，对于每个单元，我们查看八个相邻单元并计算这些邻居中有多少个“活着”。
活着的邻居的数量存储在 numNeighbors 字段中。请注意，在遍历邻居时， for (let delta of ti.ndrange(3, 3)) {...} 循环未并行化，因为它不是顶层循环。循环索引 delta 范围从 [0, 0] 到 [2, 2]，用于偏移原始单元索引 I。我们通过对 N 取模来避免越界访问。（对于拓扑倾向的读者来说，这基本上意味着游戏具有环形边界条件）。
计算每个单元的邻居数量后，我们在 updateLiveness() 内核中更新它们的活跃状态。这很简单，只需读取每个单元的活跃状态及其当前活跃邻居的数量，并根据游戏规则写回新的活跃值即可。与往常一样，此过程并行应用于所有单元。
这基本上结束了游戏模拟逻辑的实现。接下来，我们将了解如何定义 WebGPU 渲染管道以将游戏的演变绘制到网页上。
3、渲染
在 taichi.js 中编写渲染代码比编写通用计算内核稍微复杂一些，并且确实需要对顶点着色器、片段着色器和光栅化管道有一定的了解。但是，你会发现 taichi.js 的简单编程模型使这些概念非常容易使用和推理。
在绘制任何内容之前，我们需要访问要在其上绘制的一块画布。假设 HTML 中存在一个名为 result_canvas 的画布，则以下代码行将创建一个 ti.CanvasTexture 对象，该对象表示可以通过 taichi.js 渲染管道在其上渲染的一块纹理。
let htmlCanvas = document.getElementById('result_canvas');
htmlCanvas.width = 512;
htmlCanvas.height = 512;
let renderTarget = ti.canvasTexture(htmlCanvas);在我们的画布上，我们将渲染一个正方形，并将游戏的 2D 网格绘制到这个正方形上。在 GPU 中，要渲染的几何图形表示为三角形。在这种情况下，我们尝试渲染的正方形将表示为两个三角形。这两个三角形在 ti.field 中定义，它存储了两个三角形的六个顶点的坐标：
let vertices = ti.field(ti.types.vector(ti.f32, 2), [6]);
await vertices.fromArray([
[-1, -1],
[1, -1],
[-1, 1],
[1, -1],
[1, 1],
[-1, 1],
]);正如我们对 liveness 和 numNeighbors 字段所做的那样，我们需要在 taichi.js 中明确声明 renderTarget 和 vertices 变量在 GPU 内核中可见：
ti.addToKernelScope({ vertices, renderTarget });现在我们已经有了实现渲染管道所需的所有数据。以下是管道本身的实现：
let render = ti.kernel(() => {
ti.clearColor(renderTarget, [0.0, 0.0, 0.0, 1.0]);
for (let v of ti.inputVertices(vertices)) {
ti.outputPosition([v.x, v.y, 0.0, 1.0]);
ti.outputVertex(v);
}
for (let f of ti.inputFragments()) {
let coord = (f + 1) / 2.0;
let texelIndex = ti.i32(coord * (liveness.dimensions - 1));
let live = ti.f32(liveness[texelIndex]);
ti.outputColor(renderTarget, [live, live, live, 1.0]);
}
});接下来，我们定义两个顶层 for 循环，正如您所知，它们是在 WebGPU 中并行化的循环。但是，与之前迭代 ti.ndrange 对象的循环不同，这些循环分别迭代 ti.inputVertices(vertices) 和 ti.inputFragments()。这表明这些循环将被编译成 WebGPU“顶点着色器”和“片段着色器”，它们一起作为渲染管道工作。
顶点着色器有两个职责。
对于每个三角形顶点，计算其在屏幕上的最终位置（或者更准确地说，其“剪辑空间”坐标）。在 3D 渲染管道中，这通常涉及一堆矩阵乘法，将顶点的模型坐标转换为世界空间，然后转换为相机空间，最后转换为“剪辑空间”。但是，对于我们简单的 2D 正方形，顶点的输入坐标在剪辑空间中已经是其正确值，因此我们可以避免所有这些。我们要做的就是附加一个固定的 z 值 0.0 和一个固定的 w 值 1.0（如果您不知道这些是什么，请不要担心 - 这在这里并不重要！）。
ti.outputPosition([v.x, v.y, 0.0, 1.0]);对于每个顶点，生成要插值的数据，然后将其传递到片段着色器中。在渲染管道中，执行顶点着色器后，将对所有三角形执行一个称为“光栅化”的内置过程。这是一个硬件加速的过程，它计算每个三角形覆盖的像素。这些像素也称为“片段”。
对于每个三角形，程序员可以在三个顶点中的每一个顶点生成附加数据，这些数据将在光栅化阶段进行插值。对于像素中的每个片段，其对应的片段着色器线程将根据其在三角形内的位置接收插值。在我们的例子中，片段着色器只需要知道片段在 2D 方块中的位置，这样它就可以获取游戏的相应活跃度值。
为此，只需将 2D 顶点坐标传递到光栅化器中即可，这意味着片段着色器将接收像素本身的插值 2D 位置：
ti.outputVertex(v);
片段着色器的代码如下：
for (let f of ti.inputFragments()) {
let coord = (f + 1) / 2.0;
let cellIndex = ti.i32(coord * (liveness.dimensions - 1));
let live = ti.f32(liveness[cellIndex]);
ti.outputColor(renderTarget, [live, live, live, 1.0]);
}值 f 是从顶点着色器传递过来的插值像素位置。使用此值，片段着色器将查找游戏中覆盖此像素的单元格的活跃状态。首先将像素坐标 f 转换为 [0, 0] ~ [1, 1] 范围，并将此坐标存储到 coord 变量中。然后将其乘以活跃字段的尺寸，得出覆盖单元格的索引。
最后，我们获取此单元格的活跃值，如果单元格已死亡，则为 0，如果单元格还活着，则为 1。它将此像素的 RGBA 值输出到 renderTarget 上，其中 R、G、B 分量都等于活跃，A 分量等于 1，表示完全不透明。
定义渲染管道后，剩下的就是通过每帧调用模拟内核和渲染管道将所有内容放在一起：
async function frame() {
countNeighbors()
updateLiveness()
await render();
requestAnimationFrame(frame);
}
await frame();就这样！我们在 taichi.js 中完成了基于 WebGPU 的“生命游戏”实现。
如果运行该程序，您应该会看到以下动画，其中 128x128 个细胞进化了大约 1,400 代，然后融合为几种稳定的生物。
原文链接：Painless WebGPU Programming With taichi.js
BimAnt翻译整理，转载请标明出处