Como substituir um caminho quente no JavaScript do seu app pelo WebAssembly

Ele está sempre rápido, e você

Nos meus artigos anteriores, falei sobre como o WebAssembly permite que você traga o ecossistema de bibliotecas de C/C++ para a Web. Um app que faz uso extensivo de bibliotecas C/C++ é o squoosh, nosso app da Web que permite compactar imagens com vários codecs que foram compilados de C++ para WebAssembly.

O WebAssembly é uma máquina virtual de baixo nível que executa o bytecode armazenado em arquivos .wasm. Esse código de byte é fortemente tipado e estruturado de forma que pode ser compilado e otimizado para o sistema host muito mais rápido do que o JavaScript. O WebAssembly oferece um ambiente para executar um código que considera o modo sandbox e a incorporação desde o início.

Na minha experiência, a maioria dos problemas de desempenho na Web é causada por layout forçado e pintura excessiva, mas, de vez em quando, um app precisa realizar uma tarefa computacionalmente cara que leva muito tempo. O WebAssembly pode ajudar nesta situação.

O caminho quente

No squoosh, criamos uma função JavaScript que gira um buffer de imagem em múltiplos de 90 graus. Embora o OffscreenCanvas seja ideal para isso, ele não é compatível com os navegadores que estávamos segmentando e um pouco abundante no Chrome.

Essa função itera sobre cada pixel de uma imagem de entrada e a copia para uma posição diferente na imagem de saída para conseguir a rotação. Para uma imagem de 4.094 x 4.096 pixels (16 megapixels), seria necessário mais de 16 milhões de iterações do bloco de código interno, que é o que chamamos de "hot path". Apesar do grande número de iterações, dois em cada três navegadores que testamos finalizam a tarefa em dois segundos ou menos. Uma duração aceitável para esse tipo de interação.

for (let d2 = d2Start; d2 >= 0 && d2 < d2Limit; d2 += d2Advance) {
    for (let d1 = d1Start; d1 >= 0 && d1 < d1Limit; d1 += d1Advance) {
    const in_idx = ((d1 * d1Multiplier) + (d2 * d2Multiplier));
    outBuffer[i] = inBuffer[in_idx];
    i += 1;
    }
}

Um navegador, no entanto, leva mais de 8 segundos. A forma como os navegadores otimizam o JavaScript é muito complicada, e mecanismos diferentes otimizam para coisas diferentes. Alguns otimizam para execução bruta, outros para interação com o DOM. Nesse caso, encontramos um caminho não otimizado em um navegador.

Já o WebAssembly é criado inteiramente em torno da velocidade de execução bruta. Portanto, se queremos um desempenho rápido e previsível em todos os navegadores para códigos como esse, o WebAssembly pode ajudar.

WebAssembly para desempenho previsível

Em geral, JavaScript e WebAssembly podem atingir o mesmo desempenho máximo. No entanto, para JavaScript, esse desempenho só pode ser alcançado no "caminho rápido", e geralmente é difícil permanecer nesse "caminho rápido". Um dos principais benefícios que o WebAssembly oferece é o desempenho previsível, mesmo em vários navegadores. A tipagem rígida e a arquitetura de baixo nível permitem que o compilador ofereça garantias mais fortes para que o código do WebAssembly só precise ser otimizado uma vez e sempre use o "caminho rápido".

Como escrever para WebAssembly

Antes, usávamos bibliotecas C/C++ e as compilamos para WebAssembly para usar a funcionalidade na Web. Não tocamos no código das bibliotecas, apenas escrevemos pequenas quantidades de código C/C++ para formar a ponte entre o navegador e a biblioteca. Desta vez, nossa motivação é diferente: queremos escrever algo do zero com o WebAssembly em mente para podermos aproveitar as vantagens dele.

Arquitetura do WebAssembly

Ao escrever para o WebAssembly, é útil entender um pouco mais sobre o que ele realmente é.

Para citar WebAssembly.org:

Ao compilar um código C ou Rust para o WebAssembly, você recebe um arquivo .wasm que contém uma declaração de módulo. Essa declaração consiste em uma lista de "importações" que o módulo espera do ambiente, uma lista de exportações que o módulo disponibiliza para o host (funções, constantes, blocos de memória) e, é claro, as instruções binárias reais para as funções contidas nele.

