Выполнение операции уменьшения с помощью Metal

Напомним, что в Metal вычислительный проход выполняет функцию ядра в нескольких потоках, которые организованы в группы потоков, которые, в свою очередь, подразделяются на группы SIMD.

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

1. В группе SIMD
2. Внутри группы потоков, в группах SIMD
3. По группам потоков

В группе SIMD

Это основной строительный блок нашего ядра. Я позаимствую иллюстрацию из Технического блога NVIDIA, показывающую размер группы SIMD, равный 8:

Функция «перетасовки вниз», simd_shuffle_down() в Metal, позволяет потоку считывать значение переменной из другого потока в той же группе SIMD. Суммируя это с собственной частичной суммой текущего потока, мы уменьшаем количество слагаемых вдвое. Таким образом, для суммирования по группе SIMD требуется log_2(simd_size) шагов.

В примере шейдера Apple этот код находится здесь:

for (uint offset=simd_size/2; offset>0; offset/=2)
    val += simd_shuffle_down(val, offset);

Конечный результат попадает на первую полосу, поэтому доступ к нему может получить только первый поток в группе SIMD. Поэтому мы обусловливаем любой код, которому необходимо прочитать сумму SIMD-группы из val, с помощью if (simd_lane_id == 0).

Конечно, нам нужно инициализировать переменную val каждого потока другим значением из входного массива. Ой-ой, вот ОШИБКА APPLE №1:

int val = input[gid] + input[gid + lsize];  // no!
int val = input[gid];  // yes

Внутри группы потоков, в группах SIMD

Группа потоков состоит из нескольких групп SIMD, поэтому теперь, когда мы нашли частичную сумму каждой группы SIMD, мы должны их объединить.

Сначала нам понадобится способ, с помощью которого потоки смогут передавать частичные суммы своих SIMD-групп. В списке параметров ядра мы объявляем буфер хранения, общий для группы потоков, threadgroup int *ldata [[threadgroup(0)]]. Когда мы настраиваем наш вычислительный проход, мы предоставляем этому хранилищу MTLComputeCommandEncoder.setThreadgroupMemoryLength(_:index:).

После выполнения сокращения группы SIMD мы сохраняем результат в памяти группы потоков. После этого у нас есть threadgroup_barrier, чтобы дождаться сохранения частичных сумм всех SIMD-групп.

// Write per-SIMD partial reduction value to threadgroup memory.
if (simd_lane_id == 0)
    ldata[simd_group_id] = val;

// Wait for all partial reductions to complete.
threadgroup_barrier(mem_flags::mem_threadgroup);

На данный момент буфер ldata содержит все частичные суммы SIMD-групп, и теперь нам нужно их суммировать. У нас уже есть параллельное сокращение суммы внутри SIMD-группы, так почему бы не использовать его повторно? Если мы скопируем данные обратно из ldata, мы сможем просто повторить предыдущий шаг:

Функция ядра делает это с этой условной копией. Он использует индекс внутри группы потоков lid, так что первый поток считывает частичную сумму первой группы SIMD, второй поток считывает частичную сумму второй группы SIMD и так далее.

val = (lid < s) ? ldata[lid] : 0;

Сложная штука — это условное выражение lid < s. Обратите внимание, что s инициализируется количеством lsize/simd_size, то есть количеством SIMD-групп в нашей группе потоков. Это условие позволяет избежать доступа к незаполненным индексам ldata.

Когда группы SIMD загружены новыми слагаемыми, мы позволяем циклу вернуться наверх, чтобы повторить процесс:

for (uint s=lsize/simd_size; s>simd_size; s/=simd_size)  // BUG 2
{
    // compute partial sum within SIMD group
    // store partial sums in threadgroup storage ldata
    // copy ldata storage back to SIMD group
}

Каждый раз мы уменьшаем количество слагаемых в simd_size раз, поэтому идея заключается в том, что мы продолжаем до тех пор, пока не останется всего одна SIMD-группа. Однако это подводит нас к ОШИБКЕ APPLE №2. Переменная s определяется как количество SIMD-групп для обработки, поэтому завершающее условие должно быть s > 1! Мы не хотим закончить тем, что осталось суммировать simd_size групп.

for (uint s=lsize/simd_size; s>simd_size; s/=simd_size)  // no!
for (uint s=lsize/simd_size; s>1; s/=simd_size)  // yes

Когда цикл завершается, нам нужно выполнить последнее суммирование внутри группы SIMD:

// Perform final per-SIMD partial reduction to calculate
// the threadgroup partial reduction result.
for (uint offset=simd_size/2; offset>0; offset/=2)
    val += simd_shuffle_down(val, offset);

Теперь первый поток первой группы SIMD, который также является первым потоком группы потоков (lid == 0), имеет частичную сумму группы потоков в val.

По группам потоков

Наконец, нам нужно взять все частичные суммы, известные первым потокам всех групп потоков, и объединить их.

Ядро Apple делает это с помощью функции atomic_fetch_add_explicit(). Это операция чтения-изменения-записи, которая считывает содержимое буфера output, добавляет к нему val, а затем записывает сумму обратно. Атомный квалификатор гарантирует нам, что все три подэтапа будут выполнены без конфликтов с условиями гонки. (Обратите внимание, что output имеет тип atomic_int.) Поскольку суммирование является коммутативным, порядок, в котором потоки вызывают эту функцию, также не имеет значения.

Опять же, условие lid == 0 связано с тем, что только первый поток группы потоков имеет частичную сумму группы потоков.

// Atomically update the reduction result.
if (lid == 0)
    atomic_fetch_add_explicit(output, val, memory_order_relaxed);

О размерах групп потоков и групп SIMD

Кажется, не существует какого-либо API или письменной документации, в которой указано количество потоков на группу SIMD (связанный вопрос), но это число на всех устройствах Apple (по состоянию на июнь 2024 г.) составляет 32. (Вы можете только получить его программно внутри шейдера, добавив параметр с атрибутом [[threads_per_simdgroup]], как мы делаем здесь.)

Вероятно, хорошей идеей будет иметь размер группы потоков, кратный размеру группы SIMD.

Операции сокращения, кроме суммы

Должно быть возможно использовать примерно тот же алгоритм для выполнения других видов параллельных сокращений.

Операция суммирования val += simd_shuffle_down(val, offset) обобщается до val = op(val, simd_shuffle_down(val, offset)).

Однако вариантов для atomic_fetch_<op>_explicit() не так много: только add, and, max, min, or, sub и xor. Если предполагаемая операция не входит в число этих, возможно, вам придется использовать выходной буфер частичных сокращений групп потоков и позволить ЦП выполнить последний шаг сокращения. Вы должны добавить параметр для положения группы потоков в сетке, например uint tgid [[threadgroup_position_in_grid]], а затем использовать его для сохранения частичного результата:

if (lid == 0)
    output[tgid] = val;