如何玩转 WebGL 并行计算

Posted 2021-09-08 阿里云云栖号

简介： 如今在 Web 端使用 WebGL 进行高性能计算已有不少实践，例如在端智能领域中的 tensorflow.js，再比如可视化领域中的 Stardust.js。

作者 | 沧东
来源 | 阿里技术公众号

如今在 Web 端使用 WebGL 进行高性能计算已有不少实践，例如在端智能领域中的 tensorflow.js，再比如可视化领域中的 Stardust.js。在本文中，我们将介绍以下内容：

• 使用 GPU 进行通用计算（GPGPU）的历史
• 当前在 Web 端使用图形 API 实现 GPGPU 的技术原理，以及前端开发者可能遇到的难点
• 相关业界实践，包括布局计算、动画插值等
• 局限性与未来展望

一 什么是 GPGPU

由于硬件结构不同，GPU 与 CPU 擅长执行不同类型的计算任务。CPU 通过复杂的 Cache 设计实现低延迟，包含复杂的控制逻辑（分支预测），ALU 只占一小部分。而 GPU 为高吞吐量而生，包含大量 ALU。因此在单指令流多数据流（SIMD）场景下，GPU 的运算速度远超 CPU，并且这种差距还在不断拉大。

而一些现代 GPU 上甚至有专门负责张量计算、光线追踪的硬件（Tensor/RT Core），例如 Nvidia 的图灵架构。这使得在处理这些计算复杂度极高的任务时能获得更大的性能提升。

这里就需要引出一个概念，用 GPU 进行除渲染外的通用计算：General-Purpose computation on Graphics Processing Units，即 GPGPU。

自 2002 年提出以来，在实时加解密、图片压缩、随机数生成等计算领域都能看到它的身影，GPU Gems/Pro 上也有专门的章节介绍。经由 Nvidia 提出的 CUDA(Compute Unified Device Architecture) 这一统一计算架构，开发者可以使用 C、Java、Python 等语言编写自己的并行计算任务代码。

那么在 Web 端我们应该如何使用 GPU 的计算能力呢？

二 用 WebGL 实现并行计算的原理

在现代化的图形 API（Vulkan/Metal/Direct3D）中提供了 Compute Shader 供开发者编写计算逻辑。考虑到 WebGPU 仍在开发中，目前在 Web 端能使用的图形渲染 API 只有 WebGL1/2，它们都不支持 Compute Shader（WebGL 2.0 Compute 已废弃），因此只能“曲线救国”。在本文的最后一节我们将展望未来的技术手段。

我们先忽略具体的 API 用法，从 CPU 和 GPU 的角度看两者在并行计算过程中是如何协作的，前者也常被称作 host，后者为 device。第一步为数据初始化，需要从 CPU 内存中拷贝数据到 GPU 内存中，在 WebGL 中会通过纹理绑定完成。第二步 CPU 需要准备提交给 GPU 的指令和数据，完成计算程序的编译，在 WebGL 中通过调用一系列 API 实现。在第三步中将计算逻辑分配给 GPU 各个核心执行，因此这段逻辑也叫做“核函数”。最后把计算结果从 GPU 内存中拷贝回 CPU 内存，在 WebGL1 中通过读取纹理中像素值完成。

下面我们从 GPU 编程模型和执行模型入手，顺便引出线程和线程组的概念，这也是 GPU 可数据并行的关键。下图展示了网格与线程组的层次关系，并不局限于 DirectCompute。

• 通过 dispatch(x, y, z) 分配一个 3 维的线程网格（Grid），其中的线程共享全局内存空间；
• 网格中包含了许多线程组（Work Group、Thread Group、Thread Block、本地工作组不同叫法），每一个线程组中又包含了许多线程，线程组也是 3 维的，一般在 Shader 中通过 numthreads(x, y, z) 指定。它们可以通过共享内存或同步原语进行通信；
• Shader 程序最终会运行在每一个线程上。对于每一个线程，可以获取自己在线程组中的 3 维坐标，也可以获取线程组在整个线程网格中的 3 维坐标，以此映射到不同的数据上，实现数据并行的效果；

再回到硬件视角，线程对应 GPU 中的 CUDA 核心，线程组对应 SM（Streaming Multiprocessor），网格就是 GPU。

1 WebGL1 纹理映射

下图来自「GPGPU 编程技术 - 从 GLSL、CUDA 到 OpenCL」，这也是经典的 GPGPU 计算流程。

通常来说图形渲染 API 最终的输出目标就是屏幕，显示渲染结果。但是在 GPGPU 场景中我们只是希望在 CPU 侧读取最终的计算结果。因此会使用到渲染 API 提供的离屏渲染功能，即渲染到纹理，其中的关键技术就是使用帧缓存对象（Framebuffer Object/FBO）作为渲染对象。纹理用来存储输入参数和计算结果，因此在创建时我们通常需要开启浮点数扩展 OES_texture_float，该扩展在 WebGL2 中已经内置。

