译:WebGPU — All of the cores, none of the canvas

原文地址:WebGPU — All of the cores, none of the canvas

WebGPU 是一个即将推出的 Web API,让你可以使用低层的、通用的方法访问 GPU。

我在图形方面的经验不多。我通过阅读关于如何用 OpenGL 构建游戏引擎的教程来学习 WebGL 的零碎知识,并通过观看 Inigo QuilezShaderToy 上只使用着色器而不使用任何 3D 网格或模型来学习着色器。这让我足以完成像 PROXX 中的背景动画这样的工作,但我觉得 WebGL 一直不够舒适,很快我就会解释原因。

当 WebGPU 出现在我的视野中时,我想投身到其中,但许多人警告我,WebGPU 比 WebGL 更加样板。我并不气馁,但预计到最坏的情况,我把我能找到的所有教程和规范收集在一起,由于 WebGPU 还处于早期阶段,所以并没有很多。我试了试,我并没有发现 WebGPU 比 WebGL 明显地更像样板化,而实际上是一个我更喜欢的 API。

所以,我们来了。我想与大家分享我在接触 GPU 和 WebGPU 时学到的东西。这篇博文的目标是使 WebGPU 对 Web 开发者来说是可接受的。但这里要提前提醒一下。我不会使用 WebGPU 来生成图形。相反,我将使用 WebGPU 来访问 GPU 提供的原始计算能力。也许我会写一篇后续博文,介绍如何使用 WebGPU 渲染到你的屏幕上,但现在已经有 不少内容 了。我将尽可能深入地了解 WebGPU 的意义,并希望你能有效地使用它——但不一定是高效的。我不能让你成为 GPU 性能专家,主要是因为我自己也不是专家。

我写了一个很长的声明,下面正式开始!

WebGL

WebGL 出现在 2011 年,到目前为止,它是唯一可以从 Web 上访问 GPU 的底层的 API。WebGL 的 API 实际上只是 OpenGL ES 2.0,加上一些简单的封装和 helper,使其与 Web 兼容。WebGL 和 OpenGL 都是由 Khronos 组织 标准化的,该集团基本上是 3D 图形界的 W3C。

OpenGL 的 API 本身甚至可以追溯到更早以前,以今天的标准来看,它不是一个很好的 API。该设计以内部的全局状态对象为中心。因为它最大限度地减少了在任何调用中需要传入和传出到 GPU 的数据量,从这个角度来看,这种设计是有道理的。然而,它也带来了很多理解上的困难。

