WebGL 计算着色器工作分配：深入探究 GPU 线程分配

WebGL 中的计算着色器提供了一种强大的方式，可以直接在 Web 浏览器中利用 GPU 的并行处理能力执行通用计算 (GPGPU) 任务。了解如何将工作分配给各个 GPU 线程对于编写高效、高性能的计算核函数至关重要。本文全面探讨了 WebGL 计算着色器中的工作分配，涵盖了基本概念、线程分配策略和优化技术。

理解计算着色器执行模型

在深入了解工作分配之前，让我们首先通过理解 WebGL 中的计算着色器执行模型来打下基础。该模型是分层的，由几个关键组件组成：

• 计算着色器 (Compute Shader): 在 GPU 上执行的程序，包含并行计算的逻辑。
• 工作组 (Workgroup): 一组协同执行并通过共享本地内存共享数据的工作项集合。可以将其视为一个执行整体任务一部分的工作团队。
• 工作项 (Work Item): 计算着色器的一个独立实例，代表一个 GPU 线程。每个工作项执行相同的着色器代码，但操作的数据可能不同。这是团队中的个体工作者。
• 全局调用 ID (Global Invocation ID): 在整个计算分派中，每个工作项的唯一标识符。
• 本地调用 ID (Local Invocation ID): 在其工作组内，每个工作项的唯一标识符。
• 工作组 ID (Workgroup ID): 在计算分派中，每个工作组的唯一标识符。

当你分派一个计算着色器时，你需要指定工作组网格的维度。这个网格定义了将创建多少个工作组以及每个工作组将包含多少个工作项。例如，dispatchCompute(16, 8, 4) 的分派将创建一个维度为 16x8x4 的三维工作组网格。 然后，每个工作组都会填充预定义数量的工作项。

配置工作组大小

