Умножение матриц на GPU

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

Как достичь топовой производительности умножения матриц в CUDA.

Этот блог возник из внезапного осознания того, насколько я мало знаю о том, как умножение матриц работает на графическом процессоре. Проведя так много проектов по машинному обучению, я чувствую, что я должен понимать, как работает самая важная операция в машинном обучении: Что такое “Tensor Core”? Почему все говорят, что “перемещение данных является узким местом”? Насколько быстрыми могут быть графические процессоры на самом деле?

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

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

Достаточно тяготеть, давайте начнем!

• Алгоритм Дейкстры с весом, основанным на времени передвижения в сетях OSM.
• Кодирование байт-пар для начинающих
• Как расширить DataFrames в Pandas с помощью пользовательских методов для усиления функциональности и читаемости кода

Рекапитуляция архитектуры GPU

Давайте вспомним, как работают (NVIDIA) GPU. GPU достигает параллелизма, запуская множество потоков. Каждый поток выполняется на одном CUDA-ядере, но в данное время только некоторое подмножество потоков активно, поэтому количество потоков может быть намного больше, чем доступных CUDA-ядер. Каждый поток, независимо от того, активен он или нет, имеет свой набор регистров.

Группа из 32 потоков называется варпом. Все потоки в варпе должны выполняться вместе (или быть неактивными вместе). В большинстве случаев неактивных варпов намного больше, чем активных, и планировщик варпов отвечает за выбор варпов, которые будут выполняться в данное время. Это позволяет GPU скрывать задержку доступа к памяти, планируя выполнение других варпов, пока варп ждет данных.

Группа варпов называется тредблоком. Все варпы в тредблоке выполняются в одном и том же Streaming Multiprocessor (SM). У каждого тредблока есть своя собственная общая память, к которой имеют доступ все потоки в тредблоке.

Примечание: новые архитектуры

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

Архитектура Вольта также представила Tensor Cores, специализированные для решения умножений матриц определенного размера. Каждый активный варп имеет доступ к одному Tensor Core.

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

Параллелизация умножения матриц

Предположим, нам нужно вычислить:

Мы говорим, что размер проблемы в данном случае составляет (M, N, K). Чтобы параллельно решить эту задачу, мы можем разделить матрицы A и B на более маленькие, выполнить умножение матриц по отдельности и объединить результаты для получения матрицы C.

Конкретно, мы можем разделить A по строкам (то есть M на куски размера M’) и B по столбцам (то есть N на куски размера N’), чтобы получить:

Мы видим, что каждая подматрица в C независима друг от друга, поэтому мы можем легко параллельно вычислять каждую подматрицу.

На практике может быть так, что K слишком велико, чтобы загрузить его в память и вычислить. Вместо этого типичная реализация также разделит K на части размером K’ и будет выполнять итерацию по каждой части, накапливая (путем суммирования) частичные результаты. Это известно как последовательное сокращение K (в отличие от параллельного сокращения K, см. ниже). Математически это выглядит так:

Примечание: Добавление отступа

В любой момент, когда размер задачи не делится на размер разбиения, нам нужно добавить отступ. Обычно это делается неявно, когда мы загружаем разделенные данные (𝐴ᵢ,ₖ и 𝐵ₖ,ⱼ) в память более низкого уровня, где мы убеждаемся, что загруженный раздел (размером M’×K’ для 𝐴ᵢ,ₖ и K’×N’ для 𝐵ₖ,ⱼ) всегда является «полным», добавляя нули. При записи результатов обратно в глобальную память необходимо обеспечить особую осторожность, чтобы избежать ошибок выхода за границы.

На высоком уровне для параллелизации умножения матрицы на GPU используются три уровня вложенных разделений:1. Первое разделение происходит на уровне блока потоков. Каждый блок потоков отвечает за вычисление Cᵢ,ⱼ = Aᵢ Bⱼ.2. Второе разделение происходит на уровне ворпа. Проблема на уровне блока потоков Cᵢ,ⱼ дополнительно разделяется таким образом, что каждый ворп отвечает за вычисление Cᵢ,ⱼ⁽ᵐⁿ⁾ = Aᵢ⁽ᵐ⁾ Bⱼ⁽ⁿ⁾.3. Третье разделение происходит на уровне инструкции. Некоторые инструкции ожидают входные данные определенного размера. Например, тензорные ядра второго поколения работают с проблемами размером (16, 8, 8) для fp16, тогда как прямая реализация на CUDA-ядрах с использованием скалярного умножения просто работает с размером (1, 1, 1). Проблема на уровне ворпа дополнительно разделяется таким образом, что каждая часть имеет подходящий размер для инструкции: Cᵢ,ⱼ⁽ᵐⁿ⁾⁽ᵃᵇ⁾ = Aᵢ⁽ᵐ⁾⁽ᵃ⁾ Bⱼ⁽ⁿ⁾⁽ᵇ⁾.

