Понимание мозаики тензорных ядер с использованием CUDA на V100

Question

У меня есть игрушечный код, позаимствованный у NVidia simpleTensorCoreGEMM.cu . Я поменял их случайное генерирование матриц на функцию, которая считывает матрицы из файлов.

Использование этого игрушечного кода и умножение двух матриц размера [2000 x 10000] * [10000 x 3008] прекрасно работает. Вывод соответствует ожидаемому.

Когда я пытаюсь сделать намного большее умножение [20000 x 10000] * [10000 x 30000], вывод идет ужасно неправильно, и 2/3 строк равны 0.

Я убежден, что это результат того, что я не понимаю строки кода:

// blockDim.x must be a multple of warpSize
// 128x4 means we have 16 warps and a block computes a 64x64 output tile
blockDim.x = 128;
blockDim.y = 4;

gridDim.x = (MATRIX_M + (WMMA_M * blockDim.x / 32 - 1)) / (WMMA_M * blockDim.x / 32);
gridDim.y = (MATRIX_N + WMMA_N * blockDim.y - 1) / (WMMA_N * blockDim.y);

Даже если это не источник моей ошибки, я все равно должен понимать, что он делает. Я понимаю настройку blockDim.* На основу приходится 32 потока, 128 * 4/32 = 16 основ.

ВОПРОС: Может ли кто-нибудь объяснить мне логику значений и вычислений gridDim.x и gridDim.y? Правильное использование тензорных ядер, кажется, очень чувствительно к использованию правильных значений для gridDim.*.

Answer 1

Пара вступительных слов:

1. Для понимания этот код предназначен для сопровождения этой статьи блога . Последняя часть этого блога, раздел «Программный доступ к тензорным ядрам в CUDA 9.0», безусловно, полезна для понимания этого кода.

2. Как уже упоминалось в файле readme для этого кода , более простой способ получить доступ к производительности тензорных ядер (особенно для базовых операций умножения матриц, с которыми вы, кажется, играете) - просто использовать Функция CUBLAS, такая как cublasGemmEx , которая будет разумно использовать тензорные ядра при правильных обстоятельствах.

Теперь к вашему вопросу:

Может ли кто-нибудь объяснить мне логику значений и вычислений gridDim.x и gridDim.y?

Эти значения определяют размер сетки CUDA, достаточный для размера запрошенной задачи умножения матриц. Нам нужно подходить к этому иерархически.

• Прежде всего, доступ к тензорному ядру возможен на уровне деформации. В статье блога говорится, что «стратегия, которую мы будем использовать, состоит в том, чтобы иметь одну деформацию, ответственную за один раздел матрицы вывода 16 × 16». Поэтому размеры матрицы вывода будут определять размеры сетки CUDA, используемой для вычисления результата. ( Типичные наивные реализации умножения матриц также определяют размер сетки на основе размера выходной матрицы. Более конкретно, они назначают один поток на каждую выходную точку. Здесь мы назначаем один 32-ниточный перекос, чтобы отвечать за одну плитку 16x16 матрицы вывода.) В коде используются WMMA_M (то есть, сколько строк) и WMMA_N (то есть, сколько столбцов), чтобы определить, что будет обрабатывать одна тензорная базовая операция уровня деформации. Эти значения равны 16, и это определяет выбор использования плитки 16x16 в выводе для основы.

• Как часто бывает в CUDA, размеры блоков могут быть несколько произвольными, но они часто влияют на размер сетки (переменные). Деформации существуют на уровне блоков, и количество деформаций в блоке эффективно определяет, сколько плиток 16x16 в выходной матрице будет обрабатываться на блок. В этом конкретном случае код выбирает размеры блока от 128 (blockDim.x) до 4 (blockDim.y). Это бывает 4 ширины деформации, а 4 высоты - высокой, поэтому каждый блок обрабатывает набор плиток 4x4 на выходе, что означает, что каждый блок отвечает за выходные точки 64x64. Обратите внимание, что эти переменные blockDim и gridDim в коде хоста логически отделены (хотя в конечном итоге они совпадают с числовыми значениями) от встроенных переменных blockDim и gridDim в коде устройства CUDA.

• Учитывая вышеизложенное, параметры m, n и k типичной операции BLAS GEMM имеют здесь то же значение. m - количество строк левой входной матрицы. n - количество столбцов правой входной матрицы. k - количество столбцов левой матрицы, которое должно соответствовать количеству строк правой матрицы. Поэтому m, n определяют размеры выходной матрицы. Они обозначены в коде как MATRIX_M и MATRIX_N соответственно.

С учетом изложенной выше основы мы можем указать арифметику, необходимую для вычисления gridDim.x и gridDim.y в коде хоста.

1. Мы должны выбрать достаточное количество нитей в измерении x, чтобы при делении на 32 (ширина основы в измерении x) и затем умножении на WMMA_M (ответственность за ширину выходной плитки для этой основы) , у нас достаточно потоков, чтобы покрыть ширину выходной матрицы.

2. Мы должны выбрать достаточное количество нитей в измерении y, чтобы при делении на 1 («высота» основы в измерении y) затем умножить на WMMA_N (ответственность за высоту выходного элемента мозаичного изображения этой основы) ), у нас достаточно потоков, чтобы покрыть высоту выходной матрицы. Обратите внимание, что «высота» основы в измерении y в этом случае определенно равна 1, поскольку код требует, чтобы размер ширины блока был целым числом, кратным размеру основы. Поэтому любая деформация имеет постоянную составляющую threadIdx.y по всей деформации.

3. Чтобы перейти от потоков, определенных в 1 и 2 выше, к блокам в каждом измерении, мы должны масштабировать (делить) каждый на соответствующий размер блока нитей. Поэтому размер нити сетки в x должен быть разделен на blockDim.x (в коде хоста), масштабированном как в 1 выше, чтобы получить общее измерение сетки (количество блоков) в x. Эта операция деления является обычной операцией деления целых чисел «округлять вверх» в CUDA, чтобы масштабировать количество блоков, равное или превышающее необходимые потоки, чтобы учитывать размеры матриц, которые делятся неравномерно на размер блока.

Собрав все это вместе, мы получим:

gridDim.x = (MATRIX_M + (WMMA_M * blockDim.x / 32 - 1)) / (WMMA_M * blockDim.x / 32);
   ^            ^             ^                                   ^
   |            |             |                    divided by the block size scaled for the
   |            |             |                     portion of the output matrix it covers.
   |            |           rounded up
   |         the matrix size
  The grid in blocks is

И аналогично для размера сетки y. Единственное реальное отличие состоит в том, что 32 потока в x (ширина деформации) отвечают за выходной тайл 16x16, тогда как за один поток в y (высота деформации) отвечают за этот выходной тайл 16x16.