Немного о Mojo

В середине октября этого года на LLVM 2023 Developer Meeting был презентован новый язык программирования Mojo, предназначенный для программирования задач, связанных с AI/ML (кстати, кроме блог поста рекомендую к просмотру запись самой презентации). Я не буду вдаваться в вопрос зачем нужен yet another язык программирования в 2023 году, и попрошу читателей ознакомиться с главой "Why Mojo" из документации языка прежде чем оставлять комментарий такого рода. Скажу лишь, что этот компилируемый язык предназначен для того, чтобы задействовать все возможности современных акселераторов вычислений вроде GPU, TPU, различного рода AI ASICов без погружения в дебри платформо-специфичного ассемблера. Это язык как бы является естественным продолжением экосистемы MLIR, созданным для упрощения взаимодействия с MLIR в задачах, связанных с высокопроизводительными вычислениями. Для лучшего понимания предмета нашего эксперимента приведу пример многопоточной реализации фрактала Мандельброта на Mojo:

from complex import ComplexSIMD, ComplexFloat64
from math import iota
from algorithm import parallelize, vectorize
from tensor import Tensor

alias float_type = DType.float64
alias simd_width = 2 * simdwidthof[float_type]()

alias width = 960
alias height = 960
alias MAX_ITERS = 200

alias min_x = -2.0
alias max_x = 0.6
alias min_y = -1.5
alias max_y = 1.5


fn mandelbrot_kernel_SIMD[
    simd_width: Int
](c: ComplexSIMD[float_type, simd_width]) -> SIMD[float_type, simd_width]:
    """A vectorized implementation of the inner mandelbrot computation."""
    let cx = c.re
    let cy = c.im
    var x = SIMD[float_type, simd_width](0)
    var y = SIMD[float_type, simd_width](0)
    var y2 = SIMD[float_type, simd_width](0)
    var iters = SIMD[float_type, simd_width](0)

    var t: SIMD[DType.bool, simd_width] = True
    for i in range(MAX_ITERS):
        if not t.reduce_or():
            break
        y2 = y * y
        y = x.fma(y + y, cy)
        t = x.fma(x, y2) <= 4
        x = x.fma(x, cx - y2)
        iters = t.select(iters + 1, iters)
    return iters


fn main() raises:
    let t = Tensor[float_type](height, width)

    @parameter
    fn worker(row: Int):
        let scale_x = (max_x - min_x) / width
        let scale_y = (max_y - min_y) / height

        @parameter
        fn compute_vector[simd_width: Int](col: Int):
            """Each time we operate on a `simd_width` vector of pixels."""
            let cx = min_x + (col + iota[float_type, simd_width]()) * scale_x
            let cy = min_y + row * scale_y
            let c = ComplexSIMD[float_type, simd_width](cx, cy)
            t.data().simd_store[simd_width](
                row * width + col, mandelbrot_kernel_SIMD[simd_width](c)
            )

        # Vectorize the call to compute_vector where call gets a chunk of pixels.
        vectorize[simd_width, compute_vector](width)

    parallelize[worker](height, height)

Как видно, язык имеет Python-подобный синтаксис и предоставляет высокоуровневые примитивы для ускорения вычислений: SIMD (со встроенным методом fma - fuse multiply-add), ComplexSIMD , vectorize, parallelize и т.д.

На текущий момент времени язык находится в стадии closed source (с планом открытия исходного кода после окончания начальной стадии разработки) и разрабатывается компанией Modular, основанной Крисом Латтнером - одним из создателей LLVM, MLIR, компилятора Clang и языка Swift. Про Mojo уже есть пара статей на Хабре, с одной из которых я рекомендую ознакомиться тем, кому интересен обзор языка в целом.

О чем пойдет речь

Изучая сайт Modular, я наткнулся на статью с кричащим заголовком "The world's fastest unified matrix multiplication", в котором описывается как Modular AI Engine бьет по производительности библиотеки Intel MKL, OneDNN и Eigen в задаче GEMM (general matrix multiply - умножение матрицы на матрицу). Далее я ознакомился со статьей в их же блоге "AI’s compute fragmentation: what matrix multiplication teaches us" (TL DR: GEMM - фундамент современных нейронных сетей, SOTA-реализации - это километры сложнейшего кода на различных ассемблерах, и этот подход не масштабируется из-за привязки к платформе и трудоемкости разработки, поэтому нужно что-то менять, т.е. переходить на Mojo) и пришел к главе документации, которая описывает написание и оптимизацию GEMM на Mojo: используя встроенные средства языка, начинаем с наивной реализации, и постепенно улучшаем ее, добавляя векторизацию, параллелизм, тайлинг и т.п. В целом, судя по истории изменений этого конкретного примера из документации, его можно назвать условно оптимальным, так как достаточно много людей (в т.ч. разработчиков языка) языка приложило руку к его оптимизации. Учитывая все вышесказанное, мы просто обязаны сами погонять этот код и сравнить его производительность со стандартами индустрии!

