CUDA и вычисления на GPU
В главе «Скорость выполнения программ» рассматривалось ускорение умножения матриц на CPU: перестановка циклов, компиляция кода через Numba и Cython и, в завершение, вызов NumPy, давший микросекунды вместо десятков миллисекунд, затраченных наивным циклом. В главе про GIL было установлено, что потоки в обычной сборке CPython упираются в глобальную блокировку, а процессы и свободнопоточная сборка ограничены числом ядер: процессов, занятых вычислениями, имеет смысл держать столько, сколько ядер, то есть восемь-шестнадцать. Это предел CPU.
В той же машине установлена видеокарта. У современного GPU тысячи ядер, а пиковая производительность составляет десятки терафлопс против единиц у процессора. Если задача сводится к одной операции, повторяющейся над миллионами чисел (а в физике это почти всегда так), GPU даёт ускорение в 10–100 раз. Настоящая глава посвящена тому, как получить это ускорение из Python.
Видеокарты на машине, где выполнялись листинги книги, нет. Листинги CuPy и ядра Numba записаны для видеокарты; ядро в симуляторе CUDA и пример с дрейфом частиц, в котором вместо CuPy импортирован NumPy, выполнены на процессоре той же виртуальной машины под Python 3.12 с Numba 0.68 и NumPy 2.5, и их выводы приведены по этому прогону.
Слайды к главе. Устройство видеокарты, модель CUDA, ядро Numba в симуляторе и пример с дрейфом изложены также в седьмой части лекции «Оптимизация и параллельные вычисления» с демонстрациями в терминале; слайды лекции доступны на сайте книги и в PDF.
Назначение видеокарты в расчётах
GPU был создан для игр: необходимо независимо вычислить цвет каждого из миллионов пикселей 60 раз в секунду. Схема «одна формула на миллионах независимых элементов» встречается не только в пикселях:
- Моделирование частиц. Метод частиц в ячейках (PIC), молекулярная динамика, трекинг пучка в ускорителе оперируют миллионами частиц, и каждая на каждом шаге движется по уравнениям, одинаковым для всех.
- Сеточные задачи. Уравнения теплопроводности, Пуассона, Максвелла, Навье–Стокса, решаемые на разностных сетках: значение в каждом узле пересчитывается по шаблону, определяемому разностной схемой.
- Обработка данных детекторов. Реконструкция треков, фильтрация изображений, томография сводятся к конвейерам, выполняющим однотипные операции над массивами, снятыми с детектора. Триггерные фермы, установленные на LHC, давно вычисляют на GPU.
- Монте-Карло. Транспорт излучения, генерация событий: миллионы независимых историй, разыгранных по одному алгоритму.
- Нейросети. Обучение сводится к миллиардам умножений матриц; без GPU глубокое обучение не состоялось бы.
Общим для этих задач является массовый параллелизм данных. Много элементов, над каждым выполняется одинаковая и относительно простая работа, и элементы друг от друга почти не зависят. Под этот профиль GPU и предназначен.
Чем GPU отличается от CPU
CPU оптимизирован под латентность, то есть под выполнение одной цепочки инструкций за минимальное время. Ради этого в ядре предусмотрены предсказатель переходов, внеочередное исполнение и большие кеши. GPU оптимизирован под пропускную способность, то есть под выполнение максимального объёма работы в секунду, пусть каждая операция, взятая отдельно, и потребует больше времени. Ядра простые, зато их тысячи.
| CPU | GPU | |
|---|---|---|
| Ядер | 4–64 сложных | тысячи простых (у RTX 4090 их 16 384) |
| Частота | 3–5 ГГц | 1–2,5 ГГц |
| Логика ядра | предсказание переходов, внеочередное исполнение | максимально простая |
| Скрытие задержек памяти | большие кеши | тысячи нитей: пока одни ждут память, считают другие |
| Полоса памяти | ~50–100 ГБ/с (DDR5) | от сотен ГБ/с до единиц ТБ/с (GDDR6X, HBM) |
| Оптимизирован под | латентность одной задачи | суммарную пропускную способность |
| Идеальная нагрузка | ветвистая, последовательная | однотипная над большими массивами |
SIMT и warp
Ядра GPU не полностью независимы. Нити (threads) исполняются группами по 32, и группа называется warp. Все нити warp'а в каждый момент выполняют одну и ту же инструкцию над разными данными. Модель называется SIMT (Single Instruction, Multiple Threads). От SIMD, где одна инструкция применяется к вектору данных, она отличается тем, что нить здесь является самостоятельным потоком управления и может перейти по собственной ветке; отсюда возникает дивергенция, о которой сказано ниже.
Если внутри warp'а код ветвится (if), и часть нитей перешла по одной ветке, а часть по другой, GPU выполнит обе ветки по очереди, отключая нити, не принадлежащие текущей ветке. Это называется дивергенцией warp'а. Ветвление, одинаковое для всех нитей warp'а, дивергенции не вызывает, но сильно ветвящийся код, в котором соседние нити идут разными путями, ускоряется плохо; хорошей GPU-программой является та, в которой 32 соседние нити выполняют одно и то же.
Иерархия памяти
Память GPU устроена ступенями, и разница между ними составляет порядки:
- Регистры, отведённые каждой нити отдельно, доступ за ~1 такт. Локальные переменные, объявленные в ядре-функции (kernel, определение ниже), размещаются здесь.
- Разделяемая память (shared memory), общая для блока нитей, около 100 КБ, приходящихся на мультипроцессор, доступ за ~10–30 тактов. Это кеш, управляемый вручную: сюда заранее загружается фрагмент данных, читаемый блоком многократно.
- Глобальная память, гигабайты, установленные на видеокарте. Видна всем нитям, полоса огромная (сотни ГБ/с), но латентность достигает сотен тактов.
GPU скрывает латентность глобальной памяти не кешами, как CPU, а числом запущенных нитей. Пока один warp ожидает данные из памяти, мультипроцессор исполняет другие. Поэтому нитей запускается существенно больше, чем физических ядер: десятки тысяч и миллионы здесь являются нормой.
Полосу глобальной памяти определяет и порядок обращений. Обращения нитей одного warp'а к глобальной памяти объединяются (coalescing) в минимальное число 32-байтных транзакций: 32 соседним четырёхбайтным числам достаточно четырёх транзакций, а при чтении с большим шагом каждой нити требуется собственная транзакция, и полоса расходуется впустую. Поэтому данные располагают так, чтобы соседние нити обращались к соседним элементам массива.
Модель 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, полоса которой составляет 16–64 ГБ/с в зависимости от поколения (x16, с третьего по пятое). Полоса памяти самого GPU — от сотен гигабайт до нескольких терабайт в секунду: в руководстве NVIDIA по оптимизации приведены 16 ГБ/с у PCIe x16 третьего поколения против 898 ГБ/с у памяти Tesla V100. Копирование массива host → device часто обходится дороже всех вычислений над ним. Насколько дороже, определяется количеством арифметики, приходящейся на каждый переданный байт.
Оценим это на гигабайте данных, размещённом на host'е. Пересылка через PCIe третьего поколения займёт около 60 мс. Поэлементная операция вида \(y = 2x\) прочитает и запишет два гигабайта по внутренней полосе GPU и завершится примерно за 2 мс. Вычисление примерно в тридцать раз дешевле пересылки, и передавать массив туда и обратно ради одной такой операции бессмысленно. Умножение матриц представляет собой противоположный случай. Гигабайт float32 — это матрица \(16384 \times 16384\), её умножение требует \(2n^3 \approx 8{,}8 \cdot 10^{12}\) операций с плавающей точкой и при производительности около десяти терафлопс занимает около секунды, тогда как пересылка обеих матриц и результата — около 0,2 с.
Таким образом, разница определяется отношением работы к данным: умножение матриц выполняет \(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 # прогрев: первый вызов инициализирует cuBLAS
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 уже упирается в предел процессора. Умножение матриц float32 CuPy передаёт готовой библиотеке NVIDIA cuBLAS, поэтому первый вызов создаёт её контекст.
Два замечания относительно замеров на карте:
- Запуск ядер асинхронный. Вызов
A_gpu @ B_gpuставит работу в очередь GPU и возвращается немедленно. Если измерить время безcp.cuda.Stream.null.synchronize(), получатся микросекунды — время постановки в очередь, а не вычислений. Синхронизация происходит и неявно, при копировании результата на host или вызовеfloat(s), поэтому обычный код работает корректно; в бенчмарках синхронизация выполняется явно. - Первый вызов уходит на прогрев. При первом обращении CuPy компилирует поэлементные ядра и редукции под конкретную видеокарту, на что уходят доли секунды, а такие библиотеки, как cuBLAS, создают свой контекст. Замер выполняется со второго раза (как и у Numba).
Оба условия обеспечивает функция cupyx.profiler.benchmark: перед замером она выполняет прогрев, по умолчанию десять вызовов, и учитывает время работы GPU, синхронизируя потоки выполнения CUDA.
Функция cp.get_array_module(x) возвращает numpy или cupy в зависимости от того, где размещён массив x, и позволяет писать библиотечные функции, одинаково работающие и на CPU, и на GPU.
Numba CUDA: собственное ядро
CuPy покрывает всё, что выражается операциями над целыми массивами. Иногда алгоритм в них не укладывается: нестандартная поэлементная логика или собственный обход данных. В главе «Скорость выполнения программ» в такой ситуации применялся декоратор @numba.njit к функции с циклами. У Numba есть и второй бэкенд: декоратор @cuda.jit компилирует Python-функцию в CUDA-ядро.
Начиная с Numba 0.61 встроенный numba.cuda объявлен устаревшим, и поддержка CUDA перешла в пакет NVIDIA numba-cuda. В 2026 году numba-cuda переведён в режим сопровождения: до конца жизни CUDA 13 в нём исправляются только уязвимости и критические ошибки, а новые возможности NVIDIA развивает в пакете numba-cuda-mlir, совместимом с ядрами numba-cuda. Для большинства программ переход сводится к замене импорта from numba import cuda на from numba_cuda_mlir import cuda; код, использующий API расширения Numba, требует переработки, поскольку numba-cuda-mlir порождает код через MLIR, а не через LLVM IR. Пакет устанавливается командой pip install "numba-cuda-mlir[cu13]" (для CUDA 12 — [cu12]) и работает с видеокартами вычислительной способности не ниже 7.0. Примеры главы используют from numba import cuda, который предоставляют и встроенная цель Numba, и numba-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
Проверить формулу индекса и логику ядра без видеокарты позволяет симулятор CUDA, входящий в Numba. Он включается переменной окружения NUMBA_ENABLE_CUDASIM=1, заданной до импорта Numba. Ядро в симуляторе выполняет на процессоре интерпретатор Python: блоки — по одному, а нити блока — отдельными потоками. Внутри такого ядра допустимы print и отладчик pdb, а обращение за границу массива возбуждает исключение IndexError с номерами нити и блока, тогда как на видеокарте такая ошибка может остаться незамеченной. Операции уровня warp симулятор не поддерживает, типы аргументов и переменных не проверяет, поэтому в нём отлаживают ядро, которое уже компилируется для видеокарты; о скорости вычислений на видеокарте симулятор ничего не говорит.
Рассмотрим ядро сложения векторов, в котором каждая нить дополнительно записывает номер своего блока и свой номер в блоке; элементов 10, нитей в блоке 4.
# kernel.py
import numpy as np
from numba import cuda
@cuda.jit
def add_kernel(a, b, out, block, thread):
i = cuda.grid(1)
if i < out.size:
out[i] = a[i] + b[i]
block[i], thread[i] = cuda.blockIdx.x, cuda.threadIdx.x # кто обработал элемент
n, tpb = 10, 4
bpg = (n + tpb - 1) // tpb
a = np.arange(n, dtype=np.float32)
b = np.full(n, 10, dtype=np.float32)
out = np.zeros(n, dtype=np.float32)
block, thread = np.zeros(n, np.int32), np.zeros(n, np.int32)
add_kernel[bpg, tpb](a, b, out, block, thread)
print("блоков:", bpg, "| нитей:", bpg * tpb)
print("i ", *range(n))
print("блок ", *block)
print("нить ", *thread)
print("out ", *out.astype(int))
python3 -c "from numba import cuda; print(cuda.is_available())"
NUMBA_ENABLE_CUDASIM=1 python3 kernel.py
False
блоков: 3 | нитей: 12
i 0 1 2 3 4 5 6 7 8 9
блок 0 0 0 0 1 1 1 1 2 2
нить 0 1 2 3 0 1 2 3 0 1
out 10 11 12 13 14 15 16 17 18 19
На машине без видеокарты cuda.is_available() возвращает False, и запуск ядра без симулятора завершается ошибкой CudaSupportError. В симуляторе ядро отрабатывает, и таблица подтверждает формулу индекса: элемент 9 обрабатывает нить 1 блока 2, \(2 \cdot 4 + 1 = 9\), а нити 2 и 3 последнего блока не проходят проверку i < out.size и простаивают. Массивы NumPy здесь переданы в ядро напрямую: для такой проверки это допустимо, хотя в расчёте данные, как сказано выше, переносят на устройство явно.
Двумерные сетки
Сетки и блоки бывают не только одномерными. Для матриц, изображений с детекторов и разностных сеток удобны двумерные: нить получает пару индексов и обрабатывает свой узел. Ядро, вычисляющее дискретный лапласиан для уравнений теплопроводности и Пуассона:
@cuda.jit
def laplace_kernel(u, out):
j, i = cuda.grid(2) # x — столбец: соседние нити читают соседние адреса
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 = ((ny + 15) // 16, (nx + 15) // 16) # по x — ny столбцов, по y — nx строк
laplace_kernel[blocks_2d, threads_2d](d_u, d_lap)
Порядок индексов выбран не случайно: нити warp'а идут подряд по координате x, которую cuda.grid(2) возвращает первой, а в массиве NumPy подряд лежат элементы одной строки. Поэтому столбец берётся из x, и соседние нити читают соседние адреса. Каждая нить читает пять соседних узлов и записывает один, а GPU повторяет это для всех узлов сетки. Явная разностная схема для двумерной теплопроводности отсюда получается в один шаг: u += alpha * dt / h**2 * lap.
Пример из физики: дрейф частиц в скрещенных полях
Рассмотрим облако из пяти миллионов электронов в скрещенных полях \( \vec{E} = (E_0, 0, 0) \) и \( \vec{B} = (0, 0, B_0) \) и проинтегрируем их движение схемой Бориса, стандартной в вычислительной физике плазмы (родственная ей схема применялась в примере REDPIC). Согласно теории, помимо ларморовского вращения облако как целое дрейфует поперёк обоих полей со скоростью \( E_0 / B_0 \), не зависящей ни от заряда, ни от массы.
Число шагов выбрано так, чтобы расчёт длился целое число ларморовских периодов: период электрона в поле 10 мТл составляет \(T_c = 2\pi m / (|q| B_0) \approx 3{,}57\) нс, и шесть периодов при шаге 10 пс дают 2144 шага. Облако начинает движение со средней скоростью, равной нулю, поэтому его центр движется по циклоиде, и средняя скорость совпадает со скоростью дрейфа только за целое число периодов. При 2000 шагах, то есть 5,6 периода, тот же расчёт дал скорость \(-1{,}017 \cdot 10^5\) м/с, по модулю на 1,7 % больше теоретической, а при 2144 шагах расхождение составляет 0,02 %.
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, 2144 # шаг 10 пс; 21,4 нс — шесть ларморовских периодов по 3,57 нс
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
y0 = float(y.mean()) # начальный центр облака
# --- константы схемы Бориса: обычные числа, считаются один раз на 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()) - y0) / (n_steps * dt) # float() неявно синхронизирует GPU
energy = float(0.5 * m * (vx**2 + vy**2).mean())
print(f"скорость дрейфа: {v_drift:.3e} м/с (теория: {-E0 / B0:.3e})")
print(f"средняя кинетическая энергия: {energy:.3e} Дж")
скорость дрейфа: -9.998e+04 м/с (теория: -1.000e+05)
средняя кинетическая энергия: 9.105e-21 Дж
Вывод получен на процессоре: в листинге вместо CuPy импортирован NumPy, и расчёт занял 17 с. Генератор случайных чисел CuPy выдаёт другую последовательность, поэтому на видеокарте последние знаки будут иными. Скорость дрейфа отличается от теоретической на 0,02 %, а средняя кинетическая энергия совпадает с начальной энергией теплового движения \( m\sigma^2 \approx 9{,}11 \cdot 10^{-21} \) Дж, где \( \sigma = 10^5 \) м/с — разброс каждой компоненты скорости: за целое число ларморовских периодов скорость каждой частицы возвращается к начальной.
Четыре существенных момента:
- Данные создаются на GPU (
cp.random) и не покидают его все 2144 шага. На host передаются два числа диагностики, вычисленные там же; только в этот момент host дожидается GPU. - Физика умещается в несколько строк массивных операций. Каждая арифметическая операция над массивами выполняется отдельным CUDA-ядром над пятью миллионами элементов, строка цикла — одним–четырьмя ядрами. Сам цикл по шагам остаётся на CPU: 2144 итерации дёшевы, дорогая работа выполняется внутри.
- Тот же код работает на 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, необходимо выполнить профилирование (см. главу про профилирование) и оценить арифметическую интенсивность, то есть число операций, приходящихся на каждый байт данных. Умножение матриц подходит отлично, а поэлементное сложение уже упирается в память, а не в вычислители.
Работа без собственной видеокарты
Логику ядер Numba проверяет на процессоре симулятор CUDA, а код на CuPy — замена CuPy на NumPy, как в примере с дрейфом; скорость видеокарты, однако, измеряется только на видеокарте. Бесплатный GPU предоставляется в Google Colab:
- Создаётся новый блокнот (это обычный Jupyter в облаке).
- В меню Среда выполнения → Сменить среду выполнения (Runtime → Change runtime type) выбирается T4 GPU и сохраняется.
- Проверяется наличие карты.
!nvidia-smi
Команда покажет модель GPU (обычно Tesla T4 с 16 ГБ памяти), версию драйвера и загрузку. CuPy, Numba и PyTorch в Colab уже установлены, и все примеры главы запускаются там без pip install. На бесплатном тарифе сессия живёт несколько часов и при простое отключается; для учёбы и прототипов этого достаточно. Похожий бесплатный GPU предоставляют Kaggle Notebooks, а для серьёзных расчётов у института, скорее всего, есть кластер с вычислительными картами, о котором целесообразно узнать у научного руководителя. Карту на таком кластере запрашивают у планировщика (sbatch --gres=gpu:1), а расчёт в образе Apptainer получает доступ к драйверу хоста ключом --nv; работа с кластером рассмотрена в главе «Кластеры и воспроизводимое окружение для расчётов».
Экосистема: не CUDA единой
Писать ядра вручную требуется редко. Вокруг CUDA сформировался слой библиотек. PyTorch и JAX предоставляют NumPy на видеокарте в сочетании с автоматическим дифференцированием и JIT-компиляцией. Физики используют их как вычислительный бэкенд, не только для нейросетей. Перенос тензора на карту умещается в одно .to("cuda"), а jax.jit и jax.vmap компилируют и векторизуют целые модели. Библиотеки RAPIDS (cuDF, cuML) переносят на GPU pandas и scikit-learn, когда узким местом является не физика, а обработка таблиц с установки. В основе всех этих библиотек лежат те же ядра, сетки и блоки, поэтому модель CUDA пригодится в любой из них. Если пределы Python окажутся достигнуты, остаётся CUDA C++ с теми же концепциями.
Итог: лестница ускорения
Материал раздела про ускорение программ сводится к лестнице ускорения. Когда код работает медленно, достаточно пройти по ступеням сверху вниз и остановиться, как только скорость стала достаточной:
- Профилирование. Пока
line_profilerне показал, куда уходит время, оптимизация на глаз является гаданием. - Проверка алгоритма. Никакое аппаратное обеспечение не компенсирует \( O(n^2) \) там, где существует \( O(n \log n) \).
- Векторизация. NumPy вместо циклов обычно даёт первый и самый дешёвый порядок ускорения.
- Компиляция. Numba или Cython — для логики, не уложившейся в массивные операции.
- Распараллеливание на CPU. Процессы или свободнопоточная сборка Python 3.14 — для вычислений на чистом Python, потоки — для кода, освобождающего GIL (NumPy, Numba с
nogil=True), и для ожидания; асинхронность — когда одновременных ожиданий тысячи. - Перенос на GPU. Однотипная работа над большими массивами — CuPy, собственная поэлементная логика — ядра Numba CUDA (numba-cuda или numba-cuda-mlir). Два правила: данные живут на GPU, замеры — с явной синхронизацией.
- Кластер. Независимые расчёты — массивом задач планировщика, связанные — процессами MPI, окружение — в образе Apptainer; этой ступени посвящена следующая глава, «Кластеры и воспроизводимое окружение для расчётов».
Тот же векторизованный стиль, ускоривший в начале раздела код на CPU, почти без изменений переносится на видеокарту.
Полезные ссылки
- CUDA C++ Programming Guide, первоисточник: модель исполнения, иерархия памяти, все детали.
- Документация CuPy, включая таблицу соответствия функций NumPy ↔ CuPy и советы по производительности.
- Документация numba-cuda, справочник по
@cuda.jit: ядра, разделяемая память, атомарные операции, симулятор. - numba-cuda-mlir, преемник numba-cuda: установка и руководство по переходу.
- NVIDIA Deep Learning Institute, курс «Fundamentals of Accelerated Computing with CUDA Python» про Numba и CuPy.