Проект: матричная библиотека

This content is not available in your language yet.

Умножение матриц — классическая проверка любого языка «для скорости». Задача простая, три вложенных цикла, а разница между наивной и хорошей реализацией измеряется сотнями раз. В этом проекте мы напишем небольшую матричную библиотеку и пройдём весь путь: от учебных трёх циклов до кода, который на одной из наших машин идёт вровень с OpenBLAS — библиотекой, на которой работает NumPy, — а на другой отстаёт от неё на 10–20 %.

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

Как считаем скорость

Умножение матриц n×n — это n³ умножений и n³ сложений, всего 2n³ операций с плавающей точкой. Разделив их на время, получим GFLOPS — миллиарды операций в секунду. Эта мера удобна тем, что не зависит от размера матриц: если GFLOPS падают с ростом n, значит, код упирается во что-то, кроме арифметики, — чаще всего в память.

Замеры шагов ниже — для матриц 500×500 из Float64 (в NumPy тоже float64 по умолчанию), лучшее из пяти запусков, повторённых трижды. Размер 500, а не 512, выбран нарочно — почему, станет ясно на первом шаге.

Машины: облачный сервер с Intel Xeon 2,1 ГГц пятого поколения (Emerald Rapids, 2 ядра, AVX-512) и AMD Ryzen 7 9700X (виртуальная машина с 2 ядрами, AVX-512). Облако каждый раз выдаёт немного разное железо — в других главах это был Xeon 2,8 ГГц, — поэтому цифры Xeon сравнимы только внутри этой главы.

Структура `Matrix`

Матрица хранит размеры и одномерный список чисел: сначала вся первая строка, потом вся вторая и так далее. Такой порядок называется построчным (row-major), как в C и в NumPy по умолчанию.

struct Matrix(Movable, Writable):
    var rows: Int
    var cols: Int
    var data: List[Float64]

    def __init__(out self, rows: Int, cols: Int):
        self.rows = rows
        self.cols = cols
        self.data = List[Float64](length=rows * cols, fill=0.0)

    def __getitem__(self, i: Int, j: Int) -> Float64:
        return self.data[i * self.cols + j]

    def __setitem__(mut self, i: Int, j: Int, value: Float64):
        self.data[i * self.cols + j] = value

__getitem__ и __setitem__ с двумя аргументами дают привычную запись m[i, j] — как в NumPy. А метод __matmul__ даёт оператор @:

def __matmul__(self, other: Matrix) raises -> Matrix:
    return matmul(self, other)

Умножать матрицы несовместимых размеров нельзя, поэтому все функции умножения начинаются с проверки и помечены raises:

def check_shapes(a: Matrix, b: Matrix) raises:
    if a.cols != b.rows:
        raise Error(
            String("нельзя умножить ", a.rows, "×", a.cols, " на ", b.rows, "×", b.cols)
        )

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

Шаг 1. Как в учебнике

def matmul_naive(a: Matrix, b: Matrix) raises -> Matrix:
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    for i in range(a.rows):
        for j in range(b.cols):
            var acc: Float64 = 0.0
            for k in range(a.cols):
                acc += a[i, k] * b[k, j]
            c[i, j] = acc
    return c^

	Xeon	Ryzen
1. наивное	0,52 GFLOPS	1,7 GFLOPS

На Xeon с частотой 2,1 ГГц это одна операция за четыре такта. Беда во внутреннем цикле: b[k, j] при росте k прыгает по памяти через целую строку — 500 чисел, 4 КБ. Каждое обращение приносит из памяти 64 байта кэш-линии, из которых нужны 8.

А при n = 512 всё ещё хуже: 0,25 GFLOPS на Xeon и 1,3 на Ryzen. Шаг ровно в 4096 байт — степень двойки — отправляет все эти строки в одни и те же ячейки кэша, и они вытесняют друг друга. Поэтому мерить умножение матриц только на размерах-степенях двойки — значит получить заниженные цифры для наивного кода и завышенное ускорение.

Шаг 2. Другой порядок циклов

Если поменять местами циклы по j и k, внутренний цикл пойдёт вдоль строк и B, и C — подряд по памяти:

for i in range(a.rows):
    for k in range(a.cols):
        var aik = a[i, k]
        for j in range(b.cols):
            c[i, j] += aik * b[k, j]

Результат тот же, порядок сложений для каждого c[i, j] — тоже. А скорость:

	Xeon	Ryzen
1. наивное	0,52	1,7
2. порядок циклов	0,66	1,6

На Xeon чуть лучше, на Ryzen даже чуть хуже — совсем не то, что обещают учебники по кэшам. Что-то ещё держит код. Проверим догадку — соберём ту же программу с -D ASSERT=none:

500×500, GFLOPS	обычная сборка	`-D ASSERT=none`
1. наивное, Xeon	0,52	1,7
2. порядок циклов, Xeon	0,66	4,0
1. наивное, Ryzen	1,7	2,9
2. порядок циклов, Ryzen	1,6	9,5

Вот он, виновник: проверка границ. Каждое c[i, j] и b[k, j] идёт через индексацию List, а она в Mojo по умолчанию проверяет, не вышел ли индекс за пределы (подробнее — в главе «Ошибки»). Проверка — это не одно сравнение. Вот внутренний цикл шага 2 в машинном коде (mojo build --emit asm), по одному элементу j:

movq  $6, 8(%rsp)          ┐
movq  %rdi, (%rsp)         │ шесть записей в стек: заготовка
movq  %r12, 16(%rsp)       │ сообщения об ошибке — на случай,
movq  $39, 272(%rsp)       │ если проверка не пройдёт
movq  %r8, 264(%rsp)       │
movq  %r12, 280(%rsp)      ┘
cmpq  %r14, %r9            проверка индекса b[k, j]
jae   .LBB0_96
vmovsd (%r15,%r9,8), %xmm1
...                        ещё шесть записей и проверка для c[i, j]
vfmadd213sd (%r13,%rax,8), %xmm0, %xmm1     полезная работа
vmovsd %xmm1, (%r13,%rax,8)

Двенадцать записей в память и две проверки ради одного умножения-сложения. Порядок доступа к памяти на таком фоне почти не заметен.

ASSERT=none отключает проверки во всей программе — слишком грубый инструмент. Нам нужно снять их только здесь.

Шаг 3. Указатели вместо индексов

unsafe_ptr() даёт сырой указатель на данные списка, а unsafe_load / unsafe_store читают и пишут без всяких проверок (о них — в главе «Указатели и память»):

var pa = a.data.unsafe_ptr()
var pb = b.data.unsafe_ptr()
var pc = c.data.unsafe_ptr()
var n = b.cols
for i in range(a.rows):
    for k in range(a.cols):
        var aik = pa.unsafe_load(i * a.cols + k)
        for j in range(n):
            var dst = i * n + j
            pc.unsafe_store(dst, pc.unsafe_load(dst) + aik * pb.unsafe_load(k * n + j))

	Xeon	Ryzen
2. порядок циклов	0,66	1,6
3. указатели	4,1	9,1

В шесть раз — столько же, сколько давал ASSERT=none. И это не совпадение: машинный код внутреннего цикла здесь тот же, что у шага 2 без проверок, — загрузка, vfmadd213sd, запись. Но теперь проверки сняты только в одной функции. Цена та же, что всегда с unsafe_: ошибись мы с индексом — программа молча прочитает чужую память. Поэтому такой код пишут только там, где он окупается, и обкладывают тестами.

Заметьте: vfmadd213sd — скалярная инструкция, по одному числу за раз. Сам компилятор этот цикл не векторизовал.

Шаг 4. SIMD

Внутренний цикл делает одно и то же с соседними элементами строки: c[i, j] += aik * b[k, j]. Это работа ровно для SIMD — обрабатывать по W чисел за раз, где W = simd_width_of[DType.float64](): восемь на обеих наших машинах с AVX-512. Нарезку на векторы и хвост берёт на себя vectorize из главы «vectorize и parallelize»:

for i in range(a.rows):
    for k in range(a.cols):
        var aik = pa.unsafe_load(i * a.cols + k)

        def update[width: Int](j: Int) {imm pb, imm pc, imm aik, imm i, imm k, imm n}:
            var dst = i * n + j
            pc.unsafe_store(
                dst,
                pc.unsafe_load[width=width](dst)
                + aik * pb.unsafe_load[width=width](k * n + j),
            )

        vectorize[W](n, update)

aik — одно число, а pb.unsafe_load[width=width](...) — вектор; при умножении число копируется во все элементы вектора.

	Xeon	Ryzen
3. указатели	4,1	9,1
4. SIMD	9,7	33,5

Шаг 5. Потоки

Строки C считаются независимо друг от друга — их можно раздать ядрам. parallelize из пакета max делит номера строк на куски по числу потоков и вызывает функцию для каждого номера:

def one_row(i: Int) {imm pa, imm pb, imm pc, imm n, imm inner}:
    for k in range(inner):
        # ... то же обновление строки, что на шаге 4
        vectorize[W](n, update)

parallelize(one_row, a.rows)

	Xeon	Ryzen
