JavaScript 中的现代 GPU 计算和渲染

这始于 2021 年,当时我开始为 ,一个将 Python 函数编译为 CUDA、Metal 或 Vulkan 中的 GPU 内核的 Python 库。后来,我加入了元维基并开始研究,这是一种着色器语言,为Instagram和Facebook上的AR效果的跨平台GPU编程提供支持。taichi SparkSL

除了个人的快乐,我一直相信,或者至少希望,这些框架非常有用。它们使非专家更容易进行 GPU 编程,使人们无需掌握复杂的 GPU 概念即可创建引人入胜的图形内容。

在我最新一期的编译器中,我把目光转向了 WebGPU——下一代 Web 图形 API。WebGPU 承诺通过低 CPU 开销和明确的 GPU 控制带来高性能图形,这与大约七年前 Vulkan 和 D3D12 开始的趋势保持一致。

就像 Vulkan 一样,WebGPU 的性能优势是以陡峭的学习曲线为代价的。虽然我相信这不会阻止全世界有才华的程序员使用 WebGPU 构建令人惊叹的内容,但我想为人们提供一种使用 WebGPU 的方式,而不必面对它的复杂性。这就是来历。taichi.js

在编程模型下,程序员不必对 WebGPU 概念进行推理,例如设备、命令队列、绑定组等。相反,它们编写纯 JavaScript 函数,编译器将这些函数转换为 WebGPU 计算或渲染管道。这意味着任何人都可以通过 编写 WebGPU 代码,只要他们熟悉基本的 JavaScript 语法。taichi.js``taichi.js

本文的其余部分将演示通过“生命游戏”程序的编程模型。如您所见,使用少于 100 行代码,我们将创建一个完全并行的 WebGPU 程序,其中包含三个 GPU 计算管道和一个渲染管道。演示的完整源代码可以在这里找到,如果您想在不设置任何本地环境的情况下使用代码,请转到此页面taichi.js

游戏

生命游戏是细胞自动机的典型例子,细胞自动机是一种根据简单规则随着时间的推移而进化的细胞系统。它由数学家约翰·康威(John Conway)于1970年发明,此后成为计算机科学家和数学家的最爱。游戏在二维网格上进行,每个细胞可以是活的或死的。游戏规则很简单:

  • 如果一个活细胞少于两个或三个以上的活邻居,它就会死亡
  • 如果一个死细胞正好有三个活的邻居,它就会变成活的。

尽管它很简单,但生命游戏可以表现出令人惊讶的行为。从任何随机的初始状态开始,游戏通常会收敛到一些模式占主导地位的状态,就好像这些是通过进化幸存下来的“物种”一样。

模拟

让我们深入了解使用 .首先,我们在速记下导入库,并定义一个包含我们所有逻辑的异步函数。在 中,我们首先调用 ,它初始化库及其 WebGPU 上下文。taichi.js``taichi.js``ti``main()``main()``ti.init()

js
import * as ti from 'node:path/to/taichi.js'
async function main() {
  await ti.init()
}
main()

接下来,让我们定义“生命游戏”模拟所需的数据结构:ti.init()

js
const N = 128
const liveness = ti.field(ti.i32, [N, N])
const numNeighbors = ti.field(ti.i32, [N, N])
ti.addToKernelScope({ N, liveness, numNeighbors })

在这里,我们定义了两个变量,并且 都是 s。在 中,“字段”本质上是一个 n 维数组,其维数在 的第二个参数中提供。数组的元素类型在第一个参数中定义。在这种情况下,我们有 ,表示 2 位整数。但是,字段元素也可能是更复杂的类型,包括向量、矩阵和结构。liveness``numNeighbors``ti.field``taichi.js``ti.field()``ti.i32

下一行代码 确保变量 、 和 在“内核”中可见,“内核”是以 JavaScript 函数形式定义的 GPU 计算和/或渲染管道。例如,以下内核用于用初始活动值填充网格单元,其中每个单元最初有 20% 的几率处于活动状态:ti.addToKernelScope({...})``N``liveness``numNeighbors``taichi.js``init

