Автор: Чжан Сяо | Архитектор MegEngine

предисловие

В мае 2020 года Nvidia выпустила новое поколение GPU с архитектурой Ampere. Среди них TensorCore третьего поколения с высокой производительностью наиболее тесно связан с глубоким обучением Новое поколение TensorCore поддерживает более распространенные типы данных DL (Deep Learning), включая новые вычисления TesorFloat-32 (TF32), Bfloat16 (BF16). блоки и вычислительные блоки INT8, INT4 и INT1, обеспечивающие всестороннюю поддержку вывода DL. Чтобы задействовать возможности этих вычислительных блоков, операторы свертки и матричного умножения, реализованные сборкой графического процессора, были написаны старшими инженерами высокопроизводительных вычислений в прошлом для использования возможностей оборудования. Однако невозможно справиться с таким количеством типов данных, оптимизируя операторы вручную, поэтому оптимизация приложений глубокого обучения постепенно все больше и больше опирается на некоторые автоматизированные инструменты, такие как компиляторы для глубокого обучения. В соответствии с этой тенденцией Nvidia разработала библиотеку шаблонов линейной алгебры CUTLASS, которая абстрагирует ряд высокопроизводительных базовых компонентов, которые можно использовать для генерации различных типов данных, операторов свертки и матричного умножения различных вычислительных устройств. MegEngine прошла вторичную разработку на базе CUTLASS, что позволяет эффективно разрабатывать новых высокопроизводительных операторов и быстро переходить на новые архитектуры GPU. в предыдущем постестатьяВ этой статье мы кратко представили реализацию базового оператора свертки MegEngine, а в этой статье будет подробно представлен принцип реализации базового оператора свертки платформы MegEngine CUDA, а также свертка неявного GEMM Nvidia CUTLASS.ДокументацияИнтерпретировать и дополнить.

Таким образом, знания CUDA, которые читатели должны знать перед чтением этой статьи:

При доступе к глобальной памяти (Global Memory) соседние потоки в одном и том же Warp обращаются к последовательным адресам, запросы на доступ к памяти будут объединены, и комбинированный доступ к памяти может максимизировать пропускную способность Global Memory.
При доступе к глобальной памяти используйте самый широкий тип данных (float4), возможный для доступа, что максимизирует использование инструкций доступа к памяти.
Общая память CUDA разделена на один банк каждые 4 байта, и всего она разделена на 32 банка. Конфликт банков возникает, когда потоки в одном и том же варпе обращаются к разным адресам в одном и том же банке. Режим доступа к памяти без конфликта банков может максимизировать пропускную способность общей памяти.
Графический процессор имеет четыре уровня хранения: глобальная память, L2, L1 (общая память) и регистры.Задержка прямого доступа к видеопамяти очень велика.При оптимизации вычислительно-емких операторов, таких как GEMM и Convolution, необходимо
- Уменьшите количество запросов на доступ к глобальной памяти через L1 и кэши регистров.
- Скройте неизбежную задержку выборки из глобальной памяти с помощью большого количества вычислений.

Во-первых, нам нужно понять некоторые абстрактные концепции, введенные CUTLASS.

TileIterator: Используется для доступа к данным плитки в хранилище.TileIteratorДостигнутоadvance()метод, поддерживаемый вMatrix , TensorИтерация по другим типам данных.
Fragment: Тип массива, используемый для храненияTileIteratorЧтение входящих данных.FragmentДанные обычно хранятся в регистрах.

Затем мы кратко рассмотрим Pipeline высокопроизводительного оператора GEMM, разработанного CUTLASS.Оператор, реализованный в соответствии с Pipeline, может достичь более 90% производительности cublas на платформе CUDA. На следующем рисунке показан конвейерный оператор GEMM, разработанный CUTLASS:

Первая строка на рисунке демонстрируетPredicatedTileIteratorиSmemTileIteratorСотрудничайте, чтобы завершить передачу данных из глобальной памяти в общую память.
Вторая строка демонстрируетWarpTileIteratorОтвечает за перемещение данных из общей памяти вFragmentв реестре.
Третий ряд показываетWarpMmaOperatorиспользоватьFragmentДанные матрицы в регистрах выполняют операцию Matrix-Multiply-Add.

Неявный алгоритм GEMM

Свертка отображает умножение матриц

Давайте сначала посмотрим на определение оператора прямой свертки, предполагая, что входная карта объектов — это x, вес слоя свертки — w, а выход — y, где x, y и w — все 4-мерные тензоры, и четыре измерения: NxICxIHxIW, четыре измерения w: OCxICxFHxFW, а четыре измерения y: NxOCxOHxOW. Тогда математическая связь между выходом y и входом x, w может быть записана как

\text{y}( \text{n}, \text{oc}, \text{oh}, \text{ow} ) = \sum_{\text{ic}} \sum_{\text{fh}} \sum_{\text{fw}} \text{x} (\text{n}, \text{ic}, \text{ih}, \text{iw}) \cdot \text{w} ( \text{oc}, \text{ic}, \text{fh}, \text{fw} )

Строчные буквы в формуле представляют координаты тензора в каждом измерении, где отношение между ih, iw и oh, ow, fh, fw можно записать как

ih = oh * stride_h - pad_h + fh
iw = ow * stride_w - pad_w + fw

здесьstride_h, stride_w, pad_h, pad_w— параметры сверточного слоя.
Согласно принципу алгоритма im2col операция свертки, определенная в формуле, может быть преобразована в матричное умножение, то есть

C = Matmul(A, B)

Матрица A преобразуется из веса и представляет собой $\text{OC}\times\text{IC}\cdot\text{FH}\cdot\text{FW}$ матрица.
Матрица B преобразуется из карты признаков и представляет собой $\text{IC}\cdot\text{FH}\cdot\text{FW}\times\text{N}\cdot\text{OH}\cdot\text{OW}$ матрица
Матрица C представляет выходной тензор y и представляет собой $\text{OC}\times\text{N}\cdot\text{OH}\cdot\text{OW}$ матрица.

Соответствие между элементами матрицы и тензора в каждой позиции равно

где индекс матрицы $i, j, k$ Связь между координатами тензора и тензора

i = oc
j = n * OH * OW + oh * OW + ow
k = ic * FH * FW + fh * FW + fw

когда $j$ Когда известно, для вычисления координат карты объектов можно использовать следующее соотношение.

n = j / (OH * OW)
j_res = j % (OH * OW)
oh = j_res / OW
ow = j_res % OW

когда $k$ Зная координаты веса, можно рассчитать

ic = k / (FH * FW)
k_res = k % (FH * FW)
fh = k_res / FW
fw = k_res % FW

Комбинируя oh, ow, fh и fw одновременно, можно вычислить ih и iw.
Согласно приведенному выше обсуждению, мы можем записать процесс операции свертки в виде неявного матричного умножения (Implicit GEMM):

GEMM_M = OC
GEMM_N = N * OH * OW
GEMM_K = IC * FH * FW
for i in range(GEMM_M):
    oc = i
    for j in range(GEMM_N):
        accumulator = 0
        n = j / (OH * OW)
        j_res = j % (OH * OW)
        oh = j_res / OW
        ow = j_res % OW
        for k in range(GEMM_K):
            ic = k / (FH * FW)
            k_res = k % (FH * FW)
            fh = k_res / FW
            fw = k_res % FW
            ih = oh * stride_h - pad_h + fh
            iw = ow * stride_w - pad_w + fw
            accumulator = accumulator + x(n, ic, ih, iw) * w(oc, ic, fh, fw)
        y(n, oc, oh, ow) = accumulator

