В статье «Второе рождение 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, чтобы результат его работы не требовал преобразований.