前端是可以接触到GPU的 于是也是可以使用GPU的计算能力的 但是前端使用GPU计算和常规的GPU计算相差还是非常大的 咱从两个方向来分享前端的WebGL计算

0. GPU.JS

好在已经有人把相关的内容以及封装成库了 gpu.js 这个库还在开发当中的 前段时间才把核心库升级到WebGL2 目前也是还没有一个很稳定的正式版 所以很多功能都是有欠缺的 但是作为一些简单的计算绰绰有余

因为js是单线程的 所以并不适合处理CPU密集型的程序 但是GPU是有非常高的并行数 所以GPU计算的基本思想就是把计算任务分拆成N多个线程任务 每个线程返回一个结果 然后吧每个线程的结果再汇总 有点像mapreduce

咱先直接看一段代码来先感受一下

const gpu = new GPU();
const myFunc = gpu.createKernel(function() {
    return this.thread.x;
}, { output: { x: 100 } });

myFunc();  // [0,1,2,3,...98,99]

上面这个例子只取了x方向上数量100个 在函数里this.thread.x可以获取到当前执行这个函数的线程x维度上的地址

要注意的是 这个执行的函数里的语句语法是有限制的 具体的解决方法可以参考文档

0.0 矩阵平方

说了这么多 咱来写一个实际有用的东西 矩阵平方 C=AB 为了方便起见 A和B都是相同大小的方块矩阵

矩阵的乘法咱就不回顾了 只提一个非常关键的地方 就是关于矩阵C中 C(i,j)的值

这个很关键 先吧这句翻译成程序语句

// 单独封装成函数因为之后会重复使用到
let dimensions = 100; // 假设长度为100
function getValue(ma, mb, x, y) {
    let sum = 0;
    for (let i = 0; i < dimensions; i++) {
        sum += ma[x][i] * mb[i][y];
    }
    return sum;
}

于是开始写

const matrixMultiplyGPU = new GPU().createKernel(function(a, b) {
        // 计算每个坐标的值, 注意这里的x,y顺序是反的
        return getValue(a, b, this.thread.y, this.thread.x);  
    }).setOutput({ x: dimensions, y: dimensions })
    .setConstants({ dimensions }) // 添加外部变量
    .setFunctions([getValue]); // 添加外部函数, 添加外部函数的时候有非常多bug

这里注意一点 线程在两个维度上的坐标顺序是反的

然后咱们随机构造矩阵 然后开始计算

function createMatrix(dims, fn) {
    let matrix = [];
    for (let i = 0; i < dims; i++) {
        matrix[i] = [];
        for (let j = 0; j < dims; j++) {
            matrix[i][j] = fn(i, j);
        }
    }
    return matrix;
}
const randomMatrix = createMatrix(dimensions, () => Math.floor(Math.random() * 50));

const retMatrix = matrixMultiplyGPU(randomMatrix, randomMatrix);

为了方便起见 两个相同的矩阵平方好了 拿出笔和纸 用个两三个维度的值验证了一下 发现没有问题 但是这个效率和CPU的相比 有多大的优势呢

0.1 CPU矩阵平方

然后咱们利用上面的式子写一个matrixMultiplyCPU吧 因为有上面的铺垫 所以这个就非常简单了

const matrixMultiplyCPU = function(a, b) {
    return createMatrix(dimensions, getValue.bind(null, a, b));
};

0.2 验证

验证的也很简单 直接看一下运行时间对比就可以了

console.log("dimensions", dimensions)

console.time("matrixMultiplyGPU");
matrixMultiplyGPU(randomMatrix, randomMatrix);
console.timeEnd("matrixMultiplyGPU");

console.time("matrixMultiplyCPU");
matrixMultiplyCPU(randomMatrix, randomMatrix);
console.timeEnd("matrixMultiplyCPU");

电脑配置是

• CPU - 2.6 GHz Intel Core i5
• GPU - Intel Iris 1536 MB
• Browser - Chrome Version 67.0.3396.87 (64-bit)