4. SIMD	9,7	33,5
5. SIMD и потоки	13,0	36,5

Два ядра, а прирост — в 1,1–1,3 раза. Значит, упираемся не в вычисления. Посмотрим, что делает внутренний цикл на одно умножение-сложение: читает вектор B, читает и записывает вектор C — три обращения к памяти на одну полезную операцию. А главное — для каждой строки C заново читается вся матрица B: 2 МБ при n = 500. Это больше кэша второго уровня одного ядра (у этого Xeon — 2 МБ, у Ryzen — 1 МБ), так что оба ядра тянут B из общего кэша третьего уровня и стоят к нему в очереди.

Шаг 6. Упаковка и блок в регистрах

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

Возьмём полосу C шириной BLOCK = 4 × W столбцов (32 числа при AVX-512) и четыре строки сразу. Это 16 векторных регистров — сумматоры. Проходим по всем k и накапливаем:

var acc0 = SIMD[DType.float64, BLOCK](0)
# ... acc1, acc2, acc3 — ещё три строки
for k in range(inner):
    var bk = ps.unsafe_load[width=BLOCK](k * BLOCK)
    acc0 += pa.unsafe_load((i + 0) * inner + k) * bk
    acc1 += pa.unsafe_load((i + 1) * inner + k) * bk
    acc2 += pa.unsafe_load((i + 2) * inner + k) * bk
    acc3 += pa.unsafe_load((i + 3) * inner + k) * bk

Теперь на каждое k приходится одно чтение полосы B (четыре вектора) и четыре числа из A — на 16 векторных умножений-сложений, а записей в C нет вовсе: только одна в самом конце, после всех k. SIMD[DType.float64, BLOCK] шире регистра, но это не страшно: компилятор сам разложит его на четыре регистра. В машинном коде внутренний цикл — ровно 4 загрузки B, 4 размножения чисел A и 16 инструкций vfmadd231pd, без единого обращения к стеку.

Шестнадцать сумматоров помещаются в регистры только там, где регистров 32: у AVX-512 и у ARM. У AVX2 их 16, и с BLOCK = 4 × W сумматоры не влезают — компилятор выгружает их в память, и выигрыш пропадает. Проверено сборкой с --target-cpu x86-64-v3: при BLOCK = 2 × W скорость на 1024×1024 выросла с 22 до 29 GFLOPS на одном потоке и с 45 до 60 на двух. Поэтому в библиотеке ширина полосы выбирается при сборке:

comptime BLOCK = (
    2 * W if CompilationTarget.is_x86()
    and not CompilationTarget.has_avx512f() else 4 * W
)

Длина SIMD-вектора в Mojo должна быть степенью двойки, так что промежуточного 3 × W не бывает.

Остаётся проблема с B. Нужная полоса — это столбцы с j0 по j0 + 31 всех строк, а в памяти они разбросаны: между соседними строками 4 КБ. Поэтому перед счётом полосу упаковывают — копируют подряд в отдельный буфер:

var strip = List[Float64](length=inner * BLOCK, fill=0.0)
var ps = strip.unsafe_ptr()
for k in range(inner):
    for t in range(width):
        ps.unsafe_store(k * BLOCK + t, pb.unsafe_load(k * n + j0 + t))

Копирование стоит n² операций на фоне 2n³ вычислений — почти даром, зато потом все четыре строки, а следом и все остальные, читают полосу строго подряд. Полосы независимы, и parallelize раздаёт их потокам. С хвостами всё просто: последняя полоса, если она уже 32 столбцов, считается тем же кодом — недостающие столбцы буфера заполнены нулями, а в C записываются только настоящие. Отдельно обрабатываются лишь последние строки, если их число не делится на четыре.

	Xeon	Ryzen
5. SIMD и потоки	13,0	36,5
6. упаковка, 1 поток	46,2	126,3
6. упаковка, 2 потока	78,9	231,5

Один поток теперь быстрее, чем два на шаге 5, а два потока дают почти двукратный прирост: код перестал стоять в очереди к памяти.

Итог

500×500, GFLOPS	Xeon	Ryzen
1. наивное	0,52	1,7
2. порядок циклов	0,66	1,6
3. указатели	4,1	9,1
4. SIMD	9,7	33,5
5. SIMD и потоки	13,0	36,5
6. упаковка, 2 потока	78,9	231,5
ускорение	в 150 раз	в 135 раз

Заметьте, что дал каждый шаг. Больше всего — снятие проверок (×6) и правильная работа с памятью (×6). SIMD дал ×2,4–3,7, а потоки почти ничего не давали, пока код упирался в память.

А что NumPy?

