SIMD с нуля

Процессор умеет складывать не два числа, а два набора чисел — одной командой. Шестнадцать сложений за то же время, что одно. Это называется SIMD: single instruction, multiple data, одна команда над многими данными.

В Python эта возможность спрятана внутри NumPy, и добраться до неё из своего кода нельзя. В Mojo она — обычный тип языка.

Тип SIMD

def main():
    var v = SIMD[DType.float32, 4](1.0, 2.0, 3.0, 4.0)
    var w = SIMD[DType.float32, 4](10.0)

    print("вектор:", v)
    print("сумма:", v + w)
    print("произведение:", v * w)
    print("reduce_add:", v.reduce_add())
    print("элемент:", v[2])

Результат

вектор: [1.0, 2.0, 3.0, 4.0]
сумма: [11.0, 12.0, 13.0, 14.0]
произведение: [10.0, 20.0, 30.0, 40.0]
reduce_add: 10.0
элемент: 3.0

Два параметра: тип элемента и сколько их. Оба известны при сборке — иначе процессор не знал бы, какую команду выполнять.

Обратите внимание на SIMD[DType.float32, 4](10.0): одно число размножилось по всем полосам. Это частый приём — умножить весь вектор на скаляр.

`Float64` — тоже SIMD

Float64 — это SIMD[DType.float64, 1], вектор из одного элемента. Присваивание между ними работает напрямую:

var a: SIMD[DType.float64, 1] = 2.5
var b: Float64 = a          # то же самое

Отсюда и Scalar[DType.float32] — просто имя для вектора ширины 1. Поэтому арифметика над числами и над векторами выглядит одинаково: это буквально один и тот же код.

А вот Int в этот ряд не входит — это отдельный тип машинного слова, и неявно он не превращается ни во что:

error: cannot implicitly convert 'Int64' value to 'Int'

Так что «в Mojo нет скалярных типов» — красивая фраза, но неверная.

Какая ширина у вашей машины

Жёстко зашивать четвёрку не стоит: на одном процессоре за раз проходит четыре числа, на другом шестнадцать. Правильную ширину подскажет компилятор:

from std.sys import simd_width_of


def main():
    print("float32:", simd_width_of[DType.float32]())
    print("float64:", simd_width_of[DType.float64]())
    print("int8:   ", simd_width_of[DType.int8]())

На машине, где писалась эта глава (Intel Xeon 2,8 ГГц):

float32: 16
float64: 8
int8:    64

Шестнадцать чисел float32 — это 512 бит, то есть процессор с AVX-512. У вас может быть 8 или 4 — и это нормально.

Обратите внимание на закономерность: чем меньше элемент, тем больше их помещается. Регистр один и тот же, меняется только нарезка.

Имя функции поменялось в 1.0

До 1.0 она называлась simdwidthof — одним словом, — и именно так написано во всех старых статьях и примерах. Компилятор 1.0.0 такого имени не знает и ничем не помогает:

error: use of unknown declaration 'simdwidthof'

Попытка импортировать старое имя тоже ни к чему не приведёт:

error: package 'sys' does not contain 'simdwidthof'

А вот новое имя, забытое в импорте, компилятор подскажет сам — запомнить путь не придётся:

error: use of unknown declaration 'simd_width_of'; did you mean to import it from 'std.sys'? Add 'from std.sys import simd_width_of'

Чтение и запись памяти вектором

Из списка данные в вектор кладут не по одному, а сразу пачкой:

def main():
    var data = List[Float32](capacity=20)
    for i in range(20):
        data.append(Float32(i))

    var ptr = data.unsafe_ptr()

    var chunk = ptr.unsafe_load[width=4](0)
    print("загружено:", chunk)

    ptr.unsafe_store(0, chunk * 2.0)
    print("после записи:", data[0], data[1], data[2], data[3])

Результат

загружено: [0.0, 1.0, 2.0, 3.0]
после записи: 0.0 2.0 4.0 6.0

Слово unsafe в названии честное: границы никто не проверяет. Загрузить шестнадцать чисел, начиная с предпоследнего, вам никто не помешает — и прочитается чужая память.

`load` и `store` устарели

Раньше методы назывались просто load и store. Они ещё работают, но компилятор просит переименовать:

warning: 'load' is deprecated, use 'unsafe_load' instead
warning: 'store' is deprecated, use 'unsafe_store' instead

Новые имена лучше: они не дают забыть, что проверок нет.

Хвост