然后我再看下另外一个设备的结果

• CPU - 3.2 GHz Intel Core i5-6500
• GPU - NVIDIA GeForce GTX 1060 3GB
• Browser - Chrome Version 67.0.3396.87 (64-bit)

后面这个的配置就略高一些 果然整体的数据比上面的好一些

从上面的数据来看 当矩阵维度较小的时候CPU的耗时比较短 因为当数据量少的时候 时间主要花费在将数据传输上 从内存传到显存 当数据量大了之后 计算时间开始增大 GPU才开始显现出优势

你问能不能再大一点 咱试过了 再使用2048维度的计算的时候 整个浏览器已经奔溃了 说明这个库还是有一些欠缺的

1 WebGL

上面讲了个如何使用gpu.js这个库来进行简单的gpu计算 虽然简单易用 但是本身的局限也很多 这个库也不是非常完善 有待改进 那咱就从原理开始 来自己搞一个吧 当然 并不是指实现一个这个的通用的库 而是使用相关原理 完成一个是用GPU计算的demo 当然还是矩阵的乘法

前端使用GPU的能力是通过webgl实现的 更加广义的理解的可以认为是通过canvas来实现的 canvas估计对大多数前端来说并不陌生 canvas的每个像素的颜色可以有RGBA四个维度表示 每个维度范围为0-255 既8位 把RGBA表示成数值的话 那每个像素可以存32位 刚刚好就是一个int或者uint 这就是前端使用gpu计算最为核心的一点 利用像素来存储数据

WebGL涉及到的东西就非常多了 不过大致的流程是这样的

其中两个vertex shader和fragment shader为两个GLSL代码片段 当前是基于OpenGL ES 3.0 分别以像素为单位处理坐标信息和颜色数据 里面有一些全局变量和特殊类型需要注意 处理起来挺恶心的

坐标的范围是(-1, 1) 为了方便计算 我这边在vertex shader做了一个坐标转换

in vec4 g_pos;
out vec2 v_pos;

void main() {
    // 输入坐标转换  画布坐标映射到数组坐标传入fragment
    // [(1, 0), (1, 1)]
    // [(0, 0), (0, 1)]
    float curX = (1. + g_pos.x) / 2.;
    float curY = (1. + g_pos.y) / 2.;
    v_pos = vec2(curX, curY);

    // 绘制坐标不变
    gl_Position = g_pos;
}

将左下脚作为(0, 0)的一个好处是 最终读取每个像素点颜色值的时候是从这个位置开始逐行向上遍历的 得到结果的直接就可以作为最终的结果

另外 之后用到的纹理也是以这个位置作为(0, 0) 这样处理的时候可以统一一个坐标系

像素RGBA每位的范围是(0, 1) 最终输出一个RGBA颜色数据作为该像素点的颜色值

所以可以认为一个像素点里的计算 就是一个线程 至于同时会有多少像素在渲染 这个咱就不知道了 当所有像素都绘制完成之后 就是整个画布绘制完成

以原点的距离作为透明的绘制出来的结果如下, 具体咋做的先不解释

1.0 输入

WebGL的变量输入有两种方式 使用uniform关键字和使用纹理

在处理的时候要注意变量类型 由于众所周知的原因 不能直接使用js的数值进行计算 需要使用类型数组和webGL进行数据传递

1.0.0 使用 uniform 关键字传入数据

一个例子:

js中:

    function getUniformLoc(name) {
        const loc = gl.getUniformLocation(program, name);
        if (loc == null) throw `getUniformLoc ${name} err`;
        return loc;
    }
    // 获取变量
    const uniLoc = getUniformLoc("i_matrixA");
    // 设置变量值  传入一个float类型的数组
    gl.uniform1fv(uniLoc, new Float32Array([1, 2, 3, 4]));

uniform1fv()是一系列方法, 因为没有重载, 所以按照不同的类型命名方法, 支持float int max等几乎所有webgl内置类型