一个 WebGL 的内部可视化,全局状态对象。来自:[WebGL Fundamentals](https://webglfundamentals.org/webgl/lessons/resources/webgl-state-diagram.html)

内部状态对象基本上是一个指针的集合。你的 API 调用可以影响状态对象所指向的对象,但也可以影响状态对象本身。因此,API 调用的顺序是非常重要的,我总觉得这使得抽象和构建库变得很困难。你必须非常细致地理清所有可能干扰你进行 API 调用的指针和状态项,同时还要将指针和值恢复到之前的值,这样你的抽象才能正确生成。我经常发现自己盯着一块黑色的画布(因为几乎所有 WebGL 中的报错都会导致这样),并对哪个指针目前没有正确指向进行粗略的判断。老实说,我不知道 ThreeJS 是如何做到如此健壮的,但它确实做到了。我想这是大多数人使用 ThreeJS 而不是直接使用 WebGL 的主要原因之一。

这不是你的问题,而是我的问题:说白了,我不能消化 WebGL 可能是我的问题。比我聪明的人都能用 WebGL(以及 Web 之外的 OpenGL)构建出惊人的东西,但对我来说,它不适合我。

随着 ML、神经网络和加密货币的出现,GPU 已经显示出它们不仅可以在屏幕上绘制三角形,还有更大的作用。使用 GPU 进行任何形式的计算通常被称为通用 GPU(General-Purpose GPU,GPGPU),而 WebGL 1 在这方面并不出色。如果你想在 GPU 上处理任意数据,你必须将其编码为纹理,在着色器中解码,进行计算,然后将结果重新编码为纹理。WebGL 2 通过 Transform Feedback 使这个问题变得简单多了,但是 WebGL2 直到 2021 年 9 月才被 Safari 支持(而其他大多数浏览器从 2017 年 1 月开始支持 WebGL2),所以它并不是一个好的选择。即便如此,WebGL2 的某些限制仍然让人感觉有些笨拙。

WebGPU

在 Web 之外,新一代的图形 API 已经建立起来,它为显卡提供了一个更低级的接口。这些新的 API 适应了设计 OpenGL 时不存在的新用例和约束。一方面,GPU 现在几乎无处不在。甚至移动设备也内置了功能强大的 GPU。因此,现代图形编程(3D 渲染和光线追踪)和 GPGPU 的使用越来越普遍。另一方面,大多数设备都有多核处理器,因此能够从多个线程与 GPU 进行交互可能是一个重要的优化因素。当 WebGPU 的人在做这件事的时候,他们也重新审视了以前的一些设计决策,并将 GPU 必须做的很多验证工作前置,使开发者能够从 GPU 中榨取更多的性能。

下一代 GPU API 中最受欢迎的是 Khronos 集团的 Vulkan、苹果的 Metal 和微软的 DirectX 12。为了给网络带来这些新功能,WebGPU 诞生了。虽然 WebGL 只是 OpenGL 的一个简单的封装,但 WebGPU 选择了不同的道路。它引入了自己的抽象,并不直接映射原生 API。部分原因是没有一个 API 可以在所有系统上使用,同时也是因为许多概念(比如极低级别的内存管理)对于面向 Web 的 API 来说并不习惯。相反,WebGPU 的设计既让人感觉『webby』,又能舒适地位于任何原生图形 API 之上,同时抽象出它们的特性。它在 W3C 中被标准化,所有主要的浏览器厂商都有一个席位。由于其相对低级的性质及其强大的功能,WebGPU 有一点学习曲线,并且在设置上相对繁重,但我将尽可能地分解它。

适配器(Adapters)与设备(Devices)

你首先接触到的 WebGPU 抽象是适配器(adapters) 和(逻辑)设备

抽象层,从物理 GPU 到逻辑设备

物理设备是指 GPU 本身,通常区分为内置 GPU 和独立 GPU。通常情况下,大部分设备都有一个 GPU,但也有可能有两个或更多的 GPU。例如,微软的 Surface Book 有一个低功率的集成 GPU 和一个高性能的独立 GPU,操作系统会根据需要在两者之间进行切换。

驱动程序——由 GPU 制造商提供——将以操作系统理解和期望的方式向操作系统展示 GPU 的能力。操作系统反过来可以使用操作系统提供的图形 API(如 Vulkan 或 Metal)将其暴露给应用程序。
GPU 是一种共享资源。它不仅被许多应用程序同时使用,而且还控制你在屏幕上看到的内容。需要有一种东西能让多个进程同时使用 GPU,这样每个应用程序就能把自己的用户界面放在屏幕上,而不干扰其他应用程序,甚至恶意读取其他应用程序的数据。对每个进程来说,它们看起来像是对物理 GPU 拥有唯一的控制权,但这显然不是真正的情况。这种多路复用主要是由驱动程序和操作系统完成的。

适配器则是操作系统的原生图形 API 与 WebGPU 之间的转换层。由于浏览器是一个单一的操作系统级应用,它可以运行多个 Web 应用,因此又需要进行多路复用,使每个 Web 应用感觉到它对 GPU 有唯一的控制权。这在 WebGPU 中用逻辑设备的概念来模拟。
为了获得对一个适配器的访问,你可以调用 navigator.gpu.requestAdapter() 。在写这篇文章的时候,requestAdapter() 只有很少的选项。这些选项允许你请求一个高性能或低能耗的适配器。

软件渲染:一些实现为没有 GPU 或 GPU 能力不足的系统提供了一个 “后备适配器”。后备适配器实际上是一个纯粹的软件实现,它不会很快,但可以保证你应用程序的功能。

如果成功了,即返回的适配器不是空的,你可以检查适配器的能力,并使用 adapter.requestDevice() 向适配器请求一个逻辑设备。

if (!navigator.gpu) throw Error("WebGPU not supported.");

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) throw Error("Couldn’t request WebGPU adapter.");

const device = await adapter.requestDevice();
if (!device) throw Error("Couldn’t request WebGPU logical device.");

如果没有提供任何参数, requestDevice() 将返回一个不一定符合物理设备能力的设备,而是 WebGPU 团队认为合理的、所有 GPU 的最低共同标准。细节 在 WebGPU 标准中规定。例如,即使我的 GPU 能够轻松地处理高达 4GiB 的数据缓冲区,但返回的 device 只允许最多 1GiB 的数据缓冲区,并将拒绝任何更大的数据缓冲区。这看起来可能有些限制,但实际上非常有用:如果你的 WebGPU 应用能使用默认设备运行,那么它将在大多数设备上运行。如果有必要,你可以通过 adapter.limit 检查物理 GPU 的实际限制,并通过向 requestDevice() 传递一个 options 对象来请求一个具有更高限制的 decice

着色器(Shaders)

如果你使用过 WebGL,你可能对顶点着色器和像素着色器(也译作片元着色器或片段着色器)很熟悉。在不做过多深入研究的情况下,传统的准备工作是这样的:你向你的 GPU 上传一个数据缓冲区,并告诉它如何将这些数据解释为一系列的三角形。每个顶点都占据了数据缓冲区的一个块,描述该顶点在三维空间中的位置,但也可能包括颜色、纹理 ID、法线和其他数据。列表中的每个顶点都由 GPU 在顶点阶段处理,在每个顶点运行顶点着色器,它将实现平移、旋转或透视变形等操作。

着色器:『着色器』这个术语曾经让我感到困惑,因为你可以做的事情远不止是着色。但在过去的日子里(20 世纪 80 年代末),这个词是合适的:它是在 GPU 上运行的一小段代码,决定每个像素应该是什么颜色,这样你就可以对被渲染的物体进行着色,实现照明和阴影。如今,着色器泛指任何在 GPU 上运行的程序。

现在 GPU 对三角形进行光栅化处理,这意味着 GPU 会计算出每个三角形在屏幕上覆盖的像素。然后,每个像素由像素着色器处理,它可以访问像素坐标,也可以访问辅助数据来决定该像素应该是哪种颜色。如果使用得当,这个过程可以创造出令人惊叹的 3D 图形。

这种将数据传给顶点着色器,然后传给像素着色器,再直接输出到屏幕上的系统被称为管线,在 WebGPU 中,你必须明确地定义你的管线。