并行计算发生在光栅化阶段，我们将计算逻辑（核函数）写在 Fragment Shader 中，Vertex Shader 仅负责映射纹理坐标，因此 Geometry 可以使用一个 Quad（4个顶点）或者全屏三角形（3个顶点）。对于每一个像素点来说，它的工作并无变化，平时执行的渲染逻辑此时成了一种计算过程，像素值也成了计算结果。

但这种方式存在一个明显的限制，对于所有线程，纹理缓存要么是只读的，要么就是只写的，没法实现一个线程在读纹理，另一个在写纹理。本质上是由 GPU 的硬件设计决定的，如果想要实现多个线程同时对同一个纹理进行读/写操作，需要设计复杂的同步机制避免读写冲突，势必会影响到线程并行执行的效率。因此在经典 GPGPU 的实现中，通常我们会准备两个纹理，一个用来保存输入数据，一个用来保存输出数据。

除此之外，该方法并不支持线程间同步和共享内存这些特性，因此一些并行算法无法实现，例如 Bellman-Ford 单源最短路径算法。

上图中也提到了乒乓技术，很多算法需要连续运行多次，例如 G6 中使用的布局算法需要迭代多次达到稳定状态。上一次迭代中输出的计算结果，需要作为下一次迭代的输入。在实际实现中，我们会分配两张纹理缓存，每次迭代后对输入和输出纹理进行交换，实现类似乒乓的效果。

值得注意的是，由于 readPixels（在 CPU 侧读取纹理中的数据）非常慢，除了获取最终结果，过程中应当尽可能减少对它的调用，尽可能让数据留在 GPU 中。

这里我们不再展开 WebGL1 API 的实际用法，详细使用方式可以参考相关教程。

2 WebGL2 Transform Feedback

首先不得不提到已废弃的 WebGL 2.0 Compute（底层为 OpenGL ES 3.1），在草案中能看到例如用于线程间同步的 memoryBarrier 和 shared memory 这些高级特性，但最终工作组还是转向了 WebGPU。

WebGL2 中提供了另一种在 Vertex Shader 中进行并行计算的手段，即 Transform Feedback，它会跳过光栅化管线因此也不需要 Fragment Shader 参与（实际实现中提供一个空 Shader 即可）。

该方案和 WebGL1 的纹理映射方法有以下不同点：

• 不需要 Fragment Shader 参与，因此可以通过全局变量开启 gl.enable(gl.RASTERIZER_DISCARD)；
• 计算逻辑写在 Vertex Shader 中，不再需要晦涩的纹理映射，可以直接使用 Buffer 读写数据；
• 读取结果时可以直接使用 getBufferSubData。不过不变的是，该方法依然很慢；

虽然相比 WebGL1 已经有了不小进步，但依旧缺失 Compute Shader 中的一些重要特性。

同样，这里我们也不展开 WebGL2 API 的实际用法，详细使用方式可以参考相关教程。

三 实现中的难点

即使掌握了以上原理，前端开发者在具体实践中还是会遇到很大困难。除了图形 API 和 Shader 本身的学习成本，前端对于 GPU 编程模型本身也是比较陌生的。

我们遇到的第一个问题是一个算法是否可并行。有些计算任务非常耗时复杂，但并不能交给 GPU 来做，例如代码编译，因此可并行和复杂度并没有直接关系。关于是否可并行的判断并无严格标准，更多来自经验以及业界已有的实践（例如后文会提到的图布局/分析算法），通常遇到一个 "for every X do Y" 这样的任务就可以考虑是否能进行数据并行。例如下图展示了一种单源最短路径算法，不难发现里面有遍历每一个节点，针对每一条边的“松弛”操作，此时我们就可以考虑并行化，让一个线程处理一个节点。

当我们想把一个已有的可并行算法迁移到 GPU 中时，面临的第一个问题就是数据结构的设计。GPU 内存是线性的，也不存在类似对象这样结构，因此在迁移算法时不可避免的需要重新设计，如果再考虑到对 GPU 内存友好，设计难度会进一步加大。在下面应用示例「关于图布局/分析算法」一节中将看到关于图的线形表示。

下一个问题是无论 WebGL1 还是 WebGL2，都缺失了 Compute Shader 中的一些重要特性，因此一些在 CUDA 中已经实现的算法也无法直接移植。关于这个问题在本文最后一节中有详细的说明。

我们已经反复提到了共享内存和同步，这里举一个 Reduce 求和的例子帮助读者了解它们的含义。下图展示分配 16 个线程处理一个长度为 16 的数组，最终由 0 号线程将最终结果输出到共享内存的第一个元素中。该过程可分解为以下步骤：

1. 各个线程从全局内存中将数据装载到共享内存内。
2. 进行同步( barrier )，确保对于线程组内的所有线程，共享内存数据都是最新的。
3. 在共享内存中进行累加，每个线程完成后都需要进行同步。
4. 最后所有线程计算完成后，在第一个线程中把共享内存中第一个元素写入全局输出内存中。