shader中:

    uniform float i_matrixA[4];

但是这个方法有一些问题

1. 数组长度受限, 可以使用gl.getParameter(gl.MAX_FRAGMENT_UNIFORM_VECTORS)或者gl.getParameter(gl.MAX_VERTEX_UNIFORM_VECTORS)获取数组长度上限 本人实测值为1024
2. OpenGL ES 3.0不支持多维数组, 这个问题将在下个版本中得到支持, 目前无解

所以想要处理其他的数据 就需要有其他的办法了

1.0.1 使用 Texture 纹理传入数据

纹理就是另外的图案 实际使用中的含义就是把一副图叠加到另外一副图上 设计师还有PS上可能了解到多
因为图都是由像素构成的 所以可以用纹理来传入大量的数据 就不限1024了

纹理的定义有点复杂 纹理的大小非常苛刻 只能是2^n * 2^n的大小 但是数据不可能是固定的 所以 这里有个纹理进行伸缩的过程 有一大波复杂的设置

一个例子:

    function initTexture(index, tSampler, pixels) {
        // 一大波复杂的设置

        const texture = gl.createTexture();
        gl.activeTexture(gl[`TEXTURE${index}`]);
        gl.bindTexture(gl.TEXTURE_2D, texture);

        // 设置纹理 缩放渐变相关
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR);
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);

        // 将像素传入绘制到纹理上
        gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, 3, 3, 0,
            gl.RGBA, gl.UNSIGNED_BYTE, pixels, 0);

        // 将指向纹理的特殊变量使用uniform绑定到程序里
        gl.uniform1i(getUniformLoc(tSampler), index);
    }
    // 传入9个像素点, 设置中定义为3*3的纹理
    const colorMap = new Uint32Array([
        0xFF0000FF, 0x00FF00FF, 0x0000FFFF,
        0xFFFF00FF, 0xFF00FFFF, 0x00FFFFFF,
        0x000000FF, 0xFFFFFFFF, 0xF0F0F0FF,
    ]);
    // 注意, Uint32Array转换为Uint8Array时的数值顺序
    const RGBAMap = new Uint8Array(colorMap.buffer);
    initTexture(0, "samplerA", RGBAMap);

shader 中:

    uniform sampler2D samplerA;
    void main() {
        // 获取当前像素在纹理上对应点的颜色值
        vec4 color = texture(samplerA, v_pos);
        // 由于类型数组类型转换导致数值顺序变化, 需要手动修正
        o_result = color.abgr;
    }

直接将纹理的颜色输出到对应像素 直接绘制到canvas上 左下为(0,0)的颜色 即0xFF0000FF红色

1.1 输出

输出的方式只有一种 读取canvas上各个像素的值 读取像素的方法也有很多参数和重载 这里使用RGBA模式读取这四个维度的值

    // 绘制
    gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);

    let picBuf = new ArrayBuffer(canvas.width * canvas.width * 4);
    let picU8 = new Uint8Array(picBuf);
    let picU32 = new Uint32Array(picBuf);
    // 读取到Uint8Arrays类型数组中
    gl.readPixels(0, 0, canvas.width, canvas.width, gl.RGBA, gl.UNSIGNED_BYTE, picU8);
    console.log(picU32);

注意 readPixels方法必须和drawArrays方法在同一个执行队列中同步执行 否则会无法读取到数据

读取是顺序从左下开始 逐行向上 并且由于每一行中x坐标在增加 所以对应数组坐标时需要交换x与y的位置

1.2 矩阵平方

好了 搞了这么多 已经吧基本的输入输出搞定了 咱来开始试一下矩阵平方吧

因为uniform对数组长度有限制 所以只展示使用Texture的方式

const int U_LENGTH = CANVAS_SIZE;
// 计算坐标对应数组位置时偏移量, 防止出现在边界上
const float U_TEXTURE_POS_FIX = .5 / float(U_LENGTH);

