# Модуль NeoMathEngineAvx

## Введение
Модуль **NeoMathEngineAvx** предназначен для ускорения отдельных методов NeoML путём реализации их с использованием набора инструкци AVX/AVX2/FMA.
Это позволяет запускать код как на уже довольно старых процессорах x86_64 (Haswell и старше), так и на AMD процессорах, поддерживающих соответствующий набор инструкций.

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

Ядра для матричного умножения написаны с применением intrinsic вызовов для использования в классе `CMatrixMultiplier`.
По результатам замеров матричное умножение с использованием данных ядер выигрывает у MKL на AMD процессорах, но в целом проигрывает MKL на Intel даже с тем же набором инструкций AVX/AVX2/FMA.

Прямая свёртка и примитивы реализованы уже с использованием OpenSource библиотеки [xbyak](https://github.com/herumi/xbyak).
Эта библиотека позволяет генерировать машинные инструкции в рантайме, поэтому в дальнейшем будем именовать данную технологию **JIT компилятором** или просто **JIT**.
Одна из замечательных особенностей `xbyak` - это простая реализации API, которая позволяет писать код очень приближенный по виду и структуре в синтаксису ассемблера от Intel.
Основные преимущества JIT компилятора - это минимизация условных переходов, расчёт всех смещений и подстановка константных значений на этапе компиляции JIT кода, а так же разворачивание циклов для лучшей предвыборки кода конвейером.

## Описание модулей

### Matrix Multiplication Kernels
Тут чистый C++, смысла описывать этот модуль в данном документе нет.

### Forward Convolution
Изначально модуль писался с применением intrinsic вызовов, но впоследствии был переведён на JIT.
Однако реализация в виде шаблонного класса осталась, хотя в ней уже нет никакой необходимости, так как в целом время на время генерации JIT кода шаблонные параметры не влияют, а после генерации кода шаблоны уже не нужны.
*В будущем хорошо бы убрать шаблонные параметры из класса, чтобы упростить код"

JIT код для **Forward Convolution** генерируется отдельно для каждого дескриптора, так как он зависит от различных параметров дескриптора.
Это тот случай, когда ускорение достигается при помощи использования в инструкциях известных заранее констант и смещений.
Работа данного модуля уже описывалась при реализации на intrinsic-ах, и основную логику JIT не поменял.

О том, как работает JIT будет рассказано в следующем разделе.

### Primitives
JIT - очень необычный способ написания кода: с одной стороны, логика работы JIT компилятора пишется на C++, пользуясь всеми преимуществами и новыми фишками этого языка;
с другой стороны, код пишется как бы на ассемблере, получая компактный и довольно оптимальный код (хотя и не всегда читаемый для неподготовленного программиста).
Самое главное, что код пишется исходя из условий, при которых он будет запускаться: сколько раз, с какими переменными/константами и т.д.; это позволяет обеспечить максимальную производительность.
Данные 3 аспекта позволяют добиваться значительных успехов в улучшении производительности.

### Подготовка и ссылочные материалы
1. Для начала нужно ознакомится с базовыми аспектами языка Assembler и с его Intel синтаксисом.
2. Нужно почитать документацию на [xbyak](https://github.com/herumi/xbyak), там довольно небольшой объём, и он хорошо структурирован и сопровождён примерами.
3. Для поиска документации на ассемблерные инструкции я пользовался в основном следующими сайтами:
   * [x86 and amd64 instruction reference](https://www.felixcloutier.com/x86/index.html)
   * [Intel® Intrinsics Guide](https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html)
   * [X86-64_Instruction_Encoding](https://wiki.osdev.org/X86-64_Instruction_Encoding) - полезно для понимания режима адресации при формировании команд.

Вот, вроде и всё, что нужно для уверенной работы с `xbyak`.

### Проблемы или особенности, с которыми столкнулись при использовании xbyak
Считаю, что данный раздел следует поместить ДО описания работы с `xbyak`.

1. В процессе работы нельзя смешивать инструкции AVX и SSE, т.к. это приводит к фатальному замедлению, которое может быть вызвано перегрузкой модуля SIMD инструкций. 
   Более подробно про этом можно почитать [Intel® 64 and IA - 32 Architectures Optimization Reference Manual](https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf), 
   раздел 11.3 **MIXING AVX CODE WITH SSE CODE**: *Assembly/Compiler Coding Rule 72. (H impact, H generality) Add VZEROUPPER instruction after 256-bit AVX instructions are executed and before any function call that might execute SSE code.*
   *Add VZEROUPPER at the end of any function that uses 256-bit AVX instructions.*
2. `xbyak` имеет такую замечательную особенность, что очень трудно сформировать неправильную инструкцию. Внутри вызовов функций имеется множество проверок, которые не позволят вам использовать неправильные операнды внутри инструкции.
3. Получить какой-то прирост скорости при использовании `prefetch` инструкций очень сложно.
   Это связано с тем, что современные процессоры имеют очень хорошую аппаратную реализацию поиска закономерностей при доступе к памяти, и автоматически подгружают необходимые блоки.

### Описание класса `CPrimitivesJit`

#### Общие положения
В интерфейсе класса реализованы 2 механизма вызова функций примитива:
1. Непосредственно вызов метода (JIT реализация сетки LSTM),
2. Через указатель на конкретный примитив (позволяет безболезненно подменить вызовы примитивов в классе `CCpuMathEngine`).

Наибольшее ускорение можно достичь для ёмких операций с точки зрения вычисления.
По этой причине простые примитивы показывают весьма нестабильное ускорение относительно реализации с sse intrinsic-ами, т.к. основное время занимает доступ к памяти и в случае промахов кэша результат замера производительности может сильно гулять.

### Контейнер `gens`, хранящий JIT код функций примитивов
Все реализованные JIT функции перечислены в enum-е `TPrimitive`.
В классе определена структура `CGenerator`, которая по сути своей является экземпляром класса `Xbyak::CodeGenerator`, отвечающий за формирование JIT кода.
Массив **`gens`** хранит все сгенерированные JIT примитивы, их генерация происходит при первом обращении к данному примитиву в методе `GetFunctionRawPtr()`.

### Таблица констант
Таблица констант определяется при помощи ключей **`TTableKey`** и двух контейнеров **`table`** и **`tableOffsets`**.
Эта таблица позволяет хранить в одном месте все константы, которые нужны при выполнении JIT кода и иметь к ним лёгкий доступ.
Т.к. невозможно напрямую задать инициализирующее значение `ymm` регистра в коде инструкции, приходится производить загрузку их из памяти.
Локализация всех констант в одном месте благоприятно сказывается на работе с кэшем.

Таблица констант инициализируется один раз в конструкторе `CPrimitivesJit` вызовом метода `initTable`.
В каждой JIT функции адрес таблицы констант (если он используется) находится в регистре `regTablePtr` (**`r10`**).
Получить этот адрес для инициализации `ymm` регистра можно при помощи методов `getOfft` или `getAddr`, передавая при этом необходимые смещения.
Метод `getAddr` возвращает непосредственный дескриптор адреса `Xbyak::Address` реализующего механизм относительной адресации (относительно регистра `regTablePtr`).

### Структура функции примитива
JIT примитивы создаются в шаблонном методе `initPrimitive`, который 
* либо непосредственно содержит код генерации JIT функции,
* либо является обёрткой для другого шаблонного метода с именем `init...`, который обобщает реализацию похожих функций (как правило, простейших математических примитивов).

Тем не менее можно выделить следующие обязательные шаги:
* **Создание экземпляра класса `CGenerator`**.
* **Определение регистров, что должны быть сохранены перед выполнением функции `preservedReg64` и `preservedYmm`**:
  - Для Windows и Unix это разные наборы регистров, более подробно про них можно узнать в соглашении о вызовах для конкретной ОС.
* **Вызов обязательного метода `Prologue`**:
  - В котором добавляется необходимая преамбула для любой функции, на стеке сохраняются необходимые регистры, а также высчитывается дескриптор адреса, указывающий на область стека, содержащую аргументы вызываемой функции (если такие есть).
* **Определение регистров, используемых внутри функции и их инициализация**:
  - Инициализация либо регистрами, либо значениями из стека, тут тоже важно изучить соглашение о вызовах функции.
    Например, Windows передаёт через регистры только 4 аргумента в то время, как Unix передаёт 6.
    К тому же, если аргументы перемежаются значениями с плавающей точкой, то эти значения также содержатся в разных `xmm` регистрах для Windows и Unix.
    Если аргументов функции больше, чем число отведённых на это регистров, то остальные аргументы передаются через стек.
    Ещё в Windows есть такое понятие как `ShadowSpace` - зарезервированная область на стеке перед аргументами размером в 32 байта.
* **Реализация lambda функции**:
  - Лямбда выполняет основные вычислительные операции для данного примитива.
* **Пакетная обработка входных значений с разворачиванием циклов (unrolling loop)**:
  - Для этого применяется лямбда функция из предыдущего пункта.
    Количество элементов, которые можно развернуть в одном цикле, определяется тем, сколько регистров `ymm` участвует в обработке, а также тем, сколько элементов содержится в типовых случаях использования.
    Например, для сложных примитивов `Exp`, `Sigmoid`, `Tanh` и `RestOfLstm`, где задействовано много регистров `ymm`, не обрабатываем за раз более 2 значений.
* **Выполняем обработку 'хвоста' входных данных**
  - В случае, если длинна входных данных не кратна 8, производится обработка хвостовых значений отдельно.
    Она отличается только тем, что чтение и запись данных производится по маске.
* **Вызов обязательного метода `Epilogue`**:
  - В котором восстанавливаются сохранённые ранее регистры, вызывается обязательная инструкция `VZEROUPPER`, о которой говорилось выше, а также восстанавливается указатели фрейма и стека (инструкция leave).

Для простейших примитивов добавить к предыдущему описанию особо нечего.
Код, выполняющий пакетную обработку для разного количества входных данных, унифицирован и реализован в функции `insertSimpleMathFunction()`

### Описание сложных примитивов
Сложные примитивы состоят из других примитивов (**Sigmoid** и **RestOfLstm**) или вычисляются при помощи полинома (**Tanh** и **Exp**). 
Такие примитивы используют много регистров `ymm` для вычисления одного блока из 8 входных значений, а потому разворачивают за раз как правило не более 2 таких блоков.

В функции `insertPrimitive`, которая обрабатывает 2 блока (2 входных `ymm`), пришлось бы каждую строчку писать дважды, и иметь ещё такую же функцию, которая обрабатывала только 1 блок за раз. 
Подобный подход неминуемо привёл бы к расхождению кода и появлению трудно-детектируемых багов, поэтому в этих примитивах оперируем не отдельными `ymm`, а вектором `ymmVec_t`. 
При инициализации либо вручную проверяем сколько блоков должны обработать:
```
	ymmVec_t forget = wholeYmmNumber == 2 ? ymmVec_t{ ymm0, ymm1 } : ymmVec_t{ ymm0 };
```
либо используем специальную функцию `initFromAux()`, которая нарезает вектор вспомогательных регистров `ymmAux` на небольшие блоки, каждый из которых по размеру равен размеру входных данных (`ymmSrc`).

Подобное оперирование векторами `ymm` заставило нас перегрузить различные методы класса `Xbyak::CodeGenerator`, отвечающие за генерацию инструкций.
Эти перегруженные функции определены в базовом классе `CJitCommon`.
Чтобы иметь возможность перегрузить функцию в одну строчку, пришлось определить ряд дефайнов и шаблонных функций для разного количества и комбинации аргументов. 
Как итог, получили возможность использовать синтаксис, идентичный тому, что есть в `xbyak` для работы с векторами регистров.

#### Tanh
Вычисление данного примитива подробно расписано в комментариях.

Весь диапазон значений аргумента бьётся на отрезки:
1. **Линейный участок:**
   - `[0; linear_ubound]`, где `tanh(x) = x`

2. **Полиномиальные участки:**
   - `[linear_ubound; 0x1.8p-12]` - часть половины бинады
   - `[0x1.8p-12; 0x1.0p-11], ..., [0x1.8p2; 0x1.0p3]` - 29 половин бинад
   - `[0x1.0p3; saturation_ubound]`
   -  Итого, 31 интервал, где значение тангенса вычисляется полиномом 6-й степени.

3. **Участок насыщения:**
   - `[0x1.205966p3; saturation_ubound]` - участок насыщения, где `tanh(x) = 1`.

Остальные действия понятны из кода:
1. **Отбрасываем знак**, чтобы потом добавить его к ответу.
2. **Вычисляем номер бинады**, что соответствует данному аргументу.
3. **Вычисляем `tanh` при помощи полинома**, вычисляем значение гиперболического тангенса.
4. **Обработка линейного участка и участка насыщения**, подставляем значение `x` на линейном участке, и `1` и участке насыщения.
5. **Возвращаем знак** к вычисленному результату.

#### Exp
Экспонента также вычисляется полиномом.
Для вычисления экспоненты разбиваем аргумент на 2 части, поделив его на `ln(2)`.

1. **Разложим аргумент:**
  - Результат деления `x / ln(2)` будет целая часть `n` и остаток `r`, таким образом можно обратно выразить:
    ```
    x = n * ln(2) + r
    exp(x) = exp(n * ln(2) + r) = exp(ln(2))^n * exp(r) = 2^n * exp(r)
    ```
  - Затем `2^n` считается элементарно сдвигом, а `exp(r)` считается полиномом на отрезке `[0;ln(2)]`.

2. **Ограничим аргумент, чтобы избежать переполнения:**
   - Перед вычислением экспоненты нужно ограничить аргумент сверху и снизу, чтобы не получить переполнение.
     Константы `ExpFltMax` и `ExpFltMin` заданы в hex формате. Если `ExpFltMax` увеличить на 1, то при переполнении получим `inf`, что не согласуется с аналогичным результатом у MKL.

**Операции в функции:**
1. Сохраняем маску для значения меньше `ExpFltMin`, в конце функции по этой маске заполним соответствующие результаты нулями.
	```
	gen.vcmpltps( ymmMask, ymmSrc, getAddr( TTableKey::ExpFltMin ) );
	```
2. Ограничиваем аргумент `x` сверху и снизу и кэшируем значение для последующего использования.
	```
	gen.vminps( ymmSrc, ymmSrc, getAddr( TTableKey::ExpFltMax ) );
	gen.vmaxps( ymmSrc, ymmSrc, getAddr( TTableKey::ExpFltMin ) );
	gen.vmovups( ymmAux1, ymmSrc );
	```
3. Вычисляем целую часть от деления `x / ln(2)`, при этом прибавляя `0.5`, чтобы сместить остаток от деления в область нуля для более точного вычисления полинома.
	```
	n = round( x / ln(2) + 0.5 ) = round( x * log2(e) + 0.5 )
	```
4. Вычисляем остаток `r = x - n * ln(2)`
5. В случае, если `n == 128`, получим переполнение при вычислении `2^n`, поэтому поступим следующим образом:
	```
	exp(x) = 2^n * exp(r) = 2^(n-1) * exp(r) * 2
	```
	Выражение `2^(n-1) * exp(r)` будет меньше `2^127` и в итоге переполнения не будет.

	Вычислим `2^(n-1)` при помощи операции сдвига, добавив `ExpBias`, чтобы соблюсти формат числа с плавающей точкой:
	```
	gen.vsubps( ymmSrc, ymmSrc, getAddr( TTableKey::One ) );
	gen.vcvtps2dq( ymmAux2, ymmSrc );
	gen.vpaddd( ymmAux2, ymmAux2, getAddr( TTableKey::ExpBias ) );
	gen.vpslld( ymmAux2, ymmAux2, MantissaNumBits );
	```
6. Применим маску полученную в шаге 1, обнулив значения, которые меньше `ExpFltMin`
	```
	gen.vxorps( ymmSrc, ymmSrc, ymmSrc );
	gen.vblendvps( ymmAux2, ymmAux2, ymmSrc, ymmMask );
	```
7. Посчитаем `exp(r)` при помощи полинома
8. Домножим полученное значение сначала на `2^(n-1)` потом на `2`, чтобы избежать проблемы, полученной в шаге 5.

#### Sigmoid
Примитив Sigmoid по определению считается при помощи экспоненты:

$\[ \text{sigmoid}(x) = \frac{\exp(x)}{1 + \exp(x)} \]$

**В методе `insertPrimitive<Sigmoid>`:**

- **Инициализация регистров:** 
  - Последний регистр `ymmAux` уже инициализирован единицами.
    Это сделано потому, что `insertPrimitive<Sigmoid>` также активно используется в методе `insertPrimitive<RestOfLstm>`, и подобное допущение позволяем избежать лишних действий.
  
- **Функциональный объект `afterPrologue`:**
  - Чтобы не порушить концепцию унифицированных методов в `initActivationFunction()` был добавлен параметр `afterPrologue`.
    Он является функциональным объектом и вызывается после добавления пролога для инициализации последнего регистра `ymmAux` единицами для примитива Сигмоиды.

#### RestOfLstm
Выглядит реализация данного примитива очень громоздко и страшно, но всё становится более понятным, если взглянуть на схему [RestOfLstm.pdf](https://github.com/neoml-lib/neoml/NeoMathEngine/src/CPU/x86/avx/src/docs/RestOfLstm.pdf).

![image](https://github.com/user-attachments/assets/d64e8bbf-e023-40f0-ad27-7cedca41db4b)

На схеме представлено классическое изображение LSTM ячейки и её реализация в коде, разбитая на 3 блока.
Блоки выбирались из того соображения, чтобы внутри них одновременно использовалось бы максимальное количество регистров `ymm`, но не больше 16 штук. 

Каждый блок сопровождён отметкой о количестве используемых регистров **на входе блока**, **максимальное число внутри** блока и **на выходе блока**.

Основные вычисления происходят внутри lambda функции `insertCode()`, которая за раз может обрабатывать 1 или 2 входных регистра `ymm`

Слой LSTM рекуррентный, внешний цикл выполняет обработку 1 шага LSTM ячейки несколькими внутренними циклами, которые просто оптимально обрабатывают данные сначала по 16 значений, потом по 8, потом что осталось.
```cpp
	// *** Main loop ***
	gen.StartDownCountLoop( regObjectsCount, 1 );

	// ... Внутренняя обработка 1 шага LSTM

	// *** Stop Main loop ***
	gen.StopDownCountLoop();
```

После каждого внутреннего цикла обновляем указатели.
Методы `StartDownCountLoop()` и `StopDownCountLoop()` реализованы таким способом, что позволяют организовывать вложенные циклы, сохраняя лэйблы для переходов в стеке и отслеживая уровень вложенности.

## Заключение
В целом работу с JIT примитивами планировалось сделать максимально масштабируемой для возможности добавления новых примитивов, но всё равно код уже заслуживает рефакторинга:
1. Можно оставить механизм доступа к примитивам только через получение указателей.
2. Метод `insertSimpleMathFunction()` унифицирует разворачивание циклов для блоков различной длинны и обработку хвостовых значений в простейших математических примитивах, подобный механизм хорошо бы перенести и на сложные примитивы.
