Como Salvar e Restaurar Registradores SSE e AVX em Assembly NASM

Ao escrever funções ou seções de código que utilizam as extensões SSE/AVX, é uma boa prática salvar o estado dos registradores XMM/YMM antes de usá-los e restaurá-los ao final. Isso garante que o seu código não interfira com outras partes do programa que possam depender dos valores nesses registradores. O local mais comum para esse armazenamento temporário é a pilha.

Um desafio comum ao trabalhar com SSE e AVX é não saber previamente se o processador suporta as extensões AVX mais recentes. Para lidar com isso, é essencial verificar a capacidade da CPU em tempo de execução. A instrução CPUID é a ferramenta padrão para essa tarefa.

O programa a seguir para Linux demonstra como salvar e restaurar os primeiros seis registradores XMM/YMM, adaptando-se dinamicamente ao suporte de AVX.

global _start

section .data
    AVXSupport equ 0x10000000  ; Máscara para verificar o bit 28

section .text
_start:
    mov rax, 22
    movq xmm0, rax            ; Coloca um valor de teste em xmm0

    ; Reserva espaço na pilha para os registradores
    ; 6 registradores SSE (XMM) = 6 * 16 = 96 bytes
    ; 6 registradores AVX (YMM) = 6 * 32 = 192 bytes
    ; Alocamos 192 bytes para cobrir ambos os casos.
    sub rsp, 192

    ; Determina se a extensão AVX e os registradores YMM estão disponíveis
    mov rax, 1
    cpuid
    and ecx, AVXSupport       ; Verifica o bit 28 para suporte a AVX
    jnz saveAVX               ; Se AVX for suportado, salta para saveAVX

    ; Se AVX não for suportado, salva o estado dos registradores XMM0-XMM5
    movdqu [rsp + 0], xmm0
    movdqu [rsp + 16], xmm1
    movdqu [rsp + 32], xmm2
    movdqu [rsp + 48], xmm3
    movdqu [rsp + 64], xmm4
    movdqu [rsp + 80], xmm5
    jmp afterSave

saveAVX:
    ; Salva o estado dos registradores YMM0-YMM5
    vmovdqu [rsp + 0], ymm0
    vmovdqu [rsp + 32], ymm1
    vmovdqu [rsp + 64], ymm2
    vmovdqu [rsp + 96], ymm3
    vmovdqu [rsp + 128], ymm4
    vmovdqu [rsp + 160], ymm5

afterSave:
    ; Aqui, outras operações do programa seriam executadas
    mov rax, 15
    movq xmm0, rax            ; Modifica xmm0 para testar a restauração

    ; Restaura os valores dos registradores. O registrador ECX ainda contém
    ; o resultado do CPUID, então podemos reutilizá-lo.
    and ecx, AVXSupport
    jnz restoreAVX

    ; Se AVX não for suportado, restaura os registradores XMM
    movdqu xmm0, [rsp]
    movdqu xmm1, [rsp + 16]
    movdqu xmm2, [rsp + 32]
    movdqu xmm3, [rsp + 48]
    movdqu xmm4, [rsp + 64]
    movdqu xmm5, [rsp + 80]
    jmp exit

restoreAVX:
    ; Restaura os registradores YMM
    vmovdqu ymm0, [rsp + 0]
    vmovdqu ymm1, [rsp + 32]
    vmovdqu ymm2, [rsp + 64]
    vmovdqu ymm3, [rsp + 96]
    vmovdqu ymm4, [rsp + 128]
    vmovdqu ymm5, [rsp + 160]

exit:
    add rsp, 192              ; Libera o espaço alocado na pilha

    movq rdi, xmm0            ; Move o valor restaurado para rdi (deve ser 22)
    mov rax, 60
    syscall

O primeiro passo é verificar o suporte a AVX. A instrução CPUID, com RAX=1, retorna informações sobre as características do processador. O bit 28 do registrador ECX indica se AVX é suportado. Aplicamos a máscara AVXSupport para isolar este bit e, se o resultado não for zero, saltamos para a rotina que salva os registradores YMM de 256 bits. Caso contrário, executamos a rotina para os registradores XMM de 128 bits.

Para simplificar, alocamos na pilha o espaço máximo necessário (192 bytes para 6 registradores YMM). Se apenas SSE for suportado, a rotina de salvamento dos XMM simplesmente ignorará o espaço extra. Para salvar e restaurar, são usadas as instruções movdqu (para XMM) e vmovdqu (para YMM), que são adequadas para mover dados que podem não estar alinhados em fronteiras de 16 ou 32 bytes, uma situação comum na pilha.

O processo de restauração espelha o de salvamento, usando a mesma verificação para determinar quais registradores (XMM ou YMM) devem ser lidos de volta da pilha.

A seguir, a versão análoga do programa para Windows.

global _start

section .data
    AVXSupport equ 0x10000000  ; Máscara para verificar o bit 28

section .text
_start:
    mov rax, 22
    movq xmm0, rax

    sub rsp, 192

    mov rax, 1
    cpuid
    and ecx, AVXSupport
    jnz saveAVX

    movdqu [rsp + 0], xmm0
    movdqu [rsp + 16], xmm1
    movdqu [rsp + 32], xmm2
    movdqu [rsp + 48], xmm3
    movdqu [rsp + 64], xmm4
    movdqu [rsp + 80], xmm5
    jmp afterSave

saveAVX:
    vmovdqu [rsp + 0], ymm0
    vmovdqu [rsp + 32], ymm1
    vmovdqu [rsp + 64], ymm2
    vmovdqu [rsp + 96], ymm3
    vmovdqu [rsp + 128], ymm4
    vmovdqu [rsp + 160], ymm5

afterSave:
    mov rax, 15
    movq xmm0, rax

    and ecx, AVXSupport
    jnz restoreAVX

    movdqu xmm0, [rsp]
    movdqu xmm1, [rsp + 16]
    movdqu xmm2, [rsp + 32]
    movdqu xmm3, [rsp + 48]
    movdqu xmm4, [rsp + 64]
    movdqu xmm5, [rsp + 80]
    jmp exit

restoreAVX:
    vmovdqu ymm0, [rsp + 0]
    vmovdqu ymm1, [rsp + 32]
    vmovdqu ymm2, [rsp + 64]
    vmovdqu ymm3, [rsp + 96]
    vmovdqu ymm4, [rsp + 128]
    vmovdqu ymm5, [rsp + 160]

exit:
    add rsp, 192

    movq rax, xmm0    ; Move o valor restaurado para rax (código de saída)
    ret

A lógica para Windows é idêntica, alterando apenas a forma como o programa termina e retorna o valor de verificação (através do registrador RAX na instrução ret).

Resumo

É fundamental salvar e restaurar os registradores XMM/YMM ao usá-los para garantir a integridade do estado do programa.
A instrução CPUID é usada para detectar o suporte a AVX em tempo de execução, verificando o bit 28 do registrador ECX.
O espaço para os registradores é alocado na pilha com a instrução sub rsp.
As instruções movdqu (para SSE/XMM) e vmovdqu (para AVX/YMM) são usadas para mover os dados entre os registradores e a pilha.