Основы программирования GPU

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

Основные темы, рассматриваемые в руководстве:

• Понимание модели программирования CPU–GPU.
• Работа с поддержкой GPU в Mojo через стандартную библиотеку.
• Управление устройствами и контекстами GPU с помощью DeviceContext.
• Написание и выполнение функций-ядёр (kernel) для параллельных вычислений.
• Управление памятью и передача данных между CPU и GPU.
• Организация потоков и блоков потоков для достижения оптимальной производительности.

Перед тем как приступить к программированию GPU, убедитесь, что у вас есть совместимый GPU и установлена необходимая среда разработки. А если вы новичок в программировании GPU, рекомендуется сначала прочитать раздел Введение в архитектуру GPU.

Обзор программирования GPU в Mojo

Язык Mojo, включая его стандартную библиотеку и открытую библиотеку ядер MAX, позволяет разрабатывать приложения с поддержкой GPU. См. раздел документации What are the GPU requirements? для получения списка поддерживаемых в настоящее время GPU и дополнительных требований к программному обеспечению.

Поддержка GPU в стандартной библиотеке Mojo

Пакет gpu стандартной библиотеки Mojo включает несколько подпакетов для взаимодействия с GPU, при этом пакет gpu.host предоставляет большинство часто используемых API. Однако пакет sys содержит несколько базовых функций интроспекции для определения наличия поддерживаемого GPU в системе:

• has_accelerator() — возвращает True, если в системе есть ускоритель, и False в противном случае.
• has_amd_gpu_accelerator() — возвращает True, если в системе есть GPU AMD, и False в противном случае.
• has_apple_gpu_accelerator() — возвращает True, если в системе есть GPU Apple Silicon, и False в противном случае.
• has_nvidia_gpu_accelerator() — возвращает True, если в системе есть GPU NVIDIA, и False в противном случае.

Эти функции полезны для условной компиляции или выполнения кода в зависимости от того, доступен ли поддерживаемый GPU.

detect_gpu.mojo:

from sys import has_accelerator

def main():
    @parameter
    if has_accelerator():
        print("GPU detected")
        # Enable GPU processing
    else:
        print("No GPU detected")
        # Print error or fall back to CPU-only execution

Mojo требует наличия совместимой среды разработки для GPU, чтобы компилировать функции-ядра; в противном случае возникает ошибка на этапе компиляции. В этом примере используется декоратор @parameter, чтобы вычислить функцию has_accelerator() во время компиляции и скомпилировать только соответствующую ветку оператора if. В результате, если у вас нет совместимой среды разработки для GPU, при запуске программы вы увидите следующее сообщение: No GPU detected

Модель программирования GPU

Программирование GPU следует определённому шаблону, при котором работа распределяется между CPU и GPU:

• CPU (host) управляет потоком выполнения программы и координирует операции GPU.
• GPU (device) выполняет параллельные вычисления с использованием множества потоков.
• Необходимо явно управлять обменом данных между памятью host и device.

Обычно программа для GPU выполняет следующие шаги:

1. Инициализация данных в памяти host (CPU).
2. Выделение памяти на устройстве (GPU) и передача данных из памяти host в память device.
3. Запуск функции-ядра (kernel) на GPU для обработки данных.
4. Передача результатов обратно из памяти device в память host.

Этот процесс, как правило, выполняется асинхронно, позволяя CPU выполнять другие задачи, пока GPU обрабатывает данные. В любой момент, когда CPU необходимо убедиться, что GPU завершил операцию (например, перед копированием результатов из памяти устройства), он должен явно синхронизироваться с GPU, как описано в разделе Asynchronous operation and synchronizing the CPU and GPU.

Простой пример помогает понять эту модель программирования. Мы не будем сейчас подробно разбирать конкретные API, кроме комментариев в коде, однако все типы, функции и методы будут рассмотрены более детально в следующих разделах этого документа.

scalar_add.mojo:

from math import iota
from sys import exit, has_accelerator

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

comptime num_elements = 20

fn scalar_add(
    vector: UnsafePointer[Float32, MutAnyOrigin],
    size: Int,
    scalar: Float32,
):
    """
    Функция ядра для добавления скалярного значения ко всем элементам вектора.

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

    Args:
        vector: Указатель на входной вектор.
        size: Количество элементов в векторе.
        scalar: Скаляр для добавления к вектору.

    """

    # Вычисляет глобальный индекс потока во всей сетке. Каждый поток
    # обрабатывает один элемент вектора.
    #
    # block_idx.x: индекс текущего блока потока.
    # block_dim.x: количество потоков в блоке.
    # thread_idx.x: индекс текущего потока в пределах его блока.
    idx = block_idx.x * block_dim.x + thread_idx.x

    # Проверка границ: убедитесь, что мы не обращаемся к памяти, размер которой превышает размер вектора.
    # Это имеет решающее значение, когда количество потоков не совсем соответствует размеру вектора.
    if idx < UInt(size):
        # Каждый поток добавляет скаляр к соответствующему ему векторному элементу
        # Эта операция выполняется параллельно во всех потоках графического процессора
        vector[idx] += scalar


def main():
    @parameter
    if not has_accelerator():
        print("No GPUs detected")
        exit(0)
    else:
        # Инициализируйте контекст графического процессора для устройства 0 (графическое устройство по умолчанию).
        ctx = DeviceContext()

        # Создайте буфер в памяти хоста (центрального процессора) для хранения наших входных данных
        host_buffer = ctx.enqueue_create_host_buffer[DType.float32](
            num_elements
        )

        # Дожидаемся завершения создания буфера.
        ctx.synchronize()

        # Заполняем буфер хоста последовательными номерами (0, 1, 2, ..., размер-1).
        iota(host_buffer.unsafe_ptr(), num_elements)
        print("Original host buffer:", host_buffer)

        # Создаем буфер в памяти устройства (GPU) для хранения данных для вычислений.
        device_buffer = ctx.enqueue_create_buffer[DType.float32](num_elements)

        # Копируем данные из памяти хоста в память устройства для обработки на графическом процессоре.
        ctx.enqueue_copy(src_buf=host_buffer, dst_buf=device_buffer)

        # Компилируем функцию ядра scalar_add для выполнения на графическом процессоре.
        scalar_add_kernel = ctx.compile_function[
            scalar_add, scalar_add
        ]()

        # Запустите ядро графического процессора со следующими аргументами:
        #
        # - device_buffer: Память графического процессора, содержащая наши векторные данные
        # - num_elements: количество элементов в векторе
        # - Float32(20.0): скалярное значение, добавляемое к каждому элементу
        # - grid_dim=1: используем 1 блок потоков
        # - block_dim=num_elements: используем 'num_elements' процессов для каждого блока 
        #   один процесс для каждого элемента вектора
        ctx.enqueue_function(
            scalar_add_kernel,
            device_buffer,
            num_elements,
            Float32(20.0),
            grid_dim=1,
            block_dim=num_elements,
        )

        # Копируем вычисленные результаты обратно из памяти устройства в память хоста.
        ctx.enqueue_copy(src_buf=device_buffer, dst_buf=host_buffer)

        # Ожидаем завершения всех операций с графическим процессором.
        ctx.synchronize()

        # Отображаем окончательные результаты после вычисления на графическом процессоре.
        print("Modified host buffer:", host_buffer)

Это приложение выдает следующий результат:

Original host buffer: HostBuffer([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0])
Modified host buffer: HostBuffer([20.0, 21.0, 22.0, 23.0, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0])

Доступ и управление GPU с помощью DeviceContext

Пакет gpu.host включает структуру DeviceContext, которая представляет собой логический экземпляр GPU-устройства. Она предоставляет методы для выделения памяти на устройстве, копирования данных между CPU (host) и GPU, а также для компиляции и запуска функций (также называемых ядрами, kernels) на устройстве.

Создание экземпляра DeviceContext для доступа к GPU

Mojo поддерживает системы с несколькими GPU. GPU однозначно идентифицируются целочисленными индексами, начиная с 0, который считается «устройством по умолчанию». Определить количество доступных GPU можно, вызвав статический метод DeviceContext.number_of_devices().

Конструктор DeviceContext() возвращает экземпляр для взаимодействия с указанным GPU. Он принимает два необязательных аргумента:

• device_id: целочисленный индекс конкретного GPU в системе. Значение по умолчанию — 0, что означает «GPU по умолчанию».
• api: строка, указывающая API конкретного производителя. В настоящее время поддерживаются "cuda" (NVIDIA), "hip" (AMD) и "metal" (Apple).

Если в системе нет поддерживаемого GPU — или если нет GPU, соответствующего указанным device_id или api, — конструктор выбрасывает ошибку.

Асинхронная работа и синхронизация CPU и GPU

Типичное взаимодействие CPU и GPU является асинхронным, что позволяет GPU обрабатывать задачи, пока CPU занят другой работой. Каждый DeviceContext имеет связанный поток (stream) очереди операций, которые должны выполняться на GPU. Операции внутри одного потока выполняются в том порядке, в котором они были добавлены в очередь.

Метод synchronize() блокирует выполнение текущего потока CPU до тех пор, пока все операции в связанном потоке DeviceContext не будут завершены. Чаще всего его используют, чтобы дождаться, пока результат работы ядра будет скопирован из памяти gpu в память host, прежде чем обращаться к этим данным на CPU.

Функции-ядра (kernel functions)

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

Многомерные сетки и организация потоков

Как обсуждалось в разделе Модель выполнения GPU, сетка (grid) является верхнеуровневой структурой организации потоков, выполняющих ядро на GPU. Сетка состоит из нескольких блоков потоков, которые могут быть организованы в одном, двух или трёх измерениях. Каждый блок потоков далее делится на отдельные потоки, которые также могут быть организованы в одном, двух или трёх измерениях.

Размеры сетки и блоков потоков задаются с помощью именованных аргументов grid_dim и block_dim при постановке функции-ядра в очередь на выполнение с помощью метода enqueue_function(). Например:

# Enqueue the print_threads() kernel function
ctx.enqueue_function[print_threads, print_threads](
    grid_dim=(2, 2, 1),   # 2x2x1 blocks per grid
    block_dim=(4, 4, 2),  # 4x4x2 threads per block
)

Для параметров grid_dim и block_dim размеры по осям x, y и z задаются с помощью Dim или Tuple. Если измерения y и z не указаны явно, они по умолчанию принимают значение 1 (то есть (2, 2) интерпретируется как (2, 2, 1), а (8,) — как (8, 1, 1)). Также можно передать просто значение типа Int, чтобы задать только размер по оси x (например, 64 интерпретируется как (64, 1, 1)).

Рисунок 1. Организация блоков потоков и потоков внутри сетки.

Изнутри функции-ядра вы можете получить доступ к размерам сетки и блоков потоков, а также к назначенным индексам блоков и потоков для отдельных потоков, выполняющих ядро, с помощью следующих структур:

Значение времени компиляции	Описание
grid_dim	Размеры сетки по осям x, y и z (например, grid_dim.y).
block_dim	Размеры блока потоков по осям x, y и z.
block_idx	Индекс блока в сетке по осям x, y и z.
thread_idx	Индекс потока внутри блока по осям x, y и z.
global_idx	Глобальное смещение потока по осям x, y и z. То есть: global_idx.x = block_dim.x * block_idx.x + thread_idx.x, global_idx.y = block_dim.y * block_idx.y + thread_idx.y, и global_idx.z = block_dim.z * block_idx.z + thread_idx.z.

Все эти размеры и индексы имеют тип UInt.

Ниже приведён полный пример, показывающий функцию-ядро, которая просто выводит индекс блока потоков, индекс потока и глобальный индекс для каждого выполняемого потока.

print_threads.mojo

from sys import exit, has_accelerator, has_apple_gpu_accelerator

from gpu.host import DeviceContext
from gpu import block_dim, block_idx, global_idx, grid_dim, thread_idx

fn print_threads():
    """Print thread block and thread indices."""

    print(
        "block_idx: [",
            block_idx.x,
            block_idx.y,
            block_idx.z,
        "]\tthread_idx: [",
            thread_idx.x,
            thread_idx.y,
            thread_idx.z,
        "]\tglobal_idx: [",
            global_idx.x,
            global_idx.y,
            global_idx.z,
        "]\tcalculated global_idx: [",
            block_dim.x * block_idx.x + thread_idx.x,
            block_dim.y * block_idx.y + thread_idx.y,
            block_dim.z * block_idx.z + thread_idx.z,
        "]",
    )

def main():
    @parameter
    if not has_accelerator():
        print("No compatible GPU found")
    elif has_apple_gpu_accelerator():
        print(
            "Printing from a kernel is not currently supported on Apple silicon"
            " GPUs"
        )
    else:
        # Initialize GPU context for device 0 (default GPU device).
        ctx = DeviceContext()

        ctx.enqueue_function[print_threads, print_threads](
            grid_dim=(2, 2, 1),  # 2x2x1 blocks per grid
            block_dim=(4, 4, 2),  # 4x4x2 threads per block
        )

        ctx.synchronize()
        print("Done")

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

block_idx: [ 0 1 0 ] thread_idx: [ 0 0 0 ] global_idx: [ 0 4 0 ] calculated global_idx: [ 0 4 0 ]
block_idx: [ 0 1 0 ] thread_idx: [ 1 0 0 ] global_idx: [ 1 4 0 ] calculated global_idx: [ 1 4 0 ]
...
block_idx: [ 1 1 0 ] thread_idx: [ 2 3 1 ] global_idx: [ 6 7 1 ] calculated global_idx: [ 6 7 1 ]
block_idx: [ 1 1 0 ] thread_idx: [ 3 3 1 ] global_idx: [ 7 7 1 ] calculated global_idx: [ 7 7 1 ]
Done

Вывод на печать из функции-ядра в настоящее время не поддерживается на GPU Apple Silicon.

Написание функции-ядра

Функции-ядра должны быть не выбрасывающими исключения (non-raising). Это означает, что их необходимо объявлять с помощью ключевого слова fn и не использовать ключевое слово raises. (Компилятор Mojo всегда рассматривает функцию, объявленную с помощью def, как выбрасывающую исключения, даже если в теле функции нет кода, который может вызвать ошибку.)

Значения аргументов должны иметь типы, реализующие трейт DevicePassable. Кроме того, функция-ядро не может иметь возвращаемого значения. Вместо этого результат работы ядра необходимо записывать в буфер памяти, переданный в качестве аргумента. В следующих двух разделах — Passing data between CPU and GPU и DeviceBuffer and HostBuffer — более подробно объясняется, как передавать значения в функцию-ядро и получать результаты.

Как обсуждалось в разделе GPU execution model, когда GPU выполняет ядро, он назначает блоки потоков сетки различным потоковым мультипроцессорам (SM) для выполнения. Затем SM делит блок потоков на подмножества потоков, называемые варпами (warp). Размер варпа зависит от архитектуры GPU, но большинство современных GPU используют размер варпа 32 или 64 потока.

Рисунок 2. Иерархия потоков, выполняющихся на GPU, показывающая взаимосвязь между сеткой (grid), блоками потоков, варпами и отдельными потоками, на основе HIP Programming Guide ©2023–2025 Advanced Micro Devices, Inc.

Если блок потоков содержит количество потоков, не делящееся на размер варпа без остатка, SM создаёт последний частично заполненный варп, который всё равно потребляет ресурсы как полноценный варп. Например, если блок потоков содержит 100 потоков, а размер варпа равен 32, SM создаёт:

• 3 полных варпа по 32 потока каждый (всего 96 потоков);
• 1 частичный варп с только 4 активными потоками, но при этом занимающий ресурсы, эквивалентные полному варпу (32 слота потоков).

Из-за такой модели выполнения необходимо гарантировать, что потоки в вашем ядре не пытаются обращаться к данным за пределами допустимых границ. В противном случае ядро может аварийно завершиться или выдать некорректные результаты. Например, если вы передаёте в ядро вектор из 2000 элементов и запускаете его с одномерными блоками потоков по 512 потоков в каждом, причём каждый поток отвечает за обработку одного элемента, то ядро может выполнять проверку границ следующим образом, чтобы не обрабатывать элементы за пределами массива:

from gpu import global_idx

fn process_vector(vector: UnsafePointer[Float32, MutAnyOrigin], size: Int):
    if global_idx.x < size:
        # Process vector[global_idx.x] in some way

Передача данных между CPU и GPU

Все значения, передаваемые в функцию-ядро, должны иметь типы, реализующие трейт DevicePassable. Этот трейт объявляет член времени компиляции с именем device_type, который сопоставляет тип, используемый на стороне CPU (host), с соответствующим типом, используемым на стороне GPU (device).

В качестве примера, DeviceBuffer — это представление буфера на стороне host, расположенного в глобальной памяти GPU. Однако его член device_type определён как UnsafePointer, поэтому данные, представляемые DeviceBuffer, фактически передаются в функцию-ядро как значение типа UnsafePointer. В следующем разделе, DeviceBuffer and HostBuffer, более подробно описывается, как выделять буферы памяти на host и device и как обмениваться блоками данных между ними.

В следующей таблице перечислены наиболее часто используемые типы в библиотеках Mojo, которые реализуют трейт DevicePassable.

Тип на host	Тип на device	Описание
Int	Int	Знаковое целое число
SIMD[dtype, width]	SIMD[dtype, width]	Малый вектор, реализованный на основе аппаратного векторного элемента
DeviceBuffer[dtype]	UnsafePointer[SIMD[dtype, 1]]	Буфер памяти со значениями типа dtype
LayoutTensor	LayoutTensor	Мощная абстракция для многомерных данных. Подробнее см. в разделе Using LayoutTensor

DeviceBuffer и HostBuffer

В этом разделе описывается, как использовать DeviceBuffer и HostBuffer для выделения памяти соответственно на gpu и на cpu, а также для копирования данных между памятью gpu и памятью cpu.

Создание DeviceBuffer

Тип DeviceBuffer представляет собой блок памяти gpu, связанный с конкретным DeviceContext. В частности, этот буфер расположен в глобальной памяти gpu. Таким образом, буфер доступен для всех потоков всех функций-ядёр, выполняемых данным DeviceContext.

Как обсуждалось в разделе Passing data between CPU and GPU, DeviceBuffer — это тип, используемый на стороне cpu для выделения буфера и копирования данных между cpu и gpu. Однако при передаче DeviceBuffer в функцию-ядро аргумент, получаемый функцией, имеет тип UnsafePointer. Попытка использовать тип DeviceBuffer напрямую внутри функции-ядра приводит к ошибке.

Метод DeviceContext.enqueue_create_buffer() создаёт DeviceBuffer, связанный с данным DeviceContext. Он принимает тип данных в виде параметра времени компиляции DType и размер буфера как аргумент времени выполнения. Например, чтобы создать буфер для 1024 значений типа Float32, нужно выполнить:

device_buffer = ctx.enqueue_create_buffer[DType.float32](1024)

Как следует из названия метода, он является асинхронным и ставит операцию в очередь на выполнение в связанном потоке операций DeviceContext.

Создание HostBuffer

Тип HostBuffer является аналогом DeviceBuffer, но представляет собой блок памяти на стороне cpu, связанный с конкретным DeviceContext. Он поддерживает методы для передачи данных между памятью cpu и gpu, а также базовый набор методов для доступа к элементам данных по индексу и для вывода содержимого буфера.

Метод DeviceContext.enqueue_create_host_buffer() принимает тип данных в виде параметра времени компиляции DType и размер буфера как аргумент времени выполнения и возвращает объект HostBuffer. Как и все методы DeviceContext, имя которых начинается с enqueue_, этот метод является асинхронным и возвращает управление сразу, добавляя операцию в очередь на выполнение в DeviceContext. Поэтому необходимо вызвать метод synchronize(), чтобы убедиться, что операция завершена, прежде чем записывать данные в объект HostBuffer или читать из него.

device_buffer = ctx.enqueue_create_host_buffer[DType.float32](1024)

# Синхронизируется, чтобы дождаться создания буфера, прежде чем пытаться выполнить запись в него
ctx.synchronize()

# Теперь запись в буфер безопасна
for i in range(1024):
    device_buffer[i] = Float32(i * i)

Копирование данных между памятью cpu и gpu

Метод enqueue_copy() перегружен и поддерживает копирование из памяти cpu в память gpu, из памяти gpu в память cpu, а также даже из памяти одного gpu в память другого (для систем с несколькими GPU). Обычно его используют для копирования данных, подготовленных в HostBuffer, в DeviceBuffer перед выполнением ядра, а затем для копирования данных из DeviceBuffer в HostBuffer, чтобы получить результаты выполнения ядра. Пример scalar_add.mojo в разделе GPU programming model демонстрирует этот шаблон на практике. В нём функция-ядро выполняет модификацию буфера «на месте», а затем повторно использует исходный HostBuffer для копирования результатов обратно с gpu. Однако при желании можно выделить отдельные DeviceBuffer и HostBuffer для хранения результатов, если необходимо сохранить исходные данные. Помимо копирования данных между HostBuffer и DeviceBuffer, в качестве источника или назначения копирования можно использовать UnsafePointer. Однако в этом случае UnsafePointer должен указывать на память cpu. Попытка использовать UnsafePointer, ссылающийся на память gpu, приведёт к ошибке. Например, это полезно, если у вас уже есть данные в структуре данных на стороне cpu, которая может предоставить доступ к ним через UnsafePointer. В таком случае нет необходимости сначала копировать данные из этой структуры вHostBuffer перед копированием их в DeviceBuffer. И DeviceBuffer, и HostBuffer также включают методы enqueue_copy_to() и enqueue_copy_from(). Это просто удобные обёртки, которые вызывают метод enqueue_copy() соответствующего DeviceContext. Например, следующие два вызова методов эквивалентны:

ctx.enqueue_copy(src_buf=host_buffer, dst_buf=device_buffer)

# Equivalent to:
host_buffer.enqueue_copy_to(dst=device_buffer)

Наконец, в качестве удобного средства для тестирования или прототипирования можно использовать метод DeviceBuffer.map_to_host(), чтобы создать представление содержимого буфера gpu, доступное на стороне cpu. Этот метод возвращает HostBuffer в виде контекстного менеджера, который содержит копию данных из соответствующего DeviceBuffer. Кроме того, любые изменения, которые вы внесёте в HostBuffer, автоматически копируются обратно в DeviceBuffer при выходе из блока with. Например:

ctx = DeviceContext()
length = 1024

input_device = ctx.enqueue_create_buffer[DType.float32](length)

# Initialize the input

with input_device.map_to_host() as input_host:
    for i in range(length):
        input_host[i] = Float32(i)

Однако в большинстве производственного кода использовать это не рекомендуется из-за двунаправленных копирований и необходимости синхронизации. Приведённый выше пример эквивалентен следующему:

ctx = DeviceContext()
length = 1024

input_device = ctx.enqueue_create_buffer[DType.float32](length)
input_host = ctx.enqueue_create_host_buffer[DType.float32](length)

input_device.enqueue_copy_to(input_host)
ctx.synchronize()

for i in range(length):
    input_host[i] = Float32(i)

input_host.enqueue_copy_to(input_device)
ctx.synchronize()

Освобождение памяти буферов

И DeviceBuffer, и HostBuffer подчиняются стандартным механизмам владения и жизненного цикла значений в Mojo. Компилятор Mojo анализирует программу, чтобы определить последнюю точку использования владельца или ссылки на объект, и автоматически добавляет вызов деструктора объекта. Это означает, что вам не нужно явно вызывать какой-либо метод для освобождения памяти, представляемой экземпляром DeviceBuffer или HostBuffer. Более подробно о владении значениями и управлении их жизненным циклом см. в разделах Ownership и Intro to value lifecycle Руководства Mojo, а подробное объяснение уничтожения объектов — в разделе Death of a value.

Компиляция и постановка функции-ядра в очередь на выполнение

Метод compile_function() принимает функцию-ядро в качестве параметра времени компиляции и компилирует её для связанного DeviceContext. Затем скомпилированную функцию-ядро можно поставить в очередь на выполнение, передав её методу enqueue_function(). Пример в разделе GPU programming model демонстрировал этот шаблон:

scalar_add.mojo
...
scalar_add_kernel = ctx.compile_function[scalar_add, scalar_add]()

ctx.enqueue_function(
    scalar_add_kernel,
    device_buffer,
    num_elements,
    Float32(20.0),
    grid_dim=1,
    block_dim=num_elements,
)

При использовании скомпилированной функции-ядра её выполняют вызовом enqueue_function() с следующими аргументами в указанном порядке:

• Функция-ядро, которую необходимо выполнить.
• Любые дополнительные аргументы, указанные в определении функции-ядра, в порядке их объявления.
• Размеры сетки, передаваемые через именованный аргумент grid_dim.
• Размеры блока потоков, передаваемые через именованный аргумент block_dim.

Более подробно о размерах сетки и блоков потоков см. в разделе Multidimensional grids and thread organization.

Варианты API Методы compile_function_unchecked() и enqueue_function_unchecked() устарели и будут удалены в будущих версиях, так как они не выполняют проверку типов аргументов функции-ядра на этапе компиляции. При несоответствии типов аргументов возможны ошибки во время выполнения. Методы compile_function_experimental() и enqueue_function_experimental() выполняют проверку типов на этапе компиляции, но требуют передачи функции-ядра только один раз вместо двух. После того как эти API будут проверены и признаны корректными, их переименуют в compile_function() и enqueue_function(), чтобы проверка типов стала поведением по умолчанию.

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

Если приложению требуется выполнить функцию-ядро только один раз, можно использовать перегруженную версию enqueue_function(), которая одновременно компилирует ядро и ставит его в очередь на выполнение. Таким образом, следующий вызов эквивалентен отдельным вызовам compile_function() и enqueue_function(), приведённым выше (обратите внимание, что в этом случае функция-ядро передаётся как параметр времени компиляции):

ctx.enqueue_function[scalar_add, scalar_add](
    device_buffer,
    num_elements,
    Float32(20.0),
    grid_dim=1,
    block_dim=num_elements,
)