NumPy умножает матрицы через OpenBLAS — библиотеку, которую оптимизируют десятилетиями, с ядрами на ассемблере под каждое семейство процессоров. Сравним на больших матрицах (Python 3.14.7, NumPy 2.5.3, OpenBLAS 0.3.34, который на обеих машинах выбрал ядро для AVX-512):

GFLOPS	Xeon, 1024	Ryzen, 1024	Xeon, 2048	Ryzen, 2048
NumPy, 1 поток	55,2	119,6	58,3	123,1
наша библиотека, 1 поток	46,0	125,5	46,3	119,5
NumPy, 2 потока	108,5	234,7	117,7	207,4
наша библиотека, 2 потока	95,9	232,5	93,9	212,4

На Xeon мы получили 80–88 % от OpenBLAS, на Ryzen — вровень: разница в несколько процентов в обе стороны меньше разброса между запусками на этой машине. Для сотни строк на языке высокого уровня, без единой ассемблерной вставки, это отличный результат.

Чего у нас нет и что есть в OpenBLAS: блочное деление по k для больших матриц (чтобы полоса B оставалась в кэше), упаковка A, предвыборка данных, отдельные ядра под каждый процессор.

Аргумент по умолчанию вычисляется при сборке

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

def matmul(a: Matrix, b: Matrix, workers: Int = num_logical_cores()) raises -> Matrix:

Само объявление компилируется. Но как только функцию вызовут без этого аргумента — matmul(a, b), — сборка останавливается, и ошибка указывает на место вызова:

error: function instantiation failed … note: failed to compile-time evaluate function call … unable to interpret call to unknown external function: KGEN_CompilerRT_NumLogicalCores

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

Значения аргументов по умолчанию Mojo вычисляет при компиляции. А число ядер известно только при запуске — на машине, где программа будет работать. Компилятор пытается выполнить num_logical_cores() у себя и не может. То же будет с любой функцией, которой нужна работающая программа: perf_counter_ns(), argv(), случайные числа.

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

Возьмите значение-метку и разберитесь с ним внутри функции. В нашей библиотеке workers: Int = 0 означает «все ядра»: if workers > 0: parallelize(f, n, workers) else: parallelize(f, n).

Нагляднее всего это видно на функции, которая печатает:

def g() -> Int:
    print("вычисляю значение по умолчанию")
    return 4


def f(x: Int = g()) -> Int:
    return x


def main():
    print("f() =", f())

$ uv run mojo build dg.mojo
вычисляю значение по умолчанию
$ ./dg
f() = 4

Строка печатается во время сборки, а собранная программа её уже не печатает: значение 4 вшито в неё готовым. В Python аргумент по умолчанию тоже вычисляется один раз, но при выполнении строки def — то есть уже во время работы программы.

Проверка правильности

Быстрый код бесполезен, если он считает неправильно, а в шаге 6 легко ошибиться с хвостами. Поэтому библиотеку проверяют тесты (как их писать — в главе «Тестирование и отладка»):

произведение маленьких матриц, посчитанное вручную;
умножение на единичную матрицу ничего не меняет;
результат совпадает с наивной версией на «неудобных» размерах: 1, 3, 5 и 33 строки, 1, 7, 31, 33 и 65 столбцов — всё, что не делится ни на 4, ни на ширину полосы;
один поток даёт ровно тот же результат, что и все;
умножение несовместимых матриц — ошибка с понятным текстом.

С наивной версией результаты совпадают до последнего бита, хотя проверяем мы с допуском 10⁻¹². Для этого нужны два условия: каждое c[i, j] во всех версиях складывается в одном и том же порядке по k, и компилятор везде одинаково сливает умножение со сложением в одну инструкцию FMA (подробнее о ней — в главе «Тестирование и отладка»). Второе условие хрупкое: в сборке с -O0 наивная версия считает без FMA, а быстрая — с ней, и результаты расходятся примерно на 10⁻¹⁴. Поэтому тесты и сравнивают с допуском.

С NumPy интереснее. Мы сохранили наши матрицы и перемножили их в NumPy: при внутреннем размере до 384 результат OpenBLAS совпал с нашим бит в бит, а начиная с 385 разошёлся почти во всех элементах — в пределах 5·10⁻¹⁴. Похоже, ядро OpenBLAS для AVX-512 делит сумму по k на блоки по 384 — то самое блочное деление, которого нет у нас, — и порядок сложений меняется. Сравнивать свою библиотеку с NumPy можно только с допуском.

Весь код

"""Матричная библиотека: от наивного умножения до упаковки и потоков."""

