A DCT rápida: a AAN e a escala que sai de graça

A DCT rápida: a AAN e a escala que sai de graça

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

Este é o décimo nono artigo da nossa série, e eu preciso começar com um pedido de desculpas. 😅

No fim do artigo passado eu prometi que o próximo seria o julgamento: o N.148i cara a cara com JPEG, PNG e WebP.

Eu já estava com tudo preparado, os scripts prontos, as imagens escolhidas... e aí eu resolvi medir uma coisa que ainda não tinha medido.

O tempo.

E o que eu vi me fez engavetar o benchmark na hora. 😬

O susto que mudou o plano

Comprimindo uma imagem de 512×512 no modo otimizado, foi isto que o cronômetro me mostrou:

Encode: 68,76 ms
Decode: 47,93 ms

68 e 47 milísegundos é extremamente rápido, não? Sim, porém...

Pra você ter uma referência: o JPEG faz a mesma coisa em poucos milissegundos, sendo assim, é certo dizer que nós estávamos muito atrás em ordem de grandeza

Aí eu parei e pensei em duas coisas.

Primeiro: faz sentido publicar um benchmark que diz "o codec comprime bem, mas demora uma eternidade"? Não faz. Seria como entrar numa corrida com o pneu murcho de propósito, sabendo que o pneu está murcho.

Segundo, e mais importante: a Fase 5 vai empilhar predição intra-bloco, quantização adaptativa, codificação aritmética e filtros de pós-processamento. Todas essas técnicas melhoram a compressão e pioram a velocidade.

Construir tudo isso em cima de uma base lenta seria acumular dívida. Então mudei o plano: primeiro a gente conserta a velocidade, depois julga. 🔧

E não é um artigo só. Temos seis pela frente nessa jornada de otimização, e eu prometo que vale cada um:

Pega o café ☕, que hoje o assunto é velocidade! 🚀

Regra número um: meça antes de otimizar

Antes de tocar numa linha de código, existe uma regra que separa quem otimiza de quem só mexe no código torcendo pra melhorar:

Meça primeiro. Otimizar no chute é a forma mais eficiente de perder tempo consertando o que não era problema.

Então eu medi por partes. E o encode se dividia mais ou menos assim:

Faz todo sentido. A DCT é a única etapa que faz matemática pesada de verdade, representando 1024 multiplicações de ponto flutuante por bloco, e temos milhares de blocos numa imagem.

Dito isso, definimos o nosso primeiro alvo!

Mas antes de mexer na DCT, tem uma coisa muito mais fácil pra fazer primeiro. E eu confesso essa com um pouco de vergonha. 🙈

Otimização zero: a mais fácil da sua vida

Durante toda a série, a gente vem compilando todo o projeto assim:

gcc main.c header.c ppm.c tables.c dct.c huffman.c encoder.c decoder.c -o n148i -lm

Repara no que não está aí: nenhuma flag de otimização.

Sem flag, o GCC compila em modo -O0: tradução literal do que você escreveu, na ordem que você escreveu, sem nenhuma esperteza. Isso é ótimo pra depurar (o programa executa exatamente como está no papel), mas é péssimo pra velocidade.

Para melhorar isso, basta adicionar o -O2:  

gcc main.c header.c ppm.c tables.c dct.c huffman.c encoder.c decoder.c -o n148i -lm -O2

E o resultado será este:

Compilação          │  Encode  │  Decode
────────────────────┼──────────┼──────────
Sem flag (-O0)      │ 68,76 ms │ 47,93 ms
Com -O2             │ 17,85 ms │ 20,34 ms
────────────────────┼──────────┼──────────
Ganho               │   3,9x   │   2,4x

Quase 4 vezes mais rápido, sem mudar uma única linha de código. 🤯

O que o -O2 faz?

Um monte de coisa que a gente faria na mão se tivesse paciência: mantém variáveis em registradores em vez de ir na memória toda hora, desenrola laços pequenos, elimina cálculos repetidos, funde operações.

