Векторизация и параллелизм

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

Библиотека умеет делать это сама — и заодно закрывает второй резерв, который мы там намерили.

vectorize

from std.algorithm import vectorize
from std.sys import simd_width_of

comptime W = simd_width_of[DType.float32]()


def main():
    def step[width: Int](idx: Int):
        print(idx, width)

    vectorize[W](40, step)

Сорок элементов при ширине 16. Смотрите, что печатается:

0 16
16 16
32 1
33 1
34 1
...
39 1

Два полных вектора, а потом восемь вызовов с шириной 1 — это и есть хвост. vectorize не пишет для него отдельный цикл, а вызывает вашу же функцию с другой шириной. Поэтому она и объявлена как step[width: Int]: ширина приходит параметром, и один и тот же код обслуживает оба случая.

Списки захвата

Функция выше ничего не трогала снаружи. Как только понадобятся данные, всплывёт то, чего в Mojo до 1.0 не было:

error: Could not infer capture convention of the captured value ptr

Что это значит

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

Как исправить

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

def step[width: Int](idx: Int) {imm ptr}:
    ...

Конвенций пять, и они те же, что у аргументов функций:

Запись	Что значит
{imm x}	ссылка только для чтения
{mut x}	ссылка с правом изменять
{var x}	копия, снятая в момент объявления
{var x^}	значение переезжает внутрь, снаружи больше недоступно
{ref x}	изменяемость определит вызывающий

Имя можно не писать: {imm} захватит всё используемое по неизменяемой ссылке, {} — не захватит ничего.

`@parameter` ломает `vectorize`

Соблазн пометить вложенную функцию декоратором @parameter велик — именно так написано во всех примерах до 1.0. Не делайте этого: декоратор делает замыкание захватывающим в другом смысле, и vectorize его не примет.

error: no matching function in call to 'vectorize'
note: candidate not viable: value passed to 'closure' cannot be converted from 'def[width: Int](idx: Int) capturing thin -> None' to 'func'

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

В 1.1 сам декоратор вдобавок помечен устаревшим ('@parameter' is deprecated; use '@__parameter'), так что повода писать его не осталось вовсе.

unroll_factor: второй резерв

В прошлой главе мы измерили, что после векторизации остаётся ещё почти трёхкратный запас: цепочка сложений, где каждое ждёт предыдущего. Разворачивать цикл руками не нужно — у vectorize есть параметр:

vectorize[W, unroll_factor=4](len(data), step)

Библиотека выпустит четыре независимые цепочки вместо одной. Это ровно тот приём, который в прошлой главе дал 0,073 → 0,026 наносекунды на элемент.

Многопоточность: parallelize

Здесь читателя ждёт сюрприз, и лучше узнать о нём отсюда.

`parallelize` нет в стандартной библиотеке

Официальное руководство показывает parallelize в примерах, и почти все статьи о Mojo берут его из algorithm. В Mojo 1.x такого имени в стандартной библиотеке нет — ни в std.algorithm, ни где-либо ещё:

error: package 'algorithm' does not contain 'parallelize'

Функция переехала в отдельный пакет max. Ставится он рядом:

uv add max

После этого работает импорт from max.algorithm import parallelize. Проверено на 1.0.0 и 1.1.0 — других путей не нашлось.

Сама функция проста: даёт номер куска, а вы решаете, что с ним делать.

from max.algorithm import parallelize


def main():
    var results = List[Int](length=8, fill=0)
    var ptr = results.unsafe_ptr()

    def work(i: Int) {imm ptr}:
        ptr.unsafe_store(i, i * i)

    parallelize(work, 8)

    for i in range(8):
        print(results[i], end=" ")
    print()

Результат

0 1 4 9 16 25 36 49 

Сколько ядер доступно, подскажет num_physical_cores() из std.sys.

Общие данные — только осторожно

Здесь каждый вызов work пишет в свою ячейку, и это безопасно. Если бы потоки меняли одну и ту же переменную, получилась бы гонка данных, а компилятор Mojo 1.1 такое не ловит. Как делить результаты между потоками правильно — в главе «Потоки, процессы и async».

Сколько стоит запустить поток

А вот это — главное, ради чего стоит читать раздел. Одна и та же задача (корень из каждого элемента, запись в другой массив), разбитая на два куска, на машине с двумя ядрами:

Элементов	Ускорение от двух потоков
256	0,035 — в 28 раз медленнее
4 096	0,27 — вчетверо медленнее
16 384	0,83 — всё ещё медленнее
32 768	1,12
65 536	1,54

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

Почему 1,54, а не 2

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

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

Эти числа — не ваши

Замер сделан на двух физических ядрах (Intel Xeon 2,8 ГГц, облачный контейнер). Если у вас восемь или шестнадцать, ни одна цифра из таблицы вам не подойдёт, и подобрать её умножением не получится:

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

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