from std.algorithm import vectorize
from std.random import random_float64, seed
from std.sys import simd_width_of
from std.sys.info import CompilationTarget
from max.algorithm import parallelize

comptime W = simd_width_of[DType.float64]()
"""Сколько Float64 помещается в один векторный регистр."""


struct Matrix(Movable, Writable):
    """Матрица Float64, строки лежат в памяти одна за другой."""

    var rows: Int
    var cols: Int
    var data: List[Float64]

    def __init__(out self, rows: Int, cols: Int):
        self.rows = rows
        self.cols = cols
        self.data = List[Float64](length=rows * cols, fill=0.0)

    @staticmethod
    def random(rows: Int, cols: Int, seed_value: Int) -> Matrix:
        seed(seed_value)
        var m = Matrix(rows, cols)
        for i in range(rows * cols):
            m.data[i] = random_float64(-1.0, 1.0)
        return m^

    @staticmethod
    def identity(n: Int) -> Matrix:
        var m = Matrix(n, n)
        for i in range(n):
            m[i, i] = 1.0
        return m^

    def __getitem__(self, i: Int, j: Int) -> Float64:
        return self.data[i * self.cols + j]

    def __setitem__(mut self, i: Int, j: Int, value: Float64):
        self.data[i * self.cols + j] = value

    def __matmul__(self, other: Matrix) raises -> Matrix:
        """Оператор `a @ b`, как в NumPy."""
        return matmul(self, other)

    def max_abs_diff(self, other: Matrix) -> Float64:
        """Наибольшее расхождение; бесконечность, если размеры разные или есть NaN.
        """
        if self.rows != other.rows or self.cols != other.cols:
            return Float64.MAX
        var worst: Float64 = 0.0
        for i in range(len(self.data)):
            var d = abs(self.data[i] - other.data[i])
            # `not (d <= worst)` ловит и NaN: с ним любое сравнение ложно
            if not (d <= worst):
                worst = d if d == d else Float64.MAX
        return worst

    def write_to(self, mut writer: Some[Writer]):
        for i in range(self.rows):
            writer.write("[")
            for j in range(self.cols):
                if j > 0:
                    writer.write(", ")
                writer.write(self[i, j])
            writer.write("]\n")


def check_shapes(a: Matrix, b: Matrix) raises:
    if a.cols != b.rows:
        raise Error(
            String(
                "нельзя умножить ",
                a.rows,
                "×",
                a.cols,
                " на ",
                b.rows,
                "×",
                b.cols,
            )
        )


# --- Шаг 1. Как в учебнике ----------------------------------------------


def matmul_naive(a: Matrix, b: Matrix) raises -> Matrix:
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    for i in range(a.rows):
        for j in range(b.cols):
            var acc: Float64 = 0.0
            for k in range(a.cols):
                acc += a[i, k] * b[k, j]
            c[i, j] = acc
    return c^


# --- Шаг 2. Другой порядок циклов: B читается по строкам -------------------


def matmul_reordered(a: Matrix, b: Matrix) raises -> Matrix:
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    for i in range(a.rows):
        for k in range(a.cols):
            var aik = a[i, k]
            for j in range(b.cols):
                c[i, j] += aik * b[k, j]
    return c^


# --- Шаг 3. Указатели вместо индексов списка -------------------------------


def matmul_pointers(a: Matrix, b: Matrix) raises -> Matrix:
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    var pa = a.data.unsafe_ptr()
    var pb = b.data.unsafe_ptr()
    var pc = c.data.unsafe_ptr()
    var n = b.cols
    for i in range(a.rows):
        for k in range(a.cols):
            var aik = pa.unsafe_load(i * a.cols + k)
            for j in range(n):
                var dst = i * n + j
                pc.unsafe_store(
                    dst, pc.unsafe_load(dst) + aik * pb.unsafe_load(k * n + j)
                )
    return c^


# --- Шаг 4. SIMD: строка C обновляется векторами ---------------------------


def matmul_simd(a: Matrix, b: Matrix) raises -> Matrix:
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    var pa = a.data.unsafe_ptr()
    var pb = b.data.unsafe_ptr()
    var pc = c.data.unsafe_ptr()
    var n = b.cols
    for i in range(a.rows):
        for k in range(a.cols):
            var aik = pa.unsafe_load(i * a.cols + k)

            def update[
                width: Int
            ](j: Int) {imm pb, imm pc, imm aik, imm i, imm k, imm n}:
                var dst = i * n + j
                pc.unsafe_store(
                    dst,
                    pc.unsafe_load[width=width](dst)
                    + aik * pb.unsafe_load[width=width](k * n + j),
                )

            vectorize[W](n, update)
    return c^