// 
in vec2 v_pos;
// 传入指向两个纹理的特殊变量
uniform sampler2D samplerA;
uniform sampler2D samplerB;

out vec4 o_result;

// 整数与rgba转化
vec4 int2rgba(const int i) {
    vec4 v4;
    v4.r = float(i >> 24 & 0xFF) / 255.;
    v4.g = float(i >> 16 & 0xFF) / 255.;
    v4.b = float(i >>  8 & 0xFF) / 255.;
    v4.a = float(i >>  0 & 0xFF) / 255.;
    return v4;
}


// 整数与rgba转化
int rgba2int(const vec4 v) {
    int r = int(v.r * 255.) << 24;
    int g = int(v.g * 255.) << 16;
    int b = int(v.b * 255.) << 8;
    int a = int(v.a * 255.) << 0;
    return r + g + b + a;
}

// 反转rgba
vec4 reverse(const vec4 v){
    return v.abgr;
}

// 得到某坐标上的代表的数值
int getMaxtrixValue(const sampler2D sampler, const float x, const float y){
    vec4 pixel = texture(sampler, vec2(x, y));
    return rgba2int(reverse(pixel));
}

void main() {
    int sum = 0;
    float textPos = 0.0;
    for (int i = 0; i < U_LENGTH; i++) {
        // 取需要相乘的值的坐标对应的在纹理上的位置
        // 防止取到边界 增加一个位置偏移量 以取到像素块正中央
        textPos = U_TEXTURE_POS_FIX + float(i) / float(U_LENGTH);

        // 计算乘积 这里注意要交换x与y
        sum += getMaxtrixValue(samplerA, v_pos.x, textPos) * getMaxtrixValue(samplerB, textPos, v_pos.y);
    }
    // 将结果转成颜色值绘制到画布
    o_result = reverse(int2rgba(sum));
}

1.3 CPU矩阵平方

很明显 上面这个矩阵平方使用的矩阵都是一维数组 所以咱们也要实现一个基于CPU的一维数组相乘

    const matrixMultiplyCPU = function(ma, mb) {
        return createMatrix(dimensions, function(x, y) {
            let sum = 0;
            for (let i = 0; i < dimensions; i++) {
                sum += ma[x * dimensions + i] * mb[i * dimensions + y];
            }
            return sum;
        });
    };

1.4 验证

最终找了三台设备来验证

• CPU - 2.6 GHz Intel Core i5
• GPU - Intel Iris 1536 MB
• Browser - Chrome Version 67.0.3396.87 (64-bit)

然后我再看下另外一个设备的结果

• CPU - 3.2 GHz Intel Core i5-6500
• GPU - NVIDIA GeForce GTX 1060 3GB
• Browser - Chrome Version 67.0.3396.87 (64-bit)

• CPU - 3.0 GHz Intel Core i5-8500
• GPU - Nvidia GeForce GTX 1080 8GB
• Browser - Chrome Version 67.0.3396.87 (64-bit)

            

上面这些数据还是有一些偏差的 不过几个之间的数量级应该还是有明显差距的

另外

之前说过 GPU计算的主要耗时是在数据传输上 那数据传输到底占了多大比重呢 咱用上面那个1060的机器跑了另外一组数据 将代码中的任务进行拆分 把输入texImage2D 绘制drawArrays 输出readPixels 三个部分分别计时 得到一个结果

1060机器:

核显机器:

看到这个结果还是有点点惊讶到 传入数据耗时不多 考虑到js在毫秒级计时不精确 可以认为计算过程可以忽略不计

而耗时大部分都花费在读取结果 毕竟1000x1000的矩阵 就需要读取一百万个像素点的值 还是比较费时的

那么问题来了 哪有没什么计算可以做到计算量大 输入量少 同时输出量也非常少呢 有啊 挖矿

文中以上的代码在这里可以看到 gputest2

本文来自网易实践者社区，经作者徐俊鹏授权发布。