管线(Pipeline)

目前,WebGPU 允许你创建两种类型的管线。渲染管线和计算管线。顾名思义,渲染管线是指创建一个 2D 图像。该图像不需要在屏幕上,但可以直接渲染到内存中(这被称为帧缓冲区)。计算管道更通用,它返回一个缓冲区,其中可以包含任何类型的数据。在这篇博文的其余部分,我将专注于计算管线,因为我喜欢把渲染管线看作是计算管线的一种特殊优化版本。现在,这既是历史上的倒退,因为计算管线是作为专用的渲染管线的泛化而构建的,也低估了这些管线在 GPU 中是不同的物理电路。然而,就 API 而言,我发现这种心理模型非常有帮助。在未来,似乎更多类型的管道——也许是光线追踪管道——将被添加到 WebGPU 中。

在 WebGPU 中,管道由一个(或多个)可编程阶段组成,每个阶段由一个着色器和一个入口点定义。一个计算管线有一个 compute 阶段,而一个渲染管线有一个 vertex 和一个 fragment 阶段:

const module = device.createShaderModule({
    code: `
    @stage(compute) @workgroup_size(64)
    fn main() {
      // Pointless!
    }
  `,
});

const pipeline = device.createComputePipeline({
    compute: {
        module,
        entryPoint: "main",
    },
});

这是 WebGPU 着色语言 WGSL(发音为 “wig-sal”)的首次亮相。WGSL 给我的感觉是 Rust 和 GLSL 的混合。它有很多 Rust-y 语法和 GLSL 的全局函数(如 dot() , norm() , len() , …)、类型(如 vec2 , mat4x4 , …)和 符号混写(swizzling notation)(如 some_vec.xxy , …)。浏览器将把你的 WGSL 编译成底层系统期望的东西。这可能是 DirectX 12 的 HLSL,Metal 的 MSL 和 Vulkan 的 SPIR-V

SPIR-VSPIR-V 很有趣,因为它是一种开放的、二进制的、由 Khronos 集团标准化的中间格式。你可以把 SPIR-V 看作是并行编程语言编译器的 LLVM,并且支持将许多语言编译成 SPIR-V,以及将 SPIR-V 编译成许多其他语言。

在上面的着色器模块中,我们只是创建了一个名为 main 的函数,并通过使用 @stage(compute) 属性将其标记为计算阶段的入口点。你可以在一个着色器模块中把多个函数标记为入口点,因为你可以在多个管道中重复使用同一个着色器模块,并通过 entryPoint 选项选择不同的函数来调用。但那个 @workgroup_size(64) 属性是什么?

并行化(Parallelism)

GPU 是以延迟为代价来优化吞吐量的。为了理解这一点,我们必须看一下 GPU 的架构。我不想(说实话,也不能)完整地解释它。我会在觉得有必要的情况下深入研究。如果你想了解更多,Fabian Giesen 的这个由 13 部分组成的系列博文 很不错。

一些众所周知的事实是,GPU 有大量的内核,可以进行大规模的并行工作。然而,这些核心并不像你为多核 CPU 编程时那样独立。首先,GPU 的核心是分层分组的。层次结构中不同层的术语在不同制造商和 API 之间并不一致。Intel 有一份很好的 文档,对其架构进行了高层次的概述,尽管 GPU 的确切架构是受 NDA 保护的秘密,但可以简单地假设其他 GPU 也是类似的架构。

就英特尔而言,层次结构中的最低一级是 “执行单元”(EU,Execution Unit),它有多个(在这里是七个)SIMT 核心。这意味着它有七个核心,它们以锁步方式运行,并始终执行相同的指令。然而,每个内核都有自己的一组寄存器和堆栈指针。因此,虽然它们必须执行相同的操作,但它们可以在不同的数据上执行。这也是 GPU 性能专家避免分支(如 if/else 或循环)的原因。如果 EU 遇到一个 if/else ,所有的核心都必须执行这两个分支,除非所有的核心碰巧都进入相同的分支。每个核心可以被告知忽略输入的指令,但这显然浪费了本可以用于计算的宝贵周期。这一点同样适用于循环!如果一个核心提前结束了循环,它将不得不继续执行循环体,直到所有核心都完成了循环。

尽管内核的频率很高,但从内存中获取数据(或从纹理中获取像素)仍然需要相对较长的时间——Fabian 说这需要几百个时钟周期。这几百个周期本可以用在计算上。为了利用这些原本闲置的周期,每个 EU 都有大量的超额工作。每当一个 EU 闲置时(例如等待内存中的一个值),它就会切换到另一个工作项目,只有当新的工作项目需要等待时才会切换回来。这是 GPU 如何以延迟为代价优化吞吐量的关键技巧:单个工作将花费更长的时间,因为切换到另一个工作可能会使停止执行的时间超过原本需要等待的时间,但整体利用率和吞吐量更高。GPU 努力使工作总是排在队列中,以使 EU 始终保持工作。

Intel Xe 显卡芯片架构,每个 EU 有 7 个 SIMT 核心,每个子切片(SubSlice)有 8 个 EU,8 个子切片组成一个切片(Slice)