# --- Шаг 5. SIMD и потоки: строки C делятся между ядрами -------------------


def matmul_threads(a: Matrix, b: Matrix) raises -> Matrix:
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    var pa = a.data.unsafe_ptr()
    var pb = b.data.unsafe_ptr()
    var pc = c.data.unsafe_ptr()
    var n = b.cols
    var inner = a.cols

    def one_row(i: Int) {imm pa, imm pb, imm pc, imm n, imm inner}:
        for k in range(inner):
            var aik = pa.unsafe_load(i * inner + k)

            def update[
                width: Int
            ](j: Int) {imm pb, imm pc, imm aik, imm i, imm k, imm n}:
                var dst = i * n + j
                pc.unsafe_store(
                    dst,
                    pc.unsafe_load[width=width](dst)
                    + aik * pb.unsafe_load[width=width](k * n + j),
                )

            vectorize[W](n, update)

    parallelize(one_row, a.rows)
    return c^


# --- Шаг 6. Упаковка полосы B и блок 4 × BLOCK в регистрах -----------------

comptime BLOCK = (
    2 * W if CompilationTarget.is_x86()
    and not CompilationTarget.has_avx512f() else 4 * W
)
"""Ширина полосы. Сумматоров ROWS × BLOCK / W должно хватить регистров:
у AVX-512 и ARM их 32 — берём 4 вектора на строку (16 сумматоров),
у AVX2 всего 16 — только 2 (8 сумматоров), иначе они не поместятся."""
comptime ROWS = 4
"""Сколько строк C считается за один проход по полосе."""


def matmul(a: Matrix, b: Matrix, workers: Int = 0) raises -> Matrix:
    """Быстрое умножение: C = A × B.

    workers — сколько потоков использовать; 0 — все ядра.
    """
    check_shapes(a, b)
    var c = Matrix(a.rows, b.cols)
    var pa = a.data.unsafe_ptr()
    var pb = b.data.unsafe_ptr()
    var pc = c.data.unsafe_ptr()
    var m = a.rows
    var n = b.cols
    var inner = a.cols
    var strips = (n + BLOCK - 1) // BLOCK

    def one_strip(s: Int) {imm pa, imm pb, imm pc, imm m, imm n, imm inner}:
        var j0 = s * BLOCK
        var width = min(BLOCK, n - j0)

        # Столбцы j0 … j0 + BLOCK - 1 всех строк B — подряд в одном буфере.
        # Последняя полоса может быть уже: недостающее остаётся нулями.
        var strip = List[Float64](length=inner * BLOCK, fill=0.0)
        var ps = strip.unsafe_ptr()
        for k in range(inner):
            for t in range(width):
                ps.unsafe_store(k * BLOCK + t, pb.unsafe_load(k * n + j0 + t))

        var i = 0
        while i + ROWS <= m:
            # 4 строки × BLOCK столбцов C копятся в регистрах по всем k
            var acc0 = SIMD[DType.float64, BLOCK](0)
            var acc1 = SIMD[DType.float64, BLOCK](0)
            var acc2 = SIMD[DType.float64, BLOCK](0)
            var acc3 = SIMD[DType.float64, BLOCK](0)
            for k in range(inner):
                var bk = ps.unsafe_load[width=BLOCK](k * BLOCK)
                acc0 += pa.unsafe_load((i + 0) * inner + k) * bk
                acc1 += pa.unsafe_load((i + 1) * inner + k) * bk
                acc2 += pa.unsafe_load((i + 2) * inner + k) * bk
                acc3 += pa.unsafe_load((i + 3) * inner + k) * bk
            for t in range(width):
                pc.unsafe_store((i + 0) * n + j0 + t, acc0[t])
                pc.unsafe_store((i + 1) * n + j0 + t, acc1[t])
                pc.unsafe_store((i + 2) * n + j0 + t, acc2[t])
                pc.unsafe_store((i + 3) * n + j0 + t, acc3[t])
            i += ROWS

        # Оставшиеся строки, если m не делится на ROWS
        while i < m:
            var acc = SIMD[DType.float64, BLOCK](0)
            for k in range(inner):
                acc += pa.unsafe_load(i * inner + k) * ps.unsafe_load[
                    width=BLOCK
                ](k * BLOCK)
            for t in range(width):
                pc.unsafe_store(i * n + j0 + t, acc[t])
            i += 1

    if workers > 0:
        parallelize(one_strip, strips, workers)
    else:
        parallelize(one_strip, strips)
    return c^

Проверка всех шагов на небольших матрицах:

from matrix import (
    Matrix,
    matmul,
    matmul_naive,
    matmul_pointers,
    matmul_reordered,
    matmul_simd,
    matmul_threads,
)