js
const init = ti.kernel(() => {
  for (const I of ti.ndrange(N, N)) {
    liveness[I] = 0
    const f = ti.random()
    if (f < 0.2)
      liveness[I] = 1
  }
})
init()

内核是通过使用 JavaScript lambda 作为参数调用来创建的。在引擎盖下,将查看此lambda的JavaScript字符串表示,并将其逻辑编译为WebGPU代码。在这里,lambda 包含一个 -loop,其循环索引遍历 .这意味着将采用 x 个不同的值,范围从 到 。init()``ti.kernel()``taichi.js``for``I``ti.ndrange(N, N)``I``N``N``[0, 0]``[N-1, N-1]

神奇的部分来了——在内核中的所有顶级循环都将并行化。更具体地说,对于循环索引的每个可能值,将分配一个 WebGPU 计算着色器线程来执行它。在这种情况下,我们在“生命游戏”模拟中为每个单元专用一个 GPU 线程,将其初始化为随机活动状态。随机性来自一个函数,这是库中为内核使用的众多函数之一。文档中提供了这些内置实用程序的完整列表。taichi.js``for``taichi.js``ti.random()``taichi.js``taichi.js

创建游戏的初始状态后,让我们继续定义游戏的演变方式。以下是定义这种演变的两个内核:taichi.js

js
const countNeighbors = ti.kernel(() => {
  for (const I of ti.ndrange(N, N)) {
    let neighbors = 0
    for (const delta of ti.ndrange(3, 3)) {
      const J = (I + delta - 1) % N
      if ((J.x !== I.x || J.y !== I.y) && liveness[J] === 1)
        neighbors = neighbors + 1
    }
    numNeighbors[I] = neighbors
  }
})

const updateLiveness = ti.kernel(() => {
  for (const I of ti.ndrange(N, N)) {
    const 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 渲染管线,将游戏的演变绘制到网页上。

渲染

编写渲染代码比编写通用计算内核稍微复杂一些,并且通常需要对顶点着色器、片段着色器和光栅化管道有一定的了解。但是,您会发现 的简单编程模型使这些概念非常易于使用和推理。taichi.js``taichi.js

在绘制任何东西之前,我们需要访问我们正在绘制的一块画布。假设 HTML 中存在名为的画布,以下代码行将创建一个对象,该对象表示可由呈现管线呈现到的一段纹理。result_canvas``ti.CanvasTexture``taichi.js

js
const htmlCanvas = document.getElementById('result_canvas')
htmlCanvas.width = 512
htmlCanvas.height = 512
const renderTarget = ti.canvasTexture(htmlCanvas)

在我们的画布上,我们将渲染一个正方形并将游戏的 2D 网格绘制到这个正方形上。在 GPU 中,要渲染的几何图形表示为三角形。在这种情况下,我们尝试渲染的正方形将表示为两个三角形。这两个三角形在 中定义,它存储两个三角形的六个顶点中的每一个的坐标:ti.field

js
const 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],
])

正如我们对 and 字段所做的那样,我们需要显式声明 and 变量在 GPU 内核中可见:liveness``numNeighbors``renderTarget``vertices``taichi.js

    ti.addToKernelScope({ vertices, renderTarget });

现在,我们有了实现渲染管线所需的所有数据。以下是管道本身的实现:

js
const render = ti.kernel(() => {
  ti.clearColor(renderTarget, [0.0, 0.0, 0.0, 1.0])

  for (const v of ti.inputVertices(vertices)) {
    ti.outputPosition([v.x, v.y, 0.0, 1.0])
    ti.outputVertex(v)
  }

  for (const f of ti.inputFragments()) {
    const coord = (f + 1) / 2.0
    const texelIndex = ti.i32(coord * (liveness.dimensions - 1))
    const live = ti.f32(liveness[texelIndex])
    ti.outputColor(renderTarget, [live, live, live, 1.0])
  }
})

接下来,我们定义两个顶级循环,如您所知,它们是在 WebGPU 中并行化的循环。但是,与之前迭代对象的循环不同,这些循环分别迭代 和 。这表明这些循环将被编译为 WebGPU“顶点着色器”和“片段着色器”,它们作为渲染管线协同工作。for``ti.ndrange``ti.inputVertices(vertices)``ti.inputFragments()

