Параллельное умножение матриц

Урок про задачу, которую любят все параллельные системы: умножение матриц.

Умножение матриц C = A·B: элемент C[i][j] — скалярное произведение строки i матрицы A и столбца j матрицы B. Каждый элемент результата вычисляется независимо.

Почему это идеально параллельно

Каждый из n² элементов результата считается полностью независимо от других — это n² параллельных задач, не делящих состояние. Никаких зависимостей, никаких гонок при записи (каждое ядро пишет в свою ячейку C). Поэтому умножение матриц — эталонный бенчмарк параллельных машин и сердце GPU-вычислений: обучение нейросетей — это почти сплошные умножения матриц.

def matmul(A, B):
    n, m, k = len(A), len(B[0]), len(B)
    C = [[0]*m for _ in range(n)]
    for i in range(n):
        for j in range(m):
            s = 0
            for t in range(k):
                s += A[i][t] * B[t][j]    # каждый C[i][j] независим
            C[i][j] = s
    return C

A = [[1, 2], [3, 4]]
B = [[5, 6], [7, 8]]
C = matmul(A, B)
for row in C:
    print(row)
print("операций (2*n^3) =", 2 * len(A) * len(A) * len(A))

Вывод:

[19, 22]
[43, 50]
операций (2*n^3) = 16

Внешние два цикла (по i и j) можно раздать по ядрам полностью независимо: ядро отвечает за свою ячейку или блок ячеек C. Внутренний цикл по t — это редукция (сумма произведений).

Блочное умножение

Наивно раздать по одной ячейке на ядро неэффективно: данные плохо ложатся в кэш. Поэтому матрицы режут на блоки (тайлы): C делят на квадраты, и ядро вычисляет целый блок C, перемножая соответствующие блоки A и B. Блок целиком влезает в кэш, переиспользуется много раз — это резко сокращает обращения к медленной памяти. Блочность улучшает не только параллелизм, но и локальность памяти.

C делим на блоки 2x2; каждое ядро считает свой блок:
+-----+-----+
| C00 | C01 |   ядро 0 -> C00
+-----+-----+   ядро 1 -> C01
| C10 | C11 |   ядро 2 -> C10
+-----+-----+   ядро 3 -> C11
C[блок] = сумма произведений блоков A[строка] * B[столбец]

Почему масштабируется отлично

Умножение матриц имеет высокую арифметическую интенсивность: на 2n³ операций приходится лишь 3n² элементов данных. То есть на каждый загруженный из памяти байт делается много вычислений. Это значит, что задача ограничена счётом, а не памятью или сетью — а именно такие задачи параллелятся лучше всего: накладные расходы на коммуникацию малы относительно полезной работы. Поэтому умножение матриц достигает близкого к линейному ускорения даже на больших кластерах.

Как работает под капотом

На GPU умножение матриц реализуют блоками, которые загружаются в быструю разделяемую память блока, и каждый поток считает несколько элементов результата. Современные ускорители имеют специальные «тензорные ядра» именно под перемножение маленьких матриц. На кластерах используют алгоритмы вроде Cannon или SUMMA, которые организуют пересылку блоков между узлами по кольцу/решётке, минимизируя коммуникацию.

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

Частые ошибки

• Раздавать по одной ячейке на ядро — плохая локальность кэша; нужны блоки.
• Игнорировать порядок обхода (i-j-t против i-t-j) — он сильно влияет на попадания в кэш.
• Думать, что любая матричная операция так же хорошо параллелится — например, обращение матрицы имеет зависимости и сложнее.

Итоги

• Каждый элемент C[i][j] независим — n² параллельных задач без гонок.
• Блочное (тайловое) умножение улучшает локальность кэша и параллелизм.
• Высокая арифметическая интенсивность даёт почти линейное масштабирование.