Partículas GPGPU con TSL y WebGPU

Starter pack

En esta lección, crearemos cientos de miles de partículas flotantes para representar modelos 3D y texto 3D usando Three Shading Language (TSL) y WebGPU.

En lugar de usar caras, usamos toneladas de partículas, lo que nos permite hacer transiciones suaves entre diferentes modelos.

¡Un modelo 3D de un zorro, un libro y texto 3D representado con partículas GPGPU! 🚀

Sistema de Partículas GPGPU

Antes de sumergirnos en el código, tomémonos un momento para entender qué es GPGPU y cómo puede ser utilizado en Three.js.

¿Qué es GPGPU?

GPGPU (General-Purpose computing on Graphics Processing Units) es una técnica que aprovecha el poder de procesamiento paralelo de las GPUs para realizar cálculos que generalmente maneja la CPU.

En Three.js, GPGPU se utiliza a menudo para simulaciones en tiempo real, sistemas de partículas y física al almacenar y actualizar datos en texturas en lugar de depender de cálculos vinculados a la CPU.

Esta técnica permite que los shaders tengan capacidades de memoria y cálculo, habilitándolos para realizar cálculos complejos y almacenar resultados en texturas sin necesidad de intervención de la CPU.

Esto permite realizar cálculos a gran escala altamente eficientes directamente en la GPU.

Gracias a TSL, el proceso para crear simulaciones GPGPU es mucho más fácil e intuitivo. Con nodos de almacenamiento y buffer combinados con funciones de cálculo, podemos crear simulaciones complejas con un código mínimo.

Aquí hay algunas ideas de proyectos para los cuáles se puede usar GPGPU:

¡Hora de pasar de la teoría a la práctica! Vamos a crear un sistema de partículas GPGPU usando TSL y WebGPU.

Sistema de partículas

El paquete inicial es una plantilla WebGPU lista basada en la implementación de la lección WebGPU/TSL.

Paquete inicial de partículas GPGPU

Reemplacemos el mesh rosa con un nuevo componente llamado GPGPUParticles. Crea un nuevo archivo llamado GPGPUParticles.jsx en la carpeta src/components y añade el siguiente código:

import { extend } from "@react-three/fiber";
import { useMemo } from "react";
import { color, uniform } from "three/tsl";
import { AdditiveBlending, SpriteNodeMaterial } from "three/webgpu";

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  const { nodes, uniforms } = useMemo(() => {
    // uniforms
    const uniforms = {
      color: uniform(color("white")),
    };

    return {
      uniforms,
      nodes: {
        colorNode: uniforms.color,
      },
    };
  }, []);

  return (
    <>
      <sprite count={nbParticles}>
        <spriteNodeMaterial
          {...nodes}
          transparent
          depthWrite={false}
          blending={AdditiveBlending}
        />
      </sprite>
    </>
  );
};

extend({ SpriteNodeMaterial });

No hay nada nuevo aquí, estamos creando un componente GPGPUParticles que utiliza un Sprite con un SpriteNodeMaterial para renderizar las partículas.

El beneficio de usar un Sprite sobre un InstancedMesh es que es más ligero y viene con un efecto de billboard por defecto.

Agreguemos el componente GPGPUParticles al componente Experience:

import { OrbitControls } from "@react-three/drei";
import { GPGPUParticles } from "./GPGPUParticles";

export const Experience = () => {
  return (
    <>
      {/* <Environment preset="warehouse" /> */}
      <OrbitControls />
      <GPGPUParticles />
      {/* <mesh>
        <boxGeometry />
        <meshStandardMaterial color="hotpink" />
      </mesh> */}
    </>
  );
};

Podemos deshacernos de los componentes mesh y environment.

Partículas sprite blancas

Podemos ver un cuadrado en el medio de la pantalla, estas son las partículas sprite blancas. Todas en la misma posición.

¡Es momento de configurar nuestro sistema de partículas!

Buffer / Almacenamiento / Instanced Array

Para nuestra simulación GPGPU, necesitamos que nuestras partículas memoricen su posición, velocidad, edad y color sin usar la CPU.

Algunas cosas no nos requerirán almacenar datos. Podemos calcular el color basado en la edad combinada con uniforms. Y podemos generar la velocidad aleatoriamente usando un valor de seed fijo.

Pero para la posición, dado que la posición objetivo puede evolucionar, necesitamos almacenarla en un buffer. Lo mismo para la edad, queremos manejar el ciclo de vida de las partículas en el GPU.