不过,EU 只是层次结构中的最低级别。多个 EU 被分组为英特尔所谓的『SubSlice(子切片)』。一个 SubSlice 中的所有 EU 都可以访问少量的共享本地内存(SLM,Shared Local Memory),在 Intel 的 Xe 中大约是 64KiB。如果要运行的程序有任何同步命令,它必须在同一个 SubSlice 内执行,因为只有它们有共享内存用于同步。
在最后一层,多个 SubSlices 被组合成一个 Slice,形成 GPU。对于集成的 Intel GPU,总共会有 170-700 个核心。独立的 GPU 可以拥有 1500 个甚至更多的核心。同样,这里的命名来自英特尔,其他厂商可能使用不同的名称,但每个 GPU 的大体架构是类似的。

为了充分利用这种体系结构的优点,需要专门为此体系结构设置程序,以便纯程序化的 GPU 调度器可以最大限度地提高利用率。因此,图形 API 暴露了一个 线程模型,自然允许以这种方式分解工作。在 WebGPU 中,这里的重要基元是『工作组(workgroup)』。

工作组(Workgroups)

在传统的设置中,顶点着色器被每个顶点调用一次,而像素着色器被每个像素被调用一次(我知道这里忽略了一些细节)。在 GPGPU 设置中,你的计算着色器将被你安排的每个工作项调用一次。至于什么是工作项,则由你来定义。

所有工作项的集合(我称之为『工作负载(workload)』)被分解成若干个工作组(workgroups)。一个工作组中的所有工作项被安排在一起运行。在 WebGPU 中,工作负载被建模为一个 3 维网格,每个『立方体』是一个工作项,工作项组成更大的立方体以形成一个工作组。

这是一个工作负载(workload),白色边框的立方体是工作项(work item),红色边框的立方体是工作组(workgroup)

终于,我们有足够的知识来讨论 @workgroup_size(x, y, z) 属性,它甚至可以说是不言自明的:该属性允许你告诉 GPU 这个着色器的工作组应该是什么大小。或者用上图的语言来说, @workgroup_size 属性定义了红色边框立方体的大小。 x * y * z 是每个工作组的工作项数量。任何跳过的参数都被假定为 1,所以 @workgroup_size(64) 相当于 @workgroup_size(64, 1, 1)

当然,芯片上的实际 EU 并没有被安排的三维网格中。在三维网格中建立模型的目的是为了提高定位性。假设相邻的工作组可能会访问内存中的类似区域,因此当顺序运行相邻的工作组时,缓存命中的可能性会更高,从而不必从内存中获取值,从而节省几百个周期。然而,大多数硬件似乎只是以串行顺序运行工作组,因为运行 @workgroup_size(64)@workgroup_size(8, 8) 的着色器之间的差异可以忽略不计。所以这个概念可能有点过时。

然而,工作组在多个方面受到限制: device.limit 有一堆属性,值得了解一下:

// device.limits
{
    // ...
    maxComputeInvocationsPerWorkgroup: 256,
    maxComputeWorkgroupSizeX: 256,
    maxComputeWorkgroupSizeY: 256,
    maxComputeWorkgroupSizeZ: 64,
    maxComputeWorkgroupsPerDimension: 65535,
    // ...
}

工作组规模的每个维度的大小都有限制,但即使 x、y 和 z 单独在限制范围内,它们的乘积( =x * y * z )可能不在限制范围内,因为它有自己的限制。最后,每个维度最多只能有这些个工作组。

Pro tip:不要生成最大数量的线程。尽管 GPU 是由操作系统和底层调度器管理的,但你可能会因为 长时间运行的某个 GPU 程序而使整个系统停止响应

那么,什么是正确的工作组大小?这真的取决于你分配给工作项坐标的语义。我知道这并不是一个真正有用的答案,所以我想给你 Corentin 给我的同样的建议。『使用 64(工作组大小),除非你知道客户端是什么 GPU 或者你的工作负载比较与众不同。』这似乎是一个安全的数字,在许多 GPU 上表现良好,并允许 GPU 调度器保持尽可能多的 EU 忙碌。

命令(Commands)

我们已经编写了着色器并建立了流水线。剩下要做的就是实际调用 GPU 来执行。由于 GPU 可以是一个完全独立的显卡,有自己的内存芯片,你可以通过所谓的命令缓冲区或命令队列来控制它。命令队列是一大块内存,其中包含供 GPU 执行的已编码的命令。编码方式与 GPU 高度相关,所以这将由驱动程序来处理。WebGPU 暴露了一个 CommandEncoder ,来使用该功能。

const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(pipeline);
passEncoder.dispatch(1);
passEncoder.end();
const commands = commandEncoder.finish();
device.queue.submit([commands]);

commandEncoder 有多个方法,允许你将数据从一个 GPU 缓冲区复制到另一个缓冲区和操作纹理。它还允许你创建 PassEncoder 对管道的设置和调用进行编码。在这个例子中,我们有一个计算管线,所以我们必须创建一个计算通道,将其设置为使用我们预先声明的管道,最后调用 dispatch(w_x, w_y, w_z) 来告诉 GPU 沿着每个维度创建多少个工作组。换句话说,我们的计算着色器被调用的次数等于 wx * wy * wz * x * y * z 。顺便说一下,PassEncoder 是 WebGPU 的抽象,用来避免我在这篇博文开始时讨论的内部全局状态对象。所有运行 GPU 流水线所需的数据和状态都明确地通过 PassEncoder 传递。