Algo que eu não percebi até analisar isso: a pilha que faz com que o WebAssembly seja uma "máquina virtual baseada em pilha" não é armazenada no bloco de memória usado pelos módulos do WebAssembly. A pilha é completamente interna à VM e inacessível para desenvolvedores da Web, exceto pelas DevTools. Dessa forma, é possível gravar módulos WebAssembly que não precisam de memória extra e que usam apenas a pilha interna da VM.

No nosso caso, vamos precisar de mais memória para permitir o acesso arbitrário aos pixels da imagem e gerar uma versão girada dela. É para isso que WebAssembly.Memory serve.

Gerenciamento de memória

Normalmente, depois de usar memória extra, você terá que gerenciá-la de alguma forma. Quais partes da memória estão em uso? Quais são sem custo financeiro? Em C, por exemplo, você tem a função malloc(n) que encontra um espaço de memória de n bytes consecutivos. Funções desse tipo também são chamadas de "alocadores". A implementação do alocador em uso precisa ser incluída no módulo da WebAssembly e vai aumentar o tamanho do arquivo. Esse tamanho e desempenho dessas funções de gerenciamento de memória podem variar bastante dependendo do algoritmo usado. É por isso que muitos idiomas oferecem várias implementações para escolher ("dmalloc", "emmalloc", "wee_alloc" etc.).

No nosso caso, sabemos as dimensões da imagem de entrada (e, portanto, as dimensões da imagem de saída) antes de executar o módulo WebAssembly. Aqui, encontramos uma oportunidade: tradicionalmente, transmitíamos o buffer RGBA da imagem de entrada como um parâmetro para uma função do WebAssembly e retornávamos a imagem girada como um valor de retorno. Para gerar esse valor de retorno, teríamos que usar o alocador. No entanto, como sabemos a quantidade total de memória necessária (o dobro do tamanho da imagem de entrada, uma vez para entrada e outra para saída), podemos colocar a imagem de entrada na memória do WebAssembly usando JavaScript, executar o módulo do WebAssembly para gerar uma segunda imagem girada e usar o JavaScript para ler o resultado. Podemos fazer isso sem usar nenhum gerenciamento de memória.

Muitas opções

Se você analisou a função JavaScript original que queremos usar com o WebAssembly, vai notar que ela é um código puramente computacional sem APIs específicas do JavaScript. Portanto, é bastante simples portar esse código para qualquer idioma. Avaliei três linguagens diferentes que são compiladas para WebAssembly: C/C++, Rust e AssemblyScript. A única pergunta que precisamos responder para cada um dos idiomas é: como acessar a memória bruta sem usar funções de gerenciamento de memória?

C e Emscripten

O Emscripten é um compilador C para o destino do WebAssembly. O objetivo do Emscripten é funcionar como uma substituição rápida para compiladores C conhecidos, como GCC ou clang, e é compatível com a maioria das flags. Essa é uma parte importante da missão do Emscripten, que quer tornar a compilação de código C e C++ para WebAssembly o mais fácil possível.

O acesso à memória bruta está na própria natureza do C, e os ponteiros existem por esse motivo:

uint8_t* ptr = (uint8_t*)0x124;
ptr[0] = 0xFF;

Aqui, estamos transformando o número 0x124 em um ponteiro para números inteiros (ou bytes) de 8 bits sem sinal. Isso transforma a variável ptr em uma matriz que começa no endereço de memória 0x124, que pode ser usada como qualquer outra matriz, permitindo acessar bytes individuais para leitura e gravação. No nosso caso, estamos analisando um buffer RGBA de uma imagem que queremos reordenar para conseguir a rotação. Para mover um pixel, precisamos mover quatro bytes consecutivos de uma vez (um byte para cada canal: R, G, B e A). Para facilitar, podemos criar uma matriz de números inteiros não assinados de 32 bits. Por convenção, nossa imagem de entrada começará no endereço 4 e a imagem de saída começará logo após o término da imagem de entrada:

int bpp = 4;
int imageSize = inputWidth * inputHeight * bpp;
uint32_t* inBuffer = (uint32_t*) 4;
uint32_t* outBuffer = (uint32_t*) (inBuffer + imageSize);

for (int d2 = d2Start; d2 >= 0 && d2 < d2Limit; d2 += d2Advance) {
    for (int d1 = d1Start; d1 >= 0 && d1 < d1Limit; d1 += d1Advance) {
    int in_idx = ((d1 * d1Multiplier) + (d2 * d2Multiplier));
    outBuffer[i] = inBuffer[in_idx];
    i += 1;
    }
}