def same(x: Matrix, y: Matrix) -> String:
    return "да" if x.max_abs_diff(y) < 1e-12 else "НЕТ"


def main() raises:
    var a = Matrix(2, 3)
    var b = Matrix(3, 2)
    for i in range(2):
        for j in range(3):
            a[i, j] = Float64(i * 3 + j + 1)
            b[j, i] = Float64(j * 2 + i + 1)
    print(a @ b)

    # Размеры нарочно «неудобные»: 67 не делится на 4, 53 — на ширину полосы.
    var x = Matrix.random(67, 45, 1)
    var y = Matrix.random(45, 53, 2)
    var expected = matmul_naive(x, y)
    print("совпадает с наивным умножением (до 1e-12):")
    print("  переставленные циклы:", same(expected, matmul_reordered(x, y)))
    print("  указатели:           ", same(expected, matmul_pointers(x, y)))
    print("  SIMD:                ", same(expected, matmul_simd(x, y)))
    print("  SIMD и потоки:       ", same(expected, matmul_threads(x, y)))
    print("  упаковка:            ", same(expected, matmul(x, y)))
    print("  упаковка, 1 поток:   ", same(expected, matmul(x, y, workers=1)))

    try:
        _ = a @ a
    except e:
        print("ошибка:", e)

Результат

[22.0, 28.0]
[49.0, 64.0]

совпадает с наивным умножением (до 1e-12):
  переставленные циклы: да
  указатели:            да
  SIMD:                 да
  SIMD и потоки:        да
  упаковка:             да
  упаковка, 1 поток:    да
ошибка: нельзя умножить 2×3 на 2×3

Тесты:

from std.math import nan
from matrix import Matrix, matmul, matmul_naive
from std.testing import assert_equal, assert_raises, assert_true, TestSuite


def test_small_known_product() raises:
    var a = Matrix(2, 2)
    var b = Matrix(2, 2)
    a[0, 0] = 1.0
    a[0, 1] = 2.0
    a[1, 0] = 3.0
    a[1, 1] = 4.0
    b[0, 0] = 5.0
    b[0, 1] = 6.0
    b[1, 0] = 7.0
    b[1, 1] = 8.0
    var c = a @ b
    assert_equal(c[0, 0], 19.0)
    assert_equal(c[0, 1], 22.0)
    assert_equal(c[1, 0], 43.0)
    assert_equal(c[1, 1], 50.0)


def test_identity_changes_nothing() raises:
    var a = Matrix.random(37, 37, 7)
    assert_equal(a.max_abs_diff(matmul(a, Matrix.identity(37))), 0.0)
    assert_equal(a.max_abs_diff(matmul(Matrix.identity(37), a)), 0.0)


def test_matches_naive_on_awkward_sizes() raises:
    # размеры, не кратные ни 4 строкам, ни ширине полосы
    for m in [1, 3, 5, 33]:
        for k in [0, 1, 19]:
            for n in [1, 7, 31, 33, 65]:
                var a = Matrix.random(m, k, m)
                var b = Matrix.random(k, n, n)
                var diff = matmul_naive(a, b).max_abs_diff(matmul(a, b))
                assert_true(
                    diff < 1e-12,
                    msg=String(m, "×", k, " на ", k, "×", n, ": ", diff),
                )


def test_one_thread_equals_many() raises:
    var a = Matrix.random(64, 48, 1)
    var b = Matrix.random(48, 80, 2)
    assert_equal(matmul(a, b).max_abs_diff(matmul(a, b, workers=1)), 0.0)


def test_diff_notices_nan() raises:
    var a = Matrix(1, 2)
    var b = Matrix(1, 2)
    b[0, 1] = nan[DType.float64]()
    assert_true(a.max_abs_diff(b) > 1.0)


def test_shape_mismatch() raises:
    with assert_raises(contains="нельзя умножить"):
        _ = matmul(Matrix(2, 3), Matrix(2, 3))


def main() raises:
    TestSuite.discover_tests[__functions_in_module()]().run()

И замер, которым получены таблицы. Размер матриц — аргумент командной строки. Компиляция при mojo run в замер не попадает — она заканчивается до main, — но собранный бинарник удобнее запускать много раз подряд (подробнее о честных замерах — в главе «Как честно мерить скорость»):

"""Замер всех шагов: время и GFLOPS. Вывод у каждой машины свой.

Запуск: mojo build bench_matmul.mojo && ./bench_matmul 500
"""