抽象:命令缓冲区也是驱动程序或操作系统的钩子,让多个应用程序使用 GPU 而不互相干扰。当你的命令排队时,下面的抽象层会向队列中注入额外的命令,以保存前一个程序的状态并恢复你的程序的状态,这样就好像没有其它程序在使用 GPU。

运行这段代码,我们实际上是在 GPU 上生成了 64 个线程,它们不做任何事情。但它是有效的,所以这很酷。让我们来谈谈我们如何给 GPU 一些数据来工作。

交换数据

正如我承诺的那样,我不会直接使用 WebGPU 进行图形处理,因此我认为在 GPU 上运行物理模拟并使用 Canvas2D 将其可视化会很有趣。也许我称其为『物理模拟』是自欺欺人——我所做的是生成一大堆圆圈,让它们在一个平面上以随机的方向移动,并让它们发生碰撞。

为了使其发挥作用,我们需要将一些模拟参数和初始状态推送到 GPU 上,在 GPU 上运行模拟,并从 GPU 上读取模拟结果。这可以说是 WebGPU 最棘手的部分,因为有一堆数据杂技(不是说看似毫无意义的复制),但这是使 WebGPU 成为以最高性能运行的设备无关 API 的原因。

绑定组布局

为了与 GPU 交换数据,我们需要用一个绑定组的布局来扩展我们的流水线定义。绑定组是 GPU 实体(内存缓冲区、纹理、采样器等)的集合,在流水线的执行过程中可以访问。绑定组布局预先定义了这些 GPU 实体的类型、目的和用途,这使得 GPU 能够提前弄清楚如何最有效地运行管线。在这个初始步骤中简单一些,让我们的流水线访问一个单一的内存缓冲区:

const bindGroupLayout =
    device.createBindGroupLayout({
        entries: [{
            binding: 1,
            visibility: GPUShaderStage.COMPUTE,
            buffer: {
                type: "storage",
            },
        }],
    });

const pipeline = device.createComputePipeline({
    layout: device.createPipelineLayout({
        bindGroupLayouts: [bindGroupLayout],
    }),
    compute: {
        module,
        entryPoint: "main",
    },
});

binding 数字可以自由选择,用于将我们的 WGSL 代码中的变量与绑定组布局的这个槽中的缓冲区的内容联系起来。我们的 bindGroupLayout 还定义了每个缓冲区的用途,在这里是 storage 。另一个选项是 read-only-storage ,它是只读的(废话!),并允许 GPU 在永远不会被写入的基础上做进一步优化,例如不需要同步。缓冲区类型的最后一个可能的值是 uniform ,在计算管线的背景下,它在功能上主要等同于一个存储缓冲区。

绑定组布局已经到位了。现在我们可以创建绑定组本身,包含绑定组布局所期望的 GPU 实体的实例。一旦绑定组和内部的缓冲区到位,计算着色器就可以将数据填入其中,我们就可以从 GPU 上读取数据。但是还有一个障碍:暂存缓冲区(Staging Buffers)。

暂存缓冲区(Staging Buffers)

我再说一遍:GPU 在很大程度上是以延迟为代价来优化吞吐量的。GPU 需要惊人的高速度将数据输送到核心,以维持该吞吐量。Fabian 在他 2011 年的博文系列中做了一些 反推计算,得出的结论是,1280x720 分辨率的着色器下,仅纹理样本 GPU 就需要 3.3GB/s 的带宽。为了适应今天的图形需求,GPU 需要以更快的速度处理数据。只有当 GPU 的内存与核心非常紧密地结合在一起时,这才有可能实现。这种紧密的集成使其很难将相同的内存暴露给主机进行读写。

不过,GPU 有主机和 GPU 都可以访问的额外内存区块,但集成不是很紧密,读写速度也不是很快。暂存缓冲区就是在这个中间内存区分配的缓冲区,它可以 映射 到主机系统中进行读写。为了从 GPU 读取数据,我们将数据从内部的高性能缓冲区复制到一个暂存缓冲区,然后将暂存缓冲区映射到主机,这样我们就可以将数据读回主内存。对于写入,过程也是一样的,但顺序相反。

回到我们的代码:我们将创建一个可写的缓冲区,并将其添加到绑定组,这样它就可以被计算着色器写入。我们还将创建第二个同样大小的缓冲区作为暂存缓冲区。每个缓冲区的创建都有一个 usage 位掩码,你可以用它声明你打算如何使用这个缓冲区。然后,GPU 会得出满足使用情况的缓冲区的位置,如果无法实现满足标志组合的位置,则会出现错误。

const BUFFER_SIZE = 1000;

const output = device.createBuffer({
    size: BUFFER_SIZE,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
});

const stagingBuffer = device.createBuffer({
    size: BUFFER_SIZE,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});

const bindGroup = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [{
        binding: 1,
        resource: {
            buffer: output,
        },
    }],
});

注意 createBuffer() 返回一个 GPUBuffer , 不是 ArrayBuffer ,它们现在还不能读写。为此,需要将它们映射,这是另一个 API 调用,这个 API 只有在缓冲区有 GPUBufferUsage.MAP_READGPUBufferUsage.MAP_WRITE 时才会成功。

TypeScript: 我发现 TypeScript 在探索新 API 时非常有帮助。幸运的是,Chrome 的 WebGPU 团队维护着 @webgpu/types ,所以你可以享受准确的自动补全。