É o compilador aplicando décadas de conhecimento acumulado por você, de graça.  

💡 Lição prática: se você está medindo performance de código C sem -O2, os seus números não valem nada. É a primeira coisa a checar, sempre.  

Pronto, o tempo caiu em 4× e ainda nem começamos. Agora vamos ao trabalho de verdade. 💪

A promessa que ficou pendente lá no artigo 9

Lembra do bônus no fim do artigo sobre a DCT? Eu escrevi isto:

DCT naive: 4096 multiplicações por bloco DCT separável: 1024 multiplicações por bloco DCT fast (Loeffler): cerca de 176 multiplicações por bloco

E depois completei: "como o próximo passo já vai dividir cada coeficiente por um valor mesmo, dá pra embutir o fator de escala do fast DCT de graça dentro da tabela de quantização. Fica anotado esse gancho pro futuro 😉"

Pois é. O futuro chegou. 🎉

Vamos implementar a AAN — sigla de Arai-Agui-Nakajima, que são os três pesquisadores japoneses que publicaram o algoritmo em 1988.

Observação: é ela que a libjpeg usa até hoje no modo rápido.

A ideia por trás da AAN

A DCT separável que a gente escreveu faz 8 multiplicações por amostra, com 8 amostras, em 2 dimensões: 1024 por bloco. Cada uma dessas multiplicações usa um cosseno da tabela.

A AAN parte de uma observação: aqueles cossenos todos têm simetrias entre si.

Pensa no cos(π/4), que vale aproximadamente 0,7071. Ele aparece em vários lugares da conta. E vários pares de termos são, na verdade, somas e diferenças das mesmas quantidades, calculadas repetidamente.

Se você reorganizar a ordem das operações pra calcular cada quantidade intermediária uma vez só e reaproveitar, a maioria das multiplicações desaparece e vira soma, que é uma operação bem mais barata pro processador. ➕

O resultado é impressionante: a AAN faz uma DCT 1D de 8 pontos com apenas 5 multiplicações e 29 somas. Contra 64 multiplicações da versão separável.

O preço: a saída vem "torta"

Não existe almoço grátis, e o preço da AAN é peculiar.

Ela não produz exatamente os coeficientes da DCT. Ela produz os coeficientes multiplicados por um fator conhecido, e esse fator é diferente pra cada posição da matriz:

Saída da AAN[u][v] = DCT_verdadeira[u][v] × fator(u,v)

E os fatores vêm de uma tabelinha fixa, a mesma pra qualquer imagem do mundo:

Escala AAN por índice:
  1.000000  1.387040  1.306563  1.175876
  1.000000  0.785695  0.541196  0.275899

E o fator de um coeficiente (u,v) é escala[u] × escala[v] × 8.

Agora vem o problema óbvio: se a gente precisasse corrigir essa escala dividindo cada coeficiente pelo seu fator, seriam 64 divisões por bloco.

E divisão é uma das operações mais caras que existem. Boa parte do ganho iria embora pelo ralo. 😩

E aqui está a sacada

A gente não precisa corrigir.

Repara no que acontece logo depois da DCT no nosso pipeline: a quantização, que divide cada coeficiente pela tabela Q.

Ou seja, o coeficiente já vai ser dividido por alguma coisa de qualquer jeito. Então, em vez de:

1. corrigir a escala:  coef = saída_AAN / fator
2. quantizar:          resultado = arredondar(coef / Q)

A gente simplesmente embute o fator na tabela de quantização, uma vez só, na inicialização do programa:

tabela_ajustada[u][v] = 1 / (Q[u][v] × fator(u,v))

E a quantização vira uma multiplicação:

resultado = arredondar(saída_AAN × tabela_ajustada)

Olha o que ganhamos de graça nesse único movimento:

Era exatamente isso que eu tinha adiantado no artigo 9, ou seja: gancho pago. 😎

O código: a passada AAN

Vamos ao código...

Este é o coração da AAN, onde faremos uma passada 1D sobre 8 amostras. Antes de ler linha por linha, faz um exercício: conta quantas multiplicações tem aí. (Procura os números com casas decimais.)