Массивы редко бывают кратны ширине вектора. Тысяча элементов при ширине шестнадцать — это 62 полных вектора и 8 штук в остатке, которые придётся досчитать по одному:

var full = (len(data) // W) * W

for i in range(0, full, W):        # основная часть — вектором
    accumulator += ptr.unsafe_load[width=W](i)

var total = accumulator.reduce_add()

for i in range(full, len(data)):   # хвост — поштучно
    total += data[i]

Забыть про хвост — самая частая ошибка в векторном коде: результат получается почти правильным, и на маленьких тестах расхождение легко не заметить.

Редукция: из вектора обратно в число

Аккумулятор в примере выше — не число, а вектор из W независимых накопителей. Каждый складывал свою полосу. Собрать их в одно значение можно методами reduce_*:

Метод	Что делает
reduce_add()	сумма всех полос
reduce_max() / reduce_min()	наибольшее / наименьшее
reduce_mul()	произведение

Сравнение вектора — не оператором

Написать v > 2.0 для вектора нельзя, и компилятор объясняет почему:

note: constraint failed: Strict inequality is only defined for `Scalar`s; did you mean to use `SIMD.gt(...)`?

Оператор > обязан вернуть один Bool — а поэлементное сравнение даёт вектор из Bool. Поэтому для вектора зовут метод:

v.gt(2.0)     # [False, False, True, True]

Сколько это даёт на самом деле

Сумма массива из 8192 чисел float32, замер по методике из главы Как честно мерить скорость — минимум из партий, данные в куче, контрольная сумма сходится. Чтобы числу можно было верить, замер повторён на двух разных процессорах, у обоих ширина для float32 равна 16.

Машина	Поэлементно	Вектором	Ускорение
Intel Xeon 2,8 ГГц	1,254	0,073	17,1×
AMD Ryzen 7 9700X	0,367	0,022	16,4×

Наносекунд на элемент. Абсолютные числа различаются втрое с половиной — а ускорение и там, и там держится около шестнадцати, то есть ровно на ширине вектора. Именно так и должно быть, и именно поэтому этому числу можно доверять: оно свойство приёма, а не машины.

Отклонения в обе сторону — шум. На Xeon восемь прогонов дали от 15,7 до 17,5.

Ширина — не единственный резерв

Здесь легко остановиться, решив, что потолок достигнут. Это не так, и проверяется просто: добавим независимые накопители, чтобы соседние сложения не ждали друг друга.

Машина	Вектор, 1 накопитель	Вектор, 4 накопителя	Итого к наивному циклу
Intel Xeon 2,8 ГГц	0,073	0,026	48,8×
AMD Ryzen 7 9700X	0,022	0,006	60,1×

Векторная версия ускоряется ещё втрое-вчетверо. Значит, она упиралась не только в количество данных за такт, но и в длину цепочки: каждое сложение ждало результата предыдущего.

Что здесь на самом деле происходит

Сложение с плавающей точкой занимает несколько тактов. Пока результат не готов, следующее сложение к тому же накопителю начать нельзя. Получается очередь длиной во весь массив.

Вектор укорачивает её в шестнадцать раз (8192 сложения превращаются в 512), но не убирает. Четыре независимых накопителя укорачивают ещё вчетверо — и процессор наконец загружен.

Обратите внимание, что вторая половина выигрыша машинно-зависима: на Xeon она дала общий множитель 49, на Ryzen — 60. Ширина вектора предсказуема, глубина конвейера — нет.

Именно поэтому библиотечная векторизация из следующей главы не просто режет массив на векторы, но ещё и разворачивает цикл.

Векторная сумма даёт другой ответ

Сложение чисел с плавающей точкой не ассоциативно: (a + b) + c и a + (b + c) могут отличаться. Векторная редукция складывает в другом порядке — значит, и результат другой.

Проверено на массиве из 4096 элементов: одно число 1.0e8 и 4095 единиц. На ширине 16:

скаляр: 100000000.0
вектор: 100003840.0

Точный ответ — 100004095. Поэлементная сумма потеряла все единицы: к стомиллионному значению единица в float32 просто не добавляется. Векторная разложила слагаемые по шестнадцати полосам: пятнадцать из них накапливали только единицы и ничего не потеряли, и лишь одна — та, куда попало большое число, — потеряла свои 255. Отсюда и ошибка ровно в 255 вместо 4095.

На другой ширине и число будет другим: при восьми полосах выходит 100003580, при четырёх — 100003070.

Вывод не «SIMD точнее», а «результат другой». Если ваш код сверяется с эталоном до последнего бита — векторизация его сломает, и это придётся учесть.

Всё вместе

sum.mojo

# Сумма массива: поэлементно и вектором.
#
# Длина нарочно не кратна ширине вектора — так бывает почти всегда,
# и хвост приходится досчитывать отдельно.

from std.sys import simd_width_of

comptime W = simd_width_of[DType.float32]()
comptime N = 1001  # нечётное: хвост будет при любой ширине вектора


def scalar_sum(data: List[Float32]) -> Float32:
    """По одному числу за раз."""
    var total = Float32(0)
    for i in range(len(data)):
        total += data[i]
    return total


def vector_sum(data: List[Float32]) -> Float32:
    """По W чисел за раз, плюс хвост."""
    var ptr = data.unsafe_ptr()
    var accumulator = SIMD[DType.float32, W](0)

    # Основная часть: столько полных векторов, сколько влезло.
    var full = (len(data) // W) * W
    for i in range(0, full, W):
        accumulator += ptr.unsafe_load[width=W](i)

    # Складываем полосы аккумулятора в одно число.
    var total = accumulator.reduce_add()

    # Хвост: оставшиеся элементы — по одному.
    for i in range(full, len(data)):
        total += data[i]

    return total


def main():
    var data = List[Float32](capacity=N)
    for i in range(N):
        data.append(Float32((i * 37) % 100))

    var a = scalar_sum(data)
    var b = vector_sum(data)

    # W у каждой машины своё, поэтому печатаем не его, а то,
    # что от машины не зависит.
    print("элементов:", N)
    print("поэлементно:", a)
    print("вектором:   ", b)
    print("совпало:", a == b)
    print("длина нечётная, значит хвост есть при любой ширине вектора")

Результат

элементов: 1001
поэлементно: 49500.0
вектором:    49500.0
совпало: True
длина нечётная, значит хвост есть при любой ширине вектора

Программа печатает не ширину вектора и не время: и то и другое у вас будет своё. Печатает она то, что от машины не зависит, — совпадение результатов.

Проверяйте векторный код контрольной суммой

Первое, что стоит написать рядом с векторной версией, — простая поэлементная, и сравнить результаты. Она медленная, зато очевидно правильная.

Так ловятся обе главные ошибки: забытый хвост и перепутанные индексы. А если результаты разошлись в последних знаках — вы наткнулись на неассоциативность, и это тоже полезно знать заранее.

🎯 Проверь себя

Почему у `SIMD` оба параметра — тип элемента и ширина — должны быть известны при сборке?

Потому что от них зависит, какую машинную команду выпустит компилятор. Сложение восьми float32 и восьми float64 — разные команды над разными регистрами. Ширина времени выполнения означала бы, что команду надо выбирать при каждом вызове, а это уже не SIMD.

Что не так с `SIMD[DType.float32, 8]`, зашитым в код руками?

Восьмёрка верна не везде. На процессоре с AVX-512 за раз проходит шестнадцать чисел float32, и половина возможностей пропадёт; на более старом — вектор придётся эмулировать. Правильно спросить у компилятора: simd_width_of[DType.float32]().

Массив из 1000 чисел, ширина вектора 16. Что произойдёт, если написать только `for i in range(0, len(data), 16)`?

Цикл зайдёт за границу массива: на 63-м шаге i станет 992, и загрузка шестнадцати чисел прочитает 8 элементов чужой памяти — проверок у unsafe_load нет. Нужно считать до (len // 16) * 16, а оставшиеся 8 элементов досчитать по одному.

Векторная сумма массива float дала не то же число, что поэлементная. Это ошибка?

Нет, это свойство арифметики с плавающей точкой: сложение не ассоциативно, а редукция складывает в другом порядке. Причём векторная версия нередко точнее: каждая полоса накапливает свою часть и не теряет малые слагаемые. Но если результат должен совпадать с эталоном побитово — векторизовать нельзя.

Векторная версия дала 17× при ширине 16. Значит, потолок достигнут?

Нет. Проверяется добавлением независимых накопителей: векторная версия с четырьмя накопителями вместо одного ускоряется ещё втрое-вчетверо, и суммарно выходит 49× на Xeon и 60× на Ryzen. Значит, вектор снял ограничение по данным за такт, но осталось второе — цепочка сложений, где каждое ждёт предыдущего. А сами 17 вместо 16 — просто шум: прогоны давали от 15,7 до 17,5.

Что дальше

Векторизация и параллелизм: как поручить нарезку на векторы и хвост стандартной библиотеке — и как раскидать работу по ядрам.