现在我们不仅有了绑定组的布局,甚至还有了实际的绑定组本身,我们需要更新我们的调度代码来利用这个绑定组。之后映射暂存缓冲区,将结果读回 JavaScript。

const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(pipeline);
passEncoder.setBindGroup(0, bindGroup);
passEncoder.dispatch(1);
passEncoder.dispatch(Math.ceil(BUFFER_SIZE / 64));
passEncoder.end();
commandEncoder.copyBufferToBuffer(
    output,
    0, // Source offset
    stagingBuffer,
    0, // Destination offset
    BUFFER_SIZE
);
const commands = commandEncoder.finish();
device.queue.submit([commands]);

await stagingBuffer.mapAsync(
    GPUMapMode.READ,
    0, // Offset
    BUFFER_SIZE // Length
);
const copyArrayBuffer =
    stagingBuffer.getMappedRange(0, BUFFER_SIZE);
const data = copyArrayBuffer.slice();
stagingBuffer.unmap();
console.log(new Float32Array(data));

由于我们在管线中添加了一个绑定组的布局,现在没有提供绑定组的调用都会失败。在我们定义了『Pass』之后,我们通过命令编码器添加一个额外的命令,将数据从输出缓冲区复制到暂存缓冲区,并将命令缓冲区提交到队列中。GPU 将开始通过命令队列工作。我们不知道 GPU 到底什么时候会完成,但我们已经可以提交映射 stagingBuffer 的请求了。这个函数需要等待命令队列被完全处理,所以它是异步的。 stagingBuffer.getMappedRange() 让我们请求将一个区块(或整个缓冲区)作为一个 ArrayBuffer 暴露给 JavaScript。这是实际的、映射的 GPU 内存,这意味着当 stagingBuffer 被解除映射时,数据将消失( ArrayBuffer 将被『分离』),所以我使用 slice() 来创建它在 JavaScript 的副本。

虽然不是很令人兴奋,但我们从 GPU 内存中复制了这些 0

零以外的东西可能会更有说服力。在我们开始在 GPU 上进行任何高级计算之前,让我们把一些手工挑选的数据放入缓冲区,以证明我们的管道确实在按计划工作。这是我们新的计算着色器代码,为了清晰起见,增加了间距。

@group(0) @binding(1)
var<storage, write> output: array<f32>;

@stage(compute) @workgroup_size(64)
fn main(

  @builtin(global_invocation_id)
  global_id : vec3<u32>,

  @builtin(local_invocation_id)
  local_id : vec3<u32>,

) {
  output[global_id.x] =
    f32(global_id.x) * 1000. + f32(local_id.x);
}

前两行声明了一个模块范围的变量 output ,它是动态大小的 f32 数组。属性声明了数据的来源:来自我们第一个(0th)绑定组的缓冲区, binding 值为 1 的条目。数组的长度将自动反映底层缓冲区的长度(向下取整)。

变量:WGSL 与 Rust 不同的是,用 let 声明的变量是不可变的。如果你想让一个变量是可变的,要使用关键字 var

main() 函数的签名中增加了两个参数: global_idlocal_id 。我可以选择任何名字——它们的值由相关的属性决定: global_invocation_id 是一个内置值,对应该着色器调用在工作负载中的全局 x/y/z 坐标。 local_invocation_id 是这个着色器在工作组中的 x/y/z 坐标。

工作负载中标记的三个工作项 a、b 和 c 的示例

这张图片展示了 @workgroup_size(4, 4, 4) 工作负载坐标系统的一种可能情况。这取决于你为实际使用情况定义什么坐标系。如果我们接受上图中的坐标轴,则 main() 中的 a、b、c 参数如下:

  • a:
    • local_id=(x=0, y=0, z=0)
    • global_id=(x=0, y=0, z=0)
  • b:
    • local_id=(x=0, y=0, z=0)
    • global_id=(x=4, y=0, z=0)
  • c:
    • local_id=(x=1, y=1, z=0)
    • global_id=(x=5, y=5, z=0)

在着色器中,我们有 @workgroup_size(64, 1, 1) ,所以 local_id.x 会在 0~63 之间。为了能够检查这两个值,我将它们『编码』成一个数字。请注意,WGSL 是严格类型的: local_idglobal_id 都是 vec3<u32> ,所以我们必须明确地将它们转换成 f32 ,以便分配给 f32 的输出缓冲区。

由 GPU 实际填充的值。注意本地调用的 ID 在 63 后重置,而全局调用的 ID 则继续计数

这证明了我们的计算着色器确实是为输出内存中的每一个值调用的,并为其填充了一个唯一的值。我们不会知道这些数据是以何种顺序填入的,这是故意未指定的,并留给 GPU 的调度器决定。

过度调用(Overdispatching)

聪明的读者可能已经注意到,着色器调用的总数( Math.ceil(BUFFER_SIZE / 64) * 64 )将导致 global_id.x 大于数组的长度,因为每个 f32 占用 4 字节。幸运的是,对数组的访问有一个隐式的保护,所以每一次写过数组的末端都会写入数组的最后一个元素。这就避免了内存访问错误,但仍可能产生无法使用的数据。事实上,如果你检查返回的缓冲区的最后 3 个元素,你会发现数字 247055、248056 和 608032。这就需要我们在着色器代码中提前退出来防止这种情况发生:

fn main( /* ... */ ) {
    if (global_id.x >= arrayLength( & output)) {
        return;
    }
    output[global_id.x] =
        f32(global_id.x) * 100. + f32(local_id.x);
}

你可以运行这个 demo 来查看完整的源代码。

一个疯狂的结构体

现在我们的目标是让一大堆球在二维空间中移动,并进行快乐的小碰撞。为此,每个球需要有一个半径、一个位置和一个速度向量。我们可以继续使用 array<f32> ,并确定第一个浮点数是第一个球的 x 位置,第二个浮点数是第一个球的 y 位置,以此类推。这并不是我所说的符合人体工程学的做法。幸运的是,WGSL 允许我们自定义结构体,可以将多条数据放在一个整洁的袋子里。

冷知识:如果你知道什么是内存对齐,你可以跳过这一节(尽管要看一下代码示例)。如果你不知道这是什么,我不会解释原因,但会告诉你这将导致什么,以及如何解决它。

所以,使用这些组件定义一个 struct Ball ,将 array<f32> 变成 array<Ball> 是有意义的。但这一切的坏处是:我们必须谈谈 对齐问题

struct Ball {
    radius: f32;
    position: vec2 < f32 > ;
    velocity: vec2 < f32 > ;
}

@group(0) @binding(1)
var < storage, write > output: array < f32 > ;
var < storage, write > output: array < Ball > ;

@stage(compute) @workgroup_size(64)
fn main(
    @builtin(global_invocation_id) global_id: vec3 < u32 > ,
    @builtin(local_invocation_id) local_id: vec3 < u32 > ,
) {
    let num_balls = arrayLength( & output);
    if (global_id.x >= num_balls) {
        return;
    }

    output[global_id.x].radius = 999.;
    output[global_id.x].position = vec2 < f32 > (global_id.xy);
    output[global_id.x].velocity = vec2 < f32 > (local_id.xy);
}

如果你运行这个 demo,你会在控制台看见这个:

由于对齐限制,该结构体的内存布局中有一个洞(padding)

我把 999 作为结构体的第一个字段,以方便查看结构体在缓冲区中的起始位置。在下一个 999 之前总共有 6 个数字,这有点令人惊讶,因为该结构体实际上只存储 5 个数字: radiusposition.xposition.yvelocity.xvelocity.y 。这是因为对齐的问题。

每个 WGSL 数据类型都有明确的 对齐要求。如果一个数据类型的对齐方式为* N,这意味着该数据类型的值只能存储在 N *的倍数的内存地址上。 f32 的对齐是 4, 而 vec2<f32> 的对齐是 8。如果我们假设结构体从地址 0 开始,那么 radius 可以存储在地址 0,因为 0 是 4 的倍数。结构体的下一个字段是 vec2<f32> ,它的对齐是 8。 然而, radius 之后的第一个空闲地址是 4,这不是 8 的倍数。 为了解决这个问题,编译器添加了 4 个字节的 padding,来让下一个地址是 8 的倍数。 这解释了为什么我们在 DevTools 控制台看到了一个值为 0 的未使用字段。

WGSL 规范中的(缩短的)[对齐表](https://gpuweb.github.io/gpuweb/wgsl/#alignment-and-size)

现在我们知道了结构体在内存中是如何布局的,我们可以从 JavaScript 中填充它,以生成球的初始状态,也可以读回并使其可视化。

输入和输出

我们已经成功地从 GPU 中读取数据,将其带到 JavaScript 中并『解码』。现在是时候解决另一个问题了。我们需要在 JavaScript 中生成所有球的初始状态,并将其交给 GPU,以便能够在上面运行计算着色器。生成初始状态相当简单:

let inputBalls = new Float32Array(new ArrayBuffer(BUFFER_SIZE));
for (let i = 0; i < NUM_BALLS; i++) {
    inputBalls[i * 6 + 0] = randomBetween(2, 10); // radius
    inputBalls[i * 6 + 1] = 0; // padding
    inputBalls[i * 6 + 2] = randomBetween(0, ctx.canvas.width); // position.x
    inputBalls[i * 6 + 3] = randomBetween(0, ctx.canvas.height); // position.y
    inputBalls[i * 6 + 4] = randomBetween(-100, 100); // velocity.x
    inputBalls[i * 6 + 5] = randomBetween(-100, 100); // velocity.y
}

Buffer-backed-object:对于更复杂的数据结构,在 JavaScript 中操作数据可能会变得相当乏味。虽然最初是为 Web Workers 写的,但我的库 buffer-backed-object 在这里也能派上用场。

我们已经知道如何将缓冲区暴露给我们的着色器。我们只需要调整流水线绑定组的布局来加入另一个缓冲区:

const bindGroupLayout = device.createBindGroupLayout({
    entries: [{
            binding: 0,
            visibility: GPUShaderStage.COMPUTE,
            buffer: {
                type: "read-only-storage",
            },
        },
        {
            binding: 1,
            visibility: GPUShaderStage.COMPUTE,
            buffer: {
                type: "storage",
            },
        },
    ],
});

… 然后创建一个让我们可以绑定的绑定组的 GPU 缓冲区:

const input = device.createBuffer({
    size: BUFFER_SIZE,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});

const bindGroup = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [{
            binding: 0,
            resource: {
                buffer: input,
            },
        },
        {
            binding: 1,
            resource: {
                buffer: output,
            },
        },
    ],
});

