WebGL 使用纹理存储数据

[xiaoiver]

问题背景

关联 #6

在我们目前的 TS-based Shader 语法中，允许用户声明不同格式的输入输出数据，例如：

@in
data1: float[]

@in
data2: vec4[]

@in
data3: image2D

在 WebGPU Compute Shader 中，我们可以通过 buffer 声明对应类型的数组：

// float 数组
layout(std430, set = 0, binding = 0) buffer GWebGPUBuffer0 {
  float data1[];
} gWebGPUBuffer0;

// vec4 数组
layout(std430, set = 0, binding = 1) buffer GWebGPUBuffer0 {
  vec4 data2[];
} gWebGPUBuffer0;

// 纹理
layout(set = 0, binding = 2) uniform texture2D data3;
layout(set = 0, binding = 3) uniform sampler data3Sampler;

但是考虑到兼容性，当我们在 WebGL Fragment Shader 中实现时，数据就只能通过纹理存取了：

uniform sampler2D data1;
uniform sampler2D data2;
uniform sampler2D data3;

这里就涉及到不同类型数据的存取方式，我们需要考虑两点：

1. 不同类型的数据如何存储到纹理中
2. 在 Shader 中自动生成存取方法

解决方案

数据存储到纹理

纹理的尺寸是有限制的：
https://stackoverflow.com/questions/29975743/is-it-possible-to-use-webgl-max-texture-size/29985986

16384 * 16384 * 4 (RGBA) * FLOAT = 4294967296 or 4GIG!!!

因此我们不能简单创建一个数据长度 * 1 的纹理，因为数据长度很有可能超出纹理宽高限制。
另外，如果用户声明的是一个 float 数组，我们也无法充分利用纹理中每个 texel 的存储空间，即 4 个 float 塞进一个 vec4 里。原因是我们用 texel 模拟多线程的概念，即执行运算逻辑的最小单位就是 texel，这种情况下确实会出现 GPU 内存的浪费。

计算步骤如下：

1. 根据数据类型计算纹理 texel 数目。例如 vec3 需要除以3，vec4 需要除以 4 等。
2. texel 数目开根号，使用正方形纹理存储
3. 将纹理尺寸以 uniform 传入 Shader，供后续自动生成的读取纹理数据方法使用

⚠️ 暂不考虑超出 4G 的场景，后续可以拆分成多个纹理解决

Shader 中线程组相关变量实现

由于我们通过输出纹理中的每一个 texel 模拟线程组中的线程，因此 GLSL 4.5 中内置的线程组相关变量就需要我们自己模拟实现了，这样用户在 Shader 中使用时才不会报错。
需要实现的内置变量如下：

1. gl_NumWorkGroups ivec3 dispatch 的线程工作组数目
2. gl_WorkGroupSize ivec3 Shader 内声明的每一个线程工作组包含的线程数
3. gl_WorkGroupID ivec3 当前线程工作组的索引。取值范围为 (0, 0, 0) 到 (gl_NumWorkGroups.x - 1, gl_NumWorkGroups.y - 1, gl_NumWorkGroups.z - 1) 之间。
4. gl_LocalInvocationID ivec3 当前线程在自己线程组中的索引。取值范围为 (0, 0, 0) 到 (gl_WorkGroupSize.x - 1, gl_WorkGroupSize.y - 1, gl_WorkGroupSize.z - 1) 之间。
5. gl_GlobalInvocationID ivec3 当前线程在全局线程组中的索引。计算方法为 gl_WorkGroupID * gl_WorkGroupSize + gl_LocalInvocationID
6. gl_LocalInvocationIndex int 当前线程在自己线程组中的一维索引，计算方法为 gl_LocalInvocationID.z * gl_WorkGroupSize.x * gl_WorkGroupSize.y + gl_LocalInvocationID.y * gl_WorkGroupSize.x + gl_LocalInvocationID.x

⚠️ WebGL 1 中并不支持 uvec3 数据类型，因此只能使用 ivec3。详见

首先最容易实现的是线程网格和线程组的尺寸，前者和用户调用 dispatch API 时保持一致，后者和用户在 Shader 中声明的 numthreads 保持一致：

