SME в Apple M4: 65 тысяч резонаторов на коротком аудиоблоке

SME (Scalable Matrix Extension) добавляет в набор инструкций Arm операции над матрицами. Apple включила поддержку SME в чипы поколения M4. Эти инструкции доступны из нативного кода, без отправки задачи в очередь GPU.

Я попробовал применить SME к обработке звука. На Mac mini M4 расчёт банка из 65 536 резонаторов сократился с 1,59 до 0,44 мс по медиане. При буфере 64 сэмпла и частоте 48 кГц это позволило укладываться в срок: в проверенной серии не было ни одного опоздавшего вызова.

Звуковой движок с такими настройками вызывает обработчик примерно каждые 1,33 мс. Опоздание может обернуться щелчком или обрывом. Увеличение буфера даёт больше времени на расчёт, но музыкант дольше ждёт звук при игре или мониторинге. Меня интересовал выигрыш при сохранении короткого блока.

Как колебания превратились в матрицы

Банк резонаторов складывает звук из затухающих колебаний со своей частотой и амплитудой. Так можно синтезировать перкуссию и резонансные тембры. Первая версия обновляла каждый резонатор на каждом сэмпле векторными инструкциями NEON.

Между ударами поведение резонатора известно заранее. Его комплексное состояние s умножается на постоянный коэффициент z: модуль z задаёт затухание, аргумент задаёт частоту. Для монофонической суммы в точной арифметике можно записать:

s_m[t+1] = z_m * s_m[t]
y[16a+b] = Σ_m Re((s_m[0] * z_m^(16a)) * z_m^b)

m обозначает резонатор, a нумерует группы по 16 сэмплов, b задаёт позицию в группе. Состояние переношу сразу к началу группы, а её сэмплы рассчитываю через заранее подготовленные степени z от 0 до 15.

В матрице A заранее хранятся действительные части z^b и мнимые части с обратным знаком. В B находятся действительные и мнимые части состояния в начале группы, умноженные на коэффициенты панорамы. Каждому резонатору соответствуют две строки A и B. Произведение C = AᵀB складывает их вклады.

У C шестнадцать строк для позиций b. Столбцы идут парами: левый и правый канал очередной группы. Для блока 64 сэмпла использую четыре группы, то есть восемь столбцов результата. Состояние между группами переношу умножением на заранее рассчитанное z^16. Вот один шаг матричного ядра:

// В streaming-функции, ZA обнулён до цикла.
// Маска pg включает все 16 FP32-элементов M4.
svmopa_za32_f32_m(0, pg, pg,
    svld1_f32(pg, A + 16*k),  // Коэффициенты z^b.
    svld1_f32(pg, B + 16*k)); // Состояния с панорамой.
// ZA0[i,j] += A[k,i] * B[k,j].

В полном цикле чередую четыре накопителя ZA, затем складываю их. Меняющиеся данные обрабатываю порциями по 2048 резонаторов, чтобы сократить временные буферы. Если внутри аудиоблока приходит новый удар, делю обработку в этой точке; остатки короче 16 сэмплов считаю векторным кодом. Дополнительного входного сигнала ждать не требуется.

Выходы NEON, Accelerate и SME я сверил с независимой посэмпловой рекурсией double. В первой таблице сравниваются посэмпловый NEON и весь переработанный вариант с SME.

Что показали измерения

Испытания проходили на Mac mini M4 с 16 ГБ памяти, в колбэке CoreAudio при 48 кГц и буфере 64 сэмпла. Все резонаторы были возбуждены. Измерял весь расчёт стереовыхода, по три секунды на режим. p99 обозначает время, в которое уложились примерно 99% вызовов.

На 65 536 резонаторах медианное время уменьшилось примерно в 3,6 раза. Посэмпловый NEON превысил бюджет 1333 мкс во всех 1129 вызовах. Блочный вариант с SME уложился во всех 2250. На 131 072 резонаторах максимум 1290,08 мкс оставил около 43 мкс запаса. Этот размер характеризует наблюдаемую границу нагрузки в короткой серии.

