Simulando 20 milhões de partículas com JavaScript

 (dgerrells.com)
4 pontos por GN⁺ 2024-07-09 | 1 comentários | Compartilhar no WhatsApp
  • Partindo de uma simulação de partículas em JavaScript usando apenas CPU, o objetivo de 1 milhão de partículas a 60 fps no celular foi alcançado, e o exemplo final foi escalado até 20 milhões de partículas
  • O ponto central do desempenho foi evitar arrays de objetos e manter os dados das partículas em memória contígua com TypedArray e SharedArrayBuffer, dividindo o trabalho por núcleo de CPU com Web Workers
  • O gargalo real apareceu mais no desenho na tela e no acesso à memória do que no cálculo da posição das partículas, e escritas aleatórias no buffer de pixels causavam cache misses da CPU
  • O gargalo foi reduzido gradualmente passando por buffers de pixels por Worker, sincronização via postMessage, double buffering, uma grade de número de partículas por pixel e renderização de textura na GPU
  • A arquitetura final usa a CPU para simular e a GPU para renderizar uma textura de grade do tamanho da resolução da tela, processando 20 milhões de partículas a cerca de 20 fps em um Mac M1

Objetivo e abordagem básica

  • O desafio era simular 1 milhão de partículas a 60 fps em um celular usando apenas JavaScript puro e CPU, sem WebAssembly
  • Considerando que um simples array de objetos JavaScript teria limitações em desempenho single-core móvel e layout de memória, foi escolhida uma abordagem de tratar os dados como arrays contíguos
  • Em cálculos em grande escala, arrays de memória contígua e compacta, capazes de colocar o máximo de dados possível no cache da CPU, são importantes
  • Arrays e objetos comuns de JavaScript dificilmente garantem um layout tão denso quanto arrays ou structs de linguagens de baixo nível
  • TypedArray permite criar arrays de bytes contíguos de tamanho fixo, possibilitando um gerenciamento de memória relativamente baixo nível em JavaScript

Primeira implementação: SharedArrayBuffer e Web Worker

  • Os dados das partículas são armazenados colocando uma visão Float32Array sobre um SharedArrayBuffer
    • Uma partícula é representada por quatro valores de ponto flutuante de 32 bits: x, y, dx, dy
    • O buffer plano é indexado com stride = 4 e byte_stride = stride * 4
  • A simulação começa com uma atualização simples de posição em cada frame, executando x += dx, y += dy
  • O multithreading é composto por Web Workers e SharedArrayBuffer
    • A thread principal divide e entrega blocos de partículas para cada Worker
    • Cada Worker lê e escreve apenas o bloco atribuído a ele, evitando que várias threads escrevam na mesma posição de memória
  • Para usar SharedArrayBuffer no navegador, os seguintes headers são necessários
    • Cross-Origin-Embedder-Policy: require-corp
    • Cross-Origin-Opener-Policy: same-origin
  • A sincronização inicial usa um SharedArrayBuffer para sinalização separado
    • A thread principal escreve um sinal de execução para o Worker
    • O Worker verifica o sinal com setInterval(..., 1) e registra o estado de pronto quando termina o trabalho
  • A renderização é tratada como um rasterizador baseado em CPU usando ImageData
    • Cada partícula é desenhada como um pixel da tela
    • Partículas fora da tela são ignoradas para evitar acesso fora dos limites do array
    • O resultado é desenhado no Canvas com putImageData
  • No profiling inicial, o trabalho dos Workers ficava na casa de poucos ms, e a maior parte do tempo era gasta na renderização na thread principal
  • A memória era alocada uma vez no início, e nenhum objeto novo era criado no loop de simulação, gerando praticamente nenhum garbage