Depois de transferir toda a função JavaScript para C, podemos compilar o arquivo C com emcc:

$ emcc -O3 -s ALLOW_MEMORY_GROWTH=1 -o c.js rotate.c

Como sempre, o emscripten gera um arquivo de código de união chamado c.js e um módulo wasm chamado c.wasm. O módulo wasm gzips tem apenas cerca de 260 bytes, enquanto o código glue tem cerca de 3,5 KB após o gzip. Depois de algumas tentativas, conseguimos descartar o código de união e instanciar os módulos do WebAssembly com as APIs vanilla. Isso geralmente é possível com o Emscripten, desde que você não esteja usando nada da biblioteca padrão C.

Rust

O Rust é uma linguagem de programação nova e moderna com um sistema de tipos rico, sem tempo de execução e um modelo de propriedade que garante a segurança de memória e de linha de execução. O Rust também oferece suporte ao WebAssembly como um recurso principal, e a equipe do Rust contribuiu com muitas ferramentas excelentes para o ecossistema do WebAssembly.

Uma dessas ferramentas é o wasm-pack, do grupo de trabalho rustwasm (link em inglês). O wasm-pack transforma seu código em um módulo compatível com a Web que funciona pronto para uso com bundlers, como o webpack. O wasm-pack é uma experiência extremamente conveniente, mas atualmente só funciona para Rust. O grupo está considerando adicionar suporte a outros idiomas de destino do WebAssembly.

Em Rust, as fatias são o que as matrizes são em C. E, assim como no C, precisamos criar fatias que usam nossos endereços de início. Isso vai contra o modelo de segurança de memória que o Rust impõe. Portanto, para conseguirmos o que queremos, precisamos usar a palavra-chave unsafe, permitindo que escrevamos um código que não obedece a esse modelo.

let imageSize = (inputWidth * inputHeight) as usize;
let inBuffer: &mut [u32];
let outBuffer: &mut [u32];
unsafe {
    inBuffer = slice::from_raw_parts_mut::<u32>(4 as *mut u32, imageSize);
    outBuffer = slice::from_raw_parts_mut::<u32>((imageSize * 4 + 4) as *mut u32, imageSize);
}

for d2 in 0..d2Limit {
    for d1 in 0..d1Limit {
    let in_idx = (d1Start + d1 * d1Advance) * d1Multiplier + (d2Start + d2 * d2Advance) * d2Multiplier;
    outBuffer[i as usize] = inBuffer[in_idx as usize];
    i += 1;
    }
}

Compilar os arquivos Rust usando

$ wasm-pack build

produz um módulo Wasm de 7,6 KB com cerca de 100 bytes de código agrupador (ambos após o gzip).

AssemblyScript

O AssemblyScript é um projeto bastante jovem que tem como objetivo ser um compilador TypeScript para WebAssembly. No entanto, é importante observar que ele não consumirá apenas nenhum TypeScript. O AssemblyScript usa a mesma sintaxe do TypeScript, mas troca a biblioteca padrão pela própria. A biblioteca padrão deles modela os recursos do WebAssembly. Isso significa que não é possível compilar qualquer TypeScript existente no WebAssembly, mas não é preciso aprender uma nova linguagem de programação para escrever o WebAssembly.

    for (let d2 = d2Start; d2 >= 0 && d2 < d2Limit; d2 += d2Advance) {
      for (let d1 = d1Start; d1 >= 0 && d1 < d1Limit; d1 += d1Advance) {
        let in_idx = ((d1 * d1Multiplier) + (d2 * d2Multiplier));
        store<u32>(offset + i * 4 + 4, load<u32>(in_idx * 4 + 4));
        i += 1;
      }
    }

Considerando a pequena superfície de tipo que nossa função rotate() tem, foi bastante fácil portar esse código para o AssemblyScript. As funções load<T>(ptr: usize) e store<T>(ptr: usize, value: T) são fornecidas pelo AssemblyScript para acessar a memória bruta. Para compilar nosso arquivo AssemblyScript, basta instalar o pacote npm AssemblyScript/assemblyscript e executar

$ asc rotate.ts -b assemblyscript.wasm --validate -O3

O AssemblyScript vai nos fornecer um módulo Wasm de aproximadamente 300 bytes e nenhum código agrupador. O módulo só funciona com as APIs vanilla do WebAssembly.

Análise forense do WebAssembly