Есть веская причина, почему нам нужно три уровня разделения, как мы увидим в следующем разделе.

Повторное использование данных

Умножение матриц может легко стать ограниченным памятью, если мы наивно повторно загружаем данные с глобальной памяти в регистры каждый раз при выполнении вычислений. Ключевое наблюдение заключается в том, что многие подвходы Aᵢ и Bⱼ повторно используются при умножении разных подматриц. Например, Aᵢ требуется для Cᵢ,₁ , Cᵢ,₂ , … и Bⱼ требуется для C₁,ⱼ , C₂,ⱼ , … . Мы можем достичь лучшей производительности, если мы минимизируем избыточное перемещение данных и повторно используем загруженные данные в максимальной степени.

В CUDA существуют три типа доступной для пользователя памяти:

Вот общий вид того, как используется каждый тип памяти:

1. Каждый блок потоков сначала загружает необходимые данные из глобальной памяти в разделяемую память. В дальнейшем доступ к этим данным будет обслуживаться разделяемой памятью, а не медленной глобальной памятью.
2. В каждом блоке потоков каждый ворп сначала загружает необходимые данные из разделяемой памяти в регистры. В дальнейшем доступ к этим данным будет обслуживаться быстрыми регистрами непосредственно.

Глубокое погружение в детали

Уровень блока потоков

На уровне блока потоков проблема разбивается на подпроблемы размером (M’, N’, K’). Таким образом, каждый блок потоков отвечает за вычисление фрагмента C, обозначенного как:

Лишнее перемещение данных минимизируется путем загрузки подвходов Aᵢ,ₖ и Bₖ,ⱼ в разделяемую память. Когда мы закончили с вычислением Aᵢ,ₖ Bₖ,ⱼ, следующий фрагмент (Aᵢ,ₖ₊₁ и Bₖ₊₁,ⱼ) будет загружен.

Уровень warp

На уровне warp подпроблема дополнительно разбивается на подподпроблемы размером (M’‘, N’‘, K’’). Таким образом, каждый warp отвечает за вычисление фрагмента Cᵢ,ⱼ, обозначенного как Cᵢ,ⱼ⁽ᵐ ⁿ⁾:

Лишнее перемещение данных минимизируется путем загрузки подподвходов Aᵢ,ₖ⁽ᵐ ˡ⁾ и Bₖ,ⱼ⁽ˡ ⁿ⁾ в регистры. Любые обращения к Aᵢ,ₖ⁽ᵐ ˡ⁾ и Bₖ,ⱼ⁽ˡ ⁿ⁾ внутри warp будут обслуживаться быстрыми регистрами.

Примечание: Распределение данных по регистрам

Стоит отметить, что регистры это только для потока. Это означает, что входы в регистры не могут быть доступны для других потоков в warp. Конкретный способ разделения Aᵢ,ₖ⁽ᵐ ˡ⁾ и Bₖ,ⱼ⁽ˡ ⁿ⁾ между регистрами каждого потока зависит от конкретной используемой инструкции. В документации NVIDIA по инструкциям по умножению матриц на уровне warp-а приводится подробное описание для каждой инструкции.

Уровень тензорных ядер

Для фактического выполнения умножения матриц мы используем тензорные ядра на GPU. У моего GPU (RTX 2060) есть тензорные ядра второго поколения, которые специализируются на решении проблем fp16 размером (M’’‘, N’’‘, K’’’) = (16, 8, 8). Таким образом, мы еще дополнительно разбиваем Cᵢ,ⱼ⁽ᵐ ⁿ⁾, чтобы соответствовать ожидаемому размеру инструкции:

Здесь все входы уже находятся в регистрах, и поэтому издержки перемещения данных минимальны.

Примечание: Тензорные ядра

Операции с тензорными ядрами являются инструкциями на уровне warp-а, что означает, что все потоки в warp-е должны одновременно выполнять инструкцию тензорных ядер, совместно подготавливая данные для использования только одним тензорным ядром.