Segunda implementação: entrada e interação parecida com gravidade

  • Os dados de entrada são adicionados ao SharedArrayBuffer existente para o estado da simulação
    • Armazenam delta time, x/y do mouse, estado de toque e largura/altura da tela
    • O estado de toque é um valor próximo de booleano, mas usa 4 bytes
  • Ao tocar na tela, foi adicionado um efeito em que as partículas são atraídas para o ponto de toque
    • A forma 1 / r² de uma fórmula real de gravidade foi usada como referência, mas modificada visando uma interação divertida em vez de uma simulação física precisa
    • Foi adicionada uma atenuação equivalente a atrito para impedir que as partículas acelerem indefinidamente e voem para fora da tela
  • A cor da renderização foi modificada para mudar conforme a posição da partícula
    • Os valores x / width e y / height são usados para ajustar os canais RGB
    • Como armazenar a cor por partícula aumentaria o tamanho das partículas, foi usada uma cor baseada na posição para reduzir o tamanho dos dados
  • Foi possível experimentar até 2 milhões de partículas, e o número de partículas passou a poder ser alterado pelo parâmetro de query count da URL
  • Com 10 milhões de partículas, o framerate caiu, e o profiling mostrou que a parte lenta era a escrita de pixels em ImageData, mais do que os Workers
  • Ao escrever apenas os canais azul e alfa em vez de todo o RGB, em alguns ambientes o tempo de renderização caiu de cerca de 80–90 ms para 16–18 ms
    • Essa diferença variava conforme a CPU e foi maior em um desktop Ryzen do que em um MacBook Air M1

Terceira implementação: escrita do buffer de pixels nos Workers

  • Para usar mais núcleos de CPU, os Workers passaram a assumir não apenas a simulação, mas também o desenho dos pixels
  • Para evitar que vários Workers escrevessem simultaneamente no mesmo buffer de pixels, foi criado um buffer de pixels por Worker
    • É usado um SharedArrayBuffer de tamanho CPU_CORES * window.innerWidth * window.innerHeight * 3
    • Cada Worker possui um buffer RGB da tela inteira
  • A thread principal soma todos os buffers de pixels de cada Worker e depois grava em ImageData
  • Nesta versão, a tela passou a piscar
    • Enquanto um Worker limpava seu buffer de pixels com fill(0), a thread principal estava lendo o mesmo buffer
    • Antes, a renderização também podia ler dados antigos, mas a diferença entre frames não era perceptível; agora a operação de limpeza do Worker aparecia como piscadas em partes da tela
  • As possíveis soluções eram sincronização baseada em Atomics, esperar a conclusão dos Workers com postMessage ou double buffering

Quarta implementação: sincronização baseada em postMessage

  • Para eliminar as piscadas, o array de sinalização foi removido e os Workers passaram a avisar a conclusão via postMessage
  • A thread principal gerencia o valor de activeWorkers
    • Cada Worker envia uma mensagem ao terminar, e activeWorkers é decrementado
    • Quando todos os Workers terminam, a renderização é feita e o próximo requestAnimationFrame é solicitado
  • Com essa abordagem, as piscadas foram resolvidas
  • O tempo por frame não melhorou muito e podia até piorar
    • Isso acontece porque a thread principal não faz nada enquanto espera os Workers terminarem
    • Com 2 milhões de partículas, havia situações em que cerca de 7 ms poderiam ser economizados se os Workers não esperassem a renderização da thread principal
  • Para reduzir esse gargalo, foi aplicado o double buffering, comum em drivers gráficos

Quinta implementação: double buffering e gargalo de cache

  • Double buffering consiste em criar dois buffers de pixels e alternar o buffer ativo a cada frame
    • Enquanto a thread principal desenha um buffer, os Workers preparam o próximo frame no outro buffer
  • Depois da aplicação, os Workers deixaram de esperar a renderização da thread principal, e alguns frames que antes passavam de 50 ms ficaram mais próximos de 16 ms
  • O código que continuava lento eram as três linhas que somavam valores aos canais RGB do buffer de pixels
  • A causa era que, embora os dados estivessem armazenados de forma contígua, a ordem de acesso não era contígua
    • O pixelIndex calculado a partir de x, y da partícula não é alinhado com a ordem do array de partículas
    • A cada iteração, o loop salta para uma posição arbitrária do buffer de pixels, criando um padrão próximo de acesso aleatório
  • Esse padrão de acesso causa cache misses quando os dados de partículas e pixels não cabem todos no cache
  • Um exemplo de cálculo para iPhone é o seguinte
    • Fórmula: screen width * screen height * bytes per pixel + bytes per particle * total particles / number of workers
    • 2532 * 1170 * 3 + 16 * 2,000,000 / 5
    • O resultado passa de 15 MB, grande demais para caber no cache L1
  • Nesta versão, em um chip M1, a simulação de 2 milhões de partículas leva cerca de 4 ms, o acesso aos dados de pixels cerca de 3,5 ms, e a escrita no buffer de pixels cerca de 7 ms
  • Em celulares, foi possível manter 1 milhão de partículas a 60 fps, mas em desktops com mais de 20 threads a thread principal gastava 30 ms a cada frame somando mais de 20 buffers de pixels

