Descendo ao SIMD: falando a língua do processador

Descendo ao SIMD: falando a língua do processador

Olá leitor, seja muito bem vindo de volta a mais uma etapa da nossa jornada aqui no Portal da Micilini! 😊

Este é o vigésimo segundo artigo da nossa série, e o quarto da nossa empreitada de otimização.

Nos três artigos anteriores a gente ficou apertando parafusos com as ferramentas de sempre: trocamos algoritmos, guardamos resultados que estavam sendo recalculados, criamos tabelas de consulta e evitamos trabalho desnecessário.

Tudo isso segue uma mesma lógica, voltada ao objetivo de fazer menos trabalho.

Só que no fim do artigo passado a gente bateu num muro. A conversão de cor precisava mesmo calcular vermelho, verde e azul para cada pixel da imagem, e não existia uma conta ali que desse para cortar.

Dessa forma, precisamos mudar de estratégia, e em vez de fazer menos trabalho, vamos aprender a fazer vários trabalhos ao mesmo tempo. 🤯

Sendo assim, pega o café e vem comigo, porque hoje a gente desce mais um degrau na direção do hardware! ☕

A instrução que faz oito coisas de uma vez

Vou começar com uma pergunta que parece boba, mas que é a chave de tudo.

Quando você escreve a + b em C, quantas somas o processador faz?

A resposta natural seria uma, certo? Só que ela está incompleta...

Processadores modernos têm registradores especiais, bem mais largos que os normais. Um registrador comum guarda um número de 64 bits, enquanto um registrador vetorial pode guardar 256 bits, o que dá espaço para oito números de 32 bits lado a lado.

E existem instruções que operam nesses registradores inteiros de uma só vez. Uma única instrução pega oito números, soma com outros oito, e devolve oito resultados.

Soma normal:              Soma vetorial (SIMD):

  a                           a0 a1 a2 a3 a4 a5 a6 a7
+ b                         + b0 b1 b2 b3 b4 b5 b6 b7
─────                       ─────────────────────────
  c                           c0 c1 c2 c3 c4 c5 c6 c7

1 instrução, 1 resultado    1 instrução, 8 resultados

Esse nome tem uma sigla: SIMD, de Single Instruction, Multiple Data, ou seja, uma instrução operando sobre vários dados.

Cada uma dessas oito posições dentro do registrador tem um nome, e a gente chama de lane (que significa pista, em inglês).

Pense nisso como numa estrada de oito pistas onde todos os carros andam juntos, na mesma velocidade, fazendo a mesma manobra ao mesmo tempo. 🛣️

Mas eu preciso escrever assembly?

Essa é a dúvida que todo mundo tem ao ouvir falar disso pela primeira vez, e a resposta é: não.

Existem duas formas de acessar essas instruções.

A primeira é escrever assembly de verdade, colocando as instruções do processador uma a uma. É poderoso, porém trabalhoso, difícil de ler, e o código só funciona numa arquitetura.

A segunda é usar intrinsics, que são funções de C que o compilador traduz diretamente para uma instrução do processador. Você escreve algo que parece uma chamada de função normal, e o compilador gera exatamente a instrução correspondente, sem overhead nenhum.

Vou deixar abaixo um pequeno exemplo, para você ver que não tem mistério nenhum:

__m256 a = _mm256_loadu_ps(vetor_a);   // carrega 8 floats
__m256 b = _mm256_loadu_ps(vetor_b);   // carrega outros 8
__m256 c = _mm256_add_ps(a, b);        // soma os 8 de uma vez
_mm256_storeu_ps(vetor_c, c);          // guarda os 8 resultados

Traduzindo os nomes, que parecem intimidadores mas seguem um padrão bem regular:

Depois que você pega esse padrão, a maioria dos nomes fica previsível. _mm256_mul_ps multiplica oito floats, _mm256_sub_ps subtrai, e assim por diante.

É essa a rota que a libjpeg-turbo usa, e é a que a gente vai usar também. Ninguém escreve um codec inteiro em assembly, o que se faz é escrever a maior parte em C e descer ao nível do hardware apenas nos pontos que realmente importam.

O problema de usar instruções que talvez não existam

Agora aparece uma complicação bem chata.

Essas instruções vetoriais não existem em todos os processadores. Elas foram sendo adicionadas ao longo dos anos, em gerações sucessivas:

Se a gente compilar usando AVX2 e alguém rodar o programa numa máquina mais antiga, o resultado não é um erro elegante: o programa simplesmente morre com uma falha de instrução ilegal. 💥

Então o que fazer? Será que o ideal seria compilar sem AVX2 e desperdiçar o hardware moderno?

Existe uma saída bem melhor, é é o tipo de saída que os projetos sérios fazem.

Perguntando ao processador o que ele sabe fazer

A ideia é escrever duas versões de cada rotina crítica, uma portável e uma vetorial, e decidir qual delas usar no momento em que o programa começa, perguntando ao próprio processador quais instruções ele suporta.

Isso se chama despacho em tempo de execução. O resultado é um único binário que roda a nossa lógica corretamente, tanto em uma máquina antiga, quanto também em máquina moderna, sem exigir compilações diferentes.

E como se pergunta isso ao processador?

Existe uma instrução dedicada, chamada CPUID, que responde exatamente esse tipo de coisa. Aonde o compilador nos dá acesso a ela por uma função pronta:

#include <cpuid.h>

static void cpuid_count(int leaf, int sub, int regs[4]) {
    __cpuid_count(leaf, sub, regs[0], regs[1], regs[2], regs[3]);
}

A CPUID devolve um punhado de bits, cada um sinalizando um recurso. Consultar o suporte a AVX2, por exemplo, é ler um bit específico:

    int regs[4];

    cpuid_count(1, 0, regs);
    int has_sse2    = (regs[3] >> 26) & 1;
    int has_osxsave = (regs[2] >> 27) & 1;
    int has_avx     = (regs[2] >> 28) & 1;

    int has_avx2 = 0;
    if (max_leaf >= 7) {
        cpuid_count(7, 0, regs);
        has_avx2 = (regs[1] >> 5) & 1;
    }

A pegadinha que quase ninguém conta

Aqui vem um detalhe que é fácil de esquecer e que rende bugs difíceis de achar.

Perguntar ao processador não é suficiente. O sistema operacional também precisa concordar.

O motivo é o seguinte: quando o sistema troca de um programa para outro, ele precisa salvar o conteúdo dos registradores. Os registradores vetoriais são grandes, e salvá-los custa tempo, então o sistema só faz isso se souber que alguém vai usá-los.

Se o processador suporta AVX mas o sistema não habilitou o salvamento desse estado, usar as instruções mesmo assim leva a corrupção silenciosa de dados, que é justamente o pior tipo de bug possível. 😱

A verificação correta envolve uma segunda instrução, a xgetbv, que informa quais estados o sistema concordou em preservar:

// The OS must also agree to save the wide registers on a context
// switch. Checking CPUID alone is not enough: xgetbv tells us whether
// the operating system has enabled AVX state saving.
static int os_supports_avx(void) {
    unsigned int eax, edx;
    __asm__ volatile (".byte 0x0f, 0x01, 0xd0"
                      : "=a"(eax), "=d"(edx) : "c"(0));
    unsigned long long xcr = ((unsigned long long) edx << 32) | eax;
    return (xcr & 0x6) == 0x6;      // XMM and YMM state enabled
}

Repara naquele .byte 0x0f, 0x01, 0xd0, que é literalmente o código de máquina da instrução escrito à mão. Esse é o único pedacinho de assembly puro do nosso codec inteiro, e ele existe porque compiladores mais antigos não conheciam essa instrução pelo nome.

Juntando tudo, a detecção fica assim:

    if (has_avx2 && has_avx && has_osxsave && os_supports_avx())
        return N148_CPU_AVX2;
    if (has_sse2)
        return N148_CPU_SSE2;

    return N148_CPU_BASELINE;

E o resultado é guardado numa variável, porque essa pergunta só precisa ser feita uma vez na vida do programa:

static int detected = -1;
static int forced   = -1;

int n148_cpu_level(void) {
    if (forced >= 0) return forced;
    if (detected < 0) detected = detect();
    return detected;
}

void n148_cpu_force(int level) { forced = level; }

Aquele n148_cpu_force não serve para nada no uso normal, e existe puramente para testes.

Com ele a gente consegue obrigar o codec a usar o caminho escalar mesmo numa máquina moderna, e comparar as duas saídas.

Mas fique tranquilo, porque já já você vai ver por que isso é essencial. 🔬

Agora sim: a DCT com oito colunas em paralelo