Приведенный выше алгоритм неявного GEMM все еще находится в последовательной форме, и мы преобразуем его в параллельный алгоритм на CUDA. Во-первых, мы разделим всю вычислительную задачу на блоки, и пусть каждый блок потока отвечает за вычисление, а размер вывода равенTILE_MxTILE_Nматрица. Таким образом, алгоритм принимает следующий вид:

for i_out in range(GEMM_M / TILE_M):
    for j_out in range(GEMM_N / TILE_N):
        ThreadblockConvolution(x, w, y)
        
def ThreadblockConvolution(x, w, y):
    accumulate[TILE_M, TILE_N] = 0
    for i_in in range(TILE_M):
        oc = i_out * TILE_M + i_in
        for j_in in range(TILE_N):
            j = j_out * TILE_N + j_in
            n = j / (OH * OW)
            j_res = j % (OH * OW)
            oh = j_res / OW
            ow = j_res % OW
            for k in range(GEMM_K):
                ic = k / (FH * FW)
                k_res = k % (FH * FW)
                fh = k_res / FW
                fw = k_res % FW
                ih = oh * stride_h - pad_h + fh
                iw = ow * stride_w - pad_w + fw
                accumulator(i_in, j_in) = accumulator(i_in, j_in) 
                                        + x(n, ic, ih, iw) * w(oc, ic, fh, fw)
            y(n, oc, oh, ow) = accumulator(i_in, j_in)

Чтобы повысить эффективность доступа к памяти, мы можемGEMM_KБлокировка также выполняется по этому измерению, и каждый разTILE_MxTILE_Kматрица А иTILE_KxTILE_NМатрица B кэшируется в общей памяти, чтобы избежать повторных обращений к глобальной памяти. Таким образом, алгоритм принимает следующий вид:

for i_out in range(GEMM_M / TILE_M):
    for j_out in range(GEMM_N / TILE_N):
        ThreadblockConvolution(x, w, y)

def ThreadblockConvolution(x, w, y):
    accumulator[TILE_M, TILE_N] = 0
    smem_A[TILE_M, TILE_K] = 0
    smem_B[TILE_K, TILE_N] = 0
    for i_in in range(TILE_M):
        oc = i_out * TILE_M + i_in
        for j_in in range(TILE_N):
            j = j_out * TILE_N + j_in
            n = j / (OH * OW)
            j_res = j % (OH * OW)
            oh = j_res / OW
            ow = j_res % OW
            for k_out in range(GEMM_K / TILE_K):
                load_tile_to_smem(x, A_smem)
                load_tile_to_smem(w, B_smem)
                WarpGemm(A_smem, B_smem, accumulator)
            y(n, oc, oh, ow) = accumulator(i_in, j_in)

def WarpGemm(A_smem, B_smem, accumulator):
    for k_in in range(TILE_K):
        accumulator(i_in, j_in) = accumulator(i_in, j_in) 
                                + A_smem(i_in, k_in) * B_smem(k_in, j_in)

Потому что мы можем напрямую повторно использовать CUTLASS, который уже достиг высокой производительности.WarpMmaOperator, поэтому для реализации оператора свертки на основе Implicit GEMM требуется только

приспособлениеDeviceConvolution,KernelConvolutionиThreadblockConvolution, который поддерживает передачу параметров типа Tensor и Convolution Layer.
Добавить кPredicateTileIteratorПоддерживает чтение данных Tile of Tensor в общую память и неявно организует считанные данные в матрицу.
Вызывается непосредственно в основном цикле алгоритмаWarpTileIteratorПрочитайте данные из общей памяти, затемWarpGemmOperatorЗавершите операции GEMM уровня Warp.
EpilogueOperatorАдаптируйте оператор свертки и запишите данные Аккумулятора обратно в Тензор Глобальной Памяти.

Далее мы представим реализацию свертки на нижнем уровне MegEngine с оператором свертки TensorCore типа данных INT 8. В этой статье основное внимание будет уделено тому, как реализованы 2, 3 и 4. Как использовать уже написанный оператор свертки, вы можно обратиться к предыдущемустатья.

