AVX512-VNNI -- Синтез алгоритма для сверточной нейронной сети

Необходимо сделать самый быстрый алгоритм для свертки (convolution). На свертках делается современное представление о нейронных сетях, так и называется сверточные нейронные сети (CNN). Производители процессоров заявляют: мы добавили искусственный интеллект (AI) в процессоры. По сути это означает только одно - увеличили длину вектора для операции умножения. Операции свертки используются в цифровых фильтрах обработки сигналов(DSP) и в задачах связанных с обработкой звука и изображения. Есть задачи выделения контуров изображений и задачи выделения движущихся объектов на изображении. Есть задачи увеличения четкости изображения, пересчет цветности. Все подобные задачи цифровой обработки изображений требуют матричного умножения NxN элементов. При этом коэффициенты могут быть заданы числами со знаком 8бит. А значения над которыми производится операция - 8 бит без знака. Под этот класс задач компания Intel разработала систему команд AXV512-VNNI. В систему команд VNNI входит всего две инструкции. Одна из них производит покомпонентное умножение 8битных целых без знака на 8битные целые со знаком и последующее сложение с накоплением в 32 битном приемнике.
Математически выражается как-то так. Знаковые и беззнаковые компоненты перемножаются и суммируются группами по четыре, чтобы не нарушать разрядность вектора. При большом количестве таких операций может возникать переполнение разрядов. Для нейронных сетей абсолютно все выводы классификации делаются на переполнении, если есть переполнение то это и означает переход от шаманства с бубнами к логике. В данном случае операция насыщения поверх сложения не сильно оправдана. Надо выполнить более 16 тыс таких операций, чтобы сложились условия для переполнения.

VPDPBUSDS — Multiply and Add Unsigned and Signed Bytes with Saturation
p1 = ZERO_EXTEND(a.byte[4*i+0]) * SIGN_EXTEND(b.byte[4*i+0])
p2 = ZERO_EXTEND(a.byte[4*i+1]) * SIGN_EXTEND(b.byte[4*i+1])
p3 = ZERO_EXTEND(a.byte[4*i+2]) * SIGN_EXTEND(b.byte[4*i+2])
p4 = ZERO_EXTEND(a.byte[4*i+3]) * SIGN_EXTEND(b.byte[4*i+3])
DEST.dword[i] = SIGNED_DWORD_SATURATE(DEST.dword[i] + p1 + p2 + p3 + p4)

Так же выглядит вторая операция - умножение знаковых 16 битных чисел с накоплением результата в 32 битном приемнике.

VPDPWSSDS — Multiply and Add Signed Word Integers with Saturation
p1 = SIGN_EXTEND(a.word[2*i+0]) * SIGN_EXTEND(b.word[2*i+0])
p2 = SIGN_EXTEND(a.word[2*i+1]) * SIGN_EXTEND(b.word[2*i+1])
DEST.dword[i] = SIGNED_DWORD_SATURATE(DEST.dword[i] + p1 + p2)

Вот я это так и вижу. Первая операция - для умножения матриц, для обработки видео. Вторая - для умножения векторов, для обработки звука.
Я буду исследовать производительность Throughput & Latency(CPI - clock per instruction) с использованием LLVM-MCA, инструмента. Подбираю, как должен выглядеть алгоритм на языке высокого уровня, чтобы выжать из этих инструкций максимум производительности.
Под рукой у меня мануал на 5 тысяч страниц с описанием процессорных инструкций накопленных за всю историю Intel [Intel® 64 and IA-32 Architectures Software Developer’s Manual, Volume 2:Instruction Set Reference] и таблица производительности инструкций на ядрах 10th Gen Intel Core, с микро-архитектурой Ice Lake. Кроме того, практикую магию с использованием компиляторов GCC 9.2 и Clang 9.0. Вывожу результат компиляции в ассемблер и анализирую производительность выделенных фрагментов полученного кода инструментом LLVM-MCA, анализатором кода.