Com a infraestrutura pronta, vamos ao que interessa.

Lembra de como a nossa DCT funciona desde o artigo 19? Ela faz duas passadas: primeiro transforma cada uma das 8 linhas, depois transforma cada uma das 8 colunas.

Agora repara numa característica que estava lá o tempo todo, esperando ser aproveitada: as oito colunas são completamente independentes entre si. A transformada da coluna 3 não olha para nenhum valor da coluna 4.

E oito coisas independentes é exatamente o que um registrador vetorial guarda. 🎯

O truque de organizar os dados

A parte engenhosa está em como arrumar os números dentro dos registradores.

A gente carrega o bloco 8x8 em oito registradores, e cada registrador guarda uma linha inteira:

registrador r0 →  bloco[0][0]  bloco[0][1]  ...  bloco[0][7]
registrador r1 →  bloco[1][0]  bloco[1][1]  ...  bloco[1][7]
registrador r2 →  bloco[2][0]  bloco[2][1]  ...  bloco[2][7]
                              ...
registrador r7 →  bloco[7][0]  bloco[7][1]  ...  bloco[7][7]

Agora vem o pulo do gato. Olhe para a primeira lane de cada registrador, ou melhor dizendo, olhe para a posição mais à esquerda de todos eles:

r0 lane 0 →  bloco[0][0]
r1 lane 0 →  bloco[1][0]
r2 lane 0 →  bloco[2][0]
             ...
r7 lane 0 →  bloco[7][0]

isso é exatamente a COLUNA 0 do bloco

Ou seja, quando a gente faz uma operação entre os registradores r0 e r7, cada lane está processando a sua própria coluna, e as oito colunas avançam simultaneamente.

A transformada de coluna, que na versão escalar era um laço rodando oito vezes, vira uma sequência de operações entre registradores. Sem laço nenhum. 🚀

Aquela borboleta da AAN que a gente escreveu no artigo 19 traduz quase linha por linha:

    t0=_mm256_add_ps(r0,r7); t7=_mm256_sub_ps(r0,r7);
    t1=_mm256_add_ps(r1,r6); t6=_mm256_sub_ps(r1,r6);
    t2=_mm256_add_ps(r2,r5); t5=_mm256_sub_ps(r2,r5);
    t3=_mm256_add_ps(r3,r4); t4=_mm256_sub_ps(r3,r4);

Compare com a versão escalar do artigo 19 e você vai reconhecer as mesmas somas e subtrações. A diferença é que agora cada linha dessas está fazendo o trabalho de oito colunas.

E as linhas, como ficam?

Resolvemos as colunas, e agora sobra somente uma pergunta óbvia: como fica a passada das linhas?

Em vez de escrever um código diferente para as linhas, a gente gira o bloco e aplica exatamente a mesma rotina de novo.

Girar uma matriz dessa forma tem um nome e a gente chama isso transposição.

Depois de transpor, o que era linha vira coluna, e a rotina que já sabe processar colunas resolve o problema.

1. carrega o bloco nos 8 registradores
2. aplica a borboleta  → processa as 8 COLUNAS
3. transpõe o bloco    → linhas viram colunas
4. aplica a borboleta  → processa as 8 LINHAS
5. transpõe de volta   → devolve a orientação original

A transposição em si é feita com instruções que embaralham os dados entre as lanes, e é o pedaço mais críptico do código:

#define TRANSPOSE8(r0,r1,r2,r3,r4,r5,r6,r7) do {                        \
    __m256 a0,a1,a2,a3,a4,a5,a6,a7,b0,b1,b2,b3,b4,b5,b6,b7;             \
    a0=_mm256_unpacklo_ps(r0,r1); a1=_mm256_unpackhi_ps(r0,r1);         \
    a2=_mm256_unpacklo_ps(r2,r3); a3=_mm256_unpackhi_ps(r2,r3);         \
    ...

Não vou pedir que você entenda essa sequência linha por linha, porque sinceramente ela é uma daquelas coisas que a gente escreve uma vez, testa muito bem, e nunca mais mexe. 😅

O importante é o conceito: transpor permite reaproveitar a mesma rotina para as duas passadas, em vez de escrever duas versões diferentes.

Colocando tudo junto, a nossa DCT vetorial fica assim:

__attribute__((target("avx2")))
void dct_block_avx2(const float block[64], float coef[64]) {
    __m256 r0=_mm256_loadu_ps(block+0),  r1=_mm256_loadu_ps(block+8);
    __m256 r2=_mm256_loadu_ps(block+16), r3=_mm256_loadu_ps(block+24);
    __m256 r4=_mm256_loadu_ps(block+32), r5=_mm256_loadu_ps(block+40);
    __m256 r6=_mm256_loadu_ps(block+48), r7=_mm256_loadu_ps(block+56);

    // level shift: subtrai 128 das 64 amostras em 8 instrucoes
    const __m256 shift = _mm256_set1_ps(128.0f);
    r0=_mm256_sub_ps(r0,shift); r1=_mm256_sub_ps(r1,shift);
    r2=_mm256_sub_ps(r2,shift); r3=_mm256_sub_ps(r3,shift);
    r4=_mm256_sub_ps(r4,shift); r5=_mm256_sub_ps(r5,shift);
    r6=_mm256_sub_ps(r6,shift); r7=_mm256_sub_ps(r7,shift);

    AAN_FWD(r0,r1,r2,r3,r4,r5,r6,r7);
    TRANSPOSE8(r0,r1,r2,r3,r4,r5,r6,r7);
    AAN_FWD(r0,r1,r2,r3,r4,r5,r6,r7);
    TRANSPOSE8(r0,r1,r2,r3,r4,r5,r6,r7);

    _mm256_storeu_ps(coef+0,r0);  _mm256_storeu_ps(coef+8,r1);
    _mm256_storeu_ps(coef+16,r2); _mm256_storeu_ps(coef+24,r3);
    _mm256_storeu_ps(coef+32,r4); _mm256_storeu_ps(coef+40,r5);
    _mm256_storeu_ps(coef+48,r6); _mm256_storeu_ps(coef+56,r7);
}

Repara também no level shift, aquela subtração de 128 que a gente faz desde o artigo 9. Na versão escalar eram 64 subtrações num laço, e aqui viraram oito instruções.