Sexta implementação: força para voltar à posição inicial

  • O novo efeito faz com que as partículas sejam puxadas de volta com mais força quanto mais se afastam da posição inicial
  • Para isso, foram adicionados à partícula dois valores de posição inicial, sx, sy
    • Cada partícula passa a ter seis floats: x, y, dx, dy, sx, sy
    • O particleStride aumenta para 6
  • Ao código do Worker foi adicionada uma função de cálculo de força, e um objeto de cache é usado para evitar geração de garbage
  • O resultado criou um efeito visual parecido com gelatina ou fluido, com partículas formando redemoinhos como papel ou tecido amassado
  • Nesse processo surgiu a ideia de não armazenar o buffer de pixels por Worker, mas trocar para uma grade que conta apenas o número de partículas por pixel

Sétima implementação: grade de número de partículas por pixel

  • Armazenar apenas o número de partículas em cada posição de pixel, em vez dos valores RGB, reduz a memória para 1/3
  • Como a cor é calculada a partir da posição x, y, não de dados por partícula, a renderização pode ser feita apenas com a contagem por pixel
  • A implementação foi alterada para que todos os Workers compartilhem uma única grade de contagem de partículas
    • O trabalho da thread principal de somar buffers por Worker desaparece
    • O problema de o custo de soma na renderização aumentar conforme cresce o número de Workers é resolvido
  • A nova fórmula de cache é a seguinte
    • screen width * screen height + bytes per particle * total particles / number of workers
    • No iPhone: 2532 * 1170 + 16 * 2,000,000 / 5
    • Cerca de 9,3 MB, aproximadamente 30% menor que antes, mas ainda grande para o cache L1
  • A renderização percorre todos os pixels, lê o número de partículas naquele pixel, calcula os valores RGB e depois inicializa a grade com fill(0)
  • Após a simulação, o Worker calcula o pCountIndex das partículas que estão dentro da tela e executa activeGrid[pCountIndex]++
  • Essa mudança ajudou especialmente quando havia muitos Workers e permitiu usar todos os 24 threads em um desktop
  • Ainda assim, em alguns casos o MacBook Air M1 foi mais rápido que um desktop Ryzen 9 3900X
    • O Ryzen 9 3900X tem 64 MB de cache L3, mas o L1 é de 64 KB
    • O M1 de primeira geração tem 320 KB de cache L1, dos quais 128 KB são para dados e 192 KB para instruções
    • Um cache L1 de dados maior parece ter ajudado a reduzir a latência de memória

Limites da otimização somente com CPU

  • Como resultado de várias otimizações, foi obtido cerca de 2x de melhoria de velocidade em relação à primeira versão multithread
  • O objetivo de processar 1 milhão de partículas a 60 fps em um celular, com CPU e apenas JavaScript, foi alcançado
  • O resultado geral não é especialmente impressionante, e a expectativa é que uma linguagem compilada fosse 10x mais rápida
  • Haveria espaço para ficar mais rápido se fosse possível usar instruções SIMD no tight loop
  • Na etapa seguinte, a simulação permanece na CPU e apenas a renderização é movida para um pipeline tradicional de GPU para comparação