$ gcc -march=icelake-client -O3 -S -o nn.s nn.c
$ llvm-mca.exe -mcpu=icelake-client -timeline nn.s

Производительность заявленная в таблице - одна инструкция с разрядностью 512 бит на такт, т.е. 64 шт умножений. Пишу цикл свертки по вектору (по горизонтали изображения) и по вертикали.

#include <x86intrin.h>
static inline
__attribute__((__target__("avx512vnni")))
void mac_vec_512(__m512i *acc, uint8_t * src, const __m512i* weight, const size_t length)
{
    int x=0;
    for (; x<length/64;x++){
        __m512i w = _mm512_load_si512(&weight[x]);
        __m512i v = _mm512_loadu_si512(&src[x*64]);
        acc[0] = _mm512_dpbusd_epi32(acc[0], v,w);
    }
    if (length&63) {
        __mmask64 mask = (1ULL <<(length&63)) -1;
        __m512i w = _mm512_load_si512(&weight[x]);
        __m512i v = _mm512_maskz_loadu_epi8(mask, &src[x*64]);
        acc[0] = _mm512_dpbusd_epi32(acc[0], v,w);
    }
}

При разработке алгоритма я исхожу из положения - компилятор самостоятельно уберет из процедуры проверки констант и выполнит развертывание цикла. В предыдущей статье я уже объяснял основной прием работы с масками, повторяться не буду. Если размерность вектора не кратна размеру регистра, используется маска. В данном случае по маске выполняется загрузка вектора.
Ниже представлен результат анализа производительности цикла, временная диаграмма. Цикл разернут компилятором.

[0,0]     DeeeeeeeeER    .    .    .    .    vmovdqu32   (%rdx), %zmm2
[0,1]     D========eeeeeeeeeeeER   .    .    vpdpbusd    (%r8), %zmm2, %zmm0
[0,2]     D===================eER  .    .    vmovdqa64   %zmm0, (%rcx)
[0,3]     .DeeeeeeeeE-----------R  .    .    vmovdqu32   64(%rdx), %zmm3
[0,4]     .D===========eeeeeeeeeeeER    .    vpdpbusd    64(%r8), %zmm3, %zmm0
[0,5]     .D======================eER   .    vmovdqa64   %zmm0, (%rcx)
[0,6]     . DeeeeeeeeE--------------R   .    vmovdqu32   128(%rdx), %zmm4
[0,7]     . D==============eeeeeeeeeeeER.    vpdpbusd    128(%r8), %zmm4, %zmm0
[0,8]     . D=========================eER    vmovdqa64   %zmm0, (%rcx)

Видно, что ожидаемой производительности не получили, результат инструкции VPDBUSD выходит каждые четыре такта.
То что происходит дальше - шаманство с бубном. Я беру фрагмент ассемблерного кода и подбираю руками комбинацию инструкций, которая дает максимальную производительность. Я знаю что конвейер способен загружать из памяти два регистра одновременно.

[0,0]     DeeeeeeeeER    .    .    .    .    vmovdqu32  (%rdx), %zmm2
[0,1]     DeeeeeeeeER    .    .    .    .    vmovdqa64  (%r8), %zmm3
[0,2]     D========eeeeER.    .    .    .    vpdpbusd   %zmm3, %zmm2, %zmm0
[0,3]     .D===========eER    .    .    .    vmovdqa64  %zmm0, (%rcx)
[0,4]     .DeeeeeeeeE----R    .    .    .    vmovdqu32  64(%rdx), %zmm2
[0,5]     .DeeeeeeeeE----R    .    .    .    vmovdqa64  64(%r8), %zmm3
[0,6]     . D==========eeeeER .    .    .    vpdpbusd   %zmm3, %zmm2, %zmm0
[0,7]     . D==============eER.    .    .    vmovdqa64  %zmm0, (%rcx)
[0,8]     . DeeeeeeeeE-------R.    .    .    vmovdqu32  128(%rdx), %zmm2
[0,9]     .  DeeeeeeeeE------R.    .    .    vmovdqa64  128(%r8), %zmm3
[0,10]    .  D=============eeeeER  .    .    vpdpbusd   %zmm3, %zmm2, %zmm0
[0,11]    .  D=================eER .    .    vmovdqa64  %zmm0, (%rcx)