Разметка данных глобальной памяти (раскладка)

Чтобы максимизировать пропускную способность операторов свертки типа TensorCore, MegEngine использует 128-битный глобальный Инструкции по доступу к памяти, поэтому при доступе к данным Tensor адрес должен соответствовать 128-битному выравниванию. MegEngine использует формат NCHW32 для хранения тензоров, особенности формата NCHW32:

Размер канала Tensor сгруппирован по 32 каналам, и каждые 32 канала постоянно хранятся в хранилище.
Остальные измерения Тензора хранятся в хранилище в порядке W, H, C и N, от быстрого к медленному.

Из-за формата хранения, выровненного по 32 каналам, сверточный слой требует, чтобы количество каналов как для входных, так и для выходных карт объектов было кратно 32.

Смещение выборки предварительной обработки

Свертка MegEngine реализована вGEMM_KРазмер соответствует $(\text{IC}/32)\cdot \text{FH}\cdot \text{FW}\cdot32$ Порядок накопления записывается в виде псевдокода следующим образом:

kInterleaved = 32
for ic_out in range(IC//kInterleaved):
    for fh in range(FH):
        for fw in range(FW):
            for ic_in in range(kInterleaved):
                # do mma
                ......

Если это написано как слой цикла, то это должно быть записано как:

kInterleaved = 32
for k in range(GEMM_K):
    chw = k // kInterleaved
    ic_in = k % kInterleaved
    ic_out = chw // (FH * FW)
    chw_res = chw % (FH * FW)
    fh = chw_res // FW
    fw = chw_res % FW
    pointer += ic_out * C_STRIDE + fh * H_STRIDE + fw * W_STRIDE
    # do mma
    ......

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

Если можно использовать приращение адресаdeltaвыразить, тоdeltaдаFH*FWявляется периодическим, то есть:

delta(step, TILE_K) = delta(step + (FH * FW), TILE_K)

Поэтому нам нужно всего около $\text{O}\left(\text{FH}\cdot\text{FW}\right)$ места для хранения. Логика расчета смещения адреса может относиться к кодуconv2d_tile_iterator_nt_src_fprop_precomp.h. Поскольку размер буфера параметров ядра составляет 4 КБ, мы используем около 3 КБ для хранения приращения адреса, поэтому для реализации свертки MegEngine требуется уровень свертки.FH*FWРазмер не может быть слишком большим, но в целом свертки 3х3, 5х5, 7х7 справляются. Порядок итерации официальной реализации Nvidia немного отличается от описанного в этой статье:

Официальная реализация должнаICзаполнить какTILE_Kкратно , что приведет к потере некоторых вычислений, когда количество каналов невелико.
Официально реализованный блок потоков имеет большой диапазон адресов при доступе к входной карте объектов, что снижает локальность доступа к памяти и не является дружественным к кэшу.

Поэтому с точки зрения производительности реализация MegEngine будет иметь больше преимуществ, а преимущество официальной реализации в том, что ограничений на параметры Convolution Layer не слишком много, а универсальность лучше.

Инструкция Mma на уровне деформации (матрица-умножение-сложение)

cuda10.2 представляет новый уровень деформацииmmaиldmatrixинструкции, пользователь может пройтиmmaИнструкция использует TensorCore для выполнения высокоскоростных операций умножения и сложения матриц черезldmatrixДетальный контроль над тем, как Warp питает TensorCore. вmmaИнструкции используются следующим образом:

unsigned A, B;  // input matrix fragment data
int C[2], D[2]; // accumulators
asm volatile(
    "mma.sync.aligned.m8n8k16.rol.col.satfinite.s32.s8.s8.s32 {%0,$1}, {%2}, {%3}, {%4,%5};\n"
    : "=r"(D[0]), "=r"(D[1])
    : "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));

Семантика этой инструкции заключается в синхронном выполнении операции умножения и сложения матриц 8x8x16 нитями Warp 32. Она имеет три входных операнда, из которых операция умножения матриц представляет собой матрицу A 8x16 и матрицу B 16x8 соответственно. для этих двух входных матриц распределяется между 32 потоками одного и того же Warp. Схема матрицы A показана на следующем рисунке:

32 потока в одном и том же Warp разделены на 8 групп по четыре потока, отвечающих за чтение строки в матрице 8x16.
Один поток в каждой группе считывает 4 соседних int8 в каждой строке, заполняя ровно один 32-битный регистр.

Аналогичный макет для матрицы B показан ниже:

Каждая группа из 4 потоков делится на 8 групп, и каждая группа отвечает за чтение столбца в матрице 16x8.
Один поток в каждой группе отвечает за чтение 4 соседних данных в столбце.

Данные матрицы C и выходной матрицы D, участвующие в операции накопления, также распределены по 32 потокам, и их расположение показано на следующем рисунке:

Точно так же каждая группа из 4 потоков отвечает за ввод/вывод одной строки данных.
Каждый поток отвечает за вывод двух соседних данных типа int32 в строке, которые в точности составляют 64-битный регистр.

через паруmmaАнализ инструкций, если данные в Global Memory/Shared Memory хранятся в формате RowMajor или ColumnMajor, то когда один и тот же Warp выполняет две последовательные операции умножения и сложения матриц 8x8x16 в пространстве, данные, считанные каждым потоком, будут перепрыгивать, и каждый умножение может считывать в регистр только 32-разрядные данные, а инструкция загрузки с низким разрядом обычно не имеет возможности максимизировать использование пропускной способности памяти. Таким образом, Nvidia предоставляетldmatrixЭта инструкция позволяет одному и тому же варпу одновременно считывать в регистр 4 матрицы 8x16, так что каждый поток в варпе может считывать 128-битные данные за раз, максимально используя полосу пропускания.ldmarixИспользование заключается в следующем:

unsigned addr;  // shared memory pointer
int x, y, z, w; // loaded data
int4 data;      // loaded fragment
asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];"
    : "=r"(x), "=r"(y), "=r"(z), "=r"(w)
    : "r"(addr));