ivec3 numWorkGroups = ivec3(${groupX}, ${groupY}, ${groupZ});    
ivec3 workGroupSize = ivec3(${localSizeX}, ${localSizeY}, ${localSizeZ});

之前我们说过输出纹理中每个 texel 对应一个 “线程”，因此根据当前 texel 的纹理坐标就可以计算出当前线程的全局索引，注意实际上并没有 gl_GlobalInvocationIndex 这个概念，但我们可以使用该值推断剩余的变量：

int globalInvocationIndex = int(floor(v_TexCoord.x * u_OutputTextureSize.x))
  + int(floor(v_TexCoord.y * u_OutputTextureSize.y)) * int(u_OutputTextureSize.x);

例如我们可以结合线程组尺寸计算出当前线程的线程组 ID：

int workGroupIDLength = globalInvocationIndex / (workGroupSize.x * workGroupSize.y * workGroupSize.z);
ivec3 workGroupID = ivec3(workGroupIDLength / numWorkGroups.y / numWorkGroups.z, workGroupIDLength / numWorkGroups.x / numWorkGroups.z, workGroupIDLength / numWorkGroups.x / numWorkGroups.y);

根据 gl_GlobalInvocationIndex 的计算公式可以反推出 gl_LocalInvocationID：

gl_LocalInvocationIndex  = gl_LocalInvocationID.z * gl_WorkGroupSize.x * gl_WorkGroupSize.y + gl_LocalInvocationID.y * 
  gl_WorkGroupSize.x + gl_LocalInvocationID.x.

最后很容易计算出 gl_GlobalInvocationID 和 gl_LocalInvocationIndex：

ivec3 globalInvocationID = workGroupID * workGroupSize + localInvocationID;
int localInvocationIndex = localInvocationID.z * workGroupSize.x * workGroupSize.y
                + localInvocationID.y * workGroupSize.x + localInvocationID.x;

自动生成纹理读取方法

对于每个 @in 的输入，例如：

@in 
vectorA: float[]

除了 sampler2D，我们还会传入该数据纹理的尺寸，类型为 vec2，命名规则就是纹理名 + Size：

uniform sampler2D vectorA;
uniform vec2 vectorASize;

对每个数据纹理，我们需要根据数据类型生成多个参数重载的读取方法。例如上面 vectorA 为 float[] 类型，我们的返回值肯定就是 float，相应的，读取纹理数据后也需要通过 swizzling 获取正确的类型：

float getDatavectorA(vec2 address2D) {
  return float(texture2D(vectorA, address2D).r);
}
float getDatavectorA(float address1D) {
  return getDatavectorA(addrTranslation_1Dto2D(address1D, vectorASize));
}
float getDatavectorA(int address1D) {
  return getDatavectorA(float(address1D));
}

其中 1D 地址转换成 2D，来自 GPU Gem2：

vec2 addrTranslation_1Dto2D(float address1D, vec2 texSize) {
  vec2 conv_const = vec2(1.0 / texSize.x, 1.0 / (texSize.x * texSize.y));
  vec2 normAddr2D = float(address1D) * conv_const;
  return vec2(fract(normAddr2D.x), normAddr2D.y);
}

调用纹理读写方法

在编译 TS 时，我们需要转译纹理读写语法。

遇到写纹理语法（通过 AST estree 节点类型判定）时，由于我们目前只支持输出到一个纹理，左值只需要替换成 gl_FragColor 即可：

this.vectorA[globalInvocationID.x] = globalInvocationID.x;
// -> gl_FragColor = 

而右值需要固定为 vec4 类型。不用担心输出多余的数据，例如这里 vectorA 的类型为 float[]，输出纹理中每个 texel rgba 每一个分量都存储了结果，但在最后输出时会被过滤掉，只保留第一个分量结果：

gl_FragColor = vec4(globalInvocationID.x);

遇到读纹理语法时，就可以调用上一节自动生成的纹理读取方法。

const a = this.vectorA[globalInvocationID.x];
// -> float a = getDatavectorA(globalInvocationID.x);

[zuiidea]

@xiaoiver 前面提到的 “数据存储到纹理” 在当前版本中有实现吗，支持自定义纹理的宽度或者高度吗