Accelerated Computing Hub, часть 3: от NumPy к CuPy и первые ядра на Numba

Третья часть материалов NVIDIA Accelerated Computing Hub на русском. Здесь два практических текста: урок «NumPy to CuPy: ndarray Basics» из учебного курса Accelerated Python и глава 5 «CUDA Kernels with Numba» из Accelerated Python User Guide. Тексты — CC BY-NC-SA 4.0 (эта страница — на тех же условиях), код — Apache-2.0. Теория — в части 2, карта репозитория — в части 1.

Как читать вместе с курсом ChinaAI

Нужны уверенный NumPy (уроки 0.3) и понятия из части 2: ядро, поток, блок, сетка. Ядра на Numba — ступенька к уроку «Пишем ядра на Triton»: там та же идея «один экземпляр программы на кусок данных», только на уровне блоков, а не отдельных потоков.

🖥 На вашей машине

Обоим ноутбукам нужна видеокарта NVIDIA, драйвер и пакеты cupy (колесо под вашу версию CUDA, например pip install cupy-cuda12x для CUDA 12; точное имя — в инструкции по установке на cupy.dev) и numba-cuda. В оригинальном курсе всё готово в Docker-образе (tutorials/accelerated-python/brev/docker-compose.yml), его можно запустить и локально. Без видеокарты — в Google Colab с GPU-средой. Массивы в примерах — до 0,8 ГБ, на RTX 4080 всё помещается с большим запасом.

---

Урок 03. От NumPy к CuPy: основы ndarray

Оглавление

1. Создание массивов: CPU и GPU
2. Базовые операции
   • Последовательные операции и память
3. Сложные операции (линейная алгебра)
   • Код, независимый от устройства (диспетчеризация NumPy)
4. Управление устройствами
5. Упражнение — от NumPy к CuPy
   • Часть 1
   • Часть 2

---

Переключимся на высокоуровневую работу с массивами с помощью CuPy.

Что такое CuPy?

CuPy — библиотека, которая реализует знакомый API NumPy, но работает на GPU (внутри — CUDA C++).

Зачем она нужна?

• Никаких трений. Если вы знаете NumPy, вы уже знаете CuPy.
• Скорость. Она даёт ускорение на GPU для операций с массивами «из коробки».
• Простота. Часто код для CPU можно перенести на GPU, просто заменив import numpy as np на import cupy as cp.

import numpy as np
import cupy as cp
import cupyx as cpx
import matplotlib.pyplot as plt

# Вспомогательная функция, чтобы коротко выводить результаты замеров.
def print_benchmark(result):
    """Печатает результат замера по реальному времени (cpu_times) для честного сравнения."""
    avg_ms = result.cpu_times.mean() * 1000
    std_ms = result.cpu_times.std() * 1000
    print(f"{result.name}: {avg_ms:.3f} ms +/- {std_ms:.3f} ms")

1. Создание массивов: CPU и GPU

Сравним, как быстро создаётся большой трёхмерный массив (около 100 МБ) на CPU и на GPU.

Для CPU используем np.ones(), для GPU — cp.ones().

# Создание на CPU
print_benchmark(cpx.profiler.benchmark(np.ones, ((50, 500, 500),), n_repeat=10, n_warmup=1))

# Создание на GPU
print_benchmark(cpx.profiler.benchmark(cp.ones, ((50, 500, 500),), n_repeat=10, n_warmup=1))

Видно, что на GPU этот массив создаётся гораздо быстрее, чем на CPU!

(Примечание переводчика: в репозитории ноутбук сохранён без выводов, поэтому чисел мы не приводим — запустите и сравните сами. $50 \cdot 500 \cdot 500 = 12{,}5$ млн чисел типа float64 по 8 байт — ровно 100 МБ.)

О функции cupyx.profiler.benchmark()

Для замера времени операций на GPU мы используем встроенную в CuPy утилиту benchmark(). Это важно, потому что операции на GPU асинхронны: когда вы вызываете функцию CuPy, CPU ставит задачу в «список дел» GPU (поток команд, stream) и сразу идёт дальше, не дожидаясь выполнения.

Функция benchmark() берёт на себя всю сложность правильного замера времени на GPU:

• автоматически синхронизирует потоки команд GPU, чтобы замеры были точными;
• делает разогревочные прогоны, чтобы не учитывать накладные расходы «холодного старта»;
• сообщает и реальное время по часам CPU (cpu_times), и время работы ядер GPU (gpu_times). Для всех сравнений мы берём cpu_times: оно измеряет полное время от начала до конца и даёт честное сравнение кода на CPU и GPU.

Поэтому это рекомендуемый способ замерять код на CuPy: он и точный, и удобный.

Примечание переводчика: та же ловушка в PyTorch

Операции PyTorch на CUDA тоже асинхронны. Замер time.time() вокруг y = x @ x покажет только время постановки задачи в очередь. Перед остановкой секундомера вызывайте torch.cuda.synchronize() или пользуйтесь torch.utils.benchmark — подробнее в уроке «Как устроена видеокарта и куда уходит время».

2. Базовые операции

Синтаксис математических операций одинаковый. Умножим каждое значение в наших массивах на 5.

multiply_shape = (50, 500, 500)

def multiply(x):
    return x * 5

# Операция на CPU
x_cpu = np.ones(multiply_shape)
print_benchmark(cpx.profiler.benchmark(multiply, (x_cpu,), n_repeat=10, n_warmup=1))

# Операция на GPU
x_gpu = cp.ones(multiply_shape)
print_benchmark(cpx.profiler.benchmark(multiply, (x_gpu,), n_repeat=10, n_warmup=1))

GPU выполняет эту операцию заметно быстрее, а код остаётся тем же.

Последовательные операции и память

Теперь выполним несколько операций подряд.

sequential_math_shape = (50, 500, 500)

def sequential_math(x):
    x = x * 5
    x = x * x
    x = x + x
    return x

Помните: каждая из этих операций возвращает копию, а не представление (view). Это частая ловушка производительности и в NumPy, и в CuPy!

# CPU: последовательные вычисления
x_cpu = np.ones(sequential_math_shape)
print_benchmark(cpx.profiler.benchmark(sequential_math, (x_cpu,), n_repeat=10, n_warmup=1))

# GPU: последовательные вычисления
x_gpu = cp.ones(sequential_math_shape)
print_benchmark(cpx.profiler.benchmark(sequential_math, (x_gpu,), n_repeat=10, n_warmup=1))

Но даже с копиями GPU справился гораздо быстрее. Копии остаются на GPU: CuPy переносит данные с GPU на CPU, только когда это необходимо или когда об этом явно попросили.

(Примечание переводчика: каждая строка sequential_math — отдельное ядро GPU и новый массив на 100 МБ. Склеивание нескольких поэлементных операций в одно ядро называется слиянием ядер (kernel fusion); в CuPy для этого есть cupy.fuse, в PyTorch — torch.compile.)

3. Сложные операции (линейная алгебра)

GPU особенно сильны в линейной алгебре. Возьмём сингулярное разложение (SVD) — тяжёлую операцию со сложностью $O(N^3)$.

svd_shape = (3000, 1000)

# SVD на CPU
x_cpu = np.random.random(svd_shape)
print_benchmark(cpx.profiler.benchmark(np.linalg.svd, (x_cpu,), n_repeat=5, n_warmup=1))

# SVD на GPU
x_gpu = cp.random.random(svd_shape)
print_benchmark(cpx.profiler.benchmark(cp.linalg.svd, (x_gpu,), n_repeat=5, n_warmup=1))

GPU снова обходит CPU — с точно таким же API!

(Примечание переводчика: что такое SVD и зачем он нужен, объясняет урок «Ранг, SVD и низкоранговые приближения».)

Код, независимый от устройства (диспетчеризация NumPy)

Важная особенность CuPy: многие функции NumPy работают с массивами CuPy без изменения кода.

Когда вы передаёте массив CuPy на GPU (x_gpu) в функцию NumPy, которая поддерживает протокол __array_function__ (например, np.linalg.svd()), NumPy замечает вход CuPy и передаёт операцию собственной реализации CuPy, которая выполняется на GPU.

Так можно писать код в привычном синтаксисе np.*, и он будет без швов работать и на CPU, и на GPU — если в CuPy есть переопределение для этой функции.

Частый источник скрытых потерь производительности — неявные пересылки между CPU и GPU. В некоторых случаях CuPy от них защищает: например, когда NumPy пытается превратить cupy.ndarray в numpy.ndarray через протокол __array__ (скажем, np.asarray(gpu_array)), CuPy выбрасывает TypeError, а не копирует данные на хост молча.

Но в других случаях CuPy всё-таки выполняет неявные пересылки с GPU на CPU: например, при печати массива с GPU, при превращении в число Python (float, .item()) или при проверке скаляра с GPU в логическом контексте. Эти неявные пересылки мы разберём в одном из следующих ноутбуков.

# Создаём данные на GPU...
x_gpu = cp.random.random(svd_shape)

# НО вызываем обычную функцию NumPy — CuPy передаёт её выполнение на GPU!
print_benchmark(cpx.profiler.benchmark(np.linalg.svd, (x_gpu,), n_repeat=5, n_warmup=1))

4. Управление устройствами

Если у вас несколько GPU, с помощью конструкции with можно указать, на каком устройстве создавать массивы (например, GPU 0 или GPU 1).

with cp.cuda.Device(0):
   x_on_gpu0 = cp.random.random((100000, 1000))

print(f"Array is on device: {x_on_gpu0.device}")

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

(Примечание переводчика: на машине с четырьмя видеокартами попробуйте cp.cuda.Device(1), (2), (3). Массив $100\,000 \times 1000$ чисел float64 занимает 0,8 ГБ.)

---

Упражнение — от NumPy к CuPy

Часть 1

Проверим философию «замены один в один» на том же конвейере данных, что и в предыдущем ноутбуке. Один блок кода ниже делает следующее:

1. Генерирует огромный набор данных (50 миллионов элементов).
2. Обрабатывает его тяжёлой операцией (сортировкой).
3. Меняет форму и нормализует данные (broadcasting).
4. Проверяет корректность результата.

Задание (TODO):

1. Запустите ячейку ниже с xp = np (режим CPU). Запомните результат замера.
2. Поменяйте строку настройки на xp = cp (режим GPU). Запустите снова.
3. Посмотрите, как та же самая логика выполняется на GPU с CuPy значительно быстрее, сохраняя поведение NumPy.

Примечание: для замеров мы используем cupyx.profiler.benchmark() — он автоматически синхронизирует GPU.

# Шаг 1.) Настройка: выбираем, где считать
xp = np  # Переключите на 'cp', чтобы считать на GPU

print(f"Running on: {xp.__name__.upper()}")

# Шаг 2.) Генерация данных
N = 50_000_000
print(f"Generating {N:,} random elements ({N*8/2**30:.2f} GB)...")
arr = xp.random.rand(N)

# Шаг 3.) Тяжёлые вычисления (с замером)
print("Sorting data...")
# cpx.profiler.benchmark() сам синхронизирует GPU
result = cpx.profiler.benchmark(xp.sort, (arr,), n_repeat=5, n_warmup=1)
print_benchmark(result)

# Шаг 4.) Изменение формы и broadcasting
# Цель: показать, что CuPy поддерживает сложные изменения формы и правила broadcasting точно как NumPy.
# То есть логику обработки данных переписывать не нужно.

# Превращаем в матрицу из 5 столбцов
arr_new = arr.reshape((-1, 5))

# Нормализуем: делим каждую строку на её сумму с помощью broadcasting
row_sums = arr_new.sum(axis=1)
normalized_matrix = arr_new / row_sums[:, xp.newaxis]

# Шаг 5.) Проверка
# Цель: убедиться в математической корректности результата.
check_sums = xp.sum(normalized_matrix, axis=1)
xp.testing.assert_allclose(check_sums, 1.0)

print("Verification: PASSED (All rows sum to 1.0)")

Задание: когда работаете с массивами CuPy, попробуйте заменить xp.testing.assert_allclose() на np.testing.assert_allclose(). Что произойдёт и почему?

(Подсказка переводчика: вспомните абзац о протоколе __array__ и TypeError выше.)

Часть 2

Теперь создадим огромный набор данных (50 миллионов точек) — синусоиду — и посмотрим, насколько быстрее GPU сортирует его по сравнению с CPU.

Задание (TODO):

1. Сгенерируйте данные. Создайте массив NumPy (y_cpu) и массив CuPy (y_gpu), представляющие $\sin(x)$ от $0$ до $2\pi$ из 50 000 000 точек.
2. Сравните CPU и GPU. С помощью benchmark() из cupyx.profiler замерьте np.sort() и cp.sort().

# Шаг 1.) Генерация данных
N = 50_000_000
print(f"Generating {N} points...")

# TODO: создайте x_cpu с помощью np.linspace от 0 до 2*pi
# TODO: создайте y_cpu, взяв np.sin(x_cpu)

# TODO: создайте x_gpu с помощью cp.linspace от 0 до 2*pi
# TODO: создайте y_gpu, взяв cp.sin(x_gpu)


# Шаг 2.) Замер NumPy (CPU)
print("Benchmarking NumPy Sort (this may take a few seconds)...")
# TODO: используйте cpx.profiler.benchmark(function, (args,), n_repeat=5, n_warmup=1)
# Подсказка: передайте функцию `np.sort()` и аргумент `(y_cpu,)`
# Заметьте: запятая в (y_cpu,) нужна, чтобы получился кортеж!


# Шаг 3.) Замер CuPy (GPU)
print("Benchmarking CuPy Sort...")
# TODO: используйте cpx.profiler.benchmark(function, (args,), n_repeat=5, n_warmup=1)
# Подсказка: передайте функцию `cp.sort()` и аргумент `(y_gpu,)`
# Заметьте: запятая в (y_gpu,) нужна, чтобы получился кортеж!

ДОПОЛНИТЕЛЬНО: сделайте замеры для массивов разного размера и найдите размер, при котором CuPy и NumPy тратят одинаковое время. Попробуйте достать данные о времени из значения, которое возвращает cupyx.profiler.benchmark(), и вывести их по-своему. Можно даже построить график.

sizes = [5, 50, 500, 5_000, 50_000, 500_000, 5_000_000, 50_000_000]

# TODO

(Примечание переводчика: решения упражнений лежат рядом, в папке fundamentals/solutions. На маленьких массивах GPU проигрывает: запуск ядра и синхронизация стоят микросекунды, а работы почти нет. Точку пересечения на своей машине вы найдёте в дополнительном задании.)

---

Глава 5. Ядра CUDA на Numba

Оригинал: Accelerated_Python_User_Guide/notebooks/Chapter_05_CUDA_Kernels_with_Numba.ipynb.

Numba — открытый JIT-компилятор, который переводит подмножество кода на Python и NumPy в быстрый машинный код.

Numba поддерживает программирование GPU на CUDA: ограниченное подмножество Python напрямую компилируется в ядра CUDA и функции устройства по модели исполнения CUDA. Ядра, написанные на Numba, выглядят так, будто имеют прямой доступ к массивам NumPy. Массивы NumPy автоматически пересылаются между CPU и GPU.

Основы Numba

Numba поддерживает программирование GPU на CUDA: ограниченное подмножество Python напрямую компилируется в ядра CUDA и функции устройства по модели исполнения CUDA. У ядер, написанных на Numba, есть прямой доступ к массивам NumPy. Массивы NumPy автоматически пересылаются между CPU и GPU. Встроенная система компиляции Numba позволяет писать код, использующий особенности и CPU, и GPU, почти не меняя язык Python.

Установка

Прежде чем настраивать окружение для Numba, убедитесь, что выполнены условия (если вы уже ставили CuPy по инструкции, эти шаги можно пропустить):

• GPU с поддержкой CUDA (список GPU NVIDIA: https://developer.nvidia.com/cuda-gpus);
• драйверы NVIDIA с поддержкой CUDA;
• CUDA Toolkit.

Инструкция по установке: https://numba.readthedocs.io/en/stable/user/installing.html

Создаём функцию-ядро с @cuda.jit

В Numba декоратор @jit помечает функцию, которую должен оптимизировать JIT-компилятор Numba. Для GPU используется вариант @cuda.jit: он помечает функции-ядра, которые нужно оптимизировать для одновременного выполнения множеством потоков на GPU.

from numba import cuda
from numba import config as numba_config
numba_config.CUDA_ENABLE_PYNVJITLINK = True

@cuda.jit
def foo(input_array, output_array):
    # здесь будет тело ядра
    ...

Тем, кто пользовался Numba на CPU, это покажется очень знакомым.

(Примечание переводчика: в оригинале в этом фрагменте нет тела функции и закрывающих кавычек блока кода; мы добавили ..., чтобы код был синтаксически верным.)

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

Перед запуском ядра нужно указать число блоков и число потоков в блоке. Это задаёт форму сетки исполнения.

@cuda.jit
def foo(input_array, output_array):
    # Номер потока в одномерном блоке
    thread_id = cuda.threadIdx.x
    # Номер блока в одномерной сетке
    block_id = cuda.blockIdx.x
    # Ширина блока, то есть число потоков в блоке
    block_width = cuda.blockDim.x
    # Вычисляем сквозной индекс в массиве
    i = thread_id + block_id * block_width
    if i < input_array.size:  # Проверяем границы массива
        output_array[i] = input_array[i]

Чтобы вызвать функцию foo(), нужно указать размер блока и сетки.

input = np.asarray(range(10))
output = np.zeros(len(input))


block_threads = 32
grid_blocks = (input.size + (block_threads - 1)) // block_threads

foo[grid_blocks, block_threads](input, output)

flowchart LR
subgraph B0["Блок 0 (blockIdx.x = 0)"]
  T0["поток 0 → i = 0"]
  T1["поток 1 → i = 1"]
  T9["поток 9 → i = 9"]
  T10["поток 10 → i = 10:<br/>i ≥ size, ничего не делает"]
end

Сквозной индекс i = threadIdx.x + blockIdx.x · blockDim.x: для 10 элементов хватает одного блока из 32 потоков, лишние 22 потока отсекает проверка границ (схема переводчика)

Для простых примеров удобно управлять потоками, блоками и сетками функцией cuda.grid(). Весь скрипт можно переписать так:

import numpy as np
from numba import cuda
from numba import config as numba_config
numba_config.CUDA_ENABLE_PYNVJITLINK = True

input = np.asarray(range(10))
output = np.zeros(len(input))

@cuda.jit
def foo(input_array, output_array):
    i = cuda.grid(1)
    output_array[i] = input_array[i]

foo[1, len(input)](input, output)

output

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

(Примечание переводчика: во второй версии нет проверки if i < input_array.size — она не нужна только потому, что потоков запущено ровно столько, сколько элементов: foo[1, len(input)]. Для больших массивов вернитесь к первой версии с проверкой границ.)

Сколько потоков и блоков задавать

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

Практические правила для числа потоков в блоке:

• Оптимальный размер блока обычно кратен 32 (размеру варпа).
• Чтобы найти оптимальное значение, нужны профилирование и замеры.

С чего начать:

• калькулятор загрузки (Occupancy Calculator) в NSight: https://docs.nvidia.com/nsight-compute/NsightCompute/index.html#occupancy-calculator
• многие источники советуют начинать подбор с числа от 128 до 256.

Размеры блока и сетки влияют на производительность CUDA. Более крупные блоки могут лучше использовать разделяемую память и снизить накладные расходы на запуск множества мелких блоков. Но слишком большие блоки могут уменьшить число блоков, которые выполняются одновременно, и GPU будет загружен не полностью. Нужно найти баланс, чтобы использовать GPU в полную силу.

Numba вместе с CuPy

Класс cupy.ndarray реализует __cuda_array_interface__ — интерфейс обмена массивами CUDA, совместимый с Numba версии 0.39.0 и новее (подробнее — в описании CUDA Array Interface в документации Numba: https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html). Это значит, что массивы CuPy можно передавать в ядра, скомпилированные Numba.

В этом примере мы используем массивы cupy вместо numpy:

import cupy
from numba import cuda
from numba import config as numba_config
numba_config.CUDA_ENABLE_PYNVJITLINK = True

@cuda.jit
def add(x_array, y_array, output_array):
        start = cuda.grid(1)
        stride = cuda.gridsize(1)
        for i in range(start, x_array.shape[0], stride):
                output_array[i] = x_array[i] + y_array[i]

a = cupy.arange(10)
b = a * 2
out = cupy.zeros_like(a)

add[1, 32](a, b, out)

print(out)  # => [ 0  3  6  9 12 15 18 21 24 27]

(Примечание переводчика: цикл с шагом stride = cuda.gridsize(1) — приём grid-stride loop. Каждый поток обрабатывает не один элемент, а элементы start, start + stride, start + 2·stride и так далее. Так одно и то же ядро работает с массивом любого размера при фиксированной сетке. Массивы CuPy уже лежат на GPU, поэтому пересылок CPU↔GPU здесь нет — в отличие от примеров с NumPy выше.)

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

• Numba для GPU с CUDA: https://numba.readthedocs.io/en/stable/cuda/index.html
• Руководство CuPy по совместимости с другими библиотеками (включая Numba): https://docs.cupy.dev/en/stable/user_guide/interoperability.html
• Репозиторий Numba на GitHub: https://github.com/numba/numba

Примеры

В конце оригинального ноутбука те же три программы собраны в запускаемые ячейки: «Определение и запуск функции-ядра», «Упрощённое ядро с grid()» и «Numba вместе с CuPy». Их код совпадает с приведённым выше.

---

Что сделать руками

Упражнения от переводчика, в оригинале их нет.

1. Выполните урок 03 с xp = np и xp = cp и запишите время сортировки 50 млн чисел. Во сколько раз быстрее GPU?
2. Напишите на Numba ядро relu(x, out): out[i] = max(x[i], 0) с grid-stride loop. Проверьте результат против cupy.maximum(x, 0) на массиве из 10 млн элементов.
3. Запустите ядро из пункта 2 с размером блока 32, 128, 256 и 1024 потока. Как меняется время? Сравните с правилом «начинать со 128–256».
4. Перепишите то же ядро на Triton (урок «Пишем ядра на Triton») и сравните код: где в Triton «потоки», а где «блоки»?

Словарик

EN	RU
drop-in replacement	замена «один в один»
asynchronous, stream	асинхронный, поток команд (stream)
warm-up	разогрев, разогревочный прогон
wall-clock time	реальное время (по часам)
copy / view	копия / представление
dispatch	диспетчеризация (передача вызова нужной реализации)
implicit transfer	неявная пересылка
JIT compiler	JIT-компилятор (компиляция «на лету»)
kernel launch	запуск ядра
bounds check	проверка границ
grid-stride loop	цикл с шагом в размер сетки
occupancy	загрузка (заполненность) SM
kernel fusion	слияние ядер