В погоне за скоростью. Оптимизация нейросетевых вычислений на процессоре К1967ВН044 компании «Миландр»
В статье «Второе рождение DSP или запуск нейросетей на процессорах К1967ВН044 от «Миландр» мы рассмотрели в целом задачу адаптации нейросетей для DSP процессора К1967ВН044. Были вкратце описаны особенности процессора и возможные методы для эффективного его использования. В этой статье мы постараемся более детально представить один из таких методов, а именно — применение библиотеки ассемблерных функций для оптимального вычисления типичных операций, встречающихся в нейросетях.
Поскольку теперь будут появляться примеры кода на ассемблере, придётся хотя бы в общих чертах его понимать. Как было совершенно справедливо отмечено, данный процессор является развитием архитектуры TigerSHARC, так что программисты, знакомые с ним, без труда узнают этот код. Для тех, кто не имел с ним дела, можно порекомендовать «Руководство по программированию» (https://ic.milandr.ru/upload/iblock/77f/77fac90e79704374aaccc4b44f3244d6.pdf), в котором дано подробное описание всех возможностей процессора, причём с учётом многочисленных доработок, выполненных фирмой «Миландр».
Впрочем, основные идеи кочуют из одного DSP в другой, так что такие особенности, как наличие большого количества регистров, выполнение операций с данными только на регистрах, «хитрые» инструкции и т.д. не должны вызвать удивления. Кроме того, по ходу дела будут даваться краткие пояснения.
Для начала вспомним, какие преимущества мы собирались использовать для ускорения вычислений:
1. широкая шина для доступа к памяти;
2. возможность выполнять несколько инструкций за такт;
3. наличие векторных инструкций;
4. наличие «продвинутых» инструкций, способных выполнять сложные операции.
Примеры
Быстрое копирование памяти
Давайте рассмотрим простейший пример — копирование блока памяти. Причём постараемся разобрать его достаточно подробно, так как в дальнейшем уже описанные приёмы будут упоминаться вскользь.
Итак, у нас классическая функция memcpy, которая принимает на входе адреса источника и приёмника, а также количество байт. И очевидная реализация заключается в том, чтобы считать байт в регистр, записать его по нужному адресу, сдвинуть указатели и уменьшить счётчик цикла. В общем случае по-другому и не сделать, но первая же идея заключается в том, что если, например, адреса выровнены по слову, а количество байт кратно четырём, то можно оперировать словами, что даст большой выигрыш по скорости. Разумеется, эту идею немедленно хочется развить до логического завершения: использовать максимально большие порции данных, обеспечив нужное выравнивание. И в нашем случае мы можем этого добиться, поскольку наши объекты достаточно велики, чтобы пойти на потери памяти при выравнивании, если это потребуется.
В результате код функции представляется таким:
_copy_i32v8:
k5:4 = j5:4; j4 = j4 + 4;;
j5 = j5 + 4; lc0 = j6;;
loop1:
xr3:0 = q[j5 += 8]; yr3:0 = q[k5 += 8];;
if nlc0e, jump loop1; q[j4 += 8] = xr3:0; q[k4 += 8] = yr3:0;;Давайте рассмотрим его подробнее.
Во-первых, аргументы приходят на регистрах j4, j5 и j6. Как будет видно в дальнейшем, нам потребуется задействовать оба целочисленных модуля J и K, которые будут использоваться для адресации, так что мы копируем адреса источника и приёмника из J в K, причём делается это одной инструкцией. Более того, в той же линии (то есть за тот же такт) мы сдвигаем указатель приёмника (j4) на 4. Процессор позволяет читать и писать регистр в одной линии, при этом результатом чтения будет прежнее значение.
В следующей линии мы сдвигаем указатель источника (j5) и устанавливаем счётчик цикла.
Чему? Для этого смотрим на следующую линию.
В ней мы читаем в блок X сразу 4 слова по адресу j5 и 4 слова по адресу k5, причём одновременно выполняется автоинкремент обоих адресов на 8.
Теперь понятно, зачем потребовался начальный сдвиг на 4, а также ясно, что в качестве значения счётчика следует взять количество слов, поделенное на 8. При этом мы достигаем максимально возможной пропускной способности 256-битной шины.
Наконец, последняя линия записывает 8 слов по адресам j4 и k4, уменьшает счётчик цикла и, если он не равен нулю, выполняет переход в начало цикла. Этот переход помечен как предсказанный, так что сброса конвейера не происходит, и сам переход достаётся как бы «бесплатно».
Отвлечёмся на минуту и подумаем, способен ли компилятор сгенерировать такой код?
Ну, если у него будет информация, что адреса выровнены (это можно задать через атрибут в аргументах функции), а размер кратен аж 32-м…
Вот с последним обстоятельством беда, так как только программист может знать, какие размеры реально появляются. Компилятор, скорее всего, добавит к «быстрому» варианту проверку и «медленный» вариант, чтобы сохранить общность.
Но и это ещё не всё!
Процессор имеет особенность: после загрузки в модули X и Y данные становятся доступны не сразу, а только через такт. Так что в реальности внутри цикла возникает «пузырь», и один такт теряется.
Это легко обойти так:
xr3:0 = q[j5 += 8]; yr3:0 = q[k5 += 8];;
xr7:4 = q[j5 += 8]; yr7:4 = q[k5 += 8];;
q[j4 += 8] = xr3:0; q[k4 += 8] = yr3:0;;
if nlc0e, jump loop8; q[j4 += 8] = xr7:4; q[k4 += 8] = yr7:4;;Всё то же самое, только теперь мы последовательно читаем два раза по 8 слов, и первая порция готова к записи. Компилятор знает про «пузыри», и в данном случае попробовал бы в исходном варианте поставить между чтением и записью какую-то другую инструкцию, которую тут можно было бы выполнить без нарушения алгоритма, но поскольку таковой нету, он бы смирился с «пузырём». Рассчитывать на то, что он догадается выполнить такой своеобразный «unroll» было бы слишком самонадеянно.
Причём заметьте, это простейший пример — что-то будет дальше…
Однако наша задача ведь заключалась не в том, чтобы написать специализированный вариант memcpy! Пора взглянуть на Си код от TVM, который мы собираемся оптимизировать.
Он выглядит примерно так: главная функция, в которой выделяется память под временные объекты (об этом речь пойдёт в следующей статье), и вызовы функций слоёв, которых может быть очень много.
А функция слоя, в свою очередь, выделяет память под свои временные объекты, а дальше идут всякий циклы. Вот один такой:
for (int32_t i1 = 0; i1 < 25; ++i1) {
for (int32_t i2 = 0; i2 < 5; ++i2) {
for (int32_t i3 = 0; i3 < 64; ++i3) {
int32_t cse_var_1 = (((i1 * 320) + (i2 * 64)) + i3);
((int8_t*)pad_temp_let)[cse_var_1] = p0[cse_var_1];
}
}
}Опять же — первый раз остановимся на нём подольше.
Мы видим систему из трёх вложенных циклов, причём внешние два — это простейшая форма «for», где переменная пробегает значения от нуля до константной границы с шагом 1. И это не случайность, наоборот, подавляющая часть кода именно так и выглядит. Тут нет ничего удивительного, ведь нейросеть работает с тензорами, так что основная работа — это обход тензора по осям (в данном случае он двумерный) и вычисление во внутреннем цикле некой формулы. И, если присмотреться, то мы увидим, что переменная cse_var_1 во внутреннем цикле является инвариантом, к которому просто добавляется значение переменной цикла i3. Далее она используется для индексации байтового массива, причём базовый сдвиг кратен аж 64, а размер просто равен 64.
Ба, да мы можем заменить внутренний цикл на вот такой вызов нашей функции:
copy_i32v8((int32_t *)((int8_t*)pad_temp_let + i1 * 320 + i2 * 64),
(int32_t *)(p0 + i1 * 320 + i2 * 64),
64 / (2 * 8 * 4));То есть мы вычисляем два адреса, убедившись, что выравнивание не нарушается, корректируем размер, и вместо 64-х байтовых итераций получаем всего несколько инструкций в ассемблерной реализации.
Теперь, надеюсь, понятна основная идея оптимизации.
Мы полагаем, что в коде TVM в большом количестве встречаются конструкции, которые различаются только некими параметрами, а смысл вычисления повторяется. И мы будем искать их, пробуя набор шаблонов при помощи специального конвертера, и заменять внутренние циклы на своеобразные интринсики, для которых написана эффективная реализация.
Прибавление к вектору константы
Вот чуть более сложный пример:
for (int32_t ax1_2 = 0; ax1_2 < 25; ++ax1_2) {
for (int32_t ax2_2 = 0; ax2_2 < 5; ++ax2_2) {
for (int32_t ax3_2 = 0; ax3_2 < 64; ++ax3_2) {
int32_t cse_var_7 = (((ax1_2 * 320) + (ax2_2 * 64)) + ax3_2);
((int32_t*)conv2d_nhwc_let)[cse_var_7] = (((int32_t*)conv2d_nhwc_let)[cse_var_7] - 128);
}
}
}Тот, кто помнит предыдущий пример, мгновенно узнает тот же тензор, только данные теперь 32-битные, и выполняется не копирование, а вычитание 128 из каждого элемента. Это опять же типичная векторная операция, ядро которой может выглядеть так (это не самый оптимальный вариант):
xr3:0 = q[j4 + 0]; yr3:0 = q[k4 + 0];;
r1:0 = r1:0 + r5:4;;
r3:2 = r3:2 + r7:6;;
if nlc0e, jump loop6; q[j4 += 8] = xr3:0; q[k4 += 8] = yr3:0;;Здесь в прологе (я не стал его приводить) число, которое следует прибавить ко всем элементам вектора размножается в регистрах xyr7:4, далее мы опять же читаем максимальную порцию данных.
Правда, разрядность вычислителей только 64 бит, так что требуется две инструкции на цикл, но каждая из них производит 4 32-битных сложения за такт. А если ещё доработать реализацию для устранения пузырей, то опять же можно ожидать значительного ускорения. Вызов для данного случая выглядит так:
add_n_i32v8((int32_t *)conv2d_nhwc_let + ax1_2 * 320 + ax2_2 * 64,
-128,
64 / (8 * 1));Умножение со сдвигом и округлением
Давайте пропустим всякие сложения/вычитания векторов, так как уже должно быть понятно, как они обрабатываются, и рассмотрим более сложный пример (я больше не выписываю внешние циклы):
for (int32_t i3_1 = 0; i3_1 < 64; ++i3_1) {
int32_t cse_var_6 = (((i1_1 * 320) + (i2_1 * 64)) + i3_1);
((int32_t*)conv2d_nhwc_let)[cse_var_6] = ((int32_t)(((((int64_t)((int32_t*)conv2d_nhwc_let)[cse_var_6]) * ((int64_t)((int32_t*)fused_nn_conv2d_subtract_add_constant_6_let)[i3_1])) + ((int64_t)1 << ((int64_t)((((int32_t*)fused_nn_conv2d_subtract_add_constant_8_let)[i3_1] + 31) - 1)))) >> ((int64_t)(((int32_t*)fused_nn_conv2d_subtract_add_constant_8_let)[i3_1] + 31))));
}Выглядит весьма запутанно, так что лучше выписать выражение с короткими именами:
*arg1 = ((int64_t)*arg1 * (int64_t)*arg2 + (1LL << (*arg3 + 30))) >> (*arg3 + 31);Теперь становится понятно, что мы вычисляем 64-битное произведение двух векторов, причём из результатов берём только старшие биты (важно: для каждого элемента — своё количество), да ещё и округляем.
Может ли компилятор эффективно векторизовать такой цикл? Вообще-то должен, так как вся информация у него есть, но посмотрим на код, написанный вручную:
_op5_v8: // void op5_v2(int32_t *arg1, const int32_t *arg2, const int32_t *arg3, unsigned qty);
k7:4 = j7:4; j4 = j4 + 4; r13 = r13 - r13;;
j5 = j5 + 4; lc0 = j7;;
j6 = j6 + 4; r12 = (1 << 30);;
r14 = -31;;
loop5:
xr3:0 = q[k4 + 0]; yr3:0 = q[j4 + 0];;
xr7:4 = q[k5 += 8]; yr7:4 = q[j5 += 8];;
xr11:8 = q[k6 += 8]; yr11:8 = q[j6 += 8];;
r17:16 = r0 * r4 (i);;
lr19:18 = ashift r13:12 by r8;;
r8 = r14 - r8;;
lr17:16 = r17:16 + r19:18;;
lr17:16 = ashift r17:16 by r8;;
r0 = r16;;
r17:16 = r1 * r5 (i);;
lr19:18 = ashift r13:12 by r9;;
r9 = r14 - r9;;
lr17:16 = r17:16 + r19:18;;
lr17:16 = ashift r17:16 by r9;;
r1 = r16;;
r17:16 = r2 * r6 (i);;
lr19:18 = ashift r13:12 by r10;;
r10 = r14 - r10;;
lr17:16 = r17:16 + r19:18;;
lr17:16 = ashift r17:16 by r10;;
r2 = r16;;
r17:16 = r3 * r7 (i);;
lr19:18 = ashift r13:12 by r11;;
r11 = r14 - r11;;
lr17:16 = r17:16 + r19:18;;
lr17:16 = ashift r17:16 by r11;;
r3 = r16;;
if nlc0e, jump loop5; q[k4 += 8] = xr3:0; q[j4 += 8] = yr3:0;;В принципе, тут ничего особенного нет, так что даже непонятно, что комментировать.
Снова мы в прологе настраиваем указатели и готовим константы, опять выбираем из памяти максимальные порции данных. Но обратите внимание, что процесс вычислений построен таким образом, что «пузырей» (по возможности) не возникает, а главное — результаты размещаются так, чтобы их можно было в конце цикла одним махом записать в память.
Сумеет ли компилятор найти такой код? Во всяком случае, это уже не совсем тривиальная задача…
Умножение с шагом
А теперь пересмотрим список, который приводили в начале. Пункты 1–3 раскрыты, остался пункт 4 — специфические инструкции. Давайте попробуем и их.
Опять же начнём, по возможности, с простого примера. Имеем (я не стал приводить внешние циклы, их тут три, и они дают xx, yy и ff):
((int32_t*)conv2d_nhwc_let)[(((yy * 320) + (xx * 64)) + ff)] = 0;
for (int32_t rc = 0; rc < 64; ++rc) {
int32_t cse_var_3 = ((yy * 320) + (xx * 64));
int32_t cse_var_2 = (cse_var_3 + ff);
((int32_t*)conv2d_nhwc_let)[cse_var_2] = (((int32_t*)conv2d_nhwc_let)[cse_var_2] + (((int32_t)((int8_t*)pad_temp_let)[(cse_var_3 + rc)]) * ((int32_t)((int8_t*)fused_constant_2_let)[((rc * 64) + ff)])));
}Выглядит запутанно, но если присмотреться, то мы увидим, как некая 32-битная ячейка зануляется, а потом в ней суммируются некие произведения 8-битных чисел, причём один из сомножителей выбирается с постоянным шагом.
Для ассемблерного варианта проще сделать функцию, которая вернёт значение суммы, а записать её потом можно и на Си.
Итак:
_mul_i8_step_sum_i32_v8: // int32_t mul_i8_step_add_i32_v8(int8_t *arg1, int8_t *arg2, unsigned step, unsigned qty);
k7:4 = j7:4; j4 = j4 + 4;;
j0 = lshiftl j6 by 2; lc0 = j7;;
j5 = j5 + j0; r1:0 = r1:0 - r1:0; k0 = j0;;
loop2:
xr2 = pb[k4 += 8]; yr2 = pb[j4 += 8];;
xr4 = b[k5 += k6] (se); yr4 = b[j5 += j6] (se);;
xr5 = b[k5 += k6] (se); yr5 = b[j5 += j6] (se);;
xr6 = b[k5 += k6] (se); yr6 = b[j5 += j6] (se);;
xr7 = b[k5 += k6] (se); yr7 = b[j5 += j6] (se);;
sr3:2 = expand br2 (i);;
sr4 = compact r5:4 (i);;
sr5 = compact r7:6 (i);;
r7:4 = r3:2 * r5:4 (i);;
r1:0 = r1:0 + r5:4;;
if nlc0e, jump loop2; r1:0 = r1:0 + r7:6; j5 = j5 + j0; k5 = k5 + k0;;
r0 = r0 + r1;;
xr1 = yr0;;
xr0 = r0 + r1;;Полный разбор занял бы слишком много места, тем более читатель уже узнаёт привычные места. Из нового тут следующее.
Для начала, нет необходимости читать большие блоки, да это и невозможно, так как во втором байты идут «вразбивку».
Вместо этого мы читаем в регистры xyr2 восемь байт, а восемь байт второго сомножителя отдельными инструкциями в xyr7:4, куда они попадают с расширением знака до 32 бит.
Теперь главный фокус — как мы будем умножать? Процессор умеет многое, но не всё. Для нашего случая применимо умножение 16-битных чисел. А ведь у нас ни те, ни те таковыми не являются! И если сделать из 32-битного числа в дополнительном коде 16-битное ещё просто, то расширение знака без аппаратной поддержки — уже заметно сложнее.
К счастью, всё предусмотрено: инструкция expand преобразует 4 8-битных числа в 4 16-битных на двух регистрах, а две инструкции compact дают те же 4 16-битных числа на двух регистрах из 32-битных данных.
И далее следует инструкция, которая совсем потерялась в коде, хотя она выполняет 8 (!) 16-битных умножений (не забываем, что у нас два модуля X и Y), с получением восьми 32-битных результатов. Чтобы их стало поменьше, накапливаем сумму в xyr1:0, а в эпилоге получаем итог из четырёх значений.
И опять же сакраментальный вопрос: сможет ли компилятор найти такое решение?
Усечение с насыщением
Ну и последний пример, который покажет, насколько велико открывающееся пространство возможностей.
Возьмём сразу два последовательных цикла:
for (int32_t i3_2 = 0; i3_2 < 64; ++i3_2) {
int32_t cse_var_8 = (((i1_2 * 320) + (i2_2 * 64)) + i3_2);
int32_t v_ = ((int32_t*)conv2d_nhwc_let)[cse_var_8];
int32_t v__1 = (v_) < (127) ? (v_) : (127);
((int32_t*)conv2d_nhwc_let)[cse_var_8] = ((v__1) > (-128) ? (v__1) : (-128));
}
for (int32_t ax3_3 = 0; ax3_3 < 64; ++ax3_3) {
int32_t cse_var_9 = (((ax1_3 * 320) + (ax2_3 * 64)) + ax3_3);
T_cast[cse_var_9] = ((int8_t)((int32_t*)conv2d_nhwc_let)[cse_var_9]);
}Что тут хотели сказать? Так очевидно: для 32-битных чисел вычисляется 8-битное насыщение, а потом результат записывается в массив 8-битных чисел.
Теперь я даже не стану приводить всю функцию, достаточно ключевого места:
xr3:0 = q[k5 += 8]; yr3:0 = q[j5 += 8];;
sr4 = compact r1:0 (si);;
sr5 = compact r3:2 (si);;
br6 = compact sr5:4 (si);;Всего тремя инструкциями 8 слов превращаются в 8 байт, причём с насыщением, для которого даже не требуется проверять условия на значения. Вопрос про компилятор я задавать не стану…
Кто-то может сказать: вы так и будете переписывать на ассемблере весь Си?
Нет!
Ещё раз повторю основную идею: это не просто случайный код, это реализация нейросетевых операций, число которых невелико, причём полученная в результате автогенерации TVM, что обеспечивает абсолютно жёсткое форматирование. То есть рассмотренные конструкции встречаются снова и снова; меняются только размеры, возможно, типы, так что один раз написанный набор поддержки сможет применяться многократно.
Ну и не следует забывать о правиле »20/80», то есть большая часть времени тратится в малом объёме кода, так что достаточно оптимизировать самые тяжёлые операции.
Результаты
У читателя наверняка созрел вопрос: так стоила ли овчинка выделки? Давайте посмотрим…
Мы возьмём слой нейросети, из которого были взяты примеры, описанные в статье, и измерим время его выполнения (в тактах процессора). А потом будем добавлять ассемблерные вставки по одной, отмечая улучшения, возникающие на каждом шаге.
Результаты такие (эффект от вставки суть разность с предыдущим значением):
5813245 — исходный Си код, без вставок, то есть точка отсчёта.
5783619 — быстрое копирование памяти (copy_i32v8). Немного, так как размер блока очень маленький, но всё-таки.
1978783 — умножение с шагом (mul_i8_step_sum_i32_v8). Очень существенный выигрыш — тут ассемблер себя показал достойно.
1913388 — вычитание векторов (в статье не описано за тривиальностью). Немного, но курочка по зёрнышку…
1864005 — сложение векторов (аналогично).
1614949 — умножение со сдвигом и округлением (op5_v8). Заметное улучшение.
1557134 — прибавление к вектору константы (add_n_i32v8). Результат сравним со сложением/вычитанием векторов.
1444328 — усечение с насыщением (последний пример в статье, название пропущено). Заметная прибавка.
Вот как это выглядит: 
Что тут можно сказать?
С одной стороны, не стоит так уж сильно кидаться камнями в компилятор. Мы видим, что хотя ассемблер победил везде (было бы удивительно, если бы нет), но только в одном случае (заметьте, реально сложном) выигрыш разгромный (почти втрое). Остальные выглядят во всяком случае терпимо.
С другой, суммарное ускорение получилось в четыре раза! Нечего и говорить, что за такой результат стоит побороться.
Планы на будущее
Что касается планов на будущее, они вполне понятны.
Во-первых, следует искать не отдельные циклы, а их последовательности (как в последнем примере), и объединять операции, чтобы избежать промежуточного копирования. Ну и присмотреться к внешним циклам — вдруг удастся получить выигрыш, включив в оптимизацию и их.
Во-вторых, в перспективе, перенести всю эту кухню непосредственно в TVM, чтобы результат его работы не требовал преобразований.
Habrahabr.ru прочитано 1840 раз