from std.benchmark import keep
from std.sys import argv
from std.time import perf_counter_ns
from matrix import (
    Matrix,
    matmul,
    matmul_naive,
    matmul_pointers,
    matmul_reordered,
    matmul_simd,
    matmul_threads,
)


def best_seconds[
    f: def(Matrix, Matrix) thin raises -> Matrix
](a: Matrix, b: Matrix) raises -> Float64:
    """Лучшее время из пяти запусков."""
    var best = Float64.MAX
    for _ in range(5):
        var t0 = perf_counter_ns()
        var c = f(a, b)
        var t1 = perf_counter_ns()
        keep(c.data[len(c.data) - 1])
        best = min(best, Float64(t1 - t0) / 1e9)
    return best


def one_thread(a: Matrix, b: Matrix) raises -> Matrix:
    return matmul(a, b, workers=1)


def all_threads(a: Matrix, b: Matrix) raises -> Matrix:
    return matmul(a, b)


def report(name: String, n: Int, seconds: Float64):
    var gflops = 2.0 * Float64(n) ** 3 / seconds / 1e9
    print(
        name, "  ", round(seconds * 1000, 2), "мс  ", round(gflops, 1), "GFLOPS"
    )


def main() raises:
    var n = Int(argv()[1]) if len(argv()) > 1 else 500
    if n <= 0:
        raise Error("размер матриц должен быть больше нуля")
    var a = Matrix.random(n, n, 1)
    var b = Matrix.random(n, n, 2)
    print("матрицы", n, "×", n)
    if n <= 512:
        report("1. наивное             ", n, best_seconds[matmul_naive](a, b))
        report(
            "2. порядок циклов      ", n, best_seconds[matmul_reordered](a, b)
        )
    report("3. указатели           ", n, best_seconds[matmul_pointers](a, b))
    report("4. SIMD                ", n, best_seconds[matmul_simd](a, b))
    report("5. SIMD и потоки       ", n, best_seconds[matmul_threads](a, b))
    report("6. упаковка, 1 поток   ", n, best_seconds[one_thread](a, b))
    report("6. упаковка, все потоки", n, best_seconds[all_threads](a, b))

uv add max          # ради parallelize
uv run mojo build bench_matmul.mojo
./bench_matmul 500

Что можно добавить

транспонирование, сложение, умножение на число — с тем же набором приёмов;
Float32: вдвое больше чисел в регистре — в теории и вдвое больше GFLOPS; проверьте, так ли это;
блочное деление по k: найдите размер, начиная с которого GFLOPS без него падают, и сравните с размером кэша вашего процессора;
сравнение с NumPy прямо из Mojo через Python.import_module("numpy") (глава «Python из Mojo»).

Что дальше

Следующий проект — ускоряем Python-скрипт: как перенести в Mojo узкое место программы, которая уже написана.

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

Почему смена порядка циклов (шаг 2) почти ничего не дала, хотя доступ к памяти стал последовательным?

Внутренний цикл тормозила проверка границ в индексации List: на каждый элемент — две проверки и двенадцать записей в стек для будущего сообщения об ошибке. Со сборкой -D ASSERT=none та же версия ускорилась в 6 раз. Решение без отключения проверок во всей программе — указатели и unsafe_load / unsafe_store в одной горячей функции.

Почему на шаге 5 два ядра дали прирост всего в 1,1–1,3 раза?

Код упирался в память: на каждое умножение-сложение он читал вектор B и читал и записывал вектор C, а для каждой строки C заново перечитывал всю B, которая не помещается в кэш второго уровня. Второе ядро только добавляло очередь к памяти. Помогло накопление блока C в регистрах и упаковка полосы B — после этого два потока дали прирост в 1,7–1,8 раза.

Зачем копировать полосу матрицы B в отдельный буфер, если это лишняя работа?

Нужные столбцы разбросаны по памяти с шагом в целую строку. После упаковки их читают строго подряд, и так — для всех строк A. Копирование стоит n² операций на фоне 2n³ вычислений, то есть почти ничего.

Почему нельзя написать workers: Int = num_logical_cores()?

Аргументы по умолчанию в Mojo вычисляются при компиляции, а число ядер известно только при запуске. Ошибка появится при первом вызове без этого аргумента. Нужна метка вроде 0 и выбор внутри функции.

Стоит ли переписывать на Mojo numpy.dot из своей программы?

Как правило, нет: OpenBLAS внутри NumPy оптимизирован десятилетиями, и даже хороший код на Mojo в лучшем случае идёт с ним вровень. Mojo выигрывает там, где готовой библиотечной функции нет.

Примеры проверены на Mojo 1.1.0

Тексты курса — CC BY-NC-SA 4.0, код примеров — Apache 2.0