第二十四章：计算着色器(Compute Shader)：GPU通用计算，实现简单的N体模拟

说到OpenGL ES，大家第一反应肯定是画三角形、贴纹理、做特效。但你知道吗？GPU能干的事远不止渲染。计算着色器（Compute Shader）就是专门干这个的——让GPU帮你算数据，而不是画图。

我个人觉得，计算着色器是OpenGL ES里最被低估的功能。很多人学了半年OpenGL，都不知道还有这玩意儿。今天我们就拿一个经典的N体模拟来开刀，看看怎么用GPU做通用计算。

什么是计算着色器？

说白了，计算着色器就是一个跑在GPU上的小程序。它不关心渲染管线，不关心帧缓冲，只管一件事：算数。

你给它一堆数据，它用成百上千个线程同时处理。这就是所谓的GPU通用计算（GPGPU）。

我在项目中遇到过这样一个场景：需要实时计算几千个粒子的物理运动。如果用CPU算，帧率直接掉到个位数。换成计算着色器，轻轻松松60帧。这就是GPU并行计算的力量。

计算着色器的工作模型

计算着色器引入了一个新概念——工作组（Work Group）。你可以把它想象成一个线程块。

每个工作组里有一堆线程，它们共享局部内存。工作组之间是独立的，可以并行执行。

每个线程都有一个全局ID，通过这个ID可以知道自己要处理哪部分数据。比如粒子模拟中，线程ID就是粒子的索引。

N体模拟的核心逻辑

N体模拟，就是计算N个粒子之间的引力。每个粒子都受到其他所有粒子的引力影响。

公式很简单：

F = G * m1 * m2 / r²

但问题是，N个粒子就要算N²次。当N=1000时，就是100万次计算。CPU扛不住，但GPU可以。

实现步骤

我们一步步来。先看计算着色器的代码，再看怎么在OpenGL ES中调用它。

1. 计算着色器代码

#version 320 es
layout(local_size_x = 256) in;

// 粒子数据：位置和速度
layout(std430, binding = 0) buffer PosBuffer {
    vec4 positions[];
};

layout(std430, binding = 1) buffer VelBuffer {
    vec4 velocities[];
};

uniform float deltaTime;
uniform float softeningSq;  // 防止除零

void main() {
    uint i = gl_GlobalInvocationID.x;
    uint N = gl_NumWorkGroups.x * gl_WorkGroupSize.x;
    
    if (i >= N) return;
    
    vec3 pos_i = positions[i].xyz;
    vec3 acc = vec3(0.0);
    
    // 计算其他粒子对i的引力
    for (uint j = 0; j < N; j++) {
        if (i == j) continue;
        
        vec3 pos_j = positions[j].xyz;
        vec3 diff = pos_j - pos_i;
        float distSq = dot(diff, diff) + softeningSq;
        float invDist = inversesqrt(distSq);
        float invDist3 = invDist * invDist * invDist;
        
        acc += diff * invDist3;
    }
    
    // 更新速度
    velocities[i].xyz += acc * deltaTime;
    // 更新位置
    positions[i].xyz += velocities[i].xyz * deltaTime;
}

注意看，每个线程只处理一个粒子。它遍历所有其他粒子，累加引力。这就是典型的N²算法。

2. 创建计算着色器程序

在OpenGL ES中，计算着色器的创建流程和普通着色器一样：

// 创建着色器
GLuint computeShader = glCreateShader(GL_COMPUTE_SHADER);
glShaderSource(computeShader, 1, &computeSource, NULL);
glCompileShader(computeShader);

// 检查编译状态
GLint success;
glGetShaderiv(computeShader, GL_COMPILE_STATUS, &success);
if (!success) {
    char infoLog[512];
    glGetShaderInfoLog(computeShader, 512, NULL, infoLog);
    // 处理错误
}

// 创建程序并链接
GLuint computeProgram = glCreateProgram();
glAttachShader(computeProgram, computeShader);
glLinkProgram(computeProgram);

// 清理
glDeleteShader(computeShader);

3. 创建SSBO（着色器存储缓冲对象）

计算着色器不能直接用VBO，得用SSBO。SSBO可以在着色器之间共享数据。

// 创建位置缓冲
GLuint posSSBO;
glGenBuffers(1, &posSSBO);
glBindBuffer(GL_SHADER_STORAGE_BUFFER, posSSBO);
glBufferData(GL_SHADER_STORAGE_BUFFER, 
             N * sizeof(vec4), 
             initialPositions, 
             GL_DYNAMIC_DRAW);
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, posSSBO);

// 创建速度缓冲
GLuint velSSBO;
glGenBuffers(1, &velSSBO);
glBindBuffer(GL_SHADER_STORAGE_BUFFER, velSSBO);
glBufferData(GL_SHADER_STORAGE_BUFFER, 
             N * sizeof(vec4), 
             initialVelocities, 
             GL_DYNAMIC_DRAW);
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, velSSBO);