При этом производительность не изменилась, но существенно скоратилась задержка. Т.е. выигрыш будет на сравнительно небольших векторах. Видно что задержка в 4 такта определяется исключительно зависимостью от выходного результата в регистре ZMM0. В следующем приближении я уменьшаю зависимость от выходного регистра.

$ llvm-mca.exe -mcpu=icelake-client -timeline nn1.s
[0,0]     DeeeeeeeeER    .    .    .    .    vmovdqa64    (%r8), %zmm4
[0,1]     DeeeeeeeeER    .    .    .    .    movdqu32    (%rdx), %zmm5
[0,2]     D========eeeeER.    .    .    .    vpdpbusd     %zmm4, %zmm5, %zmm3
[0,3]     .DeeeeeeeeE---R.    .    .    .    vmovdqa64    64(%r8), %zmm4
[0,4]     .DeeeeeeeeE---R.    .    .    .    vmovdqu32    64(%rdx), %zmm5
[0,5]     .D========eeeeER    .    .    .    vpdpbusd     %zmm4, %zmm5, %zmm2
[0,6]     . DeeeeeeeeE---R    .    .    .    vmovdqa64    128(%r8), %zmm4
[0,7]     . DeeeeeeeeE---R    .    .    .    vmovdqu32    128(%rdx), %zmm5
[0,8]     . D========eeeeER   .    .    .    vpdpbusd     %zmm4, %zmm5, %zmm1
[0,9]     .  DeeeeeeeeE---R   .    .    .    vmovdqa64    192(%r8), %zmm4
[0,10]    .  DeeeeeeeeE---R   .    .    .    vmovdqu32    192(%rdx), %zmm5
[0,11]    .  D========eeeeER  .    .    .    vpdpbusd     %zmm4, %zmm5, %zmm0

Получил пропускную способность 64 умножения на такт. Результат операции выдается на каждый такт. В данном случае производительность ограничена пропускной способностью загрузки данных. А вот дальше происходит невозможное. Оказывается, производительность можно УДВОИТЬ, если коэффициенты для последующих умножений равны или загружены в регистры заранее. Вот пример такой оптимизации:

$ llvm-mca.exe -mcpu=icelake-client -timeline nn1.s
[0,0]     DeeeeeeeeER    .    .    vmovdqu32      (%rdx), %zmm5
[0,1]     D========eeeeER.    .    vpdpbusd       %zmm4, %zmm5, %zmm3
[0,2]     DeeeeeeeeE----R.    .    vmovdqu32      64(%rdx), %zmm5
[0,3]     D========eeeeER.    .    vpdpbusd       %zmm4, %zmm5, %zmm2
[0,4]     .DeeeeeeeeE---R.    .    vmovdqu32      128(%rdx), %zmm5
[0,5]     .D========eeeeER    .    vpdpbusd       %zmm4, %zmm5, %zmm1
[0,6]     .DeeeeeeeeE----R    .    vmovdqu32      192(%rdx), %zmm5
[0,7]     .D========eeeeER    .    vpdpbusd       %zmm4, %zmm5, %zmm0

# Еще один вариант оптимизации удвоенной производительностью
[0,0]     DeeeeeeeeeeeER .    .    vpdpbusd         (%rdx), %zmm9, %zmm0
[0,1]     DeeeeeeeeeeeER .    .    vpdpbusd       64(%rdx), %zmm9, %zmm1
[0,2]     D=eeeeeeeeeeeER.    .    vpdpbusd      128(%rdx), %zmm9, %zmm2
[0,3]     .DeeeeeeeeeeeER.    .    vpdpbusd      192(%rdx), %zmm9, %zmm3
[0,4]     .D=eeeeeeeeeeeER    .    vpdpbusd      256(%rdx), %zmm9, %zmm4
[0,5]     .D=eeeeeeeeeeeER    .    vpdpbusd      320(%rdx), %zmm9, %zmm5
[0,6]     . D=eeeeeeeeeeeER   .    vpdpbusd      384(%rdx), %zmm9, %zmm6
[0,7]     . D=eeeeeeeeeeeER   .    vpdpbusd      448(%rdx), %zmm9, %zmm7