20 milhões de partículas e renderização na GPU

  • A primeira comparação com renderização na GPU usou GPU instancing do threejs para desenhar um plane/quad por partícula
  • Essa abordagem teve desempenho menor do que o esperado
    • O instancing do threejs precisa atualizar uma matriz para cada posição de partícula, e esse trabalho é feito em uma única thread da CPU
    • Os dados de transformação das partículas precisam ser enviados para a GPU a cada frame, e essa transferência é lenta
  • Ao simular 4 milhões de partículas em um chip M1, 80% do tempo do frame era gasto enviando dados para a GPU
  • No desktop, a transferência de dados representava 30%, mas ainda era o maior gargalo
  • A alternativa é usar a grade de número de partículas criada antes como textura e renderizar um quad de tela cheia
    • A quantidade de dados enviada para a GPU fica fixa na resolução da tela, não no número de partículas
    • Em simulações grandes, a quantidade transferida diminui, mas em simulações pequenas pode ser necessário enviar mais dados; é um trade-off
  • Com uma implementação rápida usando threejs e um pixel shader customizado, a renderização na thread principal passou a ficar na casa de poucos ms, independentemente do tamanho da simulação
  • Nessa arquitetura, os Workers rodam o mais rápido possível, e a simulação fica diretamente atrelada ao número de núcleos de CPU tanto em escalas grandes quanto pequenas
  • Cada Worker usa cerca de 50% do tempo em cálculos numéricos e cerca de 50% atualizando a contagem da grade de partículas
  • O desempenho da versão final é o seguinte
    • Em um Mac M1 na bateria, executa 20 milhões de partículas a cerca de 20 fps
    • Um desktop processa cerca de 30 milhões de partículas com fps semelhante
    • No teste de um amigo com CPU de 32 núcleos, chegou a cerca de 40 milhões de partículas
  • Demos relacionadas:

Resumo

  • CPU e GPU conseguem processar cálculos numéricos com muita rapidez
  • As partes lentas são a movimentação de dados e, em especial, o acesso aleatório à memória
  • Para escrever código rápido, é preciso entender como o hardware armazena dados em cache e os movimenta
  • SharedArrayBuffer e Web Worker foram úteis para experimentos que aproveitam vários núcleos em JavaScript
  • Quando o WebGPU for adotado de forma mais ampla, há planos de testar compute shaders

Leituras relacionadas

