Привет, Хабр!
В прошлой статье я рассказывал об ускорении копирования элементов одного слайса в другой с помощью средств 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://habr.com/ru/post/720582/
Добавить комментарий