这里有个细节：glBindBufferBase的第二个参数是binding索引，要和着色器里的binding = 0对应上。

4. 调度计算

一切准备就绪后，就可以让GPU干活了：

glUseProgram(computeProgram);

// 设置uniform
glUniform1f(glGetUniformLocation(computeProgram, "deltaTime"), 0.016f);
glUniform1f(glGetUniformLocation(computeProgram, "softeningSq"), 0.01f);

// 调度计算
// 假设有1024个粒子，工作组大小256，需要4个工作组
GLuint numGroups = (N + 255) / 256;
glDispatchCompute(numGroups, 1, 1);

// 确保计算完成
glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

5. 渲染结果

计算完成后，SSBO里已经是最新的粒子位置了。怎么把它画出来？

最简单的方式：把SSBO绑定到顶点属性上，然后直接draw。

// 绑定位置SSBO到顶点属性
glBindBuffer(GL_ARRAY_BUFFER, posSSBO);
glVertexAttribPointer(0, 4, GL_FLOAT, GL_FALSE, 0, 0);
glEnableVertexAttribArray(0);

// 绘制点
glDrawArrays(GL_POINTS, 0, N);

你看，计算和渲染用的是同一份数据。这就是计算着色器的威力——数据不用来回拷贝。

性能优化技巧

N体模拟虽然简单，但想跑得快还是有点门道的。分享几个我实际用过的优化：

优化手段	说明	效果
共享内存	把粒子数据加载到共享内存，减少全局内存访问	提升2-3倍
分块计算	把粒子分成小块，每块只和附近块交互	适合大规模模拟
精度控制	用mediump代替highp	提升30%左右
减少分支	避免if语句，用数学方式处理边界	提升10-20%

我曾经做过一个测试：1024个粒子，用共享内存优化后，帧率从45fps提升到了120fps。效果非常明显。

共享内存版本示例

这里给一个使用共享内存的优化版本片段：

shared vec3 sharedPos[256];  // 工作组大小256

void main() {
    uint i = gl_GlobalInvocationID.x;
    uint localId = gl_LocalInvocationID.x;
    uint N = gl_NumWorkGroups.x * gl_WorkGroupSize.x;
    
    vec3 pos_i = positions[i].xyz;
    vec3 acc = vec3(0.0);
    
    // 分块处理
    for (uint block = 0; block < gl_NumWorkGroups.x; block++) {
        // 加载当前块到共享内存
        uint idx = block * gl_WorkGroupSize.x + localId;
        sharedPos[localId] = positions[idx].xyz;
        memoryBarrierShared();
        barrier();
        
        // 计算当前块内粒子的引力
        for (uint j = 0; j < gl_WorkGroupSize.x; j++) {
            uint globalJ = block * gl_WorkGroupSize.x + j;
            if (i == globalJ) continue;
            
            vec3 diff = sharedPos[j] - pos_i;
            float distSq = dot(diff, diff) + 0.01;
            float invDist = inversesqrt(distSq);
            acc += diff * (invDist * invDist * invDist);
        }
        
        memoryBarrierShared();
        barrier();
    }
    
    velocities[i].xyz += acc * 0.016;
    positions[i].xyz += velocities[i].xyz * 0.016;
}

调试技巧

调试计算着色器比调试普通着色器麻烦。因为你看不到中间结果。我的经验是：

• 先用小规模数据测试，比如4个粒子
• 把计算结果回读到CPU，和CPU计算结果对比
• 用glGetError()检查每一步的错误

回读数据的代码：

glBindBuffer(GL_SHADER_STORAGE_BUFFER, posSSBO);
vec4* result = (vec4*)glMapBufferRange(
    GL_SHADER_STORAGE_BUFFER,
    0,
    N * sizeof(vec4),
    GL_MAP_READ_BIT
);

// 检查结果
for (int i = 0; i < N; i++) {
    // 和CPU计算结果对比
}

glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);

嗯，这里要注意：glMapBufferRange会阻塞CPU，直到GPU完成计算。所以别在每帧都调用，只在调试时用。

总结

计算着色器让GPU不再只是画图工具。N体模拟只是冰山一角，你还可以用它做物理模拟、图像处理、甚至机器学习推理。

我个人觉得，掌握计算着色器是区分普通OpenGL开发者和高级图形工程师的分水岭。它打开了GPU通用计算的大门，让你能用GPU解决渲染之外的问题。

最后提醒一句：不同设备对计算着色器的支持程度不同。开发时记得检查GL_MAX_COMPUTE_WORK_GROUP_SIZE等参数，别想当然。