Чтобы получить максимальную производительность надо добиться, чтобы выход компилятора выглядел именно так. Я оговорюсь. Результат может быть недостоверным. Надо проверять на целевой архитектуре. Следует отметить, что при оптимизации небольших циклов практически все векторные инструкции есть смысл отделять от инструкций загрузки регистров, так код получается быстрее, хотя и выглядит немного избыточным, у планировщика появляется больще возможностей загрузить конвейер неоднородными операциями.

#include <x86intrin.h>
static inline
__attribute__((__target__("avx512vnni")))
void mac_vec_512(__m512i *acc, uint8_t * src, const __m512i* weight, const size_t length)
{
    int x=0;
    for (; x<length/64;x++){
        __m512i v = _mm512_loadu_si512(&src[x*64]);
        acc[x&7] = _mm512_dpbusd_epi32(acc[x&7], v,weight[x]);
    }
    if (length&63) {
        __mmask64 mask = (1ULL <<(length&63)) -1;
        __m512i v = _mm512_maskz_loadu_epi8(mask, &src[x*64]);
        acc[x&7] = _mm512_dpbusd_epi32(acc[x&7], v, weight[x]);
    }
}

Тут мы используем восемь аккумуляторных регистров. Потом вектора аккумуляторов надо сложить между собой, выполнить горизонтальное суммирование. Такой красивый результат, как в предыдущей "ручной" опмтимизации, получить с помощью компилятора не удается. В большинстве случаев компилятор выдает что-то не эффективное. Сказывается еще одно аппаратное ограничение - тормозит декодер иснтрукций, у него своя производительность, ограниченная длиной буфера разбора. Не всегда этой длины хватает, чтобы обрабатывать четыре инсрукции в кодировке EVEX-512 на одном такте.
Наилучший результат без ручных оптимизаций __(впишу когда протестирую, хотя бы на симуляторе)__ умножений на такт. Теоретический предел 128 умножений на такт.

Нашел у Intel идею для описания подобной операции VDPDUSD там, где ее нет, на инструкциях AVX2. # VPMADDUBSW υ8×s8 → s16 multiples, VPMADDWD broadcast1 s16 → s32, and VPADDD s32 → s32.

// цикл на операциях AVX2
    for (x=0; x<length/32;x++){
        v = _mm256_loadu_si256((void*)&src[x*32]);
        v = _mm256_maddubs_epi16(v,weight[x]);
        v = _mm256_madd_epi16(v,E);
        acc= _mm256_add_epi32(acc, v);
    }

Для AVX2 компилятор выполняет оптимизацию сносно, алгоритм модифицировать не пришлось. Могут быть расхождения в результате, может возникать насыщение в операции VPMADDUBSW. В этом случае рекомендуется снижать разрядность весовых коэффициентов. Теоретический предел для AVX2 получается 32 умножения на такт, инструкции VPMADDUBSW и VPMADDWD имеют производительность две на такт, 0.5 CPI.

Эта статья об удвоении производительности. Сначала берем хорошую новую инструкцию в 2-4 раза более производительную чем раньше, получаем результат 1/4 от заявленной производительности. Потом начинаем шаманить и водить руками, получаем 100% - теоретический предел. Потом ударяем в бубен и получаем все 200%, о такой возможности производитель умолчал. Потом замеряем реальную производительность в цикле и получаем скромный результат всего __160%__ от заявленной в документации, что тоже не плохо. Еще остается методика оптимизации.