A versão de 7,6 KB do Rust é surpreendentemente grande quando comparada a outros dois idiomas. Há algumas ferramentas no ecossistema do WebAssembly que podem ajudar a analisar os arquivos do WebAssembly (independentemente do idioma usado para criar) e informar o que está acontecendo, além de ajudar a melhorar a situação.

Twiggy

O Twiggy é outra ferramenta da equipe de WebAssembly do Rust que extrai vários dados úteis de um módulo da WebAssembly. A ferramenta não é específica do Rust e permite inspecionar itens como o gráfico de chamadas do módulo, determinar seções não utilizadas ou supérfluas e descobrir quais seções estão contribuindo para o tamanho total do arquivo do módulo. O último pode ser feito com o comando top do Twiggy:

$ twiggy top rotate_bg.wasm
Captura de tela da instalação do Twiggy

Nesse caso, podemos ver que a maioria do tamanho do arquivo vem do alocador. Isso foi surpreendente, já que nosso código não usa alocações dinâmicas. Outro fator importante é uma subseção "nomes de função".

Wasm-strip

wasm-strip é uma ferramenta do kit de ferramentas de binário do WebAssembly (ou wabt, na sigla em inglês). Ele contém algumas ferramentas que permitem inspecionar e manipular módulos do WebAssembly. wasm2wat é um desassemblizador que transforma um módulo wasm binário em um formato legível por humanos. O Wabt também contém wat2wasm, que permite transformar esse formato legível por humanos em um módulo WASM binário. Embora tenhamos usado essas duas ferramentas complementares para inspecionar nossos arquivos do WebAssembly, descobrimos que wasm-strip é a mais útil. wasm-strip remove seções e metadados desnecessários de um módulo do WebAssembly:

$ wasm-strip rotate_bg.wasm

Isso reduz o tamanho do arquivo do módulo Rust de 7,5 KB para 6,6 KB (após o gzip).

wasm-opt

wasm-opt é uma ferramenta do Binaryen. Ele pega um módulo do WebAssembly e tenta otimizar o tamanho e a performance com base apenas no bytecode. Algumas ferramentas, como o Emscripten, já executam essa ferramenta, mas outras não. É recomendável tentar salvar alguns bytes extras usando essas ferramentas.

wasm-opt -O3 -o rotate_bg_opt.wasm rotate_bg.wasm

Com wasm-opt, podemos reduzir outros bytes para deixar um total de 6,2 KB após o gzip.

#![no_std]

Depois de consultar e pesquisar, reescrevemos nosso código Rust sem usar a biblioteca padrão do Rust, usando o recurso #![no_std]. Isso também desativa totalmente as alocações de memória dinâmica, removendo o código de alocador do nosso módulo. Compilar este arquivo Rust com

$ rustc --target=wasm32-unknown-unknown -C opt-level=3 -o rust.wasm rotate.rs

produziu um módulo Wasm de 1,6 KB após wasm-opt, wasm-strip e gzip. Embora ainda seja maior do que os módulos gerados por C e AssemblyScript, ele é pequeno o suficiente para ser considerado leve.

Desempenho

Antes de tirar conclusões com base no tamanho do arquivo, fizemos essa jornada para otimizar o desempenho, não o tamanho do arquivo. Então, como medimos o desempenho e quais foram os resultados?

Como fazer a comparação

Mesmo que o WebAssembly seja um formato de bytecode de baixo nível, ele ainda precisa ser enviado por um compilador para gerar um código de máquina específico do host. Assim como o JavaScript, o compilador funciona em vários estágios. Simplificando: a primeira etapa é muito mais rápida na compilação, mas tende a gerar um código mais lento. Quando o módulo começa a ser executado, o navegador observa quais partes são usadas com frequência e as envia por um compilador mais otimizado, mas mais lento.

Nosso caso de uso é interessante porque o código para girar uma imagem será usado uma ou duas vezes. Portanto, na grande maioria dos casos, nunca vamos ter os benefícios do compilador de otimização. É importante lembrar disso ao fazer a comparação. Executar nossos módulos WebAssembly 10.000 vezes em repetição gera resultados irrealistas. Para conseguir números realistas, precisamos executar o módulo uma vez e tomar decisões com base nos números dessa única execução.

Comparação de performance

Comparação de velocidade por idioma
Comparação de velocidade por navegador

Esses dois gráficos são visualizações diferentes dos mesmos dados. No primeiro gráfico, comparamos por navegador, e no segundo, por idioma usado. Escolhi uma escala logarítmica. Também é importante que todas as comparações usem a mesma imagem de teste de 16 megapixels e a mesma máquina host, exceto um navegador, que não pode ser executado na mesma máquina.