Dito isso, Abra o src/dct.c e adicione:

// ============================================================
// SCALAR AAN — 5 multiplications per 1-D pass instead of 64
// ============================================================

static void fdct_pass(float *p, int stride) {
    float t0,t1,t2,t3,t4,t5,t6,t7,t10,t11,t12,t13,z1,z2,z3,z4,z5,z11,z13;

    t0 = p[0*stride] + p[7*stride];  t7 = p[0*stride] - p[7*stride];
    t1 = p[1*stride] + p[6*stride];  t6 = p[1*stride] - p[6*stride];
    t2 = p[2*stride] + p[5*stride];  t5 = p[2*stride] - p[5*stride];
    t3 = p[3*stride] + p[4*stride];  t4 = p[3*stride] - p[4*stride];

    t10 = t0 + t3;  t13 = t0 - t3;
    t11 = t1 + t2;  t12 = t1 - t2;

    p[0*stride] = t10 + t11;
    p[4*stride] = t10 - t11;

    z1 = (t12 + t13) * 0.707106781f;
    p[2*stride] = t13 + z1;
    p[6*stride] = t13 - z1;

    t10 = t4 + t5;
    t11 = t5 + t6;
    t12 = t6 + t7;

    z5 = (t10 - t12) * 0.382683433f;
    z2 = 0.541196100f * t10 + z5;
    z4 = 1.306562965f * t12 + z5;
    z3 = t11 * 0.707106781f;

    z11 = t7 + z3;
    z13 = t7 - z3;

    p[5*stride] = z13 + z2;
    p[3*stride] = z13 - z2;
    p[1*stride] = z11 + z4;
    p[7*stride] = z11 - z4;
}

Contou? São cinco: os 0.707106781f (duas vezes, mas é a mesma constante), 0.382683433f, 0.541196100f e 1.306562965f. Todo o resto são somas e subtrações. 🎉

E repara naquele parâmetro stride, ele é um truquezinho útil. Com stride = 1, a função percorre uma linha. Com stride = 8, ela percorre uma coluna do mesmo array.

Uma função só serve pras duas passadas:

static void dct_scalar(const float block[64], float coef[64]) {
    float t[64];
    for (int i = 0; i < 64; i++) t[i] = block[i] - 128.0f;   // level shift
    for (int r = 0; r < 8; r++) fdct_pass(&t[r*8], 1);       // linhas
    for (int c = 0; c < 8; c++) fdct_pass(&t[c],   8);       // colunas
    memcpy(coef, t, sizeof(t));
}

💡 Repara que agora usamos float em vez de double. Nesta etapa a diferença de velocidade é pequena, mas isso vai importar muito lá no artigo do SIMD: com registradores vetoriais, cabem 8 float no lugar de 4 double. É o dobro de trabalho por instrução.  

A inversa, e um aviso de quem apanhou

A IDCT rápida segue a mesma lógica, na ordem contrária:

static void idct_pass(float *p, int stride) {
    float t0,t1,t2,t3,t4,t5,t6,t7,t10,t11,t12,t13,z5,z10,z11,z12,z13;

    t0 = p[0*stride];  t1 = p[2*stride];
    t2 = p[4*stride];  t3 = p[6*stride];

    t10 = t0 + t2;
    t11 = t0 - t2;
    t13 = t1 + t3;
    t12 = (t1 - t3) * 1.414213562f - t13;

    t0 = t10 + t13;  t3 = t10 - t13;
    t1 = t11 + t12;  t2 = t11 - t12;

    t4 = p[1*stride];  t5 = p[3*stride];
    t6 = p[5*stride];  t7 = p[7*stride];

    z13 = t6 + t5;  z10 = t6 - t5;
    z11 = t4 + t7;  z12 = t4 - t7;

    t7  = z11 + z13;
    t11 = (z11 - z13) * 1.414213562f;

    z5  = (z10 + z12) * 1.847759065f;
    t10 = z5 - z12 * 1.082392200f;
    t12 = z5 - z10 * 2.613125930f;

    t6 = t12 - t7;
    t5 = t11 - t6;
    t4 = t10 - t5;

    p[0*stride] = t0 + t7;   p[7*stride] = t0 - t7;
    p[1*stride] = t1 + t6;   p[6*stride] = t1 - t6;
    p[2*stride] = t2 + t5;   p[5*stride] = t2 - t5;
    p[3*stride] = t3 + t4;   p[4*stride] = t3 - t4;
}