data = make_int4(x, y, z, w);

Приведенная выше инструкция читает ровно 4 матрицы 8x16, и каждый поток отвечает за чтение строки данных в матрице.После завершения чтения потоки будут обмениваться данными, перераспределять данные матрицы в каждый поток, читать Процесс прием показан на следующем рисунке:

В этом разделе представлены связанные с TensorCoremmaиldmatrixИнструкции, с этими двумя высокопроизводительными инструкциями, нам также необходимо разработать оригинальный формат хранения данных в общей памяти, чтобы устранить конфликт банков при чтении данных из общей памяти, тем самым повысив эффективность чтения общей памяти.

Расположение данных в общей памяти

Прежде чем представить макет данных в Shared Memory, нам нужно понять характеристики доступа к памяти Shared Memory. Общая память формирует банк каждые 4 байта и делится на банки всего 32. Когда потоки одного и того же варпа обращаются к разным адресам одного и того же банка, возникают конфликты, что приводит к снижению эффективности доступа к памяти. Когда потоки одного и того же Warp получают доступ к данным разной разрядности, поведение будет разным:

Каждый поток обращается к 32-битным данным в общей памяти, и доступ к памяти будет выполняться в один этап.
Каждый поток обращается к 64-битным данным в общей памяти, и доступ к памяти выполняется в два этапа:
- Фаза 1: первые 16 потоков извлекают 128 байт данных.
- Второй этап: последние 16 потоков извлекают 128 байт данных.
Каждый поток обращается к 128 битам данных в общей памяти, и доступ к памяти выполняется в четыре этапа:
- Каждый этап выполняется 8 потоками для выборки 128 байт данных.

