Проект: первое знакомство с GPU

У процессора несколько мощных ядер, и каждое может делать что-то своё. У видеокарты их тысячи, они проще и выполняют одну и ту же программу — каждое над своим кусочком данных. Если задача раскладывается на миллионы одинаковых независимых кусочков — пиксели картинки, элементы массива, клетки сетки, — видеокарта решит её быстрее.

Программу, которую выполняет видеокарта, называют ядром (kernel). В Mojo ядро — это обычная функция def, написанная на том же языке и лежащая в том же файле, что и остальная программа. Не нужно ни отдельного языка вроде CUDA C++, ни отдельного компилятора.

В этом проекте мы найдём видеокарту, запустим первое ядро, разберёмся, как данные попадают на GPU и обратно, нарисуем фрактал Мандельброта и заглянем в код, который получает видеокарта. А ещё выясним, почему видеокарта и процессор, считая по одной и той же формуле, иногда дают разные ответы.

Что понадобится

Пакет max. Средства для работы с GPU лежат в библиотеке MAX, в пакете max.gpu. Если вы проходили главу «Векторизация и параллелизм», он у вас уже стоит. Если нет, добавьте его в проект:

uv add max

Видеокарта. По официальному списку:

Производитель	Что поддерживается	Драйвер
NVIDIA	GeForce RTX 20xx и новее, T4, A100, H100, B200 и другие начиная с Turing	580 и новее
AMD	Radeon RX 6900, 7000- и 9000-й серий, часть встроенных (780M, 880M, 8060S), Instinct MI250X и новее	ROCm 6.3.3 и новее (для MI355X — 7.0)
Apple	M1–M5 на macOS Sequoia 15+	Xcode 16+

Если драйвер NVIDIA старше 580 (так бывает у облачных провайдеров) или видеокарта старше Turing, например GTX 10xx, поможет ptxas из CUDA Toolkit: путь к нему задают в переменной окружения MODULAR_NVPTX_COMPILER_PATH, например export MODULAR_NVPTX_COMPILER_PATH=/usr/local/cuda/bin/ptxas.

Под WSL драйвер NVIDIA ставится в Windows, а не внутрь Linux: WSL сам делает его доступным. Проверьте, что видеокарта видна изнутри WSL, командой nvidia-smi — она покажет модель и версию драйвера. Для Radeon под WSL у AMD есть отдельная инструкция по установке ROCm.

Шаг 1. Есть ли видеокарта

from std.sys import has_accelerator

from max.gpu.host import DeviceContext


def main() raises:
    comptime if has_accelerator():
        var ctx = DeviceContext()
        print("Видеокарта:", ctx.name())
        print("API:", ctx.api())
    else:
        print("Видеокарта не найдена, считать будем на процессоре")

Запуск на машине без видеокарты:

Результат

Видеокарта не найдена, считать будем на процессоре

С видеокартой вместо этого будет её название и программный интерфейс производителя: cuda у NVIDIA, hip у AMD, metal у Apple. Например:

С видеокартой

Видеокарта: NVIDIA GeForce RTX 3060
API: cuda

DeviceContext — это видеокарта с точки зрения программы. Через него выделяют память на GPU, копируют туда данные и запускают ядра.

Главное здесь — comptime if, а не просто if. Вопрос «есть ли видеокарта» решается при компиляции: компилятор должен знать, для какой видеокарты собирать ядра. Он находит её сам или берёт из флага --target-accelerator. Если видеокарты нет, ветка для GPU не собирается: компилятор проверяет её на ошибки в типах, но кода для неё не создаёт. В программу попадает только ветка для процессора.

Попробуйте убрать comptime if и оставить одну ветку для GPU. На машине без видеокарты программа не соберётся:

note: constraint failed: Unknown GPU architecture detected.

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

Ошибка начинается строкой error: function instantiation failed, а суть — в этой строке ниже. Даже простому DeviceContext() нужно знать архитектуру видеокарты, а компилятор её не нашёл. Поэтому без видеокарты дело не доходит до исключения при запуске: программа просто не компилируется.

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

Весь код, работающий с видеокартой, держите внутри comptime if has_accelerator():, а в ветке else сделайте что-то осмысленное без GPU.

Раз решение принимает компилятор, у него есть неочевидное следствие: программа, собранная на машине без видеокарты, никогда не будет считать на GPU, даже если потом запустить её там, где видеокарта есть. Собирайте программу на той машине, где её будут запускать, или явно указывайте архитектуру флагом --target-accelerator.

Шаг 2. Потоки, блоки и сетка

Ядро выполняет не один поток, а множество потоков одновременно. Потоки собраны в блоки (block), а блоки — в сетку (grid). При запуске ядра вы сами говорите, сколько будет блоков и сколько потоков в каждом. Каждый поток знает о себе четыре вещи:

Имя	Что это
`thread_idx`	номер потока внутри его блока
`block_idx`	номер блока в сетке
`block_dim`	размер блока — сколько в нём потоков
`grid_dim`	размер сетки — сколько в ней блоков

Сетка из двух блоков по четыре потока выглядит так:

             блок 0                блок 1
thread_idx:  0   1   2   3         0   1   2   3
номер:       0   1   2   3         4   5   6   7

Сквозной номер потока в сетке — block_idx.x * block_dim.x + thread_idx.x. По этому номеру поток выбирает, какой элемент данных ему обрабатывать. Запустим ядро, которое просто печатает эти числа:

from std.sys import has_accelerator

from max.gpu import block_dim, block_idx, thread_idx
from max.gpu.host import DeviceContext


def print_threads():
    var i = block_idx.x * block_dim.x + thread_idx.x
    print(block_idx.x, thread_idx.x, i, sep="\t")


def main() raises:
    comptime if has_accelerator():
        var ctx = DeviceContext()
        print("блок\tпоток\tномер")
        ctx.enqueue_function[print_threads](grid_dim=2, block_dim=4)
        ctx.synchronize()
        print("готово")
    else:
        print("Видеокарта не найдена: этому примеру нужен GPU")

enqueue_function[print_threads](grid_dim=2, block_dim=4) ставит ядро в очередь видеокарты: два блока по четыре потока, всего восемь. Ядро передаётся параметром в квадратных скобках — компилятор собирает его для видеокарты ещё при компиляции программы.

Приставка enqueue_ в названии не случайна: метод не ждёт, пока ядро отработает, а только ставит его в очередь и сразу возвращается. Процессор идёт дальше, видеокарта работает параллельно. Дождаться её помогает ctx.synchronize(): он ждёт, пока выполнится всё, что стоит в очереди.

С видеокартой вывод будет примерно таким:

С видеокартой

блок	поток	номер
1	0	4
1	1	5
1	2	6
1	3	7
0	0	0
0	1	1
0	2	2
0	3	3
готово

Порядок строк у вас может быть другим — и от запуска к запуску тоже. Потоки работают одновременно, и никто не обещает, что блок 0 закончит раньше блока 1. Не стройте программу на порядке выполнения потоков.

Почему блоки, а не просто потоки

Видеокарта состоит из потоковых мультипроцессоров — у NVIDIA они называются SM. Каждый блок целиком выполняется на одном мультипроцессоре. Поэтому потоки одного блока могут обмениваться данными через быструю общую память и ждать друг друга, а потоки разных блоков — нет. Блоки независимы, и видеокарта раскладывает их по мультипроцессорам как удобно.

В блоке может быть не больше 1024 потоков — это верно для всех видеокарт, которые поддерживает Mojo. Блоков в сетке может быть очень много: миллионы и больше.

Внутри мультипроцессора потоки выполняются группами — варпами (warp). Все потоки варпа выполняют одну и ту же инструкцию одновременно. В варпе 32 потока у NVIDIA, Radeon и Apple и 64 — у серверных AMD Instinct. Отсюда практическое правило: делайте размер блока кратным 32. Блок из 100 потоков займёт четыре варпа, и в последнем будут работать только 4 потока из 32.

Шаг 3. Данные нужно везти

У процессора и видеокарты, как правило, разная память. Процессор не может просто так прочитать память видеокарты, а видеокарта — память процессора. (У Apple M1–M5 память общая, но программа на Mojo работает с ней так же, через буферы.) Поэтому программа для GPU делает четыре шага:

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

Прибавим число к каждому элементу массива:

from std.math import ceildiv
from std.sys import has_accelerator

from max.gpu import block_dim, block_idx, thread_idx
from max.gpu.host import DeviceContext

comptime SIZE = 10
comptime BLOCK = 4


def add_scalar(
    data: Pointer[Float32, MutAnyOrigin], size: Int32, value: Float32
):
    var i = block_idx.x * block_dim.x + thread_idx.x
    if i < Int(size):
        data[unsafe_offset=i] += value


def main() raises:
    comptime if has_accelerator():
        var ctx = DeviceContext()

        # 1. Готовим данные в обычной памяти
        var host = ctx.enqueue_create_host_buffer[DType.float32](SIZE)
        ctx.synchronize()
        for i in range(SIZE):
            host[i] = Float32(i)
        print("было: ", host)

        # 2. Выделяем память на видеокарте и копируем туда данные
        var device = ctx.enqueue_create_buffer[DType.float32](SIZE)
        ctx.enqueue_copy(dst_buf=device, src_buf=host)

        # 3. Запускаем ядро: по потоку на элемент
        ctx.enqueue_function[add_scalar](
            device,
            Int32(SIZE),
            Float32(100),
            grid_dim=ceildiv(SIZE, BLOCK),
            block_dim=BLOCK,
        )

        # 4. Копируем результат обратно и ждём, пока всё выполнится
        ctx.enqueue_copy(dst_buf=host, src_buf=device)
        ctx.synchronize()
        print("стало:", host)
    else:
        print("Видеокарта не найдена: этому примеру нужен GPU")

С видеокартой

было:  HostBuffer([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
стало: HostBuffer([100.0, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0, 107.0, 108.0, 109.0])

Разберём по шагам.

Буфер в обычной памяти создаёт enqueue_create_host_buffer. Это тоже операция в очереди, поэтому перед записью в буфер стоит synchronize(): без него можно начать писать в память, которой ещё нет.

Буфер на видеокарте создаёт enqueue_create_buffer. Из программы на процессоре его элементы недоступны:

error: 'DeviceBuffer[.float32]' is not subscriptable, it does not implement the `__getitem__`/`__setitem__` methods

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

DeviceBuffer — это память видеокарты. Процессор не может прочитать или записать device[0] напрямую, поэтому индексации у типа нет.

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

Скопируйте данные в HostBuffer через enqueue_copy и работайте с ним. Для быстрых опытов есть map_to_host() — о нём ниже.

Копирование — enqueue_copy(dst_buf=..., src_buf=...): куда и откуда. Одним и тем же методом копируют в обе стороны.

Ядро получает не DeviceBuffer, а указатель Pointer[Float32, MutAnyOrigin] на ту же память — Mojo подменяет тип сам, когда передаёт аргумент на видеокарту. Элементы читаются и пишутся через data[unsafe_offset=i], как в главе «Указатели». Проверки границ здесь нет, поэтому условие if i < Int(size) обязательно. Мы запускаем ceildiv(10, 4) = 3 блока по 4 потока, то есть 12 потоков на 10 элементов, и потоки 10 и 11 должны ничего не делать. Без проверки они записали бы за конец буфера. В лучшем случае программа упадёт с ошибкой доступа к памяти, в худшем — молча испортит чужие данные.

Размер передаётся как Int32, а не Int. Это не прихоть:

note: constraint failed: Int and UInt do not conform to DevicePassable; use a fixed-width type such as Int32 or Int64 instead

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

Ширина Int зависит от устройства: на процессоре и на видеокарте она может различаться. Поэтому передавать Int в ядро нельзя. Передавать можно только типы, которые умеют переезжать на видеокарту (DevicePassable): числа фиксированной ширины, DeviceBuffer и некоторые другие.

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

Передавайте Int32(size) или Int64(size). Внутри ядра Int использовать можно — отсюда Int(size) в условии.

Эту ошибку машина без видеокарты не покажет. Типы в ветке для GPU компилятор проверит и там, но проверки внутри самого enqueue_function — а эта из таких — выполняются, только когда ветка действительно собирается. Увидеть такие ошибки без GPU можно, собрав программу под видеокарту явно, — это шаг 5.

Очередь. Копирование, запуск ядра и копирование обратно выполняются строго в том порядке, в каком поставлены в очередь. Поэтому хватает одного synchronize() в конце: к тому моменту, когда он вернётся, результат уже лежит в host.

Три правила для ядра

Функция, которую запускают на видеокарте, должна:

не бросать исключений — никакого raises;
ничего не возвращать — результат записывается в память, переданную аргументом;
принимать только DevicePassable-аргументы — как выше.

За первые два правила компилятор ругается одинаково:

error: no matching method in call to 'enqueue_function'

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

Ни один вариант enqueue_function не подошёл. Ниже в заметках компилятор печатает тип вашей функции — например, def add_scalar(...) raises thin -> None или ... thin -> Float32. Ищите в нём raises или тип результата, отличный от None.

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

Уберите raises и return со значением. Ошибочную ситуацию ядро может отметить, записав код ошибки в отдельный буфер.

Первые два правила компилятор проверяет всегда, даже без видеокарты. Третье — и вообще соответствие аргументов ядру — только при сборке под видеокарту, зато сообщения там подробные. Например, если передать буфер Float64 ядру, которое ждёт Float32, компилятор напишет, что аргумент DeviceBuffer[float64] нельзя передать как Pointer[SIMD[DType.float32, 1], ...].

Короткий путь для опытов

Для быстрых проб есть map_to_host(): он возвращает копию буфера видеокарты в обычной памяти, а при выходе из with копирует изменения обратно:

with device.map_to_host() as host:
    for i in range(SIZE):
        host[i] = Float32(i)

Удобно, но каждый раз это копирование туда и обратно плюс ожидание. В настоящей программе копируйте явно, как в add_scalar.mojo.

Шаг 4. Фрактал Мандельброта

Теперь задача, ради которой видеокарты и нужны. Возьмём точку c на комплексной плоскости и будем повторять z → z² + c, начиная с z = 0. Если |z| превысит 2, дальше последовательность уходит в бесконечность — точка «убежала». Точки, которые не убегают никогда, образуют множество Мандельброта. Раскрасив каждую точку по тому, как быстро она убежала, получаем знаменитую картинку.

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

Сначала функция, которая считает шаги для одной точки. Если записать z = x + iy, то z² + c = (x² − y² + cx) + i(2xy + cy):

def escape_time(cx: Float32, cy: Float32) -> Int32:
    var x: Float32 = 0
    var y: Float32 = 0
    for i in range(MAX_ITER):
        var x2 = x * x
        var y2 = y * y
        if x2 + y2 > 4:
            return Int32(i)
        y = 2 * x * y + cy
        x = x2 - y2 + cx
    return MAX_ITER

Это обычная функция без всяких пометок. Её можно вызвать и на процессоре, и из ядра на видеокарте — компилятор соберёт её для обоих устройств. Сравнение |z| > 2 записано как x² + y² > 4, чтобы не извлекать корень.

Ядро теперь двумерное: у картинки есть строки и столбцы, и сетку удобно сделать такой же. Номера потока по осям x и y дают столбец и строку пикселя:

def mandelbrot_kernel(
    iters: Pointer[Int32, MutAnyOrigin], width: Int32, height: Int32
):
    var col = global_idx.x
    var row = global_idx.y
    if col >= Int(width) or row >= Int(height):
        return
    var cx, cy = pixel_to_point(col, row, Int(width), Int(height))
    iters[unsafe_offset=row * Int(width) + col] = escape_time(cx, cy)

global_idx — сокращение для того же сквозного номера: global_idx.x равен block_idx.x * block_dim.x + thread_idx.x, global_idx.y — то же по оси y. Одномерный массив картинки хранится построчно, поэтому пиксель (row, col) лежит по индексу row * width + col.

Запуск — блоками 16 × 16 = 256 потоков:

ctx.enqueue_function[mandelbrot_kernel](
    device,
    Int32(WIDTH),
    Int32(HEIGHT),
    grid_dim=(ceildiv(WIDTH, BLOCK), ceildiv(HEIGHT, BLOCK)),
    block_dim=(BLOCK, BLOCK),
)

Для картинки 72 × 28 это сетка из 5 × 2 блоков: 80 × 32 = 2560 потоков на 2016 пикселей. Лишние потоки отсекает проверка в начале ядра.

Вот весь модуль — общий код для процессора и видеокарты:

"""Множество Мандельброта: общий код для процессора и видеокарты."""

from max.gpu import global_idx

comptime MAX_ITER = 200
"""Сколько шагов ждём, прежде чем решить, что точка не убегает."""


def escape_time(cx: Float32, cy: Float32) -> Int32:
    """Номер шага, на котором точка c = cx + i·cy убежала за радиус 2.

    Если за MAX_ITER шагов не убежала — считаем, что она в множестве,
    и возвращаем MAX_ITER.
    """
    var x: Float32 = 0
    var y: Float32 = 0
    for i in range(MAX_ITER):
        var x2 = x * x
        var y2 = y * y
        if x2 + y2 > 4:
            return Int32(i)
        y = 2 * x * y + cy
        x = x2 - y2 + cx
    return MAX_ITER


def pixel_to_point(
    col: Int, row: Int, width: Int, height: Int
) -> Tuple[Float32, Float32]:
    """Центр пикселя как точка прямоугольника [-2.2, 0.8] × [-1.2, 1.2]."""
    var step_x = Float32(3.0) / Float32(width)
    var step_y = Float32(2.4) / Float32(height)
    var cx = Float32(-2.2) + (Float32(col) + 0.5) * step_x
    var cy = Float32(-1.2) + (Float32(row) + 0.5) * step_y
    return (cx, cy)


def mandelbrot_kernel(
    iters: Pointer[Int32, MutAnyOrigin], width: Int32, height: Int32
):
    """Ядро: каждый поток видеокарты считает свой пиксель."""
    var col = global_idx.x
    var row = global_idx.y
    if col >= Int(width) or row >= Int(height):
        return
    var cx, cy = pixel_to_point(col, row, Int(width), Int(height))
    iters[unsafe_offset=row * Int(width) + col] = escape_time(cx, cy)


def mandelbrot_cpu(width: Int, height: Int) -> List[Int32]:
    """Та же картинка на процессоре — пиксель за пикселем."""
    var iters = List[Int32](length=width * height, fill=0)
    for row in range(height):
        for col in range(width):
            var cx, cy = pixel_to_point(col, row, width, height)
            iters[row * width + col] = escape_time(cx, cy)
    return iters^


def render(iters: Span[Int32, _], width: Int, height: Int):
    """Печатает картинку символами: чем дольше точка не убегает, тем гуще."""
    comptime shades = " .:-=+*#%"
    for row in range(height):
        var line = String()
        for col in range(width):
            var n = Int(iters[row * width + col])
            if n == MAX_ITER:
                line += "@"
            else:
                line += shades[byte=min(n // 2, shades.byte_length() - 1)]
        print(line)

И программа, которая выбирает, где считать:

from std.math import ceildiv
from std.sys import has_accelerator

from max.gpu.host import DeviceContext

from fractal import mandelbrot_cpu, mandelbrot_kernel, render

comptime WIDTH = 72
comptime HEIGHT = 28
comptime BLOCK = 16


def main() raises:
    comptime if has_accelerator():
        var ctx = DeviceContext()
        var device = ctx.enqueue_create_buffer[DType.int32](WIDTH * HEIGHT)
        ctx.enqueue_function[mandelbrot_kernel](
            device,
            Int32(WIDTH),
            Int32(HEIGHT),
            grid_dim=(ceildiv(WIDTH, BLOCK), ceildiv(HEIGHT, BLOCK)),
            block_dim=(BLOCK, BLOCK),
        )
        var host = ctx.enqueue_create_host_buffer[DType.int32](WIDTH * HEIGHT)
        ctx.enqueue_copy(dst_buf=host, src_buf=device)
        ctx.synchronize()
        render(host.as_span(), WIDTH, HEIGHT)
    else:
        var iters = mandelbrot_cpu(WIDTH, HEIGHT)
        render(iters, WIDTH, HEIGHT)

Результат

              ..............................:::::::::...................
            ...........................::::::::-*--::::::...............
           .........................::::::::::--=#+%=-::::::............
          .......................::::::::::::--=%%%#=-::::::::..........
         ......................:::::::::::---=+%@@@%%#--:::::::.........
        ....................::::::::::------==*%@@@@%+=------:::........
       ...................::::::::::-=*%#*+%%%@@@@@@@%%%%====%-:::......
       ................::::::::::----=+@@@@@@@@@@@@@@@@@@%%@@##:::......
      .............::::::::::::----=%%%@@@@@@@@@@@@@@@@@@@@%#=--:::.....
      .......:::::::--%=---==----==+#%@@@@@@@@@@@@@@@@@@@@@@@#+%-:::....
     ....:::::::::::--=*%#**%#%+==+%@@@@@@@@@@@@@@@@@@@@@@@@@@#=-:::....
     ..::::::::::::---+*%@@@@@@@%##@@@@@@@@@@@@@@@@@@@@@@@@@@@@%:::::...
     :::::::::----=+++%@@@@@@@@@@@%@@@@@@@@@@@@@@@@@@@@@@@@@@@%-:::::...
     ::---=---===+*%@@%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@%=--:::::...
     ::---=---===+*%@@%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@%=--:::::...
     :::::::::----=+++%@@@@@@@@@@@%@@@@@@@@@@@@@@@@@@@@@@@@@@@%-:::::...
     ..::::::::::::---+*%@@@@@@@%##@@@@@@@@@@@@@@@@@@@@@@@@@@@@%:::::...
     ....:::::::::::--=*%#**%#%+==+%@@@@@@@@@@@@@@@@@@@@@@@@@@#=-:::....
      .......:::::::--%=---==----==+#%@@@@@@@@@@@@@@@@@@@@@@@#+%-:::....
      .............::::::::::::----=%%%@@@@@@@@@@@@@@@@@@@@%#=--:::.....
       ................::::::::::----=+%@@@@@@@@@@@@@@@@@%%@@##:::......
       ...................::::::::::-=*%#*+%%%@@@@@@@%%%%====%-:::......
        ....................::::::::::------==*%@@@@%+=------:::........
         ......................:::::::::::---=+%@@@%%#--:::::::.........
          .......................::::::::::::--=%%%#=-::::::::..........
           .........................::::::::::--=#+%=-::::::............
            ...........................::::::::-*--::::::...............
              ..............................:::::::::...................

Справа — главная «кардиоида», слева от неё — круглый отросток, а дальше влево по оси тянется тонкая «антенна». Символ @ — точки, которые за 200 шагов так и не убежали.

Шаг 5. Проверяем ядро без видеокарты

Флаг --target-accelerator говорит компилятору, под какую видеокарту собирать. С ним has_accelerator() становится истинным, и компилятор собирает ветку для GPU целиком, вместе с ядрами:

mojo build --target-accelerator=sm_86 mandelbrot.mojo      # NVIDIA RTX 30xx
mojo build --target-accelerator=gfx1100 mandelbrot.mojo    # AMD Radeon RX 7900
mojo build --print-supported-accelerators                  # весь список

Любая ошибка в ядре или в его аргументах теперь всплывёт, даже без видеокарты. Запустить такую программу без GPU, конечно, не получится: DeviceContext() бросит исключение, потому что не найдёт драйвер — для NVIDIA это Failed to open library "libnvidia-ml.so.1".

Если добавить --emit asm, компилятор вместо программы запишет ассемблер процессора, а рядом с ним — по файлу .ptx на каждое ядро. Это код для видеокарт NVIDIA, в который превратилось ядро:

mojo build --target-accelerator=sm_86 --emit asm add_scalar.mojo -o add_scalar.s

Для add_scalar из шага 3 основная часть PTX такая:

mov.u32      %r2, %ctaid.x;        // block_idx.x
mov.u32      %r3, %ntid.x;         // block_dim.x
mul.wide.u32 %rd5, %r2, %r3;
mov.u32      %r4, %tid.x;          // thread_idx.x
cvt.u64.u32  %rd6, %r4;
add.s64      %rd7, %rd5, %rd6;     // i = block_idx.x * block_dim.x + thread_idx.x
setp.ge.s64  %p1, %rd7, %rd4;      // i >= size?
@%p1 bra     $L__BB0_2;            //   тогда сразу к выходу
...
ld.global.b32 %r5, [%rd1];         // прочитать data[i]
add.f32       %r6, %r1, %r5;       // прибавить value
st.global.b32 [%rd1], %r6;         // записать обратно
$L__BB0_2:
ret;

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

PTX — ещё не машинный код, а промежуточный язык NVIDIA. В инструкции конкретной видеокарты его переводит отдельный компилятор PTX. Mojo поставляется со своим, а вместо него можно подставить ptxas из CUDA Toolkit — через ту самую переменную MODULAR_NVPTX_COMPILER_PATH.

Для AMD флаг --emit asm создаёт файл .amdgcn, для Apple — .ll. Посмотреть их можно на любой машине, а вот полная сборка под Apple требует инструментов Metal из Xcode и работает только на macOS.

Шаг 6. Совпадает ли GPU с процессором

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

from std.math import ceildiv
from std.sys import has_accelerator
from std.testing import assert_equal, assert_true, TestSuite

from max.gpu.host import DeviceContext

from fractal import MAX_ITER, escape_time, mandelbrot_cpu, mandelbrot_kernel


def test_points_inside_never_escape() raises:
    assert_equal(escape_time(0, 0), MAX_ITER)
    assert_equal(escape_time(-0.5, 0), MAX_ITER)
    assert_equal(escape_time(-1, 0), MAX_ITER)


def test_points_outside_escape_on_known_step() raises:
    # 2.5 + 0i: уже после первого шага |z|² = 6.25 > 4
    assert_equal(escape_time(2.5, 0), 1)
    # 1 + 1i: z₁ = 1 + i, z₂ = 1 + 3i, |z₂|² = 10 > 4
    assert_equal(escape_time(1, 1), 2)


def test_gpu_matches_cpu() raises:
    comptime if has_accelerator():
        comptime width = 640
        comptime height = 480
        var ctx = DeviceContext()
        var device = ctx.enqueue_create_buffer[DType.int32](width * height)
        ctx.enqueue_function[mandelbrot_kernel](
            device,
            Int32(width),
            Int32(height),
            grid_dim=(ceildiv(width, 16), ceildiv(height, 16)),
            block_dim=(16, 16),
        )
        var gpu = ctx.enqueue_create_host_buffer[DType.int32](width * height)
        ctx.enqueue_copy(dst_buf=gpu, src_buf=device)
        ctx.synchronize()

        var cpu = mandelbrot_cpu(width, height)
        var differ = 0
        for i in range(width * height):
            if gpu[i] != cpu[i]:
                differ += 1
        print(
            "пикселей, где GPU и CPU разошлись:", differ, "из", width * height
        )
        # Совпадения бит в бит никто не обещает (см. главу), но расхождений
        # должно быть мало: доли процента пикселей, в основном у границы.
        assert_true(differ * 1000 < width * height, "слишком много расхождений")
    else:
        print("видеокарты нет — сравнение GPU и CPU пропущено")


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

Почему в последнем тесте не assert_equal, а доля расхождений? Потому что видеокарта и процессор, выполняя одну и ту же функцию, могут получить разные числа. Сравним, во что компилятор превратил escape_time для каждого из них.

Обе стороны пользуются FMA (fused multiply-add) — операцией a · b + c с одним округлением вместо двух. Например, 2xy + cy и на процессоре, и на видеокарте считается одной такой операцией. А вот проверку x² + y² > 4 компиляторы собрали по-разному. На процессоре — два умножения и сложение, каждое со своим округлением:

vmulss  %xmm10, %xmm10, %xmm9    // x * x
vmulss  %xmm12, %xmm12, %xmm8    // y * y
vaddss  %xmm8, %xmm9, %xmm11     // x² + y²

На видеокарте — умножение и FMA:

mul.f32     %r26, %r29, %r29        // y * y
fma.rn.f32  %r25, %r3, %r3, %r26    // x * x + y²

Ошибки нет ни там, ни там: FMA — такая же законная операция стандарта IEEE 754, как умножение и сложение. Само это различие почти ни на что не влияет: оно может сдвинуть выход точки на один шаг, только если |z|² отличается от 4 меньше чем на ошибку округления.

Важнее другое. PTX — не последнее слово: в x² − y² там стоят обычные mul.f32 и sub.f32 без явного режима округления, и компилятор PTX вправе слить их в одну FMA. Тогда x отличается от процессорного уже на первом шаге. Вдали от множества точка всё равно быстро убегает, а внутри — остаётся внутри. Но у границы точка держится на волоске, и крошечная разница за десятки шагов вырастает настолько, что меняет шаг выхода. Мы смоделировали такой случай на процессоре: на картинке 640 × 480 расходятся около 340 пикселей из 307 200, примерно 0,1 %. Поэтому тест допускает расхождения, но не больше 1 %.

Та же причина видна и без видеокарты. Внимательный глаз заметит, что картинка из шага 4 не совсем симметрична: в 8-й и 21-й строках отличается один символ. На бумаге эти строки симметричны относительно оси, но во Float32 их координаты получились -0.5571429 и 0.55714285 — по модулю они отличаются в последнем бите. Этого хватило: точка в 8-й строке не убежала за все 200 шагов, а её отражение в 21-й убежало на 165-м.

Отсюда правило: результаты вычислений с плавающей точкой на GPU сравнивайте с процессором с допуском, а не через ==. Для чисел это assert_almost_equal, для нашей картинки — доля несовпавших пикселей.

Запуск тестов:

mojo run test_fractal.mojo

Без видеокарты третий тест сообщит, что сравнение пропущено, и пройдёт. С видеокартой он напечатает, сколько пикселей из 307 200 разошлись.

Шаг 7. Когда видеокарта быстрее

Замер сравнивает три варианта: процессор в один поток, процессор на всех ядрах (через parallelize из главы «Векторизация и параллелизм») и видеокарту. Для видеокарты он отдельно показывает первый запуск, время ядра и ядро вместе с копированием результата в обычную память:

"""Замер: процессор (один поток и все ядра) против видеокарты.

Запуск: mojo build bench_mandelbrot.mojo && ./bench_mandelbrot 1920 1080
Вывод у каждой машины свой.
"""

from std.benchmark import keep
from std.math import ceildiv
from std.sys import argv, has_accelerator
from std.time import perf_counter_ns

from max.algorithm import parallelize
from max.gpu.host import DeviceContext

from fractal import (
    escape_time,
    mandelbrot_cpu,
    mandelbrot_kernel,
    pixel_to_point,
)


def mandelbrot_threads(width: Int, height: Int) -> List[Int32]:
    """Строки картинки делятся между ядрами процессора."""
    var iters = List[Int32](length=width * height, fill=0)
    var p = iters.unsafe_ptr()

    def one_row(row: Int) {imm p, imm width, imm height}:
        for col in range(width):
            var cx, cy = pixel_to_point(col, row, width, height)
            p[unsafe_offset=row * width + col] = escape_time(cx, cy)

    parallelize(one_row, height)
    return iters^


def report(name: String, seconds: Float64, pixels: Int):
    print(
        name,
        "  ",
        round(seconds * 1000, 2),
        "мс  ",
        round(Float64(pixels) / seconds / 1e6, 1),
        "Мпикс/с",
    )


def main() raises:
    var width = Int(argv()[1]) if len(argv()) > 1 else 1920
    var height = Int(argv()[2]) if len(argv()) > 2 else 1080
    var pixels = width * height
    print("картинка", width, "×", height)

    var best = Float64.MAX
    for _ in range(3):
        var t0 = perf_counter_ns()
        var iters = mandelbrot_cpu(width, height)
        best = min(best, Float64(perf_counter_ns() - t0) / 1e9)
        keep(iters[pixels - 1])
    report("CPU, один поток   ", best, pixels)

    best = Float64.MAX
    for _ in range(3):
        var t0 = perf_counter_ns()
        var iters = mandelbrot_threads(width, height)
        best = min(best, Float64(perf_counter_ns() - t0) / 1e9)
        keep(iters[pixels - 1])
    report("CPU, все ядра     ", best, pixels)

    comptime if has_accelerator():
        var ctx = DeviceContext()
        var device = ctx.enqueue_create_buffer[DType.int32](pixels)
        var host = ctx.enqueue_create_host_buffer[DType.int32](pixels)
        var grid = (ceildiv(width, 16), ceildiv(height, 16))
        ctx.synchronize()  # буферы созданы, дальше меряем только ядро

        # Первый запуск: сюда входят разовые расходы на загрузку ядра
        var t0 = perf_counter_ns()
        ctx.enqueue_function[mandelbrot_kernel](
            device,
            Int32(width),
            Int32(height),
            grid_dim=grid,
            block_dim=(16, 16),
        )
        ctx.synchronize()
        report(
            "GPU, первый запуск", Float64(perf_counter_ns() - t0) / 1e9, pixels
        )

        # Запуск ядра и ожидание, без копирования
        best = Float64.MAX
        for _ in range(5):
            t0 = perf_counter_ns()
            ctx.enqueue_function[mandelbrot_kernel](
                device,
                Int32(width),
                Int32(height),
                grid_dim=grid,
                block_dim=(16, 16),
            )
            ctx.synchronize()
            best = min(best, Float64(perf_counter_ns() - t0) / 1e9)
        report("GPU, только ядро  ", best, pixels)

        # Вычисления плюс копирование результата в обычную память
        best = Float64.MAX
        for _ in range(5):
            t0 = perf_counter_ns()
            ctx.enqueue_function[mandelbrot_kernel](
                device,
                Int32(width),
                Int32(height),
                grid_dim=grid,
                block_dim=(16, 16),
            )
            ctx.enqueue_copy(dst_buf=host, src_buf=device)
            ctx.synchronize()
            best = min(best, Float64(perf_counter_ns() - t0) / 1e9)
        keep(host[pixels - 1])
        report("GPU, ядро и копия ", best, pixels)
    else:
        print("видеокарты нет — замер GPU пропущен")

mojo build bench_mandelbrot.mojo
./bench_mandelbrot 1920 1080

На облачном сервере без видеокарты (Intel Xeon 2,1 ГГц, 2 ядра), картинка 1920 × 1080, лучшее из трёх запусков:

картинка 1920 × 1080
CPU, один поток       328.58 мс   6.3 Мпикс/с
CPU, все ядра         168.59 мс   12.3 Мпикс/с
видеокарты нет — замер GPU пропущен

На машине с видеокартой добавятся три строки GPU. Разницу между ними стоит понимать заранее, потому что она объясняет, когда видеокарта вообще выгодна:

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

Мандельброт — выигрышный случай: на каждый пиксель до 200 шагов вычислений, а обратно везти всего 4 байта. Противоположный пример — сложение двух массивов Float32: на каждые 12 байт, прошедших по шине, одно сложение. Видеокарта сложит мгновенно, но на копировании потеряет больше, чем выиграет. Да и запуск ядра с ожиданием результата сам по себе стоит времени. Поэтому на маленьких задачах видеокарта часто проигрывает: накладные расходы оказываются больше самих вычислений. Картинку 72 × 28 из шага 4 — 2016 пикселей — процессор считает на одном ядре примерно за треть миллисекунды, и выигрывать видеокарте тут почти не на чем.

Хорошая задача для GPU — много независимой работы на каждый байт, который нужно перевезти. Ещё лучше, если данные остаются на видеокарте между несколькими ядрами: тогда копирование платится один раз.

Типичные ошибки

Симптом	Причина	Что делать
`Unknown GPU architecture detected` при сборке	код для GPU вне `comptime if has_accelerator()`	держать всё про GPU внутри `comptime if`
`Int and UInt do not conform to DevicePassable`	`Int` передан в ядро	`Int32(n)` или `Int64(n)`
`no matching method in call to 'enqueue_function'`	ядро помечено `raises` или возвращает значение	убрать `raises`, писать результат в буфер
`'DeviceBuffer[...]' is not subscriptable`	обращение к памяти видеокарты с процессора	скопировать в `HostBuffer`
в результате нули или мусор	нет `synchronize()` перед чтением `HostBuffer`	дождаться очереди
падение или испорченные данные	нет проверки границ в ядре	`if i < Int(size)` в начале ядра
программа считает на CPU на машине с GPU	собрана на машине без видеокарты	собрать там, где запускаете, или с `--target-accelerator`
`Failed to open library "libnvidia-ml.so.1"`	драйвер NVIDIA не найден	под WSL — драйвер в Windows, проверить `nvidia-smi`

Что дальше

Мы прошли азы: ядра, потоки, блоки, память и очередь. Настоящие ускорения на GPU начинаются там, где потоки блока сотрудничают: общая память блока, синхронизация потоков, варповые операции. Куда идти дальше:

Get started with GPU programming — официальное руководство по сложению векторов. Там же тип TileTensor: многомерный взгляд на буфер вместо ручного row * width + col.
GPU programming fundamentals — справочник по DeviceContext, буферам и запуску ядер.
Mojo GPU Puzzles — задачи от простых к сложным, от сложения до свёрток и матричного умножения на GPU.

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

Почему проверку видеокарты пишут через `comptime if`, а не через обычный `if`?

Код для GPU компилируется под конкретную архитектуру видеокарты. Если её нет, компилятору не для чего собирать даже DeviceContext(), и сборка падает с Unknown GPU architecture detected. С comptime if компилятор, не найдя видеокарту, ветку для GPU не собирает. Обычный if проверил бы условие только при запуске, а собрать пришлось бы обе ветки.

Ядро запущено с `grid_dim=3, block_dim=4` на массиве из 10 элементов. Сколько потоков выполнится и что должны сделать лишние?

3 × 4 = 12 потоков. Потоки с номерами 10 и 11 должны ничего не делать — поэтому ядро начинается с проверки if i < size. Без неё они записали бы за конец буфера.

Почему в ядро нельзя передать `Int`?

Ширина Int зависит от устройства и на процессоре и видеокарте может отличаться. Передавать можно только типы фиксированной ширины, например Int32 или Int64. Внутри ядра Int использовать можно.

Что будет, если прочитать `HostBuffer` сразу после `enqueue_copy`, без `synchronize()`?

enqueue_copy только ставит копирование в очередь и сразу возвращается. Копия может быть ещё не готова, и программа прочитает старые данные. synchronize() ждёт, пока выполнится вся очередь.

Видеокарта и процессор посчитали одну и ту же картинку, и несколько пикселей не совпали. Ошибка в ядре?

Не обязательно. Компилятор для видеокарты сливает умножение и сложение в одну операцию FMA с одним округлением, а на процессоре те же вычисления могут идти с двумя. У границы множества разница в последнем бите меняет шаг выхода точки. Сравнивайте с допуском.

Почему сложение двух массивов на GPU может оказаться медленнее, чем на процессоре?

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

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

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