Afiando o encoder: o trabalho que a gente fazia duas vezes

Afiando o encoder: o trabalho que a gente fazia duas vezes

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

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

No artigo passado a gente atacou a DCT, trocamos a transformada separável pela AAN, embutimos a escala na tabela de quantização, e o encode saiu de 68 ms para 9,86 ms. Sete vezes mais rápido. 🎉

Hoje a gente ataca o resto do encoder. E o alvo principal é um desperdício que estava escondido à vista de todos desde o artigo 18. Vou te mostrar como eu o encontrei, e o porque o como é mais valioso que o quê.

Ah, e no meio do caminho vou contar a história daquela otimização que eu prometi: uma ideia que parecia genial, que eu implementei com toda a confiança do mundo... e que quebra o codec de um jeito que quase ninguém adivinharia de primeira. 🐛

Pega o café ☕ e vem comigo! 🚀

Onde está o tempo agora?

Regra número um, de novo: meça antes de otimizar. E depois de otimizar, meça de novo, porque o gargalo se move.

Isso é uma coisa que muita gente não espera. Você conserta o ponto mais lento, e aí o segundo mais lento vira o primeiro. É uma fila que nunca acaba, e a única forma de saber quem está na frente é medindo. 🎯

Dito isso, tomei a decisão de instrumentar o encode e olhei o relógio em cada parte. Numa imagem de 512×512 temos o seguinte resultado:

Etapa                        │  Tempo
─────────────────────────────┼──────────
Conversão de cor (split)     │  2,48 ms
Encode em 2 passadas         │  7,00 ms
   ...se fosse 1 passada só  │  3,97 ms
   ...dos quais a DCT        │  0,67 ms

Olha aquela terceira linha com atenção. 👀

Se o encode com duas passadas custa 7,00 ms, e com uma passada custaria 3,97 ms, então a segunda passada sozinha custa 3,03 ms, mais de 40% do tempo de encode.

E a DCT, que era a vilã do artigo passado? Caiu para 0,67 ms. Ela deixou de ser o problema.

O novo alvo apareceu sozinho, só de eu ter medido. 🎯

Porque existem duas passadas mesmo?

Vale relembrar de onde isso veio, porque não é gratuito.

Lá no artigo 18, a gente implementou as tabelas de Huffman customizadas. Em vez de usar as tabelas genéricas do JPEG, o codec constrói uma tabela sob medida para cada imagem, e isso cortou 25% do tamanho do arquivo, de graça, sem perder nada de qualidade.

Mas tinha um problema de galinha e ovo: para construir a tabela ótima, você precisa saber a frequência de cada símbolo. E para saber as frequências, precisa processar a imagem inteira.

Daí as duas passadas:

Tudo certo até aqui. O problema é o que cada passada faz.

Olha o que estava acontecendo em cada uma delas, para cada bloco:

1. Ler os 64 pixels do plano
2. Fazer a DCT          ← caro
3. Quantizar            ← caro
4. Gerar os símbolos    ← barato
5. Contar OU escrever   ← barato

Percebeu? 😲 Os passos 1, 2 e 3 — a parte cara — são feitos duas vezes, e produzem exatamente o mesmo resultado nas duas. A DCT de um bloco não muda entre a primeira e a segunda passada. Estávamos jogando fora metade do trabalho e refazendo.

A correção é quase óbvia depois que você enxerga o problema:

Faça a matemática uma vez, guarde o resultado, e deixe as duas passadas lerem dele.

Em vez de cada passada partir dos pixels, a gente cria um cache de coeficientes: um buffer com os 64 valores já quantizados de cada bloco.

A DCT roda uma vez só, e as duas passadas viram puro trabalho de contabilidade, tipo de coisa que é baratíssimo para a sua CPU.

// ============================================================
// COEFFICIENT CACHE
// ============================================================
//
// The costly part of a block is the transform, not the bookkeeping.
// Running it once and keeping the quantized result lets the counting
// pass and the writing pass share the same work instead of each
// redoing the mathematics from the pixels.

typedef struct {
    short *coefficients;    // 64 values per block, zig-zag order
    long   count;
} CoeffCache;

Repara num detalhe: os coeficientes são guardados como short (16 bits), não como float.