新的部分:发送数据到 GPU。就像读取数据一样,从技术上讲我们必须创建一个可以映射的暂存缓冲区,将数据复制到暂存缓冲区,然后发出命令将数据从暂存缓冲区复制到存储缓冲区。然而,WebGPU 提供了一个便利的功能,让我们高效的将数据送入存储缓冲区,尽管这实时创建了一个临时的暂存缓冲区:

device.queue.writeBuffer(input, 0, inputBalls);

就这样?就这样!我们甚至不需要命令编码器。我们可以将这个命令直接加入命令队列。 device.queue 也为纹理提供了一些其他类似的便利功能。
现在我们需要把这个新的缓冲区绑定到 WGSL 的一个变量上,并对它做一些处理:

struct Ball {
  radius: f32;
  position: vec2<f32>;
  velocity: vec2<f32>;
}

@group(0) @binding(0)
var<storage, read> input: array<Ball>;

@group(0) @binding(1)
var<storage, write> output: array<Ball>;

let TIME_STEP: f32 = 0.016;

@stage(compute) @workgroup_size(64)
fn main(
  @builtin(global_invocation_id)
  global_id : vec3<u32>,
) {
  let num_balls = arrayLength(&output);
  if(global_id.x >= num_balls) {
    return;
  }
  output[global_id.x].position =
    input[global_id.x].position +
    input[global_id.x].velocity * TIME_STEP;
}

我希望现在这个着色器的绝大部分代码没有让你惊讶。

每一帧都被 WebGPU 用来更新球的位置,并使用 Canvas2D 将它们绘制到屏幕上

最后,我们需要做的是将 output 缓冲区读回 JavaScript,编写一些 Canvas2D 代码来可视化缓冲区的内容,并将其全部放入 requestAnimationFrame() 循环。你可以在这个 demo 中看到效果。

性能

前面的 demo 只是沿着每个球的速度向量移动,并不惊艳或具有很复杂的计算。在看这个作品的性能特性之前,让我在着色器中加入一些适当的物理计算。我不会在这里解释——这篇博文已经够长了——但我想说我采取了最朴素的方法:每个球都检查与其他所有球的碰撞情况。如果你感到好奇,你可以看一下 最终 demo 的源代码,其中还包含我用来编写物理部分的资源链接。

... 现在有了弹性墙和弹性球!

我不想对这个实验进行精确的测量,因为我还没有优化物理算法和 WebGPU 的使用。然而,即使是这种朴素的实现也表现得非常好(在我的 M1 MacBook Air 上),这一点让我印象深刻。维持 60fps 时,大约可以处理 2500 个球。但看一下跟踪,在有 2500 个球时,很明显瓶颈是 Canvas2D 绘制场景,而不是 WebGPU 的计算。

有 14000 个球时,在 M1 MBA 上 GPU 的纯计算时间大约 16ms

为了看看到底多快,我禁用了渲染,使用 performance.measure() 来看看在耗尽每帧 16ms 的时间之前我可以模拟多少个球。在我的机器上,这是 14000 个球左右。这种未经优化的东西运行得如此之快,真让我醉心于 WebGPU 给我提供的强大计算能力。

稳定性和可用性

WebGPU 标准的制定已经有了一段时间,我认为标准小组尽快宣布该 API 是稳定的。也就是说,这个 API 只在 Chrome 和 Firefox 加入相应的 flag 后才可用。我对 Safari 支持这个 API 持乐观态度,但在写这篇文章的时候,还没有在 Safari TP 中看到。

在稳定性方面,甚至在我为这篇文章做研究的时候就有一些变化落地了。例如,属性的语法从 [[stage(compute), workgroup_size(64)]] 改为了 @stage(compute) @workgroup_size(64) 。在写这篇文章的时候,Firefox 还在使用旧的语法。 passEncoder.end() 曾经是 passEncoder.endPass() 。规范中还有一些东西还没有在任何浏览器中实现,比如 着色器常量 或在移动设备上使用 API。

我要说的基本上是:当浏览器和标准制定者在这个 API 走向✨稳定✨的过程中,预计还会有更多的重大变化。

总结

在 Web 上拥有一个现代的 API 来与 GPU 交互将是非常有趣的。在投入时间克服最初的学习曲线后,我真的感到有能力在 GPU 上使用 JavaScript 运行大规模的并行工作负载。还有 wgpu,它在 Rust 中实现了 WebGPU 的 API,允许你在浏览器外使用 API。wgpu 也支持 WebAssembly 作为编译目标,所以你可以在原生在浏览器外或通过 WebAssembl 在浏览器内运行 WebGPU 程序。有趣的是:Deno 是第一个支持 WebGPU 的运行时(多亏了 wgpu)。

如果你有问题或遇到问题,这里有一个 Matrix channel,里面有许多 WebGPU 的用户、浏览器工程师和标准人员,他们对我有很大的帮助。尝试一下吧,这非常令人激动。

感谢 Brandon Jones 对本文的校对,感谢 WebGPU Matrix channel 对我所有问题的回答。


转载规则

《译:WebGPU — All of the cores, none of the canvas》Konata 采用 知识共享署名-非商业性使用 4.0 国际许可协议 进行许可。
  目录