Если вы занимаетесь обучением крупных современных нейросетей, эта статья будет вам не совсем в тему, ведь у A100 скорость в сто раз выше (156 терафлопсов). Так что же интересного в этих полутора терафлопсах? Мы говорим не о мощных ускорителях или тензорных ядрах графических процессоров, а лишь о настоящей производительности линейной алгебры, которая отстоит от регистров процессора на один цикл. Как ни странно, Apple прячет от нас эту штуку! Но в этой статье мы рассмотрим код, который позволит нам приподнять завесу тайны. Для всего кода используется заголовок aarch.h в замечательном репозитории corsix. По сути, это SIMD на стероидах. Важная особенность в том, что отношение AMX:CPU не равно 1:1. Сопроцессор AMX есть не у каждого ядра. Вот размеры, которые можно использовать для загрузки или хранения значений:  
Размер minimum равен ширине полного регистра AVX512. Но откуда загружаются или где хранятся эти значения? Понятно, что такие размеры довольно быстро заполнили бы файл регистра NEON. Есть отдельный регистровый файл специально для AMX, выглядит он весьма странно. Регистры распределены по группам X, Y и Z. В группах X и Y хранятся входные значения для каждой инструкции, а в группе Z — выходные значения.  Группы X и Y уже не маленькие! На них ушел целый килобайт. Ну а Z вообще ни в какие ворота:  И как же попасть из X и Y в Z? Способов так много, что все они едва ли уместятся в кодировку ISA, поэтому большую часть информации об инструкциях Apple закодировали в регистре общего назначения. Оказалось, что с этим классно работать, ведь есть возможность конфигурировать код на AMX на лету, во время его выполнения. Цель этого поста — добиться от сопроцессора предельной мощности. Есть инструкции умножения векторов, которые выводят векторы одинаковой длины, но они далеко не исчерпывают вычислительные возможности этой микросхемы. Чтобы добиться реального эффекта, придется вместо этого использовать внешнее произведение. Что такое внешнее произведение? Пусть есть два вектора u и v? тогда:     Внешнее произведение этих векторов — это матрица, которая содержит произведения всех возможных попарных комбинаций их элементов. Это немного проясняет, почему группа регистров Z настолько больше групп X и Y.     В микросхеме AMX это сводится к очень простой, вот такой инструкции:  Ещё есть флаг, который можно установить, чтобы суммировать с нарастанием (accumulate) данные из прошлого результата:  Итак, у нас есть все необходимое, чтобы написать умножение матриц: перезагрузим 16 значений с плавающими точками из входных матриц и суммируем их внешние произведения с нарастанием в выводе размерностью 16x16. При этом уменьшение размерности K никакой роли не играет! Упростим задачу и неявно транспонируем матричное умножение. Как A, так и B (наши входные данные) как основную размерность будут иметь размерность уменьшения K. На практике это не имеет большого значения, но значительно упрощает код. Вот ссылка, при помощи которой предложенное решение можно проверить: А вот как проверить его на AMX: Примечательно, что мы не обращались по адресу ни к одному регистру. Вернее, обращались, но тайно. Точно так же, как reset_z кодируется битовой маской, адреса регистров кодируются в аргументах, передаваемых AMX_*. Указатели на A и B используют лишь до 56 бит, поэтому инженеры Apple запрятали данные в оставшихся восьми битах. Мы просто случайно установили их все в значение 0. Итак, в данном случае применяются регистры "0" для X и Y. Код для сохранения регистров Z в памяти немного сложнее: заполняется только первый столбец. Это означает, что нужно занять регистры 0, 4, 8 и т. д.: К сожалению, если загрузить размещенный выше код, мы увидим, что он страшно тормозит. А ведь это всего какая-то пара сотен гигафлопов. Почему же? Есть две причины. Первая — конфликт при конвейерной обработке (pipeline hazard) Каждая инструкция AMX_FMA32 зависит от предыдущей, поскольку мы суммируем с нарастанием в единое подмножество регистрового файла. В итоге используется 25% регистрового файла на полную мощность, а остальное мы оставляем незадействованным, что мешает распараллеливанию на уровне инструкций. Вторая причина — неэффективная загрузка из памяти. Одновременно можно загружать 128 байт, но приведенный выше код загружает всего 64 байта. Аналогично мы можем выполнять загрузку в другие регистры, а не загружать каждый раз в одни и те же. Это тоже позволяет немного распараллелиться на уровне инструкций. И какой же у нас план?  Мы планируем загрузить 128 байт в X и Y и затем рассчитать блок 32x32. Для этого потребуется 4 независимых расчета блоков по 16x16, что позволит добиться параллелизма на уровне инструкций, а также более эффективно использовать загруженную память (каждый 64-байтовый регистр используется дважды). Выше я добавил комментарии, но есть еще пара интересных штук с флагами для инструкций. Corsix отлично все объясняет, поэтому я просто оставлю ссылки на объяснения: Как быстро мы к этому пришли? Ну, это зависит от числа килобайт, но к 1.5TFlops мы приходим при высоких значениях.  Не удивительно, что большие задачи получают лучшую относительную производительность, ведь у кеша больше возможностей для разогрева, а у процессора — для чередования инструкций. В целом на фоне больших современных нейросетей, работающих с общим искусственным интеллектом, эти размеры просто микроскопические. Тем не менее такой тип производительности открывает двери и для меньших нейросетей, которые могут найти себе место в реальных современных вычислениях. Если прогноз можно выполнить на ноутбуке с питанием от батареи за пару десятков наносекунд, то, скорее всего, он имеет ценность больше, чем та ценность, которую мы могли бы добавить в местах, где иначе могла бы применяться эвристика. А что думаете по этому поводу вы? Спасибо за внимание!  Data Science и Machine Learning Python, веб-разработка Мобильная разработка Java и C# От основ — в глубину А также