Uso de GPU en JavaScript para mejorar el rendimiento

Actualmente la mayoría de los ordenadores cuentan tanto con CPUs como GPUs. Hasta hace poco las GPUs, a pesar de ser unos procesadores cada vez más potentes, solamente se utilizaban en juegos y otras aplicaciones que requerían realizar pesados cálculos gráficos. Actualmente la aparición de librerías como CUDA permite el uso de estos procesadores en aplicaciones de propósito general, permitiendo aprovechar la gran capacidad de paralización es estos procesadores. El uso de GPU en JavaScript se puede conseguir gracias al uso de librerías como GPU.js, gracias a la cual se puede usar estos procesadores tanto en el navegador como en Node.

GPU.js

La librería GPU.js permite llevar a cabo computación de propósito general en GPU (GPGPU, General purpose computing on GPUs) con JavaScript, pudiendo acelerar la ejecución de código tanto en los navegadores como en Node. GPU.js traduce automáticamente el código JavaScript para que este pueda ser ejecutado de forma eficiente en la GPU. En el caso de que no exista un GPU en el equipo el programa seguirá funcionando normalmente en la CPU.

GPU.js se puede instalar en Node tanto a través de npm como de yarn y en el navegador solamente requiere la importación de un archivo JavaScript.

Multiplicación de matrices con la GPU en JavaScript

Una tarea en la que las GPU funcionan mejor que las CPU es la multiplicación de matrices, ya que las operaciones necesarias se pueden ejecutar de forma paralela fácilmente. Así se puede escribir una función para multiplicar dos matrices cuadradas en JavaScript y su equivalente con GPU.js

import { GPU } from 'gpu.js';

function multiplyMatrix(a, b) {
    const result = Array(a.length).fill([]).map(() => Array(a.length).fill(0));

    for (let i = 0; i < a.length; i++) {
        for (let j = 0; j < a.length; j++) {
            for (let k = 0; k < a.length; k++) {
                result[i][j] += a[i][k] * b[k][j];
            }
        }
    }

    return result;
}

function multiplyMatrixGpu(a, b) {
    const gpu = new GPU();
    
    const multiplyMatrix = gpu.createKernel(function (a, b, matrixSize) {
        let sum = 0;

        for (let i = 0; i < matrixSize; i++) {
            sum += a[this.thread.y][i] * b[i][this.thread.x];
        }

        return sum;
    }).setOutput([a.length, a.length]);

    return multiplyMatrix(a, b, a.length);
}

const A = [[1, 2], [3, 4]];
const B = [[1, 1], [0, 1]];

console.log(multiplyMatrix(A, B));
console.log(multiplyMatrixGpu(A, B, 2));

[ [ 1, 3 ], [ 3, 7 ] ]
[ Float32Array(2) [ 1, 3 ], Float32Array(2) [ 3, 7 ] ]

La función en JavaScript pura es posiblemente la más fácil de entender, simplemente se ha implementado el proceso para multiplicar dos matrices.

En el caso de la versión creada para usar GPU.js es algo más complicado. Inicialmente se tiene que crear un objeto GPU y dentro de este se crea un kernel que tiene como parámetro una función en la que se implementa el proceso de cálculo. Lo primero que se puede notar es que en la función usada en esta caso se han eliminado los dos bucles externos, los cuales han sido reemplazados por this.thread.x y this.thread.y. Esto es así porque el kernel ya itera sobre la matriz que se define con la propiedad setOutput(), por lo que solamente se tiene que definir la función más interna del proceso.

Una vez hecho esto se puede comprobar cómo el resultado es el correcto en ambas implementaciones. Quizás la única diferencia es que en la GPU se han usado un tipo de dato float de 32 bits, mientras que en la GPU el tipo de dato es el number de JavaScript, esto es float de 64 bits. Una precisión que puede no ser suficiente en algunos casos.

Comparación del rendimiento de la CPU y GPU en JavaScript

Ahora que sabemos que las funciones realizar el cálculo correctamente se puede comparar el tiempo que tarda una implementación y la otra. Para lo que se puede crear dos matrices cuadradas de 1024 elementos y comparar el tiempo que tarda en llevar a cabo cada una de las funciones. Para lo que se puede usar el siguiente código.

const size = 1024;
const A = Array(size).fill([]).map(() => Array(size).fill(0).map(() => Math.random()));
const B = Array(size).fill([]).map(() => Array(size).fill(0).map(() => Math.random()));

console.time("CPU");
multiplyMatrix(A, B);
console.timeEnd("CPU");

console.time("GPU");
multiplyMatrixGpu(A, B, size);
console.timeEnd("GPU");

CPU: 9.692s
GPU: 458.771ms

En este caso se puede apreciar que la implementación con la CPU ha tardado más de 9 segundos frente a los 0,5 segundos de la GPU. Lo que supone casi un factor 20 de mejora. A pesar de que la GPU usada en la prueba es una Intel Iris Plus, una tarjeta gráfica de gama baja. Posiblemente al usar una tarjeta más potente se puede mejorar este rendimiento aún más.

Conclusiones

En esta ocasión hemos visto una librería interesante con la que es posible ejecutar código JavaScript y TypeScript en la GPU de nuestros equipos. Algo que puede suponer importantes mejoras de rendimiento, en el ejemplo de la entrada se ha observado un casi un factor con una GPU de gama baja. Si las tareas que necesitamos llevar a cabo pueden aprovechar la ventaja de estos procesadores es una buena idea darle una oportunidad a esta librería.

Imagen de Nana Dua en Pixabay