Apple Silicon на Swift ускоряет умножение матриц в 10 раз

Автор пытается запустить собственный код умножения матриц в Swift на Apple Silicon с максимальной скоростью для обучения больших языковых моделей (LLM). Его цель — разобраться в возможностях компонентов системы: CPU, SIMD, AMX и GPU. В качестве эталона берётся реализация Andrej Karpathy'а llm.c. Начальный вариант кода на Swift работал в 15–20 раз медленнее C из-за проверок безопасности и функций _ArrayBuffer.beginCOWMutation(). Решение пришло с появлением MutableSpan в Swift 6.2, что ускорило циклы обучения более чем в три раза. Следующим шагом стала работа с числами через Relaxed.multiplyAdd, которая задействовала инструкции fmla (fused-multiply-add) и добавила почти 10-кратный прирост к производительности вывода токенов. В дальнейшем автор перешёл на работу с InlineArray для имитации кэшированных массивов C и применил DispatchQueue.concurrentPerform для многопоточности. Этот подход позволил достичь производительности обучения, сравнимой с C, но код стал громоздким из-за необходимости использования unsafe-блоков. Финальный прорыв дал AMX — Apple Matrix Coprocessor, доступный через обратную инженерируемые инструкции типа AMX_MATFP, AMX_LDX и AMX_LDY. Реализация на этих командах ускорила процесс ещё в 1,67 раза, несмотря на то, что официальный фреймворк Accelerate работает немного быстрее. Завершает серию код на Metal, где внутренняя часть ядра написана на Metal/C++, а внешний слой управления остаётся на Swift. Это решение разделяет вычисления над данными и вызов вычислений, используя 2D-параллелизм для обработки батчей и колонок. Автор признаёт, что чистый код без библиотек уступает оптимизированным ядрам производителей, но демонстрирует путь к созданию высокопроизводительных вычислений без привязки к внешним зависимостям.