E note o __attribute__((target(\"avx2\"))) logo acima da função. Ele avisa o compilador que apenas essa função pode usar instruções AVX2, o que permite compilar o projeto inteiro sem flags de arquitetura e ainda assim ter o caminho rápido disponível. 🎯

Ligando os dois caminhos

Com as duas versões prontas, a escolha entre elas fica num único lugar:

void dct_block_fast(const float block[64], float coefficients[64]) {
#if defined(__x86_64__) || defined(__i386__)
    if (n148_cpu_level() >= N148_CPU_AVX2) {
        dct_block_avx2(block, coefficients);
        return;
    }
#endif
    dct_scalar(block, coefficients);
}

Repare que existem duas camadas de proteção aqui, e cada uma resolve um problema diferente.

O #if defined(__x86_64__) é resolvido na compilação, e serve para que o código nem exista quando o projeto for compilado para outra arquitetura, como um processador ARM.

Já o if (n148_cpu_level() >= ...) é resolvido durante a execução, e cuida do caso de um x86 antigo que não tem AVX2.

O teste que não dá para pular

Quando você tem duas implementações da mesma coisa, aparece um risco novo que não existia antes: elas podem discordar.

E uma discordância aqui seria catastrófica de um jeito bem particular, porque o arquivo gerado numa máquina moderna abriria diferente numa máquina antiga. O bug só apareceria em produção, na máquina de outra pessoa. 😰

Por isso aquele n148_cpu_force existe, e com ele, dá para rodar a mesma imagem pelos dois caminhos de modo a fazer a comparação:

=== escalar vs AVX2 (devem bater) ===
portrait:
  scalar    26542 bytes  PSNR 32.0612 dB
  AVX2      26542 bytes  PSNR 32.0612 dB
texture:
  scalar    54180 bytes  PSNR 27.1186 dB
  AVX2      54180 bytes  PSNR 27.1186 dB

Mesmos bytes, mesmo PSNR. Os dois caminhos são intercambiáveis. ✅

E na transformada isolada, estamos ainda comparando coeficiente por coeficiente:

precisao AVX2 vs escalar: erro max 0.000244
roundtrip AVX2:           erro max 0.000023

Aquela diferença minúscula é apenas arredondamento de ponto flutuante, do tipo que aparece quando as mesmas somas são feitas numa ordem levemente diferente.

Ou seja, não é aproximação, é a mesma matemática. 💎

Os resultados, e uma surpresa

Medindo a transformada isolada, nos 6144 blocos de uma imagem 512x512, nós temos os seguintes resultados:

Transformada │ Escalar  │  AVX2    │ Ganho
─────────────┼──────────┼──────────┼───────
DCT direta   │ 0,63 ms  │ 0,17 ms  │  3,7x
IDCT inversa │ 0,62 ms  │ 0,16 ms  │  3,9x

Quase quatro vezes mais rápido. Isso é excelente. 🎉

Agora tendo a visão do codec completo:

Imagem    │ Artigo 21 │ Artigo 22 │ Ganho
──────────┼───────────┼───────────┼───────
portrait  │  6,08 ms  │  5,80 ms  │ 1,05x
smooth    │ 17,67 ms  │ 16,01 ms  │ 1,10x
texture   │  6,90 ms  │  6,55 ms  │ 1,05x

Isso é cinco por cento. 😐

Quatro vezes mais rápido numa parte, e cinco por cento no total. Parece frustrante, mas na verdade é a coisa mais instrutiva deste artigo.

Por que o ganho quase sumiu

A explicação é aquela mesma lei que já apareceu algumas vezes aqui, e que vale a pena nomear: a Lei de Amdahl.

Ela diz algo simples e implacável: o ganho total que você consegue é limitado pela fração do tempo que a parte otimizada representava.

Vamos aos números. Perfilando o encode hoje:

Etapa                     │ portrait │ smooth
──────────────────────────┼──────────┼──────────
Conversão de cor          │  2,48 ms │  7,61 ms
Encode (com DCT dentro)   │  3,40 ms │  8,83 ms
   ...dos quais a DCT     │  0,17 ms │  0,53 ms

A DCT representa agora cerca de 3% do tempo de encode. Deixar ela quatro vezes mais rápida economiza uns dois pontos percentuais, e é matematicamente impossível ganhar mais que isso mexendo só nela.

Repara na ironia: no artigo 19 a DCT era a vilã absoluta e otimizá-la deu 7x. Três artigos depois, a mesma otimização, aplicada com uma técnica bem mais poderosa, rende quase nada.

Não é que a otimização ficou pior. É que o alvo mudou de lugar.

E o alvo atual está escancarado naquela tabela: a conversão de cor, que sozinha come 42% do encode. No decode a situação é ainda mais gritante:

Etapa                        │ portrait │ smooth
─────────────────────────────┼──────────┼──────────
decode_image (entropia+IDCT) │  2,14 ms │  4,75 ms
merge_channels               │  7,27 ms │ 19,17 ms

O merge_channels, que é onde vive a conversão de cor, responde por 77% do decode.

A boa notícia é que agora a gente tem a ferramenta certa na mão. 🔧

O que vem por aí

Neste artigo a gente construiu a fundação: aprendeu o que é SIMD, entendeu que intrinsics resolvem sem precisar de assembly, montou a detecção de processador em tempo de execução, e aplicou tudo isso na DCT.

O ganho foi pequeno porque escolhemos de propósito começar pela parte que a gente já conhecia bem. A DCT é território familiar desde o artigo 9, e isso deixou o foco na técnica nova em vez de em duas coisas novas ao mesmo tempo.

Com a fundação pronta, o próximo artigo aponta a mesma arma para onde o tempo realmente está.

Repositório no GitHub

E sim, você pode acompanhar a evolução desse CODEC de imagens em um repositório no GitHub, commit por commit 👋

Segue o link abaixo do commit desse artigo, onde implementamos a detecção de processador em tempo de execução e a transformada vetorial:

Resumo

Recapitulando:

Aquele resultado de cinco por cento é, na minha opinião, a lição mais valiosa deste artigo.

É muito fácil se apaixonar por uma otimização bonita e esquecer de perguntar se ela mexe no que importa.

No próximo artigo, a gente aponta o SIMD para a conversão de cor, que é onde o tempo está de verdade.

E aí vou contar três histórias de armadilha, incluindo uma que me pegou de jeito: uma instrução criada exatamente para o problema que eu tinha, que parecia perfeita no papel, e que na prática se revelou duas vezes e meia mais lenta que a solução ingênua. 🪤

Prepara o café ☕, porque agora é onde o SIMD mostra a que veio.

Até lá! 👋