Если на каждом этапе описанного выше процесса не возникает конфликта банков, может быть достигнута максимальная эффективность доступа к общей памяти. Обычно, чтобы избежать конфликта банков общей памяти, мы дополняем данные общей памяти, чтобы данные, к которым обращается поток, располагались в шахматном порядке, чтобы избежать попадания в один и тот же банк. Однако проблема заключается в том, что размер Shared Memory, требуемый ядром, станет больше, но кэш L1 (Shared Memory) на SM ограничен, поэтому заполнение уменьшит занятость ядра, что, в свою очередь, снизить производительность ядра. Поэтому CUTLASS разработала ступенчатую структуру Shared Memory, которая может сделать адрес доступа к потоку без конфликта банков без заполнения. Далее мы возьмем матрицу 64x64 в качестве примера, чтобы подробно описать расположение данных в Shared Memory. Прежде всего, гранулярность считываемых потоком данных составляет 128 бит, то есть 16 частей данных типа INT8, поэтому мы всегда используем 16 частей данных как группу при демонстрации размещения данных. Если матрица организована в построчном формате (RowMajor), то логическая компоновка выглядит следующим образом:

Как видно из рисунка

Каждая группа из 16 элементов, называемая вектором, окрашена в свой цвет.
Смежные 32 элемента каждой строки называются Crosswise, которые представляют собой данные группы каналов в формате NCHW32.

В физическом хранилище Shared Memory данные матрицы переупорядочиваются, как показано на следующем рисунке:

Мы видим, что физическое расположение Shared Memory имеет следующие характеристики:

Данные одного Crosswise из каждых 4 строк сохраняются как группа в общей памяти, а затем будут сохраняться данные следующего Crosswise из этих 4 строк.
Каждый набор данных содержит 8 векторов, занимающих 128 байтов, что составляет 32 различных банка в общей памяти.
Каждая группа данных чередуется в расположении, гарантируя, чтоldmatrixникакого конфликта с банком не произойдет.

Видеопамять -> передача данных Shared Memory

В этом разделе мы представим передачу данных из видеопамяти (глобальной памяти) в общую память. Перенос данных из видеопамяти в Shared Memory осуществляетсяConv2dTileSrcIteratorFpropPrecompВ завершение в этой статье не будет подробно интерпретироваться реализация кода, но будет описан процесс обработки данных потоком, чтобы помочь вам создать интуитивно понятное впечатление и лучше понять код. Если взять в качестве примера логическую структуру Shared Memory из предыдущего раздела, логическая структура данных, считываемых каждым потоком в одном и том же Warp, показана на рисунке ниже.Каждый поток считывает 16 фрагментов данных типа INT8, что в точности составляет Вектор.

В реальной физической видеопамяти распределение данных доступа потоков показано на следующем рисунке:

Мы видим, что каждый поток считывает 128 бит данных.
Данные, считываемые соседними потоками, физически непрерывны.

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

Каждые 8 потоков записывают 128 байт данных в разделяемую память, которые попадают в 32 различных банка разделяемой памяти.
Доступ к памяти одного и того же варпа выполняется в четыре этапа, и на каждом этапе не возникает конфликта банков.

На следующем рисунке показан процесс записи Warp в Shared Memory:

Общая память -> перемещение данных регистров

Перенос данных из разделяемой памяти в регистры осуществляетсяMmaTensorOpMultiplicandTileIteratorЗаконченный. Тот же Warp будет считывать четыре матрицы 8x16 в регистры на каждой итерации, и каждый поток будет считывать строку данных. Например, в первом раунде итерации логическая структура данных, считанных потоком, показана на следующем рисунке:

На самом деле физическое расположение данных в Shared Memory выглядит следующим образом:

можно увидеть:

Каждый поток считывает 128 бит данных, поэтому выборка памяти выполняется в четыре этапа.
Данные, считанные 8 потоками на каждом этапе, попадают в 32 банка общей памяти, и между данными, полученными потоками, нет конфликта.