1 comentários

 
GN⁺ 2024-07-09
Comentários do Hacker News

  • Seria bom incorporar a simulação diretamente no blog
    Tive que rolar bastante para cima e para baixo até encontrar o link da simulação real
    Daria para escolher valores que rodem razoavelmente até em celulares antigos, ou ajustar de acordo com a taxa de quadros; ou então apenas colocar alguns links no topo do artigo
    https://ciechanow.ski/ é um exemplo de primeiríssimo nível mundial de como pode ser incrível colocar simulações dentro de um artigo
    Antigamente, todos os sites costumavam incorporar elementos interativos legais assim
    E acho que dá para rodar simulações de partículas na GPU mesmo sem WebGPU
    Ex.: https://news.ycombinator.com/item?id=19963640

    • É um dos melhores exemplos de blog explicativo que já vi
      Eu gostaria que esse fosse o padrão para compartilhar informação. Se algo vale a pena ser compartilhado, também vale a pena torná-lo fácil de entender
    • Aquela simulação de fluidos é realmente impressionante
    • Concordo que incorporar seria bom
      Pensei em incorporar cada versão, mas fiquei preocupado com a possibilidade de muitos workers rodarem ao mesmo tempo
      Pretendo atualizar o texto para incorporar a versão final no fim
      Aquele blog é excelente; cada exemplo tem um nível de acabamento enorme
      Edit: tentei adicionar uma versão incorporada, mas os headers necessários não se encaixavam bem com outros embeds, e as versões antigas ainda estão presas ao codesandbox
    • Basta avisar: “pule para o final se quiser brincar com o app final”

  • Testei diretamente a demo mobile [0] e é bem impressionante
    [0] https://dgerrells.com/sabby

    • Funciona com vários dedos também
      É surpreendente que dê para fazer isso em JavaScript puro; como o atraso aumenta conforme há mais dedos, parece que mais dados estão indo e voltando entre threads
    • Estranhamente, parece orgânico, como estruturas lipídicas em uma sopa primordial
    • No meu celular, o Firefox roda melhor que o Chrome, o que é satisfatório

  • Por pura curiosidade, fico pensando se, em vez de representar cada partícula com quatro números de ponto flutuante de 32 bits — x, y, dx, dy —, seria possível codificar esses dados em um único número JS
    O intervalo de inteiros seguros do JS é 2^53 - 1 = 9,007,199,254,740,991, e o intervalo do Float32Array usado no texto vai de -3.4e38 a 3.4e38
    Por exemplo, se a posição na tela for em um canvas 1000x1000, ela pode ser representada por números de 0 a 1.000.000; mesmo adicionando 10 níveis de subpixel, seriam 100.000.000, ainda bem dentro do intervalo do JS
    As velocidades dx e dy também são calculadas como decimais aleatórios entre -10 e +10 com (Math.random()*2-1)*10, mas, se limitar a uma casa decimal for suficiente, daria para convertê-las para uma faixa inteira de -100 a +100, e bastariam 10.000 valores possíveis
    Combinando os dois, 10,000 * 100,000,000 = 1,000,000,000,000, ou seja, um trilhão de números, daria para representar uma partícula dentro do MAX_SAFE_INTEGER do JS
    Então parece que os dados de uma única partícula poderiam caber em um único MAX_SAFE_INTEGER ou em um único elemento de Float32Array, o que eliminaria a necessidade de stride e talvez tornasse a consistência dos dados mais garantida
    Mas o custo de codificar/decodificar pode ser maior que a economia de memória, então talvez não valha nada em termos de desempenho

  • Muito legal, mas fiquei curioso sobre esta parte
    Javascript does support an Atomics API but it uses promises which are gross. Eww sick.
    Exceto por waitAsync[1], a Atomics API parece não usar promises
    Já usei Atomics antes, mas não precisei lidar com código assíncrono/promises
    Fiquei curioso se ela usa promises internamente ou se há alguma parte que eu deixei passar
    [1] https://developer.mozilla.org/en-US/docs/Web/JavaScript/Refe...

  • Também existe uma simulação que lida com quase 20 milhões de partículas
    Dá para executar e programar: https://tinlizzie.org/~ohshima/shadama2/
    https://tinlizzie.org/~ohshima/shadama2/live2017/
    https://www.youtube.com/watch?v=tqBO_5GMfWw

  • O vídeo é legal, mas os links do CodeSandbox em try it out here não funcionam no desktop Chrome no MacOS
    Aparecem erros de CORS como Uncaught ReferenceError: SharedArrayBuffer is not defined e ERR_BLOCKED_BY_RESPONSE.NotSameOriginAfterDefaultedToSameOriginByCoep

    • Como o editor inline do CodeSandbox impede a configuração de headers, é preciso abrir a prévia em uma aba dedicada
      Se você usa um navegador focado em privacidade, isso também pode ser bloqueado
      Daqui em diante vou tentar incluir exemplos embedados
      A versão final pode ser testada aqui: https://dgerrells.com/sabby

  • Excelente
    Passei uma hora entendendo e mexendo no código, e deixei aqui uma implementação em tempo real: https://particules.kapochamo.com/index.html

    • Levei um susto quando abri pela primeira vez
      É muito bonito, e é difícil acreditar que tantas partículas estejam sendo simuladas em JS
    • Legal
      Impressiona que você tenha se aprofundado tanto no código

  • Vale conferir se chrome://tracing traz mais insights
    Descobri por aqui: https://youtu.be/easvMCCBFkQ?t=114

  • Quero muito mandar isso para a equipe de UI que faz apps React lentos
    JS é realmente rápido, especialmente quando é bem escrito, é muito rápido

    • O problema é que JS idiomático e JS rápido, na prática, são quase opostos
      JS rápido parece um dialeto ruim de C, e não se deve fazer alocação de objetos com coleta de lixo, mas a linguagem carece de boas alternativas sem alocação
      Em apps JS complexos, quase ninguém vai tratar todas as alocações de memória como ponteiros dentro de um ArrayBuffer gigante; nesse ponto é mais fácil partir para WebAssembly
    • JavaScript provavelmente é a linguagem que recebeu a maior quantidade de horas humanas investidas em otimizações de engines
      O lamentável é que não dá para depender só de JS, porque é preciso misturar muita manipulação do DOM, e nas UIs que fazemos a parte lenta costuma ser justamente essa manipulação do DOM