实现向量加法 - antvis/g-webgl-compute GitHub Wiki

先来看一个简单的例子，如何实现两个向量的加法。类似 tensorflow 中 tf.add 的效果。在 tensorflow 中是通过 API 提供的，用户完全不需要关心内部后端（WebGL/WASM/WebGPU）的实现：

// https://js.tensorflow.org/api/latest/#add
const a = tf.tensor1d([1, 2, 3, 4]);
const b = tf.tensor1d([10, 20, 30, 40]);

a.add(b).print();  // or tf.add(a, b)

下面我们通过两步完成该计算任务的创建：

创建 Compute Pipeline
用 TS 语法编写 Compute Shader

通过这个例子也能看出相比 API 组合调用，用户通过 TS 来编写自己的并行计算任务显然能满足更多的场景。不过相同的是，我们的 Parser 会生成适合不同目标平台的 Shader 代码，用户只需要知道这份代码最终会运行在 GPU 侧。

创建 Compute Pipeline

首先调用 API 完成 Compute Pipeline 的创建：

// 获取 HTMLCanvasElement
const canvas = document.getElementById('application');

const world = new World(canvas, {
  engineOptions: {
    supportCompute: true,
  },
});

const compute = world.createComputePipeline({
  shader: `
    //...
  `, // 下一节的 Shader 文本
  dispatch: [1, 1, 1], // 线程网格
  onCompleted: (result) => {
    console.log(result); // [2, 4, 6, 8, 10, 12, 14, 16]
    world.destroy(); // 计算完成后销毁相关 GPU 资源
  },
});

// 绑定输入到 Compute Shader 中的两个参数
world.setBinding(compute, 'vectorA', [1, 2, 3, 4, 5, 6, 7, 8]);
world.setBinding(compute, 'vectorB', [1, 2, 3, 4, 5, 6, 7, 8]);

编写计算任务代码

完整 Shader 代码如下，前端开发者应该会相当熟悉，稍后我们会详细解释每一行代码的含义：

import { numthreads, in, out, main, globalInvocationID } from 'g-webgpu';

@numthreads(8, 1, 1)
class Add2Vectors {
  @in @out
  vectorA: float[];

  @in
  vectorB: float[];

  sum(a: float, b: float): float {
    return a + b;
  }

  @main
  compute() {
    // 获取当前线程处理的数据
    const a = this.vectorA[globalInvocationID.x];
    const b = this.vectorB[globalInvocationID.x];
  
    // 输出当前线程处理完毕的数据，即两个向量相加后的结果
    this.vectorA[globalInvocationID.x] = this.sum(a, b);
  }
}

还记得上一节在创建 Compute Pipeline 时我们绑定的两个参数吗？在 Shader 的 main 函数中我们在参数中声明。这里使用了 TS 中的参数装饰器语法：

@in @out vectorA: vec4[] // 表示既作为输入又作为输出

在并行计算中，最重要的概念就是数据并行，即分配多个线程运行同一份程序处理不同的数据（个人理解）。类似 CUDA 中的 kernel 函数（GPU.js 也沿用了这种叫法）。

首先在 API 中我们声明了线程网格的大小，在这个简单的例子中整个网格只有一个线程组：

dispatch: [1, 1, 1], // 线程网格

然后通过函数装饰器语法声明每个线程组的大小，这里包含了 8 个线程：

@numthreads(8, 1, 1)

⚠️ numthreads 中 x * y * z 的值最大为 1024。

最后对于每一个线程来说，必须要在程序中获取自己关心的部分数据。因此需要知道自身在线程网格中的位置，通过 import 语法可以获取 dispatchThreadID，这是一个网格中的三维坐标。在我们的例子中，网格中包含 1 * 1 * 1 个线程组，每个线程组包含 8 * 1 * 1 个线程，因此对于第一个线程，dispatchThreadID.x 在执行时值为 0，第二个线程执行时这个值为 1：

import { globalInvocationID } from 'g-webgpu';

const a = vectorA[globalInvocationID.x];

计算函数主体十分简单，就不多做介绍了。

预编译 Shader 代码

目前在运行编译用户编写的 TS 代码到目标平台会消耗很多时间，这还是在我们选择了 Pegjs 这样相对较轻（相比 Antlr）的 Parser 的情况下。在绝大部分场景下，用户都不会在运行时修改编译好的 Shader 代码，因此预编译是一个不错的选择。

我们提供了 API 便于用户获取编译好的 Shader 代码。这样用户可以在开发过程中试运行成功后保存下编译结果，在实际运行代码中直接传入：

// 试运行代码
const compute = world.createComputePipeline({
  shader: `...`,
  onCompleted: (result) => {
    // 获取 Shader 的编译结果，用户可以输出到 console 中并保存
    console.log(world.getPrecompiledBundle(compute));
  },
});

// 实际运行代码，使用试运行过程中获得的编译结果创建
const compute = world.createComputePipeline({
  precompiled: true,
  shader: compiledBundle, // JSON 字符串，包含 WebGL & WebGPU 所需的 Shader 代码以及上下文环境
});

事实上 GPU.js 也是这么做的： https://github.com/gpujs/gpu.js#precompiled-and-lighter-weight-kernels

后续改进

目前我们有两个优化方向：

提升开发体验
WebGPU 输出多份数据

VS Code 插件

毕竟我们只是借用了 TS 的语法，在原生数据类型上还是有很大差别。我们可以通过提供 VS Code 插件（例如我们可以叫 g 语言，处理 .g 文件）完成语法高亮（通过 TextMate 实现）、自动提示等语言特性。

输出多份数据

受限于 WebGL 的 GPGPU 实现（不支持 Compute Shader 因此我们只能通过常规渲染管线模拟），目前我们只能输出一份结果。即只能声明一个 @out。

WebGPU 因为原生支持 Compute Shader 并没有这个限制。如果想在 WebGL 中也支持，可能需要考虑生成两份 Shader 代码实现输出两次的目的。

总结

如果你对数据类型、内置函数还不熟悉，可以参考：如何使用 TS 语法写 Compute Shader。

API 说明

另外你可能觉得这个计算任务太过简单，下面我们会介绍一个更为复杂的并行算法：Fruchterman 布局算法。