Para almacenar datos en el GPU, podemos usar el storage node. Nos permite almacenar grandes cantidades de datos estructurados que pueden ser actualizados en el GPU.

Para usarlo con el mínimo código, utilizaremos la función TSL InstancedArray que se basa en el storage node.

Esta parte de Three.js nodes aún no está documentada, es sumergiéndose en los ejemplos y el código fuente como podemos entender cómo funciona.

Vamos a preparar nuestro buffer en el useMemo donde colocamos nuestros shader nodes:

// ...
import { instancedArray } from "three/tsl";

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  const { nodes, uniforms } = useMemo(() => {
    // uniforms
    const uniforms = {
      color: uniform(color("white")),
    };

    // buffers
    const spawnPositionsBuffer = instancedArray(nbParticles, "vec3");
    const offsetPositionsBuffer = instancedArray(nbParticles, "vec3");
    const agesBuffer = instancedArray(nbParticles, "float");

    return {
      uniforms,
      nodes: {
        colorNode: uniforms.color,
      },
    };
  }, []);

  // ...
};

// ...

instancedArray es una función TSL que crea un buffer del tamaño y tipo especificado.

El mismo código usando el storage node se vería así:

import { storage } from "three/tsl";
import { StorageInstancedBufferAttribute } from "three/webgpu";

const spawnPositionsBuffer = storage(
  new StorageInstancedBufferAttribute(nbParticles, 3),
  "vec3",
  nbParticles
);

Con estos buffers, podemos almacenar la posición y edad de cada partícula y actualizarlos en el GPU.

Para acceder a los datos en los buffers, podemos usar .element(index) para obtener el valor en el índice especificado.

En nuestro caso, usaremos el instancedIndex de cada partícula para acceder a los datos en los buffers:

// ...
import { instanceIndex } from "three/tsl";

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  const { nodes, uniforms } = useMemo(() => {
    // ...

    // buffers
    const spawnPositionsBuffer = instancedArray(nbParticles, "vec3");
    const offsetPositionsBuffer = instancedArray(nbParticles, "vec3");
    const agesBuffer = instancedArray(nbParticles, "float");

    const spawnPosition = spawnPositionsBuffer.element(instanceIndex);
    const offsetPosition = offsetPositionsBuffer.element(instanceIndex);
    const age = agesBuffer.element(instanceIndex);

    return {
      uniforms,
      nodes: {
        colorNode: uniforms.color,
      },
    };
  }, []);

  // ...
};

// ...

instanceIndex es una función TSL integrada que devuelve el índice de la instancia actual siendo procesada.

Esto nos permite acceder a los datos en los buffers para cada partícula.

No lo necesitaremos para este proyecto, pero al poder acceder a los datos de otra instancia, podemos crear interacciones complejas entre partículas. Por ejemplo, podríamos crear una bandada de pájaros que se sigan entre ellos.

Cálculo inicial

Para configurar la posición y edad de las partículas, necesitamos crear una función de cálculo que se ejecutará en la GPU al inicio de la simulación.

Para crear una función de cálculo con TSL, necesitamos usar el nodo Fn para llamarla, y utilizar el método compute que retorna con el número de partículas:

// ...
import { Fn } from "three/src/nodes/TSL.js";

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  const { nodes, uniforms } = useMemo(() => {
    // ...
    const spawnPosition = spawnPositionsBuffer.element(instanceIndex);
    const offsetPosition = offsetPositionsBuffer.element(instanceIndex);
    const age = agesBuffer.element(instanceIndex);

    // inicializar Fn
    const lifetime = randValue({ min: 0.1, max: 6, seed: 13 });

    const computeInit = Fn(() => {
      spawnPosition.assign(
        vec3(
          randValue({ min: -3, max: 3, seed: 0 }),
          randValue({ min: -3, max: 3, seed: 1 }),
          randValue({ min: -3, max: 3, seed: 2 })
        )
      );
      offsetPosition.assign(0);
      age.assign(randValue({ min: 0, max: lifetime, seed: 11 }));
    })().compute(nbParticles);

    // ...
  }, []);

  // ...
};

// ...

Creamos una función computeInit que asigna a nuestros buffers valores aleatorios.

La función randValue no existe, necesitamos crearla nosotros mismos.

Las funciones que tenemos a nuestra disposición son:

  • hash(seed): Para generar un valor aleatorio basado en una semilla entre 0 y 1.
  • range(min, max): Para generar un valor aleatorio entre min y max.

Más información en el Wiki de Three.js Shading Language.