Что и как будем сравнивать

• Мы будем сравнивать производительность слегка модифицированного примера Matmul на Mojo 0.6.0, C/C++ библиотек Intel MKL 2020.4.304, OpenBLAS 0.3.25 и Eigen 3.4.0, а также наивной реализации умножения матриц на C++.

• Сравнение будем проводить на CPU (в текущей версии Mojo 0.6.0 CPU - это единственный доступный таргет компиляции) Intel Xeon Platinum 8124M (16 ядер, инстанс c5.4xlarge AWS EC2) на ОС Ubuntu 22.04.

• Для бенчмарков возьмем размеры матриц использованные самими Modular в их публикации. Сравним мы производительность как однопоточных версий кода, так и многопоточных.

• Для компиляции бенчмарков на C++ будем использовать компилятор Clang с флагами -O3 -march=native - последний флаг означает, что наша наивная реализация автоматически становится векторизованной.

• Сравнивать мы будем GFLOPS (Giga Floating Point Operations Per Second), рассчитывая их по формуле GLOPS = 2 * M * N * K / time / 1e9, где M, N, K - размеры матриц (MxN=MxK*KxN), а time - время в секундах, затраченное на выполнение операции умножения матрицы на матрицу.

Весь использованный код с инструкциями по запуску доступен в этом репозитории.

Сравнение

Многопоточная версия

Однопоточная версия

* - Код на Mojo падает с сегфолтом на кейсе 1024x352x512, поэтому в этом случае для Mojo указано 0 GFLOPS

Выводы

Итак, что мы видим? Во-первых, в случае однопоточного выполнения векторизованный код на Mojo показывает ту же производительность, что и код на C++, векторизованный с помощью Clang. Это вполне ожидаемо, но не слишком интересно и впечатляюще: кому нужно однопоточное выполнение в эпоху, когда закон Мура больше не работает?

Куда занятнее результаты для многопоточной версии. Здесь мы видим, что:

• Intel MKL ожидаемо дает фору всем остальным вариантам (ожидаемо потому, что эта библиотека для процессоров Intel разрабатывается самими Intel).

• Eigen проигрывает вообще всем, кроме наивной реализации, что несколько странно и неожиданно.

• Mojo показывает производительность, которая сопоставима с производительностью OpenBLAS.

Последний пункт особенно важен. Почему? А потому что OpenBLAS - это монстр, состоящий из ассемблерного кода, который годами оптимизировали красноглазые гении в подвалах под каждую существующую микроархитектуру Intel, в то время как программа на Mojo - это пара сотен строк условно-платформонезависимого питоноподобного кода, что очень впечатляет.

Заключение

Учитывая абсолютно звездный состав команды Modular, и производительность, которую код на Mojo показывает на текущем этапе, можно со сдержанным оптимизмом предположить, что Mojo успешно займет планируемую нишу: быть языком сравнительно высокого уровня для реализации кросс-плафторменных производительных ML-специфичных вычислений.

Эта статья (а скорее даже заметка), как и этот набор бенчмарков, не претендует ни на научную точность, ни на то, чтобы быть исчерпывающей, по крайней мере, по следующим причинам:

• Тесты проводились на одном единственном CPU.

• Mojo имеет закрытый исходный код и он удивительно нестабилен и попросту крив: просто посмотрите на обилие лежащих на поверхности багов и ошибок, которые в нем есть на данный момент.

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

• Последнее, но от того не менее важное: возможно я где-то не дожал и/или использовал одну из библиотек не совсем правильно, поэтому производительность оказалась плачевной.

В любом случае, основной целью этой небольшой статьи (которой, она, я надеюсь, достигла) являлось рассказать сообществу еще раз о Mojo как о потенциальном будущем стандарте индустрии, которым он вполне может стать в ближайшие годы, и немного подогреть интерес русскоязычного сообщества к нему.