Partículas GPGPU com TSL & WebGPU

Starter pack

Nesta lição, criaremos centenas de milhares de partículas flutuantes para renderizar modelos 3D e texto 3D usando Three Shading Language (TSL) e WebGPU.

Em vez de usar faces, usamos toneladas de partículas, permitindo uma transição suave entre diferentes modelos.

Um modelo 3D de uma raposa, um livro e texto 3D renderizados com partículas GPGPU! 🚀

Sistema de Partículas GPGPU

Antes de mergulharmos no código, vamos reservar um momento para entender o que é GPGPU e como ele pode ser usado no Three.js.

O que é GPGPU?

GPGPU (Computação de Propósito Geral em Unidades de Processamento Gráfico) é uma técnica que aproveita o poder de processamento paralelo das GPUs para realizar cálculos geralmente manejados pela CPU.

No Three.js, GPGPU é frequentemente usado para simulações em tempo real, sistemas de partículas e física ao armazenar e atualizar dados em texturas em vez de depender de cálculos restritos à CPU.

Esta técnica permite que os shaders tenham capacidades de memória e computação, permitindo a realização de cálculos complexos e o armazenamento de resultados em texturas sem a necessidade de intervenção da CPU.

Isso possibilita cálculos em grande escala e altamente eficientes diretamente na GPU.

Graças ao TSL, o processo de criação de simulações GPGPU torna-se muito mais fácil e intuitivo. Com nós de armazenamento e buffer combinados a funções de computação, podemos criar simulações complexas com código mínimo.

Aqui estão algumas ideias de projetos nos quais o GPGPU pode ser usado:

Hora de mudar da teoria para a prática! Vamos criar um sistema de partículas GPGPU usando TSL e WebGPU.

Sistema de Partículas

O starter pack é um modelo pronto para WebGPU baseado na implementação da lição WebGPU/TSL.

GPGPU particles starter pack

Vamos substituir o pink mesh por um novo componente chamado GPGPUParticles. Crie um novo arquivo chamado GPGPUParticles.jsx na pasta src/components e adicione o seguinte 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 });

Nada novo aqui, estamos criando um componente GPGPUParticles que usa um Sprite com um SpriteNodeMaterial para renderizar as partículas.

O benefício de usar um Sprite em vez de InstancedMesh é que ele é mais leve e vem com um efeito de billboard por padrão.

Vamos adicionar o componente GPGPUParticles ao 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 nos livrar dos componentes mesh e environment.

White sprite particles

Podemos ver um quadrado no meio da tela, estas são as partículas de white sprite. Todas na mesma posição.

É hora de configurar nosso sistema de partículas!

Buffer / Armazenamento / Instanced Array

Para a nossa simulação GPGPU, precisamos que nossas partículas memorizem suas position, velocity, age, e color sem usar a CPU.

Algumas coisas não exigirão que armazenemos dados. Podemos calcular a color com base na age combinada a uniforms. E podemos gerar a velocity aleatoriamente usando um valor seed fixo.

Mas para a position, como a posição alvo pode evoluir, precisamos armazená-la em um buffer. O mesmo para a age, queremos lidar com o ciclo de vida das partículas na GPU.

Para armazenar dados na GPU, podemos usar o storage node. Ele nos permite armazenar grandes quantidades de dados estruturados que podem ser atualizados na GPU.

Para usá-lo com o mínimo de código, utilizaremos a função TSL InstancedArray que se baseia no storage node.

Esta parte dos nodes do Three.js ainda não está documentada, é mergulhando nos exemplos e no código-fonte que conseguimos entender como funciona.

Vamos preparar nosso buffer no useMemo onde colocamos nossos 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 é uma função TSL que cria um buffer do tamanho e tipo especificados.

O mesmo código usando o storage node ficaria assim:

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

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

Com esses buffers, podemos armazenar a position e age de cada partícula e atualizá-los na GPU.

Para acessar os dados nos buffers, podemos usar .element(index) para obter o valor no índice especificado.

No nosso caso, usaremos o instancedIndex de cada partícula para acessar os dados nos 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 é uma função TSL embutida que retorna o índice da instância atual sendo processada.

Isso nos permite acessar os dados nos buffers para cada partícula.

Não precisaremos disso para este projeto, mas ao poder acessar os dados de outra instância, podemos criar interações complexas entre partículas. Por exemplo, poderíamos criar um bando de pássaros que se seguem.

Cálculo Inicial

Para configurar a posição e a idade das partículas, precisamos criar uma função de compute que será executada na GPU no início da simulação.

Para criar uma função compute com TSL, precisamos usar o nó Fn, chamá-lo e usar o método compute que ele retorna com o 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);

    // init 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);

    // ...
  }, []);

  // ...
};

// ...

Criamos uma função computeInit que atribui aos nossos buffers valores aleatórios.

A função randValue não existe, precisamos criá-la.

As funções à nossa disposição são:

  • hash(seed): Para gerar um valor aleatório com base em uma seed entre 0 e 1.
  • range(min, max): Para gerar um valor aleatório entre min e max.

Mais informações no Wiki Three.js Shading Language.

Mas a função range define um atributo e armazena o valor dele. Não é o que queremos.

Vamos criar uma função randValue que retornará um valor aleatório entre min e max com base em uma seed:

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 }) => {
  // ...
};
// ...

A função randValue recebe valores min, max e seed, e retorna um valor aleatório entre min e max baseado na seed.

/*#__PURE__*/ é um comentário usado para tree-shaking. Ele informa ao bundler para remover a função se ela não for usada. Mais detalhes aqui.

Agora precisamos chamar nossa função computeInit. Esta é uma tarefa para o renderer. Vamos importá-lo com useThree e chamá-lo logo após sua declaração:

// ...
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 podermos visualizá-lo, precisamos mudar o positionNode do SpriteNodeMaterial para usar os buffers spawnPosition e offsetPosition.

// ...

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

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

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

  // ...
};

// ...

Definimos o positionNode para a soma dos vetores spawnPosition e offsetPosition.

Está funcionando? Vamos conferir!

Partículas com posições aleatórias totalmente brancas

Mayday! Está tudo branco! ⬜️

Afastar um pouco?

Partículas com posições aleatórias afastadas

Ufa, conseguimos ver as partículas, elas são apenas grandes demais e pintaram a tela inteira! 😮‍💨

Vamos corrigir isso definindo o scaleNode com um valor aleatório:

// ...
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>
    </>
  );
};

// ...

Neste cenário, podemos usar a função range para gerar um valor aleatório entre 0.001 e 0.01.

Perfeito, temos nossas partículas com diferentes tamanhos e posições! 🎉

Está um pouco estático, entretanto, precisamos adicionar algum movimento a elas.

Atualizar cálculo

Como fizemos para a função de inicialização de cálculo, vamos criar uma função de atualização de cálculo que será executada em cada frame.

Nesta função, iremos atualizar a posição e a idade das 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.