Agora o aviso, e eu falo por experiência dolorosa: a AAN inversa tem sinais traiçoeiros. 🐛

Na minha primeira tentativa, a DCT direta batia perfeitamente com a versão lenta, onde eu conferi coeficiente por coeficiente, erro na ordem de 10⁻⁵. Perfeito. Mas o roundtrip (ida e volta) devolvia lixo puro.

Eu tinha invertido dois sinais na parte ímpar, e trocado a ordem de dois resultados na saída. Detalhe suficiente pra destruir tudo, e sutil o bastante pra não saltar aos olhos.

Se você for implementar, teste a ida e a volta separadamente:

Assim, quando quebrar, você sabe exatamente em qual metade está o problema. Testar as duas juntas é como procurar uma agulha em dois palheiros ao mesmo tempo. 🪡

Embutindo a escala no encoder

Agora a parte que paga o gancho. No src/encoder.c, criamos uma tabela com os recíprocos já ajustados:

// The fast transform leaves every coefficient multiplied by a known
// constant. Rather than dividing it back out on every block, the
// constant is folded into the quantization table once, at startup.
// Quantizing then becomes a multiplication instead of a division.
typedef struct {
    float reciprocal[64];
} ScaledQuant;

static void build_scaled_quant(int quantization[8][8], ScaledQuant *out) {
    for (int row = 0; row < 8; row++) {
        for (int column = 0; column < 8; column++) {
            out->reciprocal[row * 8 + column] = (float)(1.0 /
                (quantization[row][column] * aan_scale_factor(row, column)));
        }
    }
}

E a quantização, que antes dividia, agora multiplica:

static int tokenize_block(const float block[64],
                          const ScaledQuant *quantization,
                          int *previous_dc, Token tokens[64]) {
    float coefficients[64];
    dct_block_fast(block, coefficients);

    int zigzag[64];
    for (int i = 0; i < 64; i++) {
        int index = ZIGZAG[i];
        zigzag[i] = (int)lrintf(
            coefficients[index] * quantization->reciprocal[index]);
    }
    /* ... o resto continua igual ... */
}

Repara como ficou limpo: nenhum vestígio da escala da AAN aparece no caminho quente. Ela vive escondida na tabela. 😌

E no decoder, o que faremos?

O mesmo raciocínio, espelhado. A IDCT rápida espera receber os coeficientes já multiplicados pelas constantes AAN, então a gente embute isso na tabela de dequantização da seguinte forma:

// The fast inverse transform expects its input already multiplied by
// the AAN constants, so those are folded into the dequantization
// table exactly as they were folded into the encoder's table.
typedef struct {
    float multiplier[64];
} ScaledDequant;

static void build_scaled_dequant(int quantization[8][8], ScaledDequant *out) {
    for (int row = 0; row < 8; row++) {
        for (int column = 0; column < 8; column++) {
            out->multiplier[row * 8 + column] = (float)(
                quantization[row][column] *
                aan_scale_factor(row, column) / 64.0);
        }
    }
}

E a dequantização fica assim:

    // Undo quantization and zig-zag ordering in one pass.
    float coefficients[64];
    for (int i = 0; i < 64; i++) {
        int index = ZIGZAG[i];
        coefficients[index] = zigzag[i] * quantization->multiplier[index];
    }

    idct_block_fast(coefficients, block);

💡 Aquele / 64.0 no meio da conta é o ajuste final de normalização da AAN inversa. Se você tirar, a imagem sai com o contraste completamente errado.