Чтобы проверить выбор матричного ядра, провёл отдельный прогретый тест без CoreAudio. Блочный алгоритм, матрицы и подготовка данных одинаковы; меняется только умножение. На вариант пришлось 4 прогрева и 40 измерений, порядок вызовов чередовался. Измерял весь стереорасчёт с подготовкой B и переносом состояний.

На 65 536 резонаторах медианы составили 838,96 мкс для блочного NEON, 425,27 мкс для Accelerate и 343,77 мкс для SME. При одной блочной формулировке путь SME оказался быстрее в 2,44 раза относительно моего NEON-ядра и в 1,24 раза относительно Accelerate. Это отдельная серия; её числа не заменяют результаты аудиоколбэка.

Точность проверил на 21 635 временных позициях стереовыхода относительно независимого эталона double. Отношение сигнала к численной ошибке, SNR, составило 98,73 дБ, максимальная абсолютная ошибка 2,39 × 10⁻⁶. В эталоне использовались те же квантованные коэффициенты резонаторов.

Для рабочего режима оставил 8192 резонатора. За десять секунд на блоке 64 получил медиану 41,12 мкс и максимум 264,25 мкс. Медиана занимает около 3,1% интервала обработки. Собственная алгоритмическая задержка синтезатора составляет 0 сэмплов: дополнительной очереди внутри него нет. Задержка аудиоустройства сохраняется.

Audio Unit прошли штатную валидацию. Во всех 40 проверенных наборах параметров выход через API Audio Unit побитно совпал с выходом процессора напрямую. Функциональные тесты прошли при 32, 44,1, 48, 96 и 192 кГц и размерах блоков от 1 до 511 сэмплов.

Что получилось со свёртками

Следующий тест относился к слою аудионейросети: случайные веса, блок 64 сэмпла, ядро длиной 3 и dilation 2, то есть шаг между элементами ядра в два сэмпла. Постоянные веса упакованы заранее; подготовка входа входит в измерение.

После прогрева на 128 каналах SME оказался примерно в 12 раз быстрее реализации на Eigen из стенда и в 1,4 раза быстрее Accelerate. Eigen складывал три матричных произведения. Для Accelerate и SME вход собирался в матрицу методом im2col: элементы для свёртки копировались в строки. Таблица сравнивает реализации слоя вместе с этой подготовкой.

Для проверки целой сети взял обученный вокодер GeneratorNSF с выходом 48 кГц и закреплённой конфигурацией. Заменил 76 свёрток кодом с SME и Accelerate; остальные операции оставил в PyTorch.

На одном потоке CPU после 3 прогревов и 20 измерений медиана синтеза порции длительностью 20 мс сократилась с 10,18 до 6,38 мс, примерно в 1,6 раза. На одинаковых синтетических входных признаках SNR относительно исходного расчёта составил от 113,5 до 119,7 дБ. Это время вокодера целиком, без анализа входа и аудиоустройства.

Что из этого можно использовать

Для резонансного эффекта, который обрабатывает входящий звук, к свободному затуханию я добавил причинную свёртку входа и поправку состояния. Дополнительная буферная задержка не понадобилась. На 1024 входных сэмплах максимальное расхождение с независимой посэмпловой рекурсией составило 5,13 × 10⁻⁸ по всем 12 проверенным настройкам.

Вспомогательные аудиопотоки можно объединять через Audio Workgroups, чтобы macOS учитывала их общий срок готовности.

Выигрыш зависит от размеров матриц и общего матричного ресурса кластера. Серии по 3 и 10 секунд проверяют выполнимость задачи; устойчивость длительного мониторинга с несколькими экземплярами ещё предстоит проверить. Ускорение также не отменяет ожидание будущих данных, заложенное в модель.

Банк резонаторов стал успевать к следующему вызову, а обученный вокодер быстрее считал ту же порцию синтеза. Полученный запас можно потратить на более сложный тембр, больше одновременно звучащих голосов или дополнительную обработку при прежнем размере буфера.