Sem analisar muito esses gráficos, fica claro que resolvemos nosso problema de desempenho original: todos os módulos do WebAssembly são executados em cerca de 500 ms ou menos. Isso confirma o que estabelecemos no início: o WebAssembly oferece um desempenho previsível. Não importa qual idioma escolhermos, a variação entre navegadores e idiomas é mínima. Para ser exato: a variação padrão do JavaScript em todos os navegadores é de cerca de 400 ms, enquanto a variação padrão de todos os módulos do WebAssembly em todos os navegadores é de cerca de 80 ms.

Esforço

Outra métrica é a quantidade de esforço que tivemos que fazer para criar e integrar nosso módulo do WebAssembly ao squoosh. É difícil atribuir um valor numérico ao esforço. Por isso, não vou criar gráficos, mas gostaria de apontar algumas coisas:

O AssemblyScript foi fácil. Além de permitir que você use o TypeScript para escrever o WebAssembly, facilitando a revisão de código para meus colegas, ele também produz módulos do WebAssembly sem cola que são muito pequenos e têm desempenho decente. As ferramentas do ecossistema do TypeScript, como as mais bonitas e as tslint, provavelmente funcionarão.

Rust em combinação com wasm-pack também é extremamente conveniente, mas se destaca mais em projetos maiores do WebAssembly em que vinculações e gerenciamento de memória são necessários. Tivemos que divergir um pouco do caminho da felicidade para atingir um tamanho de arquivo competitivo.

O C e o Emscripten criaram um módulo WebAssembly muito pequeno e de alto desempenho imediatamente, mas sem a coragem de pular para o código de união e reduzi-lo às necessidades básicas, o tamanho total (módulo WebAssembly + código de união) acaba sendo muito grande.

Conclusão

Então, qual linguagem você deve usar se tiver um caminho de acesso rápido do JS e quiser torná-lo mais rápido ou mais consistente com o WebAssembly. Como sempre com perguntas de performance, a resposta é: depende. Então, o que enviamos?

Gráfico de comparação

Comparando a troca de tamanho / desempenho do módulo das diferentes linguagens que usamos, a melhor escolha parece ser C ou AssemblyScript. Decidimos enviar o Rust. Há vários motivos para essa decisão: todos os codecs enviados no Squoosh até agora são compilados usando o Emscripten. Queríamos ampliar nossos conhecimentos sobre o ecossistema do WebAssembly e usar uma linguagem diferente na produção. O AssemblyScript é uma boa alternativa, mas o projeto é relativamente novo e o compilador não é tão maduro quanto o Rust.

Embora a diferença no tamanho do arquivo entre Rust e os outros idiomas pareça bastante drástica no gráfico de dispersão, ela não é tão grande na realidade: carregar 500B ou 1,6 KB, mesmo com 2 GB, leva menos de um décimo de segundo. E esperamos que o Rust feche a lacuna em termos de tamanho do módulo em breve.

Em termos de desempenho em tempo de execução, o Rust tem uma média mais rápida em todos os navegadores do que o AssemblyScript. Especialmente em projetos maiores, o Rust tem mais probabilidade de produzir um código mais rápido sem precisar de otimizações manuais. Mas isso não impede que você use aquilo com que se sente mais confortável.

O AssemblyScript foi uma grande descoberta. Ele permite que desenvolvedores da Web produzam módulos WebAssembly sem precisar aprender uma nova linguagem. A equipe do AssemblyScript tem respondido muito bem e está trabalhando ativamente para melhorar a cadeia de ferramentas. Vamos ficar de olho no AssemblyScript no futuro.

Atualização: Rust

Após a publicação deste artigo, Nick Fitzgerald da equipe Rust nos indicou o excelente livro Rust Wasm, que contém uma seção sobre como otimizar o tamanho do arquivo. Seguir as instruções (principalmente, ativar otimizações de tempo de link e tratamento manual de pânico) nos permitiu escrever um código Rust "normal" e voltar a usar Cargo (o npm do Rust) sem aumentar o tamanho do arquivo. O módulo Rust termina com 370B após o gzip. Para mais detalhes, confira a solicitação de correção que abri no Squoosh.

Agradeço a Ashley Williams, Steve Klabnik, Nick Fitzgerald e Max Graey por toda a ajuda nesta jornada.