工作组大小在计算着色器源代码中使用 layout 限定符定义：

            layout (local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
            

              
                Copy
              
            
          

此声明指定每个工作组将包含 8 * 8 * 1 = 64 个工作项。local_size_x、local_size_y 和 local_size_z 的值必须是常量表达式，并且通常是 2 的幂。最大工作组大小取决于硬件，可以使用 gl.getParameter(gl.MAX_COMPUTE_WORK_GROUP_INVOCATIONS) 进行查询。此外，工作组的各个维度也有限制，可以使用 gl.getParameter(gl.MAX_COMPUTE_WORK_GROUP_SIZE) 进行查询，该函数返回一个包含三个数字的数组，分别表示 X、Y 和 Z 维度的最大大小。

示例：查找最大工作组大小

            const maxWorkGroupInvocations = gl.getParameter(gl.MAX_COMPUTE_WORK_GROUP_INVOCATIONS);
const maxWorkGroupSize = gl.getParameter(gl.MAX_COMPUTE_WORK_GROUP_SIZE);

console.log("Maximum workgroup invocations: ", maxWorkGroupInvocations);
console.log("Maximum workgroup size: ", maxWorkGroupSize); // Output: [1024, 1024, 64]

            

              
                Copy
              
            
          

选择合适的工作组大小对于性能至关重要。较小的工作组可能无法充分利用 GPU 的并行性，而较大的工作组可能会超出硬件限制或导致低效的内存访问模式。通常，需要通过实验来确定特定计算核函数和目标硬件的最佳工作组大小。一个好的起点是尝试 2 的幂次的工作组大小（例如 4、8、16、32、64），并分析它们对性能的影响。

GPU 线程分配和全局调用 ID

当计算着色器被分派时，WebGL 实现负责将每个工作项分配给一个特定的 GPU 线程。每个工作项都由其全局调用 ID 唯一标识，这是一个 3D 向量，表示它在整个计算分派网格中的位置。在计算着色器中，可以使用内置的 GLSL 变量 gl_GlobalInvocationID 访问此 ID。

gl_GlobalInvocationID 是使用以下公式从 gl_WorkGroupID 和 gl_LocalInvocationID 计算得出的：

            gl_GlobalInvocationID = gl_WorkGroupID * gl_WorkGroupSize + gl_LocalInvocationID;
            

              
                Copy
              
            
          

其中 gl_WorkGroupSize 是在 layout 限定符中指定的工作组大小。此公式突出了工作组网格与单个工作项之间的关系。每个工作组被分配一个唯一的 ID (gl_WorkGroupID)，该工作组内的每个工作项被分配一个唯一的本地 ID (gl_LocalInvocationID)。然后通过组合这两个 ID 来计算全局 ID。

示例：访问全局调用 ID

            #version 450

layout (local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

layout (binding = 0) buffer DataBuffer {
    float data[];
} outputData;

void main() {
    uint index = gl_GlobalInvocationID.x + gl_GlobalInvocationID.y * gl_NumWorkGroups.x * gl_WorkGroupSize.x;
    outputData.data[index] = float(index);
}
            

              
                Copy
              
            
          

在这个例子中，每个工作项使用 gl_GlobalInvocationID 计算其在 outputData 缓冲区中的索引。这是在大数据集上分配工作的常见模式。行 `uint index = gl_GlobalInvocationID.x + gl_GlobalInvocationID.y * gl_NumWorkGroups.x * gl_WorkGroupSize.x;` 至关重要。让我们分解一下： * `gl_GlobalInvocationID.x` 提供工作项在全局网格中的 x 坐标。 * `gl_GlobalInvocationID.y` 提供工作项在全局网格中的 y 坐标。 * `gl_NumWorkGroups.x` 提供 x 维度中工作组的总数。 * `gl_WorkGroupSize.x` 提供每个工作组 x 维度中的工作项数量。 这些值共同使每个工作项能够计算其在扁平化输出数据数组中的唯一索引。如果你正在处理 3D 数据结构，你还需要将 `gl_GlobalInvocationID.z`、`gl_NumWorkGroups.y`、`gl_WorkGroupSize.y`、`gl_NumWorkGroups.z` 和 `gl_WorkGroupSize.z` 纳入索引计算中。

内存访问模式和合并内存访问

工作项访问内存的方式会显著影响性能。理想情况下，工作组内的各个工作项应该访问连续的内存位置。这被称为合并内存访问，它允许 GPU 高效地大块抓取数据。当内存访问是分散或不连续时，GPU 可能需要执行多个较小的内存事务，这可能导致性能瓶颈。

为了实现合并内存访问，仔细考虑内存中数据的布局以及工作项如何分配给数据元素非常重要。例如，在处理 2D 图像时，将工作项分配给同一行中的相邻像素可以实现合并内存访问。

示例：图像处理的合并内存访问

            #version 450

layout (local_size_x = 16, local_size_y = 16, local_size_z = 1) in;

layout (binding = 0) uniform sampler2D inputImage;
layout (binding = 1) writeonly uniform image2D outputImage;

void main() {
    ivec2 pixelCoord = ivec2(gl_GlobalInvocationID.xy);
    vec4 pixelColor = texture(inputImage, vec2(pixelCoord) / textureSize(inputImage, 0));

    // Perform some image processing operation (e.g., grayscale conversion)
    float gray = dot(pixelColor.rgb, vec3(0.299, 0.587, 0.114));
    vec4 outputColor = vec4(gray, gray, gray, pixelColor.a);

    imageStore(outputImage, pixelCoord, outputColor);
}
            

              
                Copy
              
            
          

在这个例子中，每个工作项处理图像中的一个像素。由于工作组大小为 16x16，同一工作组中的相邻工作项将处理同一行中的相邻像素。这在从 inputImage 读取和写入 outputImage 时促进了合并内存访问。

然而，考虑一下如果你转置图像数据，或者你以列优先而非行优先顺序访问像素会发生什么。你很可能会看到性能显著下降，因为相邻的工作项将访问非连续的内存位置。

共享本地内存

共享本地内存，也称为本地共享内存 (LSM)，是工作组内所有工作项共享的一小块快速内存区域。它可以通过缓存频繁访问的数据或促进同一工作组内工作项之间的通信来提高性能。共享本地内存使用 GLSL 中的 shared 关键字声明。

示例：使用共享本地内存进行数据归约

            #version 450

layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

layout (binding = 0) buffer InputBuffer {
    float inputData[];
} inputBuffer;

layout (binding = 1) buffer OutputBuffer {
    float outputData[];
} outputBuffer;

shared float localSum[gl_WorkGroupSize.x];

void main() {
    uint localId = gl_LocalInvocationID.x;
    uint globalId = gl_GlobalInvocationID.x;

    localSum[localId] = inputBuffer.inputData[globalId];

    barrier(); // Wait for all work items to write to shared memory

    // Perform reduction within the workgroup
    for (uint i = gl_WorkGroupSize.x / 2; i > 0; i /= 2) {
        if (localId < i) {
            localSum[localId] += localSum[localId + i];
        }
        barrier(); // Wait for all work items to complete the reduction step
    }

    // Write the final sum to the output buffer
    if (localId == 0) {
        outputBuffer.outputData[gl_WorkGroupID.x] = localSum[0];
    }
}
            

              
                Copy