试想如果没有共享内存和同步机制，最终的结果显然不会是正确的，有点类似并发编程中的 mutex，如果没有读写锁会得到意想不到的混乱结果。

最后，GPU 编程中的优化空间很大程度依赖开发者对硬件本身的了解，还是以上面 Reduce 求和为例，在 DirectCompute Optimizations and Best Practices 中能找到基于该版本 5 个以上的优化版本。

另外，GPU 在执行 Shader 时无法中断，这也带来了代码难以调试的问题，很多渲染引擎也同样面临这样的问题，Unity 有 RenderDoc 这样的工具，WebGL 暂无。

四 应用示例介绍

下面我们着重介绍一些 GPGPU 在可视化领域的应用，它们分别来自图算法、高性能动画以及海量数据并行处理场景。

既然是通用计算，我们必然无法覆盖所有领域的计算场景，我们尝试分析以下计算任务的设计实现思路，希望能给读者一些启发，当遇到特定场景的可并行算法时，可以尝试使用 GPU 加速这个过程。

1 图算法

布局和分析是图场景中常见的两类算法。CUDA 有 nvGRAPH 这样的高性能图分析算法库，包含类似最短路径、PageRank 等，支持多达 20 亿条边的规模。

在实现具体算法前，我们首先需要思考一个问题，即如何用线性结构表示一个图。最直观的数据结构是邻接矩阵，如下图所示。如果我们有 6 个节点，就可以用一个 6 x 6 的矩阵表示，有连接关系的就在对应元素上 + 1。下图来自维基百科对于邻接矩阵的展示。

但这样的数据结构存在一个明显的问题，过于稀疏导致空间浪费，尤其当节点数增多时。邻接表是更好的选择，该线性结构分成节点和边两部分，充分考虑 GPU 内存的顺序读，尽可能压缩（例如每一个 Edge 的 rgba 分量都存储了临接节点的 index）。以斥力计算（G6 的实现）为例，需要遍历除自身外的全部其他节点，这全部都是顺序读操作。同样的，在计算吸引力时，遍历一个节点的所有边也都是顺序读。随机读只会出现在获取端点坐标时才会出现。

这里不展开具体算法实现，迁移 G6 已有布局算法的过程详见。最终效果依不同算法实现差距很大，效果最好的是 Fruchterman 布局，节点数过千后 GPU 版本有百倍以上的提升，但 GForce 布局在少量节点的情况下甚至不如 CPU 版本。

2 SandDance

SandDance 提供了多维数据在多种布局下流畅切换的效果，它扩展了 Vega 规范，在 2D 场景中增加了深度信息，同时使用 Deck.gl 做渲染。具体到布局切换使用的技术，Luma.gl（Deck.gl 的底层渲染引擎） 提供了基于 WebGL2 Transform Feedback 的高级封装，用于在 GPU 中完成动画和数据变换的插值。对比传统的在 CPU 中做插值动画性能要高很多。在通用渲染引擎中，该技术也常用于粒子特效的实现。

3 P4: Portable Parallel Processing Pipelines

P4 致力于海量数据的处理和渲染，在运行时生成数据聚合和渲染的 Shader 代码，前者有点类似 tfjs 中的一些 op。值得一提的是通过 WebGL 的 Blending 操作实现了一些 Reduce 操作（例如最大最小值、计数、求和、平均值）。例如在实现 Reduce 求和时使用到的 blendEquation 为 gl.ADD。

五 当前局限性与未来展望

我们可以看出 WebGL 受限于底层 API 能力，在很多计算相关的特性上有不同程度的缺失，导致很多可并行算法无法实现。另一方面，可视化领域又缺失有不少适合的场景，我们迫切需要下一代能力更强的 Web API。

WebGPU 作为 WebGL 的继任者，底层依赖各个操作系统上更现代化的图形 API，提供了更低级的接口，这意味着开发者对 GPU 有更多直接控制以及更少的驱动资源消耗，渲染计算一视同仁。目前已经可以在 Chrome/Edge/Safari 的预览版本中使用它。

WebGPU 在 Shader 语言的选择上抛弃了 WebGL 使用的 GLSL，转向新的 WGSL。在我们关心的计算相关特性上，它提供了 storage/workgroupBarrier 同步方法。有了这些特性，一些算法就可以移植到 Web 端了，例如单源最短路径等其他图算法的 CUDA 开源实现，笔者尝试用 WGSL 实现它。

目前一个更成熟的实践是，Apache TVM 社区加入了 WebAssembly 和 WebGPU 后端支持。在 MacOS 上可以获得和直接本地运行 native metal 几乎一样的效率。

总之在可预见的未来，这无疑是 Web 端 GPGPU 的最佳选择。

原文链接
本文为阿里云原创内容，未经允许不得转载。