Когда достигается второй раунд итерации, физическое расположение данных, к которым обращается каждый поток, выглядит следующим образом:

На каждом этапе одной и той же выборки не возникает конфликта банков.

Аккумулятор выполняет обратную запись в глобальное хранилище

В случае с int8 за вывод результатов 64x64 отвечает тот же Warp, а ядро будет делиться на 8 раз для записи обратно в Global Memory, каждый раз обратно записывается матрица 32x8. Это гарантирует, что каждый раз, когда Tensor записывается обратно в видеопамять в соответствии с форматом NCHW32, 32 потока одного и того же Warp записывают ровно 256 байтов данных, которые физически непрерывны, и каждый поток записывает обратно 8 байтов, гарантируя, что 64 байта Типы битовых данных выполняют операции записи в видеопамять, чтобы максимизировать использование полосы пропускания. так какmmaХарактеристики инструкции, данные выходной матрицы распределяются по каждому потоку, а для того, чтобы иметь возможность объединять доступ к памяти, то есть: чтобы адреса, записываемые обратно соседними потоками, были непрерывными, мы используем Shared Память до 32 потоков в том же Warp.change. После обмена данными каждый поток имеет 8 последовательных каналов данных, а адреса, записанные потоком, являются последовательными, что гарантирует, что обратная запись в глобальную память соответствует требованиям комбинированного доступа к памяти. Процесс обмена данными между потоками показан на следующем рисунке:

На каждой итерации 32 потока в Warp записывают данные матрицы 32x16 в общую память. Затем, как показано на рисунке ниже, каждый поток будет считывать в регистр данные 8 последовательных каналов.

Обмен данными Shared Memory осуществляется следующими двумяIteratorЗаконченный

InterleavedTileIteratorTensorOpВыполнена каждая итерация для записи данных 32x8 в общую память.
InterleavedSharedLoadIteratorTensorOpОтвечает за чтение данных из последовательных 8 каналовFragmentв реестре.

Когда поток считывает замененные данные вFragmentПосле регистрации он будет заменен наEpilogueOp, сделанный на основе сверткиBiasAddоперация. отBiasAddLinearCombinationReluНапример, на самом деле он делает следующее:

accumulator = conv(x, w)
y = alpha * accumulator + beta * bias + gamma * z

где предвзятостьPerChannelТензор представляет смещение каждого выходного канала, z — это тензор того же размера, что и выход свертки, используемый дляConvolutionиElemwiseAddслияние. НаконецEpilogueOpРезультат будет предоставленTensorPredicatedTileIteratorTensorOpНа самом деле отпишитесь в глобальную память. Данные, записываемые обратно каждым потоком, показаны на следующем рисунке:

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

Суммировать

В этой статье представлен принцип реализации оператора свертки в нижней части MegEngine.Производительность оператора может достигать более 80% cudnn.Результаты измерения скорости можно найти встатья.

MegEngine продолжит оптимизировать реализацию свертки для дальнейшего повышения производительности оператора.В настоящее время можно выполнить две оптимизации:

Изучите официальную свертку CUTLASS ImplicitGEMM от Nvidia для обработки масок, улучшенияTileIteratorДля эффективности оценки маски.
Текущая реализация свертки использует Shared Memory для обмена данными при обратной записи в видеопамять, и возникает конфликт банков. Две оптимизации будут рассмотрены позже
- Изучите расположение данных в общей памяти, устраните конфликты банков и оптимизируйте эффективность обмена данными в общей памяти.
- Изучите расположение тензора весов в глобальной памяти, улучшите расположение аккумулятора в каждом потоке и избегайте обмена данными в общей памяти.

использованная литература

Фрагмент матрицы варп-уровня Mma PTX Documentation
Официальная документация CUTLASS по неявной свертке GEMM
Volta architecture and performance optimization
Developing CUDA kernels to push Tensor Cores to the absolute limit on Nvidia A100