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

Ele é sempre rápido

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 códigos que consideram o 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.

Hot path

No squoosh, escrevemos uma função JavaScript que gira um buffer de imagem em múltiplos de 90 graus. Embora OffscreenCanvas seja ideal para isso, ele não tem suporte para todos os navegadores que estávamos segmentando e é um pouco instável 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 4094 x 4096 px (16 megapixels), seriam necessárias mais de 16 milhões de iterações do bloco de código interno, que é o que chamamos de "caminho rápido". Apesar do grande número de iterações, dois em cada três navegadores que testamos concluíram a tarefa em 2 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. Neste 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 quisermos 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, o JavaScript e o WebAssembly podem alcançar o mesmo pico de desempenho. No entanto, para JavaScript, essa performance só pode ser alcançada no "caminho rápido", e muitas vezes é 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 programar para o 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 é.

Citando 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 esse 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.

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. Assim, é possível escrever módulos do WebAssembly que não precisam de nenhuma memória adicional e 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. WebAssembly.Memory serve para isso.

Gerenciamento de memória

Normalmente, depois de usar mais memória, você vai precisar gerenciar essa memória 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 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 transformar em 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, a imagem de entrada vai começar no endereço 4, e a imagem de saída vai 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 portar 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 é de 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 segurança de memória e de linha de execução. O Rust também oferece suporte à WebAssembly como um recurso principal, e a equipe do Rust contribuiu com muitas ferramentas excelentes para o ecossistema da WebAssembly.

Uma dessas ferramentas é o wasm-pack, do grupo de trabalho rustwasm. 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

gera um módulo wasm de 7,6 KB com cerca de 100 bytes de código de união (ambos após o gzip).

AssemblyScript

O AssemblyScript é um projeto jovem que tem como objetivo ser um compilador TypeScript para WebAssembly. No entanto, é importante observar que ele não consome apenas 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 você não pode simplesmente compilar qualquer TypeScript que tenha para o WebAssembly, mas significa que você não precisa 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 fornecer um módulo wasm de aproximadamente 300 bytes e nenhum código de união. O módulo só funciona com as APIs vanilla do WebAssembly.

Análise forense do WebAssembly

Os 7,6 KB de Rust são surpreendentemente grandes em comparação com os outros dois idiomas. Há algumas ferramentas no ecossistema do WebAssembly que podem ajudar a analisar seus arquivos do WebAssembly (independentemente do idioma usado para criá-los) e informar o que está acontecendo, além de ajudar a melhorar a situação.

Twiggy

O Twiggy é outra ferramenta da equipe do Rust WebAssembly que extrai vários dados úteis de um módulo 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 desassemblador 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. O 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 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

gerou 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. Como medimos a performance e quais foram os resultados?

Como fazer a comparação

Apesar de o WebAssembly ser um formato de bytecode de baixo nível, ele ainda precisa ser enviado por um compilador para gerar o código de máquina específico do host. Assim como o JavaScript, o compilador funciona em vários estágios. Em termos simples, 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. Na grande maioria dos casos, nunca teremos os benefícios do compilador de otimização. É importante lembrar disso ao fazer a comparação. A execução dos módulos do WebAssembly 10.000 vezes em um loop daria resultados não realistas. Para conseguir números realistas, é necessário 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 todos os comparativos de mercado 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 original de desempenho: todos os módulos do WebAssembly são executados em cerca de 500 ms ou menos. Isso confirma o que dissemos no início: o WebAssembly oferece 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 no ecossistema do TypeScript, como prettier e tslint, provavelmente vão 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 desviar um pouco do caminho ideal para alcançar 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. 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 lançar 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 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 deve impedir você de usar o que for 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 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 especialmente a Ashley Williams, Steve Klabnik, Nick Fitzgerald e Max Graey por toda a ajuda nesta jornada.