如何使用 TS 语法写 Compute Shader - antvis/g-webgl-compute GitHub Wiki

问题背景

很多可并行算法适合放在 GPU 中执行，但对于前端开发者来说，迁移这些算法成本很高，需要学习 WebGL/WebGPU API 以及 Shader 语法。并且 GLSL 语言本身存在多个版本（WebGL 1 使用 GLSL 1.0，WebGPU 使用 GLSL 4.5），在兼容性场景下还得考虑语法层面的差异。

为了降低前端开发者的学习成本，我们决定使用一种类 TS 语言，通过 Pegjs 生成 Parser，将源码转成 AST 后根据不同目标（WebGL/WebGPU）输出不同版本的 GLSL 代码。

对于前端开发者来说，只需要了解该语言的一些限制就能轻松上手编写 Shader 代码了。

语言特性

语法限制

受制于 GLSL 的语法限制：

  • 无法使用 JS 中的原生类型例如 String Date RegExp 等等
  • 无法使用 JS 中的 Math 方法，请使用对应的 GLSL 原生函数（后续会考虑提供转译，例如将 Math.sin() 转译成 sin()），详见原生函数
  • 无法使用解构。
  • 无法使用箭头函数，后续可能会支持。
  • 循环无法使用 for in。循环长度必须为定长。
  • 单文件，不支持类似 ESModule 引用依赖，import 语法仅用于引用工具方法。未来可能会提供简单的 Shader 模块化功能。

DataType

GLSL 是一种类 C 的强类型语言，因此 TS 开发者会很熟悉。需要注意的是两者的原生类型有很大差异，例如 TS 中的 number 类型会细分成标量、向量、矩阵等。因此我们只能使用 GLSL 的原生类型，包括：

标量

  • float
  • int
  • bool

向量

  • vec2 每个元素都是 float
  • vec3
  • vec4
  • ivec2 每个元素都是 int
  • ivec3
  • ivec4
  • bvec2 每个元素都是 bool
  • bvec3
  • bvec4

矩阵

  • mat3 每个元素都是 float
  • mat4

swizzling

在 GLSL 中有一种“特殊”的向量操作，通过 rgba/xyzw 可以代替下标访问向量中的元素：

const v: vec4 = [1, 2, 3, 4];

// v.r/v.x 等价于 a[0]
// v.g/v.y 等价于 a[1]
// v.b/v.z 等价于 a[2]
// v.a/v.w 等价于 a[3]

const b = a.rrr;
// -> vec3 b = vec3(1, 1, 1);

原生函数

来自 GLSLangSpec.4.30 第八节。

首先需要介绍 component-wise 的概念。大部分内置的三角函数、数学函数都是针对标量（float、int、uint）进行，但也可以传入向量，此时函数会作用于向量的每一个分量，例如：

max(1.0, 2.0) // 返回 float 2.0
max(int(1.0), int(2.0)) // 返回 int 1
max(uint(1.0), uint(2.0)) // 返回 uint 2
max(vec2(1.0, 2.0), vec2(2.0, 1.0)) // 返回 vec2(2.0, 2.0)

由于考虑到 WebGL 1 的兼容性，部分不支持的函数就不列出来了。兼容性支持可以查看对应函数说明下方的支持表。

三角函数，全部都是 component-wise：

指数函数，全部都是 component-wise：

常用数学函数，全部都是 component-wise：

uniform 变量

uniform 也是 Shader 中的一种变量类型，和常量不同，通常通过在 CPU 侧计算完成后传值，例如相机矩阵。 当然我们在 Shader 中也不需要修改这些变量值。在声明时有以下注意点：

  • 作为类属性声明
  • 需要声明类型
  • 不可以在此直接定义值，运行时从 CPU 侧传入，如何传值详见 ComputePipeline API
class MyProgram {
  @in
  param1: float;

  @in
  param2: vec3;
}

// -> WebGL 1 GLSL 1.0
// uniform float param1;
// uniform vec3 param2;

// -> WebGPU GLSL 4.5
// layout(std140, set = 0, binding = 1) uniform Params {
//   float param1;
//   vec3 param2;
// } params;

常量

通过在全局作用域使用全大写变量名，就可以声明一个常量，后续在自定义函数和 main 函数中都可以引用：

