CUDA и вычисления на GPU

В главе «Скорость выполнения программ» рассматривалось ускорение умножения матриц на CPU: перестановка циклов, компиляция кода через Numba и Cython и, в завершение, вызов NumPy, давший микросекунды вместо десятков миллисекунд, затраченных наивным циклом. В главе про GIL было установлено, что потоки в Python упираются в глобальную блокировку, а процессов, занятых вычислениями, имеет смысл держать столько, сколько ядер, то есть восемь-шестнадцать. Это предел CPU.

В той же машине установлена видеокарта. У современного GPU тысячи ядер, а пиковая производительность составляет десятки терафлопс против единиц у процессора. Если задача сводится к одной операции, повторяющейся над миллионами чисел (а в физике это почти всегда так), GPU даёт ускорение в 10–100 раз. Настоящая глава посвящена тому, как получить это ускорение из Python.

Назначение видеокарты в расчётах

GPU был создан для игр: необходимо независимо вычислить цвет каждого из миллионов пикселей 60 раз в секунду. Схема «одна формула на миллионах независимых элементов» встречается не только в пикселях:

• Моделирование частиц. Метод частиц в ячейках (PIC), молекулярная динамика, трекинг пучка в ускорителе оперируют миллионами частиц, и каждая на каждом шаге движется по уравнениям, одинаковым для всех.
• Сеточные задачи. Уравнения теплопроводности, Пуассона, Максвелла, Навье–Стокса, решаемые на разностных сетках: значение в каждом узле пересчитывается по шаблону, определяемому разностной схемой.
• Обработка данных детекторов. Реконструкция треков, фильтрация изображений, томография сводятся к конвейерам, выполняющим однотипные операции над массивами, снятыми с детектора. Триггерные фермы, установленные на LHC, давно вычисляют на GPU.
• Монте-Карло. Транспорт излучения, генерация событий: миллионы независимых историй, разыгранных по одному алгоритму.
• Нейросети. Обучение сводится к миллиардам умножений матриц; без GPU глубокое обучение не состоялось бы.

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

Чем GPU отличается от CPU

CPU оптимизирован под латентность, то есть под выполнение одной цепочки инструкций за минимальное время. Ради этого в ядре предусмотрены предсказатель переходов, внеочередное исполнение и большие кеши. GPU оптимизирован под пропускную способность, то есть под выполнение максимального объёма работы в секунду, пусть каждая операция, взятая отдельно, и потребует больше времени. Ядра простые, зато их тысячи.

SIMT и warp

Ядра GPU не полностью независимы. Нити (threads) исполняются группами по 32, и группа называется warp. Все нити warp'а в каждый момент выполняют одну и ту же инструкцию над разными данными. Модель называется SIMT (Single Instruction, Multiple Threads). От SIMD, где одна инструкция применяется к вектору данных, она отличается тем, что нить здесь является самостоятельным потоком управления и может перейти по собственной ветке; отсюда возникает дивергенция, о которой сказано ниже.

Если внутри warp'а код ветвится (if), и часть нитей перешла по одной ветке, а часть по другой, GPU выполнит обе ветки по очереди, отключая нити, не принадлежащие текущей ветке. Это называется дивергенцией warp'а. Ветвистый код на GPU не ускоряется, а деградирует; хорошей GPU-программой является та, в которой 32 соседние нити выполняют одно и то же.

Иерархия памяти

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

• Регистры, отведённые каждой нити отдельно, доступ за ~1 такт. Локальные переменные, объявленные в ядре-функции (kernel, определение ниже), размещаются здесь.
• Разделяемая память (shared memory), общая для блока нитей, около 100 КБ, приходящихся на мультипроцессор, доступ за ~10–30 тактов. Это кеш, управляемый вручную: сюда заранее загружается фрагмент данных, читаемый блоком многократно.
• Глобальная память, гигабайты, установленные на видеокарте. Видна всем нитям, полоса огромная (сотни ГБ/с), но латентность достигает сотен тактов.

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

Модель CUDA

CUDA является платформой NVIDIA для вычислений общего назначения на GPU. Её терминология используется во всех инструментах для GPU.

• Host объединяет CPU и оперативную память, а Device объединяет GPU и его память. У них разные адресные пространства; данные из одного необходимо явно копировать в другой.
• Ядром (kernel) называется функция, которую host отправляет на исполнение в device. Она исполняется одновременно тысячами нитей, и каждая нить выполняет один и тот же код, но обрабатывает свой элемент данных.
• Нити организованы в двухуровневую иерархию: сетка (grid) → блоки (blocks) → нити (threads).

Grid (сетка — весь запуск ядра)
├── Block 0 ── [нить 0][нить 1] ... [нить 255]
├── Block 1 ── [нить 0][нить 1] ... [нить 255]
├── Block 2 ── [нить 0][нить 1] ... [нить 255]
└── ...

глобальный индекс нити = blockIdx.x * blockDim.x + threadIdx.x

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

Внутри ядра нить в первую очередь определяет свой индекс. У неё есть номер блока blockIdx, размер блока blockDim и номер внутри блока threadIdx. Глобальный индекс собирается из этих трёх чисел, которые предоставляет среда исполнения.

i = blockIdx.x * blockDim.x + threadIdx.x

Например, нить номер 5 в блоке номер 3 при блоках по 256 нитей получит i = 3 * 256 + 5 = 773 и будет обрабатывать 773-й элемент. Цикл for i in range(n), привычный по CPU-коду, на GPU исчезает: вместо него запускается n нитей, и каждая выполняет одну итерацию.

Главное правило: данные должны жить на GPU

У GPU собственная память, и соединена она с host'ом шиной PCIe, обеспечивающей полосу порядка 10–30 ГБ/с. Полоса памяти самого GPU составляет около терабайта в секунду. Копирование массива host → device часто обходится дороже всех вычислений над ним. Насколько дороже, определяется количеством арифметики, приходящейся на каждый переданный байт.

Рассмотрим гигабайт данных на host'е. Пересылка через PCIe займёт около 50 мс. Поэлементная операция вида \(y = 2x\) прочитает и запишет два гигабайта по внутренней полосе GPU и завершится за 2 мс. Вычисление в двадцать пять раз дешевле пересылки, и передавать массив туда и обратно бессмысленно. Умножение матриц представляет собой противоположный случай. Гигабайт float32 — это матрица \(16384 \times 16384\), её умножение требует \(2n^3 \approx 8{,}8\) терафлоп и занимает сотни миллисекунд даже на быстрой карте, то есть на порядок дороже пересылки.

Разница состоит в том, что умножение матриц выполняет \(O(n^3)\) работы над \(O(n^2)\) данных, а поэлементная операция выполняет \(O(n)\) работы над \(O(n)\) данных. Чем выше отношение работы к данным, тем легче окупается копирование в начале расчёта.

Отсюда следует антипаттерн номер один:

плохо:  [CPU: данные] → GPU → шаг → [CPU] → GPU → шаг → [CPU] → ...
хорошо: [CPU: данные] → GPU → шаг → шаг → шаг → ... → [CPU: результат]

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

CuPy: NumPy на видеокарте

Самым коротким путём к GPU из Python является CuPy. Это реализация интерфейса NumPy поверх CUDA, предоставляющая те же ndarray, sum, dot, fft, linalg, только массивы размещаются в памяти видеокарты, а операции над ними выполняет GPU.

import cupy as cp

x = cp.arange(10**7, dtype=cp.float32)   # массив рождается сразу в памяти GPU
y = cp.sin(x) ** 2 + cp.cos(x) ** 2      # считает GPU, host только раздаёт команды
s = y.sum()                              # редукция тоже на GPU, результат — 0-мерный массив

y_host = cp.asnumpy(y)                   # явное копирование device → host (в NumPy-массив)
x_gpu = cp.asarray(y_host)               # и обратно, host → device

Разработчик, владеющий векторизованным NumPy-кодом, уже владеет программированием GPU. В большинстве скриптов достаточно заменить import numpy as np на import cupy as cp.

Проверим это на умножении матриц.

import numpy as np
import cupy as cp
from time import perf_counter

n = 4096
A = np.random.rand(n, n).astype(np.float32)
B = np.random.rand(n, n).astype(np.float32)

# --- CPU ---
t0 = perf_counter()
C = A @ B
t_cpu = perf_counter() - t0

# --- GPU ---
A_gpu = cp.asarray(A)                  # копируем данные на видеокарту один раз
B_gpu = cp.asarray(B)

C_gpu = A_gpu @ B_gpu                  # прогрев: первый вызов компилирует ядра
cp.cuda.Stream.null.synchronize()

t0 = perf_counter()
C_gpu = A_gpu @ B_gpu
cp.cuda.Stream.null.synchronize()      # дожидаемся, пока GPU реально досчитает
t_gpu = perf_counter() - t0

print(f"CPU: {t_cpu:.3f} c, GPU: {t_gpu:.4f} c, ускорение x{t_cpu / t_gpu:.0f}")

На бесплатной T4, предоставляемой в Colab, это даёт ускорение в один-два порядка относительно NumPy, притом что сам NumPy уже упирается в предел процессора.

Два замечания относительно замеров на карте:

• Запуск ядер асинхронный. Вызов A_gpu @ B_gpu ставит работу в очередь GPU и возвращается немедленно. Если измерить время без cp.cuda.Stream.null.synchronize(), получатся микросекунды — время постановки в очередь, а не вычислений. Синхронизация происходит и неявно, при копировании результата на host или вызове float(s), поэтому обычный код работает корректно; в бенчмарках синхронизация выполняется явно.
• Первый вызов уходит на прогрев. При первом обращении CuPy компилирует ядра под конкретную видеокарту, на что уходят доли секунды. Замер выполняется со второго раза (как и у Numba).

Функция cp.get_array_module(x) возвращает numpy или cupy в зависимости от того, где размещён массив x. Так пишутся библиотечные функции, одинаково работающие и на CPU, и на GPU.

Numba CUDA: собственное ядро

CuPy покрывает всё, что выражается операциями над целыми массивами. Иногда алгоритм в них не укладывается: нестандартная поэлементная логика или собственный обход данных. В главе про профилирование в такой ситуации применялся декоратор @numba.jit к функции с циклами. У Numba есть и второй бэкенд: декоратор @cuda.jit компилирует Python-функцию в CUDA-ядро.

Каноническим примером является сложение векторов.

import numpy as np
from numba import cuda

@cuda.jit
def add_kernel(a, b, out):
    i = cuda.grid(1)              # глобальный индекс нити в одномерной сетке
    if i < out.size:              # нитей запущено чуть больше, чем элементов
        out[i] = a[i] + b[i]      # каждая нить обрабатывает ровно один элемент

n = 10_000_000
a = np.random.rand(n).astype(np.float32)
b = np.random.rand(n).astype(np.float32)

# явно переносим данные на device и заводим там массив под результат
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_out = cuda.device_array_like(d_a)

# конфигурация запуска: сколько нитей в блоке и сколько блоков в сетке
threads_per_block = 256
blocks_per_grid = (n + threads_per_block - 1) // threads_per_block  # округление вверх

add_kernel[blocks_per_grid, threads_per_block](d_a, d_b, d_out)
cuda.synchronize()                # дожидаемся завершения (запуск, как всегда, асинхронный)

out = d_out.copy_to_host()        # забираем результат на CPU

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

• Тело ядра содержит код одной нити. Циклов по элементам нет: параллелизм создаётся миллионами запущенных нитей. cuda.grid(1) является сокращением для cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x, формулы глобального индекса.
• Проверка if i < out.size обязательна. Число запущенных нитей кратно размеру блока, поэтому для n = 10^7 при блоках по 256 запустится 39 063 блока, то есть 10 000 128 нитей, на 128 больше, чем элементов. Лишние нити без проверки обратились бы за границу массива.
• threads_per_block обычно выбирается в диапазоне 128–512 (максимум 1024, предпочтительно кратно 32, размеру warp'а). Точное значение подбирается замерами; для начала 256 является хорошим выбором.
• blocks_per_grid вычисляется из размера задачи: классическая формула округления вверх (n + tpb - 1) // tpb.
• Конфигурация в квадратных скобках kernel[blocks, threads](...) сообщает CUDA через Numba, сколько нитей запустить.
• Ядро ничего не возвращает, а способно только записывать в переданные ему массивы, поэтому d_out создаётся заранее.

Если передать в ядро обычные NumPy-массивы, Numba скопирует их на GPU и обратно самостоятельно, что удобно для первых проб, но в цикле по шагам это антипаттерн с копированием. Явные cuda.to_device / copy_to_host позволяют контролировать трафик через PCIe. Ядра Numba принимают и массивы CuPy.

Внутри @cuda.jit-функции доступно ограниченное подмножество Python: числа, арифметика, циклы, условия, обращение к массивам, математика из math. Списки, словари и создание массивов недоступны, как и в nopython-режиме обычной Numba, только ограничения строже.

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

@cuda.jit
def laplace_kernel(u, out):
    i, j = cuda.grid(2)                   # у нити теперь два индекса: строка и столбец
    if 0 < i < u.shape[0] - 1 and 0 < j < u.shape[1] - 1:   # только внутренние узлы
        out[i, j] = (u[i + 1, j] + u[i - 1, j] +
                     u[i, j + 1] + u[i, j - 1] - 4 * u[i, j])

threads_2d = (16, 16)                     # блок 16x16 = 256 нитей
blocks_2d = ((nx + 15) // 16, (ny + 15) // 16)   # покрываем сетку nx на ny узлов
laplace_kernel[blocks_2d, threads_2d](d_u, d_lap)

Каждая нить читает пять соседних узлов и записывает один, а GPU повторяет это для всех узлов сетки. Явная разностная схема для двумерной теплопроводности отсюда получается в один шаг: u += alpha * dt / h**2 * lap.

Пример из физики: дрейф частиц в скрещенных полях

Рассмотрим облако из пяти миллионов электронов в скрещенных полях \( \vec{E} = (E_0, 0, 0) \) и \( \vec{B} = (0, 0, B_0) \) и проинтегрируем их движение схемой Бориса, стандартной в вычислительной физике плазмы (родственная ей схема применялась в примере REDPIC). Согласно теории, помимо ларморовского вращения облако как целое дрейфует поперёк обоих полей со скоростью \( E_0 / B_0 \), не зависящей ни от заряда, ни от массы.

import cupy as cp   # замени на "import numpy as cp" — и тот же код посчитает CPU

# --- параметры задачи (СИ) ---
q, m = -1.602e-19, 9.109e-31          # заряд и масса электрона
E0, B0 = 1e3, 1e-2                    # поля: В/м и Тл
dt, n_steps = 1e-11, 2000             # шаг 10 пс, всего 20 нс (~5 ларморовских периодов)
N = 5_000_000                         # пять миллионов частиц

# --- начальное состояние: облако ~1 мм с тепловым разбросом скоростей ---
# все массивы создаются сразу в памяти GPU и живут там до конца расчёта
rng = cp.random.default_rng(42)
x = rng.standard_normal(N, dtype=cp.float32) * 1e-3    # м
y = rng.standard_normal(N, dtype=cp.float32) * 1e-3
vx = rng.standard_normal(N, dtype=cp.float32) * 1e5    # м/с
vy = rng.standard_normal(N, dtype=cp.float32) * 1e5

# --- константы схемы Бориса: обычные числа, считаются один раз на CPU ---
h = q * dt / (2 * m)                  # q·dt/2m
t = h * B0                            # tan(θ/2) поворота в магнитном поле за шаг
s = 2 * t / (1 + t * t)

for step in range(n_steps):
    vx += h * E0                      # полшага разгона электрическим полем
    px = vx + vy * t                  # поворот скорости вокруг B:
    py = vy - vx * t                  #   промежуточный вектор v'
    vx, vy = vx + py * s, vy - px * s #   точный поворот на угол ω_c·dt
    vx += h * E0                      # ещё полшага разгона
    x += vx * dt                      # сдвигаем частицы
    y += vy * dt

# --- диагностика: на host уходят только два числа, не массивы ---
v_drift = float(y.mean()) / (n_steps * dt)   # float() неявно синхронизирует GPU
energy = float(0.5 * m * (vx**2 + vy**2).mean())
print(f"скорость дрейфа: {v_drift:.3e} м/с (теория: {-E0 / B0:.3e})")

Четыре существенных момента:

• Данные создаются на GPU (cp.random) и не покидают его все 2000 шагов. На host передаются два числа диагностики, вычисленные там же; только в этот момент host дожидается GPU.
• Физика умещается в несколько строк массивных операций. Каждая строка цикла разворачивается в одно-два CUDA-ядра над пятью миллионами элементов. Сам цикл по шагам остаётся на CPU: 2000 итераций дёшевы, дорогая работа выполняется внутри.
• Тот же код работает на CPU. Замена первой строки даёт NumPy-версию, пригодную для отладки на ноутбуке без видеокарты. На GPU она выполняется в десятки раз быстрее.
• float32. Для такой задачи одинарной точности достаточно, а на игровых и облачных картах она в разы быстрее двойной (об этом сказано ниже).

Заготовку можно развивать: неоднородные поля (поле зависит от координаты частицы, но это по-прежнему поэлементные операции), потери и рождение частиц (маски), самосогласованное поле (гистограмма зарядов + Пуассон через cp.fft). Всё это остаётся в рамках CuPy.

Границы применимости GPU

Ситуации, в которых GPU бесполезен или вреден:

• Малые данные. Запуск ядра стоит ~5–10 микросекунд, а копирование через PCIe обходится ещё дороже. Если работы меньше чем на миллисекунду, накладные расходы поглотят выигрыш. Массивы из тысячи элементов NumPy сложит быстрее.
• Ветвистая логика. Из-за дивергенции warp'ов код, в котором каждый элемент обрабатывается по-своему (деревья решений, разбор текста, сложные if-каскады), параллелится плохо.
• Последовательные зависимости. Если шаг i+1 невозможно начать без результата шага i, а внутри шага работы мало, параллелить нечего. Одна частица на миллиард шагов остаётся задачей для CPU и Numba. Однако миллион независимых траекторий по миллиарду шагов является идеальной задачей для GPU; в физике часто распараллеливают по ансамблю или по параметрам, а не по времени.
• Двойная точность на игровых картах. GeForce вычисляет float64 в 32–64 раза медленнее float32, так как блоки FP64 там почти отсутствуют. Полноценный FP64 (половина от FP32) присутствует только у вычислительных карт наподобие A100/H100. Если задаче (небесная механика, длинное интегрирование, плохо обусловленная алгебра) требуется двойная точность, её необходимо измерить отдельно.
• Задачи про ожидание, а не про вычисления. Если программа ожидает сеть или диск, имеет смысл обратиться к главе про асинхронность; GPU не поможет.

Прежде чем переписывать на GPU, необходимо выполнить профилирование (см. главу про профилирование) и оценить арифметическую интенсивность, то есть число операций, приходящихся на каждый байт данных. Умножение матриц подходит отлично, а поэлементное сложение уже упирается в память, а не в вычислители.

Работа без собственной видеокарты

Бесплатный GPU предоставляется в Google Colab:

1. Создаётся новый блокнот (это обычный Jupyter в облаке).
2. В меню Среда выполнения → Сменить среду выполнения (Runtime → Change runtime type) выбирается T4 GPU и сохраняется.
3. Проверяется наличие карты.

!nvidia-smi

Команда покажет модель GPU (обычно Tesla T4 с 16 ГБ памяти), версию драйвера и загрузку. CuPy, Numba и PyTorch в Colab уже установлены, и все примеры главы запускаются там без pip install. На бесплатном тарифе сессия живёт несколько часов и при простое отключается; для учёбы и прототипов этого достаточно. Похожий бесплатный GPU предоставляют Kaggle Notebooks, а для серьёзных расчётов у института, скорее всего, есть кластер с вычислительными картами, о котором целесообразно узнать у научного руководителя.

Экосистема: не CUDA единой

Писать ядра вручную требуется редко. Вокруг CUDA сформировался слой библиотек. PyTorch и JAX предоставляют NumPy на видеокарте в сочетании с автоматическим дифференцированием и JIT-компиляцией. Физики используют их как вычислительный бэкенд, не только для нейросетей. Перенос тензора на карту умещается в одно .to("cuda"), а jax.jit и jax.vmap компилируют и векторизуют целые модели. Библиотеки RAPIDS (cuDF, cuML) переносят на GPU pandas и scikit-learn, когда узким местом является не физика, а обработка таблиц с установки. В основе всех этих библиотек лежат те же ядра, сетки и блоки, поэтому модель CUDA пригодится в любой из них. Если пределы Python окажутся достигнуты, остаётся CUDA C++ с теми же концепциями.

Итог: лестница ускорения

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

1. Профилирование. Пока line_profiler не показал, куда уходит время, оптимизация на глаз является гаданием.
2. Проверка алгоритма. Никакое аппаратное обеспечение не спасёт \( O(n^2) \) там, где существует \( O(n \log n) \).
3. Векторизация. NumPy вместо циклов обычно даёт первый и самый дешёвый порядок ускорения.
4. Компиляция. Numba или Cython — для логики, не уложившейся в массивные операции.
5. Распараллеливание на CPU. Процессы — для вычислений (GIL!), асинхронность — когда узким местом является ожидание, а не вычисления.
6. Перенос на GPU. Однотипная работа над большими массивами — CuPy, собственная поэлементная логика — Numba CUDA. Два правила: данные живут на GPU, замеры — с явной синхронизацией.

Тот же векторизованный стиль, ускоривший в начале раздела код на CPU, почти без изменений переносится на видеокарту.

Полезные ссылки

• CUDA C++ Programming Guide, первоисточник: модель исполнения, иерархия памяти, все детали.
• Документация CuPy, включая таблицу соответствия функций NumPy ↔ CuPy и советы по производительности.
• Numba for CUDA GPUs, справочник по @cuda.jit: ядра, разделяемая память, атомарные операции.
• NVIDIA Deep Learning Institute, курс «Fundamentals of Accelerated Computing with CUDA Python» про Numba и CuPy.