Выбор размеров разделения

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

Размер разделения блока потоков

Асимптотически, с ростом размера проблемы, да, мы хотим использовать как можно больше разделяемой памяти и регистров. Однако для небольших размеров проблемы мы можем столкнуться с двумя проблемами:

1. Большой размер разделения означает, что у нас будет меньше блоков потоков. Поскольку каждый блок потоков может выполняться только на одном SM, это может означать, что мы не сможем использовать все SM.
2. Для проблемных размеров, которые не делятся на размер разделения, нам придется добавлять больше заполнения входам. Это снижает эффективность, так как будет выполнено меньше вычислений на значимых входах.

Обычная реализация может использовать размер разбиения (M’, N’, K’) = (128, 256, 32).

Размер разбиения варпов

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

Обычная реализация может использовать размер разбиения (M”, N”, K”) = (64, 64, 32).

Размер разбиения инструкций

Это полностью определяется инструкциями, поддерживаемыми вашей графической картой. Для моей RTX 2060 инструкцией ptx для умножения матрицы Tensor Core fp16 (с накоплением fp16) является mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16, которая ожидает размеры входных данных (16, 8, 8).

Еще больше оптимизаций

Вышеуказанные методы позволяют приблизиться к теоретической максимальной производительности графического процессора при большом размере проблемы. Однако для малых размеров проблем они не так эффективны. Существуют две общие методики для дальнейшего улучшения производительности умножения матриц: параллельное сокращение К и программная конвейеризация.

Параллельное сокращение К

В случаях, когда M и N малы, у нас может быть только несколько блоков потоков. Например, если (M’, N’) = (128, 256), и исходный размер проблемы имеет M ≤ 128 и N ≤ 256, у нас будет только один блок потоков, и мы будем использовать только часть вычислительной мощности графического процессора! (Например, моя RTX 2060 имеет 30 SM, поэтому для максимального использования нам нужно как минимум 30 блоков потоков).

В случаях, когда K большое (несмотря на малые M и N), мы можем использовать больше параллелизма с помощью параллельного сокращения К. Напомним, что в последовательном сокращении К каждый блок потоков выполняет следующую сумму:

и накапливает промежуточные результаты в Cᵢ,ⱼ. В параллельном сокращении К каждый блок потоков вычисляет только один элемент суммы (т.е. Aᵢ,ₖ Bₖ,ⱼ). Это позволяет увеличить количество блоков потоков в K/K’ раз, таким образом, используя больше SM.

Однако теперь нам нужно выделить больше памяти для сохранения результатов каждого блока потоков и вызвать второе ядро для выполнения окончательного сокращения частичных результатов, чтобы получить Cᵢ,ⱼ.

Программная конвейеризация

Обычно CUDA скрывает задержку доступа к памяти, планируя выполнение других варпов, пока один варп ожидает данных. Для этого нам нужно иметь достаточное количество варпов, чтобы скрыть задержку.

Однако количество варпов обычно относительно невелико при выполнении умножения матриц. Это связано с тем, что количество варпов ограничено “Количество доступных регистров на блок потоков, деленное на количество требуемых регистров на варп”. Для умножения матриц мы используем много регистров на варп, чтобы максимизировать повторное использование данных. В результате у нас может не быть достаточного количества варпов для скрытия задержки.

“Элементы аккумулятора обычно занимают не менее половины общего бюджета регистров потока”. — Документация CUTLASS

Для смягчения этого эффекта мы можем использовать программную конвейеризацию. В основном, мы можем (вручную) предварительно загрузить входные данные для следующей итерации цикла асинхронно, используя специальные инструкции. Во время загрузки входных данных мы можем продолжать вычисления на текущей итерации. Это показано на следующей диаграмме:

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

Умножение матриц в действии

Если вы хотите увидеть, как все эти концепции объединяются в реальной реализации, ознакомьтесь с моей реализацией обучения MNIST с нуля с использованием CUDA. Там я обучал многослойный перцептрон на наборе данных MNIST с использованием CUDA и достиг ускорения в 6 раз по сравнению с оптимизированным PyTorch для скрытого размера 128:

GitHub – andylolu2/cuda-mnist

Присоединяйтесь к разработке andylolu2/cuda-mnist, создав аккаунт на GitHub.

github.com

Ссылки

1. Документация CUTLASS2. Документация CUDA3. Примеры CUTLASS