Depois da quantização os valores são inteiros pequenos, então 16 bits dão e sobram, e ocupam quatro vezes menos memória que um float. 🧮

O preenchimento do cache é direto, vejamos:

static int cache_fill(CoeffCache *cache, Plane *plane,
                      const ScaledQuant *quantization) {
    int blocks_x = (plane->width + 7) / 8;
    int blocks_y = (plane->height + 7) / 8;
    long total = (long)blocks_x * blocks_y;

    cache->count = total;
    cache->coefficients = (short *)malloc((size_t)total * 64 * sizeof(short));
    if (!cache->coefficients) {
        return 0;
    }

    long index = 0;
    for (int block_y = 0; block_y < blocks_y; block_y++) {
        for (int block_x = 0; block_x < blocks_x; block_x++, index++) {
            float block[64];
            for (int row = 0; row < 8; row++) {
                for (int column = 0; column < 8; column++) {
                    block[row * 8 + column] = (float)plane_sample(
                        plane, block_x * 8 + column, block_y * 8 + row);
                }
            }
            quantize_block(block, quantization,
                           &cache->coefficients[index * 64]);
        }
    }
    return 1;
}

E a função que antes fazia tudo junto (tokenize_block) agora recebe os coeficientes prontos e só gera os símbolos.

Quanto de memória isso custa? Numa imagem de 512×512 em 4:2:0, são 6144 blocos × 64 valores × 2 bytes ≈ 786 KB. Nada demais para um computador moderno.

É a clássica troca de memória por tempo: gastamos um pouco de RAM para não refazer trabalho caro. 💰

A ideia genial que quebra o codec

Agora a história que eu prometi.

Antes de chegar no cache, eu tive uma ideia que me pareceu muito mais esperta, porque pensa comigo:

A primeira passada existe só para estimar as frequências dos símbolos. E para estimar uma distribuição estatística, você não precisa olhar todos os dados — uma amostra representativa basta!

Faz todo sentido, né? É assim que funciona pesquisa eleitoral: ninguém entrevista o país inteiro.

Então implementei: a passada de contagem passaria a olhar 1 de cada 3 blocos. A mudança é de uma linha:

long step = frequencies ? 3 : 1;
    for (long block = 0; block < cache->count; block += step) {

Compilei, rodei, e o encoder ficou mais rápido. 🎉

E aí eu rodei o teste de validação, e olha o que apareceu:

COM amostragem:
  portrait    encode_image FALHOU (retornou 0)
  texture     encode_image FALHOU (retornou 0)
  smooth      encode_image FALHOU (retornou 0)

O encoder se recusou a funcionar. Em todas as imagens. 😱

Mas dai você deve estar se perguntando... porque quebra?

Leva um tempo para a ficha cair, e a explicação é sutil e engraçada rs

Pensa no que acontece com um símbolo que aparece apenas nos blocos que a amostragem pulou:

O erro não está na estatística. A amostra até estima muito bem as proporções. O erro é de outra natureza:

Amostrar é seguro para estimar uma distribuição. Mas é perigoso quando o resultado precisa cobrir todos os casos possíveis. Uma tabela de Huffman não pode ter buracos. 🕳️

É uma diferença conceitual bonita: a tabela não é um resumo estatístico da imagem, na verdade ela é um dicionário completo, ou seja, se falta uma palavra, o texto não pode ser escrito.

Repara que o nosso codec falhou de forma limpa, com uma mensagem clara., e isso não foi apenas sorte: é uma guarda que existe no código desde o artigo 18.

            HuffTable *table = &tables[table_index];
            int length = table->length[token->symbol];
            if (length == 0) {
                writer->failed = 1;
                return -1;
            }

Aquele if (length == 0) é a diferença entre um bug que grita e um bug que se esconde. 🔊

Porque veja o que aconteceria sem essa verificação: o codec escreveria o código do símbolo, que é... zero bits. Nada.

O símbolo simplesmente sumiria do fluxo. O arquivo sairia com o tamanho quase certo, o encoder não reclamaria de nada, e só a imagem decodificada mostraria o estrago, mostrando um borrão colorido a partir do ponto onde o decoder dessincronizou.

Esse é o pior tipo de bug que existe: silencioso, e com sintoma distante da causa. Você olharia para o decoder por horas, procurando um erro que está no encoder.

💡 A lição que eu levo dessa: quando uma tabela é indexada por algo que "sempre deveria existir", verifique mesmo assim. O custo é uma comparação; o benefício é transformar um bug de horas num erro de segundos.

Otimização 2: escrevendo bits em blocos

Com o cache resolvido, sobrou olhar para como a gente escreve os bits. E aqui tinha outro desperdício.

Assim estava o nosso bw_write_bits:

    for (int i = count - 1; i >= 0; i--) {
        writer->accumulator =
            (writer->accumulator << 1) | ((value >> i) & 1);
        writer->bit_count++;

        if (writer->bit_count == 8) {
            /* ...grava um byte... */
        }
    }

Um laço bit a bit. Para escrever um código Huffman de 12 bits, são 12 iterações, cada uma com um deslocamento, uma máscara, uma soma e um teste. 😩

Mas repara: os bits já vêm juntos. O código Huffman é um número. Por que desmontar em bits e remontar?

A solução é usar um acumulador maior. Em vez de 8 bits (um byte), usamos um de 64 bits, onde o código inteiro entra de uma vez:

// Pushing one bit at a time costs a shift, a mask and a test for
// every single bit: a twelve bit code becomes twelve iterations. Here
// a 64 bit accumulator swallows whole codes in one shift, and memory
// is only touched once at least four bytes are ready to leave.
static void bw_write_bits(BitWriter *writer, unsigned int value, int count) {
    if (writer->failed || count <= 0) {
        return;
    }

    writer->accumulator = (writer->accumulator << count)
                        | (value & ((count >= 32) ? 0xFFFFFFFFu
                                                  : ((1u << count) - 1u)));
    writer->bit_count += count;

    if (writer->bit_count >= 32) {
        /* ...garante espaço no buffer... */

        writer->bit_count -= 32;
        unsigned int word =
            (unsigned int)(writer->accumulator >> writer->bit_count);
        word = __builtin_bswap32(word);
        memcpy(writer->buffer + writer->size, &word, 4);
        writer->size += 4;
    }
}

Duas mudanças importantes aqui:

Um deslocamento em vez de doze. O código inteiro entra no acumulador com uma única operação. 🎯

Quatro bytes por vez na memória. Quando pelo menos 32 bits estão prontos, eles saem juntos, como uma palavra de 4 bytes. Antes eram quatro idas separadas à memória.

O bswap32 e a ordem dos bytes

Aquele __builtin_bswap32 merece explicação, porque ele conecta com um assunto do artigo 14.

Lembra da conversa sobre endianness?

Quando a gente escreve um número de 32 bits na memória em x86, ele é gravado com o byte menos significativo primeiro (little-endian).

Mas o nosso fluxo de bits é lido do bit mais alto para o mais baixo, ou seja, big-endian.

Se a gente gravasse a palavra direto, os bytes sairiam na ordem trocada. O bswap32 inverte a ordem dos 4 bytes numa única instrução do processador, deixando tudo no lugar certo.

É um bom exemplo de como conhecer o hardware rende: a alternativa "portável" seria montar o valor com deslocamentos e máscaras, custando várias operações.

E o compilador nos dá isso de graça. 🎁

Otimização 3: contando bits com uma instrução

Essa é pequena, mas é bonita.

A função category() responde "quantos bits esse número precisa?". E ela é chamada uma vez para cada coeficiente sobrevivente, nas duas passadas. É bastante coisa:

    int absolute = abs(value);
    int size = 0;
    while (absolute) {
        size++;
        absolute >>= 1;
    }
    return size;

Um laço que desloca o número um bit por vez até zerar. Para o valor 200, são 8 iterações.

Acontece que o processador já sabe responder isso. Existe uma instrução chamada count leading zeros, que conta quantos bits zero existem antes do primeiro bit 1. E o número de bits necessários é simplesmente 32 menos isso:

// How many bits a value needs. The loop version reads well but runs
// once per surviving coefficient; counting leading zeros gets the
// same answer in a single instruction on any modern processor.
static inline int category(int value) {
    unsigned int magnitude = (unsigned int)(value < 0 ? -value : value);
    if (magnitude == 0) {
        return 0;
    }
#if defined(__GNUC__) || defined(__clang__)
    return 32 - __builtin_clz(magnitude);
#else
    int size = 0;
    while (magnitude) {
        size++;
        magnitude >>= 1;
    }
    return size;
#endif
}

Repara que mantive o laço original dentro de um #else. O __builtin_clz é uma extensão do GCC e do Clang; em outro compilador, o código continua funcionando pelo caminho antigo.

Lembre-se dessa frase: Otimizar não deveria custar portabilidade. 🌍  

Os resultados

Medindo cada otimização isoladamente, no tempo de encode, nós temos:

Imagem    │ Artigo 19 │ + bit writer │ + cache  │ Ganho
──────────┼───────────┼──────────────┼──────────┼───────
portrait  │ 11,47 ms  │   8,81 ms    │ 6,50 ms  │ 1,76x
smooth    │ 27,49 ms  │  25,51 ms    │ 17,44 ms │ 1,58x
texture   │ 11,65 ms  │   9,61 ms    │  7,06 ms │ 1,65x
micro     │ 11,63 ms  │   9,30 ms    │  6,65 ms │ 1,75x

E aquela segunda passada que custava 3,03 ms? Agora custa 0,52 ms. Sobrou só o trabalho de contabilidade, que era o que ela deveria fazer desde o começo. 🎯

Somando com o artigo passado, nós temos o seguinte:

Imagem    │ Artigo 18 │ Artigo 19 │ Artigo 20 │ Total
──────────┼───────────┼───────────┼───────────┼────────
portrait  │ 17,67 ms  │ 10,61 ms  │  6,62 ms  │  2,7x
smooth    │ 49,92 ms  │ 27,51 ms  │ 17,79 ms  │  2,8x

E contando desde o ponto de partida real (o artigo 18 sem -O2), o encode saiu de 68,76 ms para 6,25 ms: onze vezes mais rápido. 🚀  

O verdadeiro teste que importa...

A pergunta de sempre: a saída mudou?

Imagem    │ Artigo 19 │ Artigo 20 │ PSNR
──────────┼───────────┼───────────┼───────
portrait  │ 26.708 B  │ 26.708 B  │ igual
smooth    │ 30.805 B  │ 30.805 B  │ igual
texture   │ 54.318 B  │ 54.318 B  │ igual
micro     │ 35.492 B  │ 35.492 B  │ igual
colorful  │ 25.933 B  │ 25.933 B  │ igual

Byte por byte idênticos. 🎉

E isso faz todo sentido: diferente do artigo passado (onde a troca de double por float mudou alguns arredondamentos), aqui a gente não mexeu em nenhuma conta. Só reorganizou quando e como o trabalho é feito.

Essa é a assinatura de uma otimização bem feita: o resultado é exatamente o mesmo, só chegou mais cedo. 💎

O que sobrou?

Perfilando de novo o encode, agora:

Etapa                     │  Tempo
──────────────────────────┼──────────
Conversão de cor          │  2,52 ms
Encode (2 passadas)       │  3,86 ms
   ...a 2ª passada agora  │  0,52 ms
   ...a DCT               │  0,67 ms

A conversão de cor virou 40% do tempo de encode. Ela era irrelevante quando o encode custava 68 ms; agora que custa 6, ela aparece.

É a Lei de Amdahl em ação: conforme você acelera uma parte, o que sobra domina cada vez mais.

E tem outra coisa incomodando: o decode praticamente não mudou nestes dois artigos. Ele está em ~16 ms, enquanto o encode já está em 6. Virou o lado lento do codec.

E obviamente que esses dois são assunto dos próximos artigos. 😉

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 o cache de coeficientes, o escritor de bits de 64 bits e a contagem de bits por instrução:

Resumo

Recapitulando:

No próximo artigo, a atenção vai para o decoder, que ficou para trás.

Vamos construir uma tabela de consulta que resolve a maioria dos códigos Huffman num único passo, ensinar o codec a reconhecer blocos que reconstroem para uma cor chapada e pular a transformada inteira, e trocar o upsampling de croma por aritmética inteira.

E tem mais uma história de bug esperando: uma otimização que funcionou, foi mais rápida, produziu arquivos idênticos... e mesmo assim piorou a imagem em 0,15 dB.

Descobrir de onde vinha essa perda foi um dos momentos mais instrutivos da série. 🔍

Te espero no próximo artigo 👋