Свои числа — только замером, как в главе Как честно мерить скорость.

Порядок действий

Если код тормозит, браться за инструменты стоит в этом порядке:

1. Замерить. Без этого всё остальное — гадание.
2. Векторизовать — vectorize. Даёт кратно ширине вектора, работает всегда, ничего не стоит.
3. Развернуть — unroll_factor. Ещё в два-три раза, тоже даром.
4. Распараллелить — parallelize. Только если данных много и работа над каждым элементом не пустячная.
5. Замерить снова. Особенно после четвёртого пункта.

Первые три шага почти безопасны. Четвёртый — единственный, который может сделать хуже, и на маленьких данных делает это уверенно.

Всё вместе

Одна и та же операция — умножение массива на число — написана дважды: руками и через vectorize. Результаты сверяются.

auto.mojo

# vectorize делает то же, что мы писали руками в главе про SIMD:
# режет массив на векторы и досчитывает хвост. Только сам.
#
# Обратите внимание на список захвата {imm ptr} — без него вложенная
# функция не увидит переменных снаружи.

from std.algorithm import vectorize
from std.sys import simd_width_of

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


def scale_by_hand(mut data: List[Float32], k: Float32):
    """Ручная версия: основная часть плюс хвост."""
    var ptr = data.unsafe_ptr()
    var full = (len(data) // W) * W

    for i in range(0, full, W):
        ptr.unsafe_store(i, ptr.unsafe_load[width=W](i) * k)

    for i in range(full, len(data)):
        ptr.unsafe_store(i, ptr.unsafe_load(i) * k)


def scale_with_vectorize(mut data: List[Float32], k: Float32):
    """То же самое, но нарезку и хвост берёт на себя библиотека."""
    var ptr = data.unsafe_ptr()

    # width приходит параметром: для основной части он равен W,
    # для хвоста библиотека вызовет эту же функцию с width = 1.
    def step[width: Int](idx: Int) {imm ptr, imm k}:
        ptr.unsafe_store(idx, ptr.unsafe_load[width=width](idx) * k)

    vectorize[W](len(data), step)


def main():
    var a = List[Float32](length=N, fill=0)
    var b = List[Float32](length=N, fill=0)
    for i in range(N):
        a[i] = Float32(i)
        b[i] = Float32(i)

    scale_by_hand(a, 2.0)
    scale_with_vectorize(b, 2.0)

    var same = True
    for i in range(N):
        if a[i] != b[i]:
            same = False

    print("элементов:", N)
    print("первый и последний:", a[0], a[N - 1])
    print("руками и через vectorize совпало:", same)

Результат

элементов: 1001
первый и последний: 0.0 2000.0
руками и через vectorize совпало: True

Ручная версия — восемь строк с двумя циклами и вычислением границы. Версия с vectorize — четыре строки, и хвост в ней не упомянут вовсе.

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

Почему функция для `vectorize` объявлена как `step[width: Int](idx: Int)`, а не просто `step(idx: Int)`?

Потому что vectorize вызывает её с разной шириной: для основной части — с полной, а для хвоста — с шириной 1, по одному элементу. Один и тот же код обслуживает оба случая, и отдельный цикл для хвоста писать не нужно.

Компилятор говорит `Could not infer capture convention of the captured value ptr`. Что это значит?

Вложенная функция обратилась к переменной снаружи, а как её захватывать — компилятор не угадывает. Нужен список захвата в фигурных скобках после аргументов: {imm ptr} для чтения, {mut ptr} если требуется менять.

В старом примере вложенная функция помечена `@parameter`. Можно ли передать её в `vectorize`?

Нет. Декоратор делает замыкание capturing, а vectorize такое не принимает — причём ошибка появится, даже если функция ничего не захватывает. Нужна обычная вложенная def со списком захвата.

Массив из 256 элементов, два ядра. Стоит ли распараллеливать?

Категорически нет. На таком размере два потока оказались в 28 раз медленнее одного: запустить поток дороже, чем выполнить всю работу. Окупаться параллелизм начал только после тридцати тысяч элементов — и то дав 1,1, а не 2. На машине с большим числом ядер порог будет ещё выше: потоков больше, запускать их дороже.

Статья предлагает `from algorithm import parallelize`. Почему не собирается?

Потому что в Mojo этой функции в стандартной библиотеке нет: package 'algorithm' does not contain 'parallelize'. Она переехала в отдельный пакет max — нужно uv add max и импорт from max.algorithm import parallelize. Статья написана до 1.0.

Что дальше

Указатели и память: что такое Pointer, чем он отличается от ссылки и зачем в его методах так часто встречается слово unsafe.