Pero la función range define un atributo y almacena su valor. No es lo que queremos.

Vamos a crear una función randValue que devuelva un valor aleatorio entre min y max basado en una semilla:

import { hash } from "three/tsl";

const randValue = /*#__PURE__*/ Fn(({ min, max, seed = 42 }) => {
  return hash(instanceIndex.add(seed)).mul(max.sub(min)).add(min);
});

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  // ...
};
// ...

La función randValue toma valores min, max, y seed, y devuelve un valor aleatorio entre min y max basado en la semilla.

/*#__PURE__*/ es un comentario usado para tree-shaking. Indica al empaquetador que elimine la función si no se utiliza. Más detalles aquí.

Ahora necesitamos llamar a nuestra función computeInit. Esto es trabajo para el renderer. Vamos a importarlo con useThree y llamarlo justo después de su declaración:

// ...
import { useThree } from "@react-three/fiber";

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  const gl = useThree((state) => state.gl);

  const { nodes, uniforms } = useMemo(() => {
    // ...
    const computeInit = Fn(() => {
      // ...
    })().compute(nbParticles);

    gl.computeAsync(computeInit);

    // ...
  }, []);

  // ...
};

// ...

Para poder visualizarlo, necesitamos cambiar el positionNode del SpriteNodeMaterial para usar los buffers spawnPosition y offsetPosition.

// ...

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  // ...

  const { nodes, uniforms } = useMemo(() => {
    // ...

    return {
      uniforms,
      nodes: {
        positionNode: spawnPosition.add(offsetPosition),
        colorNode: uniforms.color,
      },
    };
  }, []);

  // ...
};

// ...

Establecemos el positionNode a la suma de los vectores spawnPosition y offsetPosition.

¿Funciona? ¡Vamos a comprobarlo!

Partículas con posiciones aleatorias completamente blancas

¡Mayday! ¡Todo es blanco! ⬜️

¿Hacemos un poco de zoom hacia afuera?

Partículas con posiciones aleatorias alejadas

Uf, ahora podemos ver las partículas, ¡simplemente son demasiado grandes y pintaron toda la pantalla! 😮‍💨

Arreglemos eso estableciendo el scaleNode con un valor aleatorio:

// ...
import { range } from "three/tsl";

// ...

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  // ...

  const { nodes, uniforms } = useMemo(() => {
    // ...

    const scale = vec3(range(0.001, 0.01));

    return {
      uniforms,
      nodes: {
        positionNode: spawnPosition.add(offsetPosition),
        colorNode: uniforms.color,
        scaleNode: scale,
      },
    };
  }, []);

  return (
    <>
      <sprite count={nbParticles}>
        <spriteNodeMaterial
          {...nodes}
          transparent
          depthWrite={false}
          blending={AdditiveBlending}
        />
      </sprite>
    </>
  );
};

// ...

En este escenario, podemos usar la función range para generar un valor aleatorio entre 0.001 y 0.01.

Perfecto, ¡tenemos nuestras partículas con diferentes tamaños y posiciones! 🎉

Sin embargo, es un poco estático, necesitamos añadirle algo de movimiento.

Actualizar computación

Como hicimos para la función de computación init, vamos a crear una función de actualización de computación que se ejecutará en cada frame.

En esta función, actualizaremos la posición y la edad de las partículas:

// ...
import { deltaTime, If } from "three/tsl";

export const GPGPUParticles = ({ nbParticles = 1000 }) => {
  // ...

  const { nodes, uniforms } = useMemo(() => {
    // ...

    const instanceSpeed = randValue({ min: 0.01, max: 0.05, seed: 12 });

    // update Fn
    const computeUpdate = Fn(() => {
      age.addAssign(deltaTime);

      If(age.greaterThan(lifetime), () => {
        age.assign(0);
        offsetPosition.assign(0);
      });

      offsetPosition.addAssign(vec3(instanceSpeed));
    })().compute(nbParticles);

    // ...
  }, []);

  // ...
};

// ...
Three.js logoReact logo

React Three Fiber: The Ultimate Guide to 3D Web Development

✨ You have reached the end of the preview ✨

Go to the next level with Three.js and React Three Fiber!

Get full access to this lesson and the complete course when you enroll:

  • 🔓 Full lesson videos with no limits
  • 💻 Access to the final source code
  • 🎓 Course progress tracking & completion
  • 💬 Invite to our private Discord community
Unlock the Full Course – Just $85

One-time payment. Lifetime updates included.