顶点着色器有两个职责:

  • 对于每个三角形顶点,计算其在屏幕上的最终位置(或者更准确地说,计算其“剪辑空间”坐标)。在 3D 渲染管线中,这通常涉及一堆矩阵乘法,将顶点的模型坐标转换为世界空间,然后转换为相机空间,最后转换为“剪辑空间”。但是,对于简单的 2D 正方形,顶点的输入坐标在剪辑空间中已经处于正确的值,因此我们可以避免所有这些。我们所要做的就是附加一个固定值 0.0 和一个固定值 (如果您不知道它们是什么,请不要担心 - 这在这里并不重要!z``w``1.0
js
ti.outputPosition([v.x, v.y, 0.0, 1.0])
  • 对于每个顶点,生成要插值的数据,然后传递到片段着色器中。在渲染管线中,执行顶点着色器后,将在所有三角形上执行称为“光栅化”的内置进程。这是一个硬件加速的过程,用于计算每个三角形覆盖的像素。这些像素也称为“片段”。 对于每个三角形,程序员可以在三个顶点中的每一个生成额外的数据,这些数据将在光栅化阶段进行插值。对于像素中的每个片段,其相应的片段着色器线程将根据其在三角形中的位置接收插值。在我们的例子中,片段着色器只需要知道片段在 2D 正方形中的位置,因此它可以获取游戏的相应活动值。 为此,将 2D 顶点坐标传递到光栅器中就足够了,这意味着片段着色器将接收像素本身的插值 2D 位置:
js
ti.outputVertex(v)

片段着色器的代码如下所示:

js
for (const f of ti.inputFragments()) {
  const coord = (f + 1) / 2.0
  const cellIndex = ti.i32(coord * (liveness.dimensions - 1))
  const live = ti.f32(liveness[cellIndex])
  ti.outputColor(renderTarget, [live, live, live, 1.0])
}

该值是从顶点着色器传递的插值像素位置。使用此值,片段着色器将在覆盖此像素的游戏中查找单元格的活动状态。这是通过首先将像素坐标转换为范围,然后将此坐标存储到变量中来完成的。然后将其乘以字段的维度,从而生成覆盖单元格的索引。f``f``[0, 0] ~ [1, 1]``coord``liveness

最后,我们获取这个单元格的值,即它是否死亡,以及它是否还活着。它将此像素的 RGBA 值输出到 ,其中 R、G、B 分量都等于 ,A 分量等于 ,以获得完全不透明度。live``0``1``renderTarget``live``1

定义渲染管线后,剩下的就是通过每帧调用模拟内核和渲染管线将所有内容放在一起:

    async function frame() {        countNeighbors()        updateLiveness()        await render();        requestAnimationFrame(frame);    }    await frame();

就是这样!我们已经在 中完成了基于 WebGPU 的“生命游戏”实现。taichi.js

如果运行该程序,则应看到以下动画,其中 128x128 个细胞进化了大约 1,400 代,然后汇聚到几种稳定的生物。

习题

我希望你觉得这个演示很有趣!如果你这样做了,我有一些额外的练习和问题,我邀请你尝试和思考。(顺便说一句,为了快速试验代码,请转到此页面)

  1. [简单]在演示中添加一个 FPS 计数器!使用当前设置可以获得什么 FPS 值?尝试增加 的值,看看帧速率如何变化。你能写一个普通的JavaScript程序,在没有WebGPU的情况下获得这个帧率吗?N = 128``N``taichi.js
  2. [中]如果我们合并并合并到单个内核中并将计数器保留为局部变量会发生什么?该程序是否仍然始终正常工作?countNeighbors()``updateLiveness()``neighbors
  3. [硬]在 中,始终生成一个函数,无论它包含计算管线还是渲染管线。如果你必须猜测,这个 -ness 是什么意思?调用这些电话的意义是什么?最后,在上面定义的函数中,为什么我们只为函数而没有为其他两个函数放?taichi.js``ti.kernel(..)``async``async``await``async``frame``await``render()

最后两个问题特别有趣,因为它们涉及编译器的内部工作原理和框架的运行时,以及GPU编程的原理。让我知道你的答案!taichi.js

最后更新:2024-05-08