Assembler в Go: техники ускорения и оптимизации / forpes.ru

06.03.2023 07:56

svistunov 1 3200 Источник

Привет, Хабр!

В прошлой статье я рассказывал об ускорении копирования элементов одного слайса в другой с помощью средств Go. В этот раз я решил пойти дальше и посмотреть, что можно достичь, начав разговаривать с процессором на его языке. Я выбрал одну из оптимизированных версий функции Copy в качестве объекта исследования из решения задачи VK Cup'22/23, которая копирует только синий компонент RGBA в Paletted картинку:

func Copy(srcRGBA, dstPaletted []uint8) {
	srcPtr := unsafe.Add(unsafe.Pointer(&srcRGBA[0]), 2)
	dstPtr := unsafe.Pointer(&dstPaletted[0])
	for range dstPaletted {
		*(*uint8)(dstPtr) = *(*uint8)(srcPtr)
		dstPtr = unsafe.Add(dstPtr, 1)
		srcPtr = unsafe.Add(srcPtr, 4)
	}
}

Ниже я расскажу как её ускорить почти в 10 раз.

Disclaimer

С вероятностью 99.9% приведённое ниже никогда не пригодится в реальной разработке, более того, я абсолютно не гарантирую, что данный код является 100% надёжным. Вообще последний раз я игрался с ассемблером почти 20 лет назад и всё описанное здесь было сделано чисто в исследовательских целях для саморазвития.

Ассемблер в Go

Знакомство с Ассемблером в Go я начал с документа "A Quick Guide to Go's Assembler". В нём кратко описано что это такое и как оно устроено, но, если честно, с первого раза мало что понятно, кроме того, что это не привычный ассемблер, а Plan 9 Assembler. Поэтому первое что приходит в голову, это посмотреть как выглядит приведённая выше функция Copy на нём. К счастью, Go предоставляет флаг компиляции -gcflags "-S", который выводит программу в виде машинного кода. Всю работу я делал в *_test.go файле, поэтому для его превращения в код на Ассемблере надо выполнить go test -c -gcflags "-S -B" ./fastcopy, где fastcopy - директория с исходниками. Я также отключил Bounds Check с помощью флага -B, чтобы упростить код. На выходе будет довольно много текста, но самое интересное, это как выглядит функция Copy:

0x0000 00000     TEXT    round-3/fastcopy.Copy(SB), NOSPLIT|ABIInternal, $0-48
0x0000 00000     MOVQ    AX, round-3/fastcopy.srcRGBA+8(FP)
0x0005 00005     MOVQ    DI, round-3/fastcopy.dstPaletted+32(FP)
0x000a 00010     FUNCDATA        $0, gclocals·cNGUyZq94N9QFR70tEjj5A==(SB)
0x000a 00010     FUNCDATA        $1, gclocals·J5F+7Qw7O7ve2QcWC7DpeQ==(SB)
0x000a 00010     FUNCDATA        $5, round-3/fastcopy.Copy.arginfo1(SB)
0x000a 00010     FUNCDATA        $6, round-3/fastcopy.Copy.argliveinfo(SB)
0x000a 00010     PCDATA  $3, $1
0x000a 00010     ADDQ    $2, AX
0x000e 00014     XORL    CX, CX
0x0010 00016     JMP     33
0x0012 00018     MOVBLZX (AX), DX
0x0015 00021     MOVB    DL, (DI)
0x0017 00023     INCQ    CX
0x001a 00026     INCQ    DI
0x001d 00029     ADDQ    $4, AX
0x0021 00033     CMPQ    SI, CX
0x0024 00036     JGT     18
0x0026 00038     RET

Строка 1 - это описание функции, состоящее из названия функции, флагов и размера стека. Стек равен 48 байтам, так как функция получает 2 слайса в качестве параметров, а каждый слайс состоит из структуры с тремя полями (указатель на массив с данными, длина и капасити), каждое из которых равно 8 байт. В строках 2 и 3 в регистры помещаются их длины.

Строки 4-8 сообщают информацию для Garbage Collector, а далее уже идёт само тело функции.

Из необычного в глаза бросается отсутствие регистров RAX, EAX, ..., везде используется AX. Первая мысль была: почему всё 16 битное? На самом деле размер регистров задаётся командой, например MOVQ $1, AX - это MOVQ $1, %RAX, а MOVD $1, AX - это MOVD $1, %EAX.

Первая функция на Ассемблере.

Прежде чем создавать функцию на Ассемблере, сначала надо описать её на Go, для этого я в файле fastcopy.go рядом с функцией Copy добавил функцию CopyAsm без тела:

func CopyAsm(srcRGBA, dstPaletted []uint8)

Реализацию на ассемблере надо положить в файл с расширением .s, в моём случае это fastcopy_amd64.s. Суффикс _amd64 нужен для того, чтобы не дать собраться приложению на других архитектурах, ведь для них нужна будет другая реализация. Первая версия получилась такая:

TEXT ·CopyAsm(SB),$0-48                 // Имя функции должно начинаться с ·, флагов нет, размер стека 48 байт
        MOVQ    srcRGBA+0(FP), AX       // В AX кладём адрес массива srcRGBA
        ADDQ    $2, AX                  // И смещаем его на 2, т.к. нужен только синий компонент
        MOVQ    dstPaletted+24(FP), CX  // В CX кладём адрес массива dstPaletted
        MOVQ    CX, BX                  // Для выхода из цикла нужен адрес последнего элемента dstPaletted
        ADDQ    $512*512, BX            // Берём начальный адрес и добавляем к нему 512*512

LOOP:   MOVBLZX (AX), DX                // Копируем 1 байт по адресу который хранится в AX (srcRGBA) в DX
        MOVB    DL, (CX)                // И затем копируем его в ячейку массива dstPaletted.

        INCQ    CX                      // Увеличиваем адрес ткущей ячейки dstPaletted на 1
        ADDQ    $4, AX                  // Увеличиваем адрес ткущей ячейки srcRGBA на 4

        CMPQ    BX, CX                  // Если не достигли конца
        JGT     LOOP                    // То идём на строку 8
        RET                             // Иначе выход

В данной реализации я избавился от хранения переменной цикла в отдельном регистре, а так она очень похожа на то, что генерирует Go. Если сравнить их производительность, то, как и ожидалось, они примерно на одном уровне:

cpu: AMD Ryzen 7 5800H with Radeon Graphics         
BenchmarkCopy
BenchmarkCopy-16                   	   93292	    128104 ns/op
BenchmarkCopyAsm-16                	   94081	    127012 ns/op

1. Уменьшаем количество чтений из памяти

В строке 8 функции CopyAsm читается только 1 байт из памяти в регистр имеющий размер 8 байт. Первое что приходит в голову, а не прочитать ли за раз все 8 байт, а дальше уже с помощью SHRQ достать второй элемент:

Код получился таким:

TEXT ·CopyAsm1R2W(SB),$0-48
        MOVQ    srcRGBA+0(FP), AX
        ADDQ    $2, AX
        MOVQ    dstPaletted+24(FP), CX
        MOVQ    CX, BX
        ADDQ    $512*512, BX

LOOP:
        MOVQ (AX), DX      // Копируем 8 байт в DX
        MOVB    DL, (CX)   // Сохраняем первый синий компонент
        SHRQ    $32, DX    // Смещаем на 4 байта
        MOVB    DL, 1(CX)  // Сохраняем второй синий компонент

        ADDQ    $2, CX     // Адрес получателя смещаем уже на 2, т.к. сохранили 2 элемента
        ADDQ    $8, AX     // А адрес источника на прочитанные за раз 8 байт

        CMPQ    BX, CX
        JGT     LOOP
        RET

С точки зрения производительности всё стало сильно лучше:

cpu: AMD Ryzen 7 5800H with Radeon Graphics         
BenchmarkCopy
BenchmarkCopy-16                   	   93292	    128104 ns/op
BenchmarkCopyAsm-16                	   94081	    127012 ns/op
BenchmarkCopyAsm1R2W-16            	  185229	     64904 ns/op

А если ещё развернуть цикл, то можно добиться ещё большей производительности:

cpu: AMD Ryzen 7 5800H with Radeon Graphics         
BenchmarkCopy
BenchmarkCopy-16                   	   93292	    128104 ns/op
BenchmarkCopyAsm-16                	   94081	    127012 ns/op
BenchmarkCopyAsm1R2W-16            	  185229	     64904 ns/op
BenchmarkCopyAsm1R2WUnrolled-16    	  240812	     49406 ns/op

TEXT ·CopyAsm1R2WUnrolled(SB),$0-48
        MOVQ    srcRGBA+0(FP), AX
        ADDQ    $2, AX
        MOVQ    dstPaletted+24(FP), CX
        MOVQ    CX, BX
        ADDQ    $512*512, BX

LOOP:
        // 1
        MOVQ    (AX), DX
        MOVB    DL, (CX)
        SHRQ    $32, DX
        MOVB    DL, 1(CX)

        // 2
        MOVQ    8(AX), DX
        MOVB    DL, 2(CX)
        SHRQ    $32, DX
        MOVB    DL, 3(CX)

        ADDQ    $16, AX
        ADDQ    $4, CX

        CMPQ    BX, CX
        JGT     LOOP
        RET

Увеличение развёртки профита не дало, по крайней мере на моём CPU.

2. Уменьшаем количество записей в память

Анализируя код, появляется желание избавиться от двух MOVB DL, (CX) и превратить их в один MOVW BX, (CX), где регистр BX будет хранить 2 байта:

TEXT ·CopyAsmAcc2(SB),$0-48
        MOVQ    srcRGBA+0(FP), AX
        ADDQ    $2, AX
        MOVQ    dstPaletted+24(FP), CX
        MOVQ    CX, DI
        ADDQ    $512*512, DI

LOOP:   MOVQ    (AX), DX   // Копируем 8 байт в DX
        MOVB    DL, BL     // Сохраняем первый синий компонент в BL
        SHRQ    $32, DX    // Смещаем на 4 байта
        MOVB    DL, BH     // Сохраняем второй синий компонент в BH

        MOVW    BX, (CX)   // Копируем 2 байта в dstPaletted

        ADDQ    $2, CX
        ADDQ    $8, AX

        CMPQ    DI, CX
        JGT     LOOP
        RET

И тут меня ждал сюрприз, этот вариант работает сильно медленнее предыдущего:

cpu: AMD Ryzen 7 5800H with Radeon Graphics         
BenchmarkCopy
BenchmarkCopy-16                   	   93292	    128104 ns/op
BenchmarkCopyAsm-16                	   94081	    127012 ns/op
BenchmarkCopyAsm1R2W-16            	  185229	     64904 ns/op
BenchmarkCopyAsm1R2WUnrolled-16    	  240812	     49406 ns/op
BenchmarkCopyAsmAcc2-16            	  155895	     76325 ns/op

Причём я случайно заметил, что если на первом шаге цикла делать XORQ BX, BX, то время выполнения немного уменьшается. Но это уже шаманство.

Следующая попытка была собрать в регистре BX сразу 8 байт и только после этого их записать в память:

TEXT ·CopyAsmAcc8(SB),$0-48
        MOVQ    srcRGBA+0(FP), AX
        ADDQ    $2, AX
        MOVQ    dstPaletted+24(FP), CX
        MOVQ    CX, DI
        ADDQ    $512*512, DI

LOOP:   XORQ    BX, BX
        MOVQ    (AX), DX
        MOVB    DL, BL
        SHLQ    $8, BX
        SHRQ    $32, DX
        MOVB    DL, BL

        MOVQ    8(AX), DX
        SHLQ    $8, BX
        MOVB    DL, BL
        SHLQ    $8, BX
        SHRQ    $32, DX
        MOVB    DL, BL

        MOVQ    16(AX), DX
        SHLQ    $8, BX
        MOVB    DL, BL
        SHLQ    $8, BX
        SHRQ    $32, DX
        MOVB    DL, BL

        MOVQ    24(AX), DX
        SHLQ    $8, BX
        MOVB    DL, BL
        SHLQ    $8, BX
        SHRQ    $32, DX
        MOVB    DL, BL

        BSWAPQ  BX           // Байты получились в обратном порядке, пожтому переворачиваем
        MOVQ    BX, (CX)     // Записываем 8 байт в dstPaletted

        ADDQ    $32, AX
        ADDQ    $8, CX

        CMPQ    DI, CX
        JGT     LOOP
        RET

Этот вариант победил CopyAsm1R2W, но в Unrolled версии немного проиграл:

cpu: AMD Ryzen 7 5800H with Radeon Graphics         
BenchmarkCopy
BenchmarkCopy-16                   	   93292	    128104 ns/op
BenchmarkCopyAsm-16                	   94081	    127012 ns/op
BenchmarkCopyAsm1R2W-16            	  185229	     64904 ns/op
BenchmarkCopyAsm1R2WUnrolled-16    	  240812	     49406 ns/op
BenchmarkCopyAsmAcc2-16            	  155895	     76325 ns/op
BenchmarkCopyAcc8-16               	  203373	     57072 ns/op
BenchmarkCopyAcc8Unrolled-16       	  218060	     51953 ns/op

SIMD

Все предыдущие попытки упирались в то, что команды выполняются последовательно, ну почти все, а хочется обрабатывать данные блоками. Для этого в CPU есть векторные инструкции. Для C++ есть Intrinsics, которые их используют, в нашем случае придётся использовать инструкции напрямую. К счастью на моей любимой странице про Intrinsics есть описание ассемблерных команд. Например

Вообще есть несколько поколений векторных инструкций, начиная от совсем старых типа SSE, MMX, до современных - типа AVX/AVX2 позволяющих работать с регистрами размером 256 бит. Или даже AVX-512, но поддержка этих инструкций мало где есть. Мой AMD Ryzen поддерживает инструкции AVX2, поэтому буду работать с ними.

Идея: копируем 32 байта из памяти в AVX регистр, а их нам доступно 16 штук (Y0-Y15), перемещаем за 1 команду в нужные позиции 8 компонент синего и записываем их в память. Но, к сожалению, нет такой команды, с помощью которой можно переместить каждый четвёртый байт в первые 8 байт, но есть замечательная команда _mm256_shuffle_epi8, которая позволяет переместить 4 нужных байта в первых 16ти, и 4 - во вторых 16. Если прочитать 4*32 байт в четыре регистра, расставить в них байты в нужном порядке и объединить с помощью операции OR, то получим нужный результат:

Для VPSHUFB и для VPERMD нужны 32 байтные маски, я их определил как глобальные переменные в Go в файле fastcopy.go:

var (
	shuffleMask1 = [32]byte{
		2, 6, 10, 14, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
		2, 6, 10, 14, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
	}

	shuffleMask2 = [32]byte{
		128, 128, 128, 128, 2, 6, 10, 14, 128, 128, 128, 128, 128, 128, 128, 128,
		128, 128, 128, 128, 2, 6, 10, 14, 128, 128, 128, 128, 128, 128, 128, 128,
	}

	shuffleMask3 = [32]byte{
		128, 128, 128, 128, 128, 128, 128, 128, 2, 6, 10, 14, 128, 128, 128, 128,
		128, 128, 128, 128, 128, 128, 128, 128, 2, 6, 10, 14, 128, 128, 128, 128,
	}

	shuffleMask4 = [32]byte{
		128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 2, 6, 10, 14,
		128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 2, 6, 10, 14,
	}

	permutateMask = [8]uint32{0, 4, 1, 5, 2, 6, 3, 7}
)

VPSHUFB в результат вставляет 0, если старший бит байта в маске равен 1, иначе вставляет число из позиции, на которую указывает байт в маске. VPERMD просто переставляет 4 байтовые слова в указанном в маске порядке.

Из Ассемблера к этим переменным можно добраться с помощью конструкции ·имяПеремнной(SB), символ "·" в начале очень важен. В итоге код получился такой:

TEXT ·CopyAsmAvx(SB),$0-48
        MOVQ    srcRGBA+0(FP), AX
        MOVQ    dstPaletted+24(FP), CX
        MOVQ    CX, DI
        ADDQ    $512*512, DI

        // Копируем маски в регистры Y10-Y14 
        VMOVDQU ·shuffleMask1(SB), Y10
        VMOVDQU ·shuffleMask2(SB), Y11
        VMOVDQU ·shuffleMask3(SB), Y12
        VMOVDQU ·shuffleMask4(SB), Y13
        VMOVDQU ·permutateMask(SB), Y14

LOOP:
        // Копируем 128 байт из памяти в регистры Y0-Y4
        VMOVDQU (AX), Y0
        VMOVDQU 32(AX), Y1
        VMOVDQU 64(AX), Y2
        VMOVDQU 96(AX), Y3

        // Применяем маски, чтобы расставить синюю компоненту в нужные позиции
        VPSHUFB Y10, Y0, Y0
        VPSHUFB Y11, Y1, Y1
        VPSHUFB Y12, Y2, Y2
        VPSHUFB Y13, Y3, Y3

        // Объединяем полученные данные в регистр Y0
        VPOR Y0, Y1, Y0
        VPOR Y2, Y3, Y2
        VPOR Y0, Y2, Y0

        // Исправляем порядок
        VPERMD Y0, Y14, Y0

        // Сохраняем 32 числа в dstPaletted
        VMOVDQU Y0, (CX)

        ADDQ    $32, CX
        ADDQ    $128, AX

        CMPQ    DI, CX
        JGT     LOOP
        RET

Этот вариант работает уже сильно быстрее, чем все предыдущие:

cpu: AMD Ryzen 7 5800H with Radeon Graphics         
BenchmarkCopy
BenchmarkCopy-16                   	   93292	    128104 ns/op
BenchmarkCopyAsm-16                	   94081	    127012 ns/op
BenchmarkCopyAsm1R2W-16            	  185229	     64904 ns/op
BenchmarkCopyAsm1R2WUnrolled-16    	  240812	     49406 ns/op
BenchmarkCopyAsmAcc2-16            	  155895	     76325 ns/op
BenchmarkCopyAcc8-16               	  203373	     57072 ns/op
BenchmarkCopyAcc8Unrolled-16       	  218060	     51953 ns/op
BenchmarkCopyAsmAvx-16             	  823080	     14464 ns/op

Заключение

Можно ли ещё ускорится? Думаю что да, например если работать с выровненными данными, возможно есть более оптимальные способы как решить эту задачу, ... Если есть идеи, то пишите в комментариях или проверяйте их в похожей задаче на HighLoad.Fun.

Полезные ссылки:

Assembler в Go: техники ускорения и оптимизации +28