Сохранение состояния регистров SSE/AVX

При работе с инструкциями расширений SSE/AVX нам может потребоваться сохранить значения регистров XMM/YMM. Рассмотрим на примере сохранения первых 6 регистров XMM/YMM:

.globl _start

.data
AVXSupport = 0x10000000  # проверяем бит 28

.text
_start: 
    movq $22, %rax
    movq %rax, %xmm0    # для теста помещаем некоторое число в xmm0

    # Резервируем в стеке место  для регистров AVX/SSE
    # Для 6 регистров SSE необходимо только 96 байт
    # Для 6 регистров AVX 196 байт 
    # Выделяем одно пространство для двух типов расширений - 196 байт
    # Тогда SSE просто игнорируют дополнительные 64 байта
    
    subq $192, %rsp
    # Определяем, доступны ли расширения AVX и соответственно регистры YMM
    movq $1, %rax
    cpuid
    andl $AVXSupport, %ecx     # проверяем на поддержку AVX бит 28
    jnz saveAVX         # если AVX доступно, переходим к метке preserveAVX
    # Если поддержка расширений AVX отсутствует, сохраняет состояние регистров XMM0-XMM5
    movdqu %xmm0, 0(%rsp) 
    movdqu %xmm1, 16(%rsp) 
    movdqu %xmm2, 32(%rsp) 
    movdqu %xmm3, 48(%rsp) 
    movdqu %xmm4, 64(%rsp) 
    movdqu %xmm5, 80(%rsp) 
    jmp afterSave
    # сохраняем состояние регистров YMM0-YMM5
saveAVX: 
    vmovdqu %ymm0, 0(%rsp)
    vmovdqu %ymm1, 32(%rsp)
    vmovdqu %ymm2, 64(%rsp)
    vmovdqu %ymm3, 96(%rsp)
    vmovdqu %ymm4, 128(%rsp)
    vmovdqu %ymm5, 160(%rsp)

afterSave:      # здесь какие-нибудь действия в программе
    movq $15, %rax      # например, поместим в регистр RAX какое-нибудь число
    movq %rax, %xmm0    # и из регистра rax в xmm0, чтобы он изменился

# восстанавливаем значения регистров - в регистре ECX по прежнему значение для AVXSupport
# можно его использовать, чтобы заново не вызывать инструкцию cpuid
    andl $AVXSupport, %ecx     # проверяем на поддержку AVX бит 28
    jnz restoreAVX         # если AVX доступно, переходим к метке restoreAVX

# если не доступно AVX, восстанавливаем регистры XMM
    movdqu 0(%rsp), %xmm0
    movdqu 16(%rsp), %xmm1
    movdqu 32(%rsp), %xmm2
    movdqu 48(%rsp), %xmm3
    movdqu 64(%rsp), %xmm4
    movdqu 80(%rsp), %xmm5
    jmp exit
# восстанавливаем регистры YMM
restoreAVX:
    vmovdqu 0(%rsp), %ymm0
    vmovdqu 32(%rsp), %ymm1
    vmovdqu 64(%rsp), %ymm2
    vmovdqu 96(%rsp), %ymm3
    vmovdqu 128(%rsp), %ymm4
    vmovdqu 160(%rsp), %ymm5
exit:
    addq $192, %rsp

    movq %xmm0, %rdi    # проверяем восстановленное значение
    movq $60, %rax 
    syscall

Как и в общем случае, проблема может заключаться в том, что мы можем не знать точно, какие именно расширения поддерживаются на текущей машине. Поэтому сначала с помощью инструкции cpuid проверяем наличие поддержки регистров YMM. Для этого к результату инструкции в регистре ECX применяем маску AVXSupport, которая выявляет установку 28 бита (который собственно и говорит о поддержке AVX):

movq $1, %rax
cpuid
andl $AVXSupport, %ecx     # проверяем на поддержку AVX бит 28
jnz saveAVX 

Если бит установлен, соответственно после выполнения инструкции AND результат ненулевой, то переходим к метке saveAVX для сохранения состояния регистров YMM. Если бит не установлен, то сохраняем регистры XMM.

Для сохранения 6 регистров XMM нужно 16 байт * 6 = 96 байт. На сохранение 6 регистров YMM - 196 байт. В данном случае для простоты для обоих типов регистров в стеке выделяется 196 байт. Затем последовательно сохраняем регистры:

    movdqu %xmm0, 0(%rsp) 
    movdqu %xmm1, 16(%rsp) 
    movdqu %xmm2, 32(%rsp) 
    movdqu %xmm3, 48(%rsp) 
    movdqu %xmm4, 64(%rsp) 
    movdqu %xmm5, 80(%rsp) 
    jmp afterSave
    # сохраняем состояние регистров YMM0-YMM5
saveAVX: 
    vmovdqu %ymm0, 0(%rsp)
    vmovdqu %ymm1, 32(%rsp)
    vmovdqu %ymm2, 64(%rsp)
    vmovdqu %ymm3, 96(%rsp)
    vmovdqu %ymm4, 128(%rsp)
    vmovdqu %ymm5, 160(%rsp)

При завершении программы аналогично восстанавливаем регистры.