const CONST = 100;
// -> #define CONST 100

但是有一类特殊的常量，只能在运行时确定，因此无法像上面一样直接写在 Shader 中，但又没法以 uniform 变量形式传入。 例如 GLSL 1.0 中的循环变量只能和常量比较：

const loopLength;
// -> uniform float loopLength;

for (let i = 0; i < loopLength; i++)
// 报错:
// Loop index cannot be compared with non-constant expression

面对这种情况，我们可以写成全大写，然后像 uniform 一样在运行时传值：

const LOOP_LENGTH;
// -> #define LOOP_LENGTH 100

for (let i = 0; i < LOOP_LENGTH; i++)

声明函数

需要注意两点：

  • 只能通过类方法定义，不可以在一个方法内声明另一个函数
  • 一定要声明返回值和参数类型
class MyProgram {
  myFunc(p1: float, p2: float): float {
    //...
  }
}
// ->
// float myFunc(float p1, float p2) {...}

主函数

在 GLSL 中只能有一个 main 函数作为程序入口，为了和其他自定义函数区分开，我们使用 main 类方法装饰器描述：

class MyProgram {
  @main
  compute() {
   //...
  }
}
// -> void main() {...}

引用工具方法

我们未来会内置一些工具方法，例如 random() noise() 等，开发者可以通过 import 引入，例如：

import { noise } from 'g-webgpu';

class MyProgram {
  @main
  compute() {
    noise();
  }
}

类型推导

在声明变量时，可以尽量写明类型，以减少自动推导成本：

const a = [true, true, true];
// 会自动推导成 bvec3 a = bvec3(true);
// 也可写上类型：
// const a: bvec3 = [true, true, true];

但为了提高开发者的效率，我们针对以下情况做了自动类型推导，此时就可以省略类型的声明。

类型比较与运算

GLSL 在比较、运算时需要遵循类型的匹配，例如：

float a = 2.0;
if (a > 1) {}
// 报错，因为 float 不能和 int 比较，需要写成：
// if (a > 1.0) {}

为了减轻前端开发者的负担，我们会进行类型推导与转换，因此不必担心下面的代码会编译失败：

const a = 2.0;
if (a > 1) {}
// ->
float a = 2.0;
if (a > float(1)) {}

在运算时有些类型可以自动转换，例如 float 与 vec3 相加时我们会自动推导结果的类型：

const a = 1;
const b = [1, 2, 3];
const c = a + b; // 可以不手动声明 c 的类型为 vec3
// ->
float a = 1;
vec3 b = vec3(1.0, 2.0, 3.0);
vec3 c = a + b;
内置函数返回值类型推导

https://stackoverflow.com/questions/12085403/whats-the-logic-for-determining-a-min-max-vector-in-glsl https://www.khronos.org/registry/OpenGL/specs/gl/GLSLangSpec.4.30.pdf

另外很多 GLSL 内置函数都是可以接受多种数据类型的参数的，例如：

max(1.0, 2.0) // 返回 float 2.0
max(int(1.0), int(2.0)) // 返回 int 1
max(uint(1.0), uint(2.0)) // 返回 uint 2
max(vec2(1.0, 2.0), vec2(2.0, 1.0)) // 返回 vec2(2.0, 2.0)

这就给我们的返回值类型自动推导带来了困难，我们希望用户可以直接写：

const a = max(a, b);
// -> float a = max(a, b); 如果 ab 为 float
// -> int a = max(a, b); 如果 ab 为 int
swizzling

https://www.khronos.org/opengl/wiki/Data_Type_(GLSL)#Swizzling

const a = [1, 2, 3];
const b = a.rg;
// ->
vec3 a = vec3(1.0, 2.0, 3.0);
vec2 b = a.rg;
必须声明类型的情况

目前针对先声明再赋值的情况还做不到自动类型推导，因此不带初始值的变量声明必须带上类型：

const a; // 目前还无法自动推断类型
a = [1, 1, 1];
// 需要声明类型
const a: vec3;
a = [1, 1, 1];

后续针对这种情况会考虑更新 AST（通过赋值语句向前查找未声明类型的变量，更新该变量的类型）后二次生成代码。

示例[WIP]