Foi um dos meus bugs, aliás, eu cheguei a aplicar a divisão duas vezes (uma na tabela e outra no código) e o PSNR despencou de 43 dB para 13 dB. 😵  

Os resultados

Primeiro, a transformada isolada. Rodando nos 6144 blocos de uma imagem 512×512 em 4:2:0, teve o seguinte resultado:

Transformada │ Separável │   AAN   │ Ganho
─────────────┼───────────┼─────────┼───────
DCT direta   │  3,34 ms  │ 0,69 ms │  4,8x
IDCT inversa │  5,08 ms  │ 0,66 ms │  7,6x

E o mais importante: os coeficientes são matematicamente os mesmos. Conferi contra a versão lenta e o erro máximo ficou em 0,000041, puro arredondamento de ponto flutuante. Não é uma aproximação; é a mesma transformada, só que calculada de forma mais esperta.

Agora vamos ao pipeline completo, medindo a jornada deste artigo:

Estágio                  │  Encode  │  Decode
─────────────────────────┼──────────┼──────────
Artigo 18 (sem -O2)      │ 68,76 ms │ 47,93 ms
Só ligando -O2           │ 17,85 ms │ 20,34 ms
+ DCT AAN e escala       │  9,86 ms │ 15,81 ms
─────────────────────────┼──────────┼──────────
Ganho total              │   7,0x   │   3,0x

Sete vezes mais rápido no encode, três vezes no decode, e olha que ainda estamos no primeiro artigo da série de otimização. 🎉

O teste que realmente importa...

De que adianta ser rápido se a saída mudou? Essa é a pergunta que qualquer otimização precisa responder, e ela vale mais que qualquer número de velocidade.

Rodei as cinco imagens de teste comparando a saída antes e depois, e se liga no resultado:

Imagem    │ Artigo 18       │ Artigo 19       │ PSNR
──────────┼─────────────────┼─────────────────┼───────
portrait  │ 26.711 B        │ 26.708 B        │ igual
smooth    │ 30.867 B        │ 30.805 B        │ igual
texture   │ 54.317 B        │ 54.318 B        │ igual
micro     │ 35.497 B        │ 35.492 B        │ igual
colorful  │ 25.934 B        │ 25.933 B        │ igual

O PSNR ficou idêntico em todas. Os tamanhos variam em alguns bytes, no smooth, 62 bytes em 30 mil, ou 0,2%.

De onde vem essa diferencinha? Do arredondamento.

A versão antiga usava round() sobre double; a nova usa lrintf() sobre float. Quando um coeficiente cai exatamente no meio do caminho entre dois inteiros, as duas funções podem escolher lados diferentes.

Isso muda um coeficiente aqui e ali, e o Huffman gera alguns bits a mais ou a menos.

É uma diferença real e eu não vou varrer pra debaixo do tapete.

Mas é irrelevante na prática: o formato é o mesmo, a qualidade é a mesma, e o tamanho oscila em fração de por cento, às vezes pra cima, às vezes pra baixo.

O que ainda incomoda?

Terminamos o artigo 7× mais rápidos, mas sendo honesto sobre o que sobrou:

Cada um desses é um pedaço do próximo artigo. 😉

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 transformada rápida AAN e as tabelas de quantização pré-escaladas:

Resumo

Recapitulando o primeiro passo da nossa jornada de otimização:

E olha que ainda estamos no começo. 😄

No próximo artigo, vamos afiar o encoder: matar aquela segunda passada da DCT guardando os coeficientes, trocar a escrita bit a bit por um acumulador de 64 bits, e ensinar o codec a pular direto de um coeficiente sobrevivente ao próximo em vez de varrer os 63.

E vou contar a história de uma otimização que parecia genial, produzia arquivos do tamanho certo, não gerava nenhum aviso... e destruía a imagem completamente.

E não vou mentir, eu levei um bom tempo pra entender o porquê, e a explicação é uma das mais bonitas da série. 🐛

Até a próxima! 👋