Прежде чем перейти к статье, хочу вам представить, экономическую онлайн игру Brave Knights, в которой вы можете играть и зарабатывать. Регистируйтесь, играйте и зарабатывайте!
Привет, Хабр!
В прошлой статье я рассказывал об ускорении копирования элементов одного слайса в другой с помощью средств 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/op1. Уменьшаем количество чтений из памяти
В строке 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/opTEXT ·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/opSIMD
Все предыдущие попытки упирались в то, что команды выполняются последовательно, ну почти все, а хочется обрабатывать данные блоками. Для этого в 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.
Полезные ссылки:
https://go.dev/doc/asm
https://www.amd.com/system/files/TechDocs/24592.pdf
https://www.agner.org/optimize/instruction_tables.pdf
https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html
