Давайте рассмотрим реализацию конвеевской игры «Жизнь» при помощи графической карты. Я хочу поэкспериментировать с разными библиотеками и методиками, чтобы понять, как обеспечить наилучшую производительность. Начнём мы с простого и постепенно будем повышать сложность.
Игра «Жизнь» — это простой клеточный автомат, поэтому она должна хорошо поддаваться GPU-ускорению. Правила просты: каждая ячейка в двухмерной сетке или жива, или мертва. На каждом шаге мы подсчитываем живых соседей ячейки (включая диагонали). Если ячейка жива, она остаётся живой, если живы два или три её соседа. В противном случае она умирает. Если клетка мертва, она оживает, если живы ровно три соседа. Из этих простых правил возникает потрясающий объём сложности, о котором написаны подробные статьи.
Для простоты я буду рассматривать только сети N×N и пропущу вычисления на краях. Всё будет работать на Nvidia A40, а бенчмарк производительности я буду проводить при N=216. Пока мы будем хранить каждую ячейку в виде 1 байта, поэтому весь массив займёт 4 ГБ.
Весь код выложен в репозитории GitHub.
Теоретический предел
Пропускная способность памяти A40 составляет 696 ГБ/с. Это скорость передачи данных из DRAM в ядра, выполняющие вычисления. Для вычисления ячейки GPU нужно загрузить как минимум один байт, а после завершения сохранить один байт. То есть, если учитывать только передачу через память, нам понадобится как минимум 4 ГБ * 2 / 696 ГБ/с = 11,5 мс.
Так как правила игры относительно просты, мы редко будем исчерпывать вычислительные мощности GPU, которые обычно тоже становятся ограничивающим фактором. То есть 11,5 мс — это та цель, к которой нам нужно максимально приблизиться.
На Pytorch
Pytorch — это популярная Python-библиотека для машинного обучения, имеющая полезный набор операций. В Torch есть встроенная поддержка свёрток, но они предназначены для чисел с плавающей запятой и содержат ненужные умножения. Вместо них я использую разделяемый box blur для вычисления количества живых клеток в каждом прямоугольнике 3×3. Дальше простая логика будет реализовывать правила игры.
def gol_torch_sum(x: torch.Tensor) -> torch.Tensor: y = x[2:] + x[1:-1] + x[:-2] z = y[:, 2:] + y[:, 1:-1] + y[:, :-2] z = torch.nn.functional.pad(z, (1, 1, 1, 1), value=0) return ((x == 1) & (z == 4)) | (z == 3).to(torch.int8)
В своём бенчмарке я получил 223 мс. Torch имеет большой оверхед, потому что каждая операция передаёт GPU отдельную операцию, из-за чего происходит множество необязательных передач данных.
Можно использовать torch.compile для комбинирования и оптимизации максимально большого количества операций, что даст нам намного более приемлемые 38,1 мс, или 30% от максимальной производительности.
На CUDA
CUDA — это традиционный способ максимизации вычислительных мощностей GPU. Это программный интерфейс, проприетарный для GPU Nvidia. Разработчики пишут на C/C++ ядра, которые компилируются в программы, выполняемые непосредственно в графической карте. CUDA обеспечивает высокую степень контроля за вычислениями, но для начала мы напишем простое ядро. Оно похоже на функцию C, вычисляющую значение одной ячейки. CUDA будет параллельно выполнять этот код одновременно для миллионов ячеек.
__global__ void gol_kernel_i8(const int8_t* __restrict__ x_ptr, int8_t* __restrict__ out_ptr, int64_t rowstride, int64_t n) { int64_t x = blockIdx.x * blockDim.x + threadIdx.x; int64_t y = blockIdx.y * blockDim.y + threadIdx.y; if (x >= n - 2 || y >= n - 2) return; int8_t r00 = x_ptr[y * rowstride + x + 0 * rowstride + 0]; int8_t r01 = x_ptr[y * rowstride + x + 0 * rowstride + 1]; int8_t r02 = x_ptr[y * rowstride + x + 0 * rowstride + 2]; int8_t r10 = x_ptr[y * rowstride + x + 1 * rowstride + 0]; int8_t r11 = x_ptr[y * rowstride + x + 1 * rowstride + 1]; int8_t r12 = x_ptr[y * rowstride + x + 1 * rowstride + 2]; int8_t r20 = x_ptr[y * rowstride + x + 2 * rowstride + 0]; int8_t r21 = x_ptr[y * rowstride + x + 2 * rowstride + 1]; int8_t r22 = x_ptr[y * rowstride + x + 2 * rowstride + 2]; int8_t sum = r00 + r01 + r02 + r10 + r12 + r20 + r21 + r22; int8_t result = (r11 > 0) ? ((sum == 2) || (sum == 3) ? 1 : 0) : (sum == 3 ? 1 : 0); out_ptr[(y + 1) * rowstride + (x + 1)] = result; } void gol(torch::Tensor x, torch::Tensor out, int block_size_row, int block_size_col) { const long n = x.size(0); const int row_blocks = (n - 2 + block_size_row - 1) / block_size_row; const int col_blocks = (n - 2 + block_size_col - 1) / block_size_col; auto stream = at::cuda::getCurrentCUDAStream(); dim3 grid(col_blocks, row_blocks); dim3 block(block_size_col, block_size_row); gol_kernel_i8<<<grid, block, 0, stream>>>( x.data_ptr<int8_t>(), out.data_ptr<int8_t>(), x.stride(0), n); TORCH_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed"); }
В этом коде есть настраиваемые параметры block_size_row и block_size_col. Ядро разбивает наши ячейки N×N на сетку блоков потоков размером block_size_row×block_size_col. Исполнение каждого блока назначается одному из 84 потоковых мультипроцессоров (streaming multiprocessor, SM). Все вычисления в блоке потоков выполняются в одном SM, а потому имеют общий кэш L1.
Кэширование очень важно для производительности нашего ядра. Если бы не было кэширования, то показанное выше ядро наивно бы считывало 9 ячеек и записывало 1. И такой объём передаваемых данных занимал бы больше 55 мс. Однако благодаря множественным слоям внутреннего кэширования обычно достаточно считать ячейки только один раз, сохранить их в кэш, а затем использовать для всех 9 соседних ячеек, которым нужно это значение.
Зависимости между размером блоков и кэшированием — сложная тема. Если блоки квадратные, то мы минимизируем периметр, то есть ячейки, которые нужно загружать в различные блоки потоков. С другой стороны, столбцы выстроены в памяти последовательно и для этого есть оптимизации, поэтому, возможно, стоит предпочесть широкие прямоугольники. Из-за внутренних ограничений GPU площадь блока не может быть больше 1024 и должна быть кратной 32, чтобы потоки не тратились впустую. Если блок слишком большой, он занимает множество регистров и управлять большим их количеством конкурентно становится сложно, что снижает заполненность.
То есть рассуждать об этом довольно сложно. Обычно просто проводят бенчмарки и выбирают наилучшее сочетание. В моём случае наилучшие результаты были у блока размером 1×128, это обычный размер для ядер с ограниченной пропускной способностью.
Таким образом мы получаем приличные 26 мс, или 44% от пиковой производительности.
На Triton
Как мы видели выше, torch.compile обеспечил рост производительности в 5 раз. torch.compile сканирует программу на torch и преобразует её в ядро Triton. Triton — это специализированный язык программирования, который и упрощает написание ядер GPU, и позволяет управлять многими улучшениями производительности, которые было бы слишком трудоёмко реализовывать на CUDA вручную. Код на Triton работает подобно коду CUDA, но написан на Python и позволяет выполнять пакетные векторные операции при помощи triton.language.tensor. Мы можем использовать его напрямую вместо работы с torch.compile и добиться ещё большей производительности.
В Triton нет возможностей создавать срезы, но можно выполнять смещение базового указателя, поэтому я вычисляю три смещения строк и столбцов, чтобы получить 9 пересекающихся тензоров.
Примечание: это ядро довольно сложно читать. Возможно, стоит начать с одномерного ядра, в котором принимается BLOCK_SIZE_ROW=1.
@triton.jit def gol_triton_2d_kernel(x_ptr, out_ptr, row_stride: tl.int64, N: tl.int64, BLOCK_SIZE_ROW: tl.constexpr, BLOCK_SIZE_COL: tl.constexpr): row_id = tl.program_id(1) col_id = tl.program_id(0) col_offsets0 = (col_id * BLOCK_SIZE_COL + tl.arange(0, BLOCK_SIZE_COL) + 0)[None, :] col_offsets1 = (col_id * BLOCK_SIZE_COL + tl.arange(0, BLOCK_SIZE_COL) + 1)[None, :] col_offsets2 = (col_id * BLOCK_SIZE_COL + tl.arange(0, BLOCK_SIZE_COL) + 2)[None, :] row_offsets0 = (row_id * BLOCK_SIZE_ROW + tl.arange(0, BLOCK_SIZE_ROW) + 0)[:, None] row_offsets1 = (row_id * BLOCK_SIZE_ROW + tl.arange(0, BLOCK_SIZE_ROW) + 1)[:, None] row_offsets2 = (row_id * BLOCK_SIZE_ROW + tl.arange(0, BLOCK_SIZE_ROW) + 2)[:, None] col_mask0 = (col_offsets0 >= 0) & (col_offsets0 < N) col_mask1 = (col_offsets1 >= 0) & (col_offsets1 < N) col_mask2 = (col_offsets2 >= 0) & (col_offsets2 < N) row_mask0 = (row_offsets0 >= 0) & (row_offsets0 < N) row_mask1 = (row_offsets1 >= 0) & (row_offsets1 < N) row_mask2 = (row_offsets2 >= 0) & (row_offsets2 < N) row00 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets0, mask=row_mask0 & col_mask0, other=0) row01 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets1, mask=row_mask0 & col_mask1, other=0) row02 = tl.load(x_ptr + row_offsets0 * row_stride + col_offsets2, mask=row_mask0 & col_mask2, other=0) row10 = tl.load(x_ptr + row_offsets1 * row_stride + col_offsets0, mask=row_mask1 & col_mask0, other=0) row11 = tl.load(x_ptr + row_offsets1 * row_stride + col_offsets1, mask=row_mask1 & col_mask1, other=0) row12 = tl.load(x_ptr + row_offsets1 * row_stride + col_offsets2, mask=row_mask1 & col_mask2, other=0) row20 = tl.load(x_ptr + row_offsets2 * row_stride + col_offsets0, mask=row_mask2 & col_mask0, other=0) row21 = tl.load(x_ptr + row_offsets2 * row_stride + col_offsets1, mask=row_mask2 & col_mask1, other=0) row22 = tl.load(x_ptr + row_offsets2 * row_stride + col_offsets2, mask=row_mask2 & col_mask2, other=0) sum = row00 + row01 + row02 + row10 + row12 + row20 + row21 + row22 result = tl.where(row11 > 0, (sum == 2) | (sum == 3), sum == 3).to(tl.int8) tl.store(out_ptr + row_offsets1 * row_stride + col_offsets1, result, mask=row_mask1 & col_mask1) def gol_triton_2d(x: torch.Tensor, BLOCK_SIZE_ROW, BLOCK_SIZE_COL): output = torch.empty_like(x) grid = ( triton.cdiv(x.shape[1] - 2, BLOCK_SIZE_COL), triton.cdiv(x.shape[0] - 2, BLOCK_SIZE_ROW), ) gol_triton_2d_kernel[grid](x, output, x.stride(0), x.shape[0], BLOCK_SIZE_ROW=BLOCK_SIZE_ROW, BLOCK_SIZE_COL=BLOCK_SIZE_COL) return output
Теперь нам снова надо замерить производительность при разных размерах блоков. Размер блока в Triton не совсем эквивалентен этому понятию в CUDA. Triton использует группу потоков размера 128; если создать блок большего размера, то он начнёт вычислять больше элементов в одном потоке.
При 1024 элементах в блоке потока и 128 потоках на блок каждый поток будет вычислять 8 ячеек. Triton выравнивает нагрузки так, чтобы они были последовательными, чтобы можно было использовать широкие или векторизированные загрузки/сохранения. Кроме того, Triton может пользоваться преимуществами общей памяти1. Я проверил, что сгенерированный код2 делает что-то подобное.
В итоге мы получили время исполнения 22,5 мс, или 51% от пиковой производительности.
Сгруппированные ядра CUDA
Первоначальное ядро CUDA вычисляло по одной ячейке на поток. Одна из главных оптимизаций Triton заключается в снижении количества потоков на блок и в вычислении большего количества ячеек на блок. Для этого в CUDA просто нужно добавить дополнительные циклы вокруг всего ядра.
#define COL_GROUP_SIZE 4 #define ROW_GROUP_SIZE 2 __global__ void gol_kernel_grouped(const int8_t* __restrict__ x_ptr, int8_t* __restrict__ out_ptr, int64_t rowstride, int64_t n) { for (int j = 0; j < ROW_GROUP_SIZE; j++) { int64_t y = (blockIdx.y * blockDim.y + threadIdx.y) * ROW_GROUP_SIZE + j; if (y >= n - 2) break; for (int i = 0; i < COL_GROUP_SIZE; i++) { int64_t x = (blockIdx.x * blockDim.x + threadIdx.x) * COL_GROUP_SIZE + i; if (x >= n - 2) break; int8_t r00 = x_ptr[y * rowstride + x + 0 * rowstride + 0]; int8_t r01 = x_ptr[y * rowstride + x + 0 * rowstride + 1]; int8_t r02 = x_ptr[y * rowstride + x + 0 * rowstride + 2]; int8_t r10 = x_ptr[y * rowstride + x + 1 * rowstride + 0]; int8_t r11 = x_ptr[y * rowstride + x + 1 * rowstride + 1]; int8_t r12 = x_ptr[y * rowstride + x + 1 * rowstride + 2]; int8_t r20 = x_ptr[y * rowstride + x + 2 * rowstride + 0]; int8_t r21 = x_ptr[y * rowstride + x + 2 * rowstride + 1]; int8_t r22 = x_ptr[y * rowstride + x + 2 * rowstride + 2]; int8_t sum = r00 + r01 + r02 + r10 + r12 + r20 + r21 + r22; int8_t result = (r11 > 0) ? ((sum == 2) || (sum == 3) ? 1 : 0) : (sum == 3 ? 1 : 0); out_ptr[(y + 1) * rowstride + (x + 1)] = result; } } } void gol(torch::Tensor x, torch::Tensor out, int block_size_row, int block_size_col) { TORCH_CHECK(block_size_col % COL_GROUP_SIZE == 0, "block_size_col must be divisible by COL_GROUP_SIZE"); TORCH_CHECK(block_size_row % ROW_GROUP_SIZE == 0, "block_size_row must be divisible by ROW_GROUP_SIZE"); TORCH_CHECK(x.size(0) % COL_GROUP_SIZE == 0, "n must be divisible by COL_GROUP_SIZE"); TORCH_CHECK(x.size(1) % ROW_GROUP_SIZE == 0, "n must be divisible by ROW_GROUP_SIZE"); const long n = x.size(0); const int row_blocks = (n - 2 + block_size_row - 1) / block_size_row; const int col_blocks = (n - 2 + block_size_col - 1) / block_size_col; auto stream = at::cuda::getCurrentCUDAStream(); dim3 grid(col_blocks, row_blocks); dim3 block(block_size_col / COL_GROUP_SIZE, block_size_row / ROW_GROUP_SIZE); gol_kernel_grouped<<<grid, block, 0, stream>>>( x.data_ptr<int8_t>(), out.data_ptr<int8_t>(), x.stride(0), n); TORCH_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed"); }
После профилирования я выяснил, что наилучшие результаты получаются при ROW_GROUP_SIZE=1, BLOCK_SIZE_ROW=1, COL_GROUP_SIZE=4, BLOCK_SIZE_COL=512. Так можно получить 14,7 мс, или 78% от пиковой производительности. Обычно при программировании CUDA 60% считается хорошим результатом, а 80% — замечательным!
Честно говоря, меня удивила такая хорошая производительность этого ядра. Оно по-прежнему считывает и записывает по одному байту за раз3. Предположительно, сочетания из кэширования развёртывания циклов и исполнения вне порядка достаточно для максимального роста производительности.
Я провёл ещё несколько экспериментов с CUDA: загрузка 4 байт за одну команду вместо COL_GROUP_SIZE=4 дало эквивалентные результаты. Также я поэкспериментировал с использованием общей памяти, но по сравнению с более наивными решениями производительность была ужасной. Наверно, лучшего нам уже не добиться. Напишите мне, если у вас есть какие-то предложения.
Ядра Triton с упаковкой бит
Когда я начинал этот проект, то хотел поэкспериментировать с 1 байтом на ячейку, поскольку это ближе к типичному ядру машинного обучения, использующего дискретизированные числа с плавающей запятой, поэтому проигнорировал очевидное решение: на самом деле, для игры «Жизнь» достаточно одного бита на ячейку. Так как мы ограничены пропускной способностью памяти, уменьшение хранимых данных в 8 раз даст гораздо больший выигрыш, чем любые инкрементные улучшения, вносимые в ядро.
Я создал ядро Triton, которое распаковывает 8 бит (и загружает дополнительные биты из соседних байт) и вычисляет 8 ячеек. Также я создал 32-битный и 64-битный варианты, использующие трюки с битами для вычисления множества ячеек в одной команде.
Теоретический предел | 1,4 мс |
Triton с 8-битной упаковкой бит | 14,9 мс |
Triton с 32-битной упаковкой бит | 5,21 мс |
Triton с 64-битной упаковкой бит | 6,44 мс |
Как видите, эти ядра намного эффективнее, чем теоретический максимум без упаковки бит. Но на самом деле они далеки от восьмикратного ускорения. Подозреваю, дело в том, что мы всё-таки упёрлись в ограничения вычислительных мощностей. В этих ядрах активно используются побитовые вычисления. В идеале я бы использовал Nsight Systems или Nsight Compute, чтобы подтвердить это, но не смог запустить их на своём виртуальном сервере.
Ядра CUDA с упаковкой бит
Естественно, я попробовал повторить это на CUDA. Не знаю, почему, но это сработало гораздо лучше, чем на Triton. Затем я попросил ChatGPT написать мне ядро с упаковкой бит, обрабатывающее все 64 бита за раз при помощи побитовых операций с uint64.
Теоретический предел | 1,4 мс |
CUDA с 8-битной упаковкой бит | 8,04 мс |
CUDA с 64-битной упаковкой бит | 1,84 мс |
Эти ядра практически не зависят от размера блоков. Подозреваю, что из-за того, что отдельные потоки уже выполняют достаточно широкую загрузку, чтобы воспользоваться преимуществами линий кэша и тому подобного. Поэтому неважно, последовательны ли два потока, их кэширование практически не пересекается. Группирование тоже не обеспечивает роста производительности — каждый такой поток уже обрабатывает по 64 ячейки, что достаточно объёмно.
Результаты
Приведу результаты в удобной таблице:
Длительность | % от пиковой пропускной способности | |
223 мс | 5% | |
38,1 мс | 30% | |
26 мс | 44% | |
22,5 мс | 51% | |
14,7 мс | 78% | |
14,9 мс | 9% (от пика восьмикратного ускорения) | |
5,21 мс | 27% | |
1,84 мс | 76 % |
В итоге мы увеличили производительность в 120 раз! И это при том, что начинали мы с ускоряемого GPU кода, который уже на порядки величин быстрее кода на CPU.
Меня сильно впечатлил стек ускорителей/компиляторов и Triton, и CUDA. У них внутри таится достаточно магии для того, чтобы даже если делать что-то наивно, они могли бы самостоятельно разобраться со всеми подробностями. Основные улучшения, которых мне удалось добиться, связаны с настройкой размеров блоков и групп, а не с какими-то очень хитрыми трюками.
Сам по себе Triton довольно противоречив. С ним определённо немного удобнее работать, чем с CUDA, особенно если не знаешь C/C++. В нём изначально много чего есть, поэтому можно добиться большей производительности меньшими усилиями, чем в CUDA. Однако в конечном итоге ему не удалось сравниться по производительности с относительно простым в написании ядре CUDA, а когда что-то пойдёт не так, возможности его настройки ограничены. Откровенно говоря, частично это связано с моей нагрузкой — я работаю на GPU старого поколения, а в проекте нет операций с плавающей запятой, поэтому многие из лучших операций Triton я не задействую. Но я всё равно считаю, что если есть время для масштабных трудозатрат, всё равно стоит выбирать CUDA.
У Triton есть и другие трюки, например, использование асинхронных операций копирования и тензоров, но они или недоступны A40, или неактуальны для конкретно этого ядра.
Чтобы извлечь скомпилированный вывод, я использовал команду
k = gol_triton_2d_kernel.warmup(x, out, x.stride(0), x.shape[0],BLOCK_SIZE_ROW=4, BLOCK_SIZE_COL=256, grid=(1, 1))print(k.asm['llir'])print(k.asm['ptx'])LLIR — это промежуточное представление LLVM, которое не очень легко читать, а PTX — это байт-код, используемый GPU Nvidia.
Исследовать скомпилированный PTX и SASS можно при помощи
cuobjdump --dump-ptx <file>
