Полное руководство по оптимизации умножения матриц CUDA

искусственный интеллект глубокое обучение алгоритм задняя часть

Автор: Ма Джун | Архитектор MegEngine

предисловие

Умножение матриц с одинарной точностью (SGEMM) — это случай, которого не может избежать почти каждый изучающий CUDA.Этот классический случай с интенсивными вычислениями может хорошо продемонстрировать методы оптимизации, обычно используемые в программировании на GPU, и возможность написать эффективное ядро ​​SGEMM. отличный тестовый вопрос, который отражает понимание программистом CUDA архитектуры графического процессора. В этой статье подробно представлены методы оптимизации CUDA SGEMM, подходящие для серьезного чтения."Руководство по программированию на CUDA C++", Студенты, у которых есть определенные основы программирования CUDA, могут прочитать его, надеясь вдохновить студентов, стремящихся к максимальной производительности.

Подробное объяснение методов оптимизации умножения матриц CUDA

Анализ реализации Naive: в чем отличие?

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

__global__ void matrixMul(const float *A, const float *B, float *C, 
                          int M, int N, int K) {
    int tx = blockIdx.x * blockDim.x + threadIdx.x;
    int ty = blockIdx.y * blockDim.y + threadIdx.y;
    if(ty < M && tx < N) {
        float c = 0;
        for(int i = 0; i < K; ++i){
            c += A[ty * K + i] * B[i * N + tx];
        }
        C[ty * N + tx] = c;
    }
}

Конечно, такое Наивное Ядро — это не то, на что рассчитывает автор, потому что производительность этого Ядра в принципе можно сделать вывод, что она не составляет даже 1/10 от производительности cublas, что явно не отвечает нашему стремлению к высокой производительности. Так что же не так с реализацией этого Наива?

Анализируя код, мы видим, что перед вычислением FMA (умножение и накопление) необходимо один раз прочитать A и B. FMA обычно требуется всего несколько циклов, и много времени тратится на доступ к памяти. Возможно, студенты с активным мышлением сразу подумали, что матрицы A и B можно переместить в Shared Memory (память на кристалле с малой задержкой в ​​SM, разделяемая потоками в блоке, со схемой структуры памяти NVIDIA GPU), чтобы уменьшить доступ к памяти. overhead, Это действительно хорошая идея, но она может лишь снизить стоимость доступа к памяти с сотен тактов до десятков тактов, а сути проблемы не меняет. Суть проблемы в том, что основной цикл состоит из двух инструкций Load и одной инструкции FMA.Инструкции расчета составляют только 1/3 от общего количества.Коэффициент доступа к памяти вычислений слишком низкий, что в конечном итоге приводит к обращению к памяти задержка не может быть скрыта, что приводит к неудовлетворительной производительности.

Давайте откроем наши умы.Если поток вычисляет не только один результат, но результаты 4x4 и использует оптимизацию Shared Memory, как будет выглядеть Hot Loop?Псевдокод выглядит следующим образом:

    float c[4][4] = {{0}};
    float a_reg[4];
    float b_reg[4];
    for(int i = 0; i < K; i += TILE_K){
        __syncthreads();
        // transfer tile from global mem to shared mem
        load_gmem_tile_to_smem(A, i, smemA);
        load_gmem_tile_to_smem(B, i, smemB);
        __syncthreads();
    #pragma unroll
        for(int j = 0; j < TILE_K; ++j) {
            // load tile from shared mem to register 
            load_smem_tile_to_reg(smemA, j, a_reg);
            load_smem_tile_to_reg(smemB, j, b_reg);
            // compute matrix multiply accumulate 4x4
            mma4x4(a_reg, b_reg, c)；
        }
    }

Анализ показывает, что для чтения из smemA в регистр a_reg требуется 4 операции доступа к памяти Аналогично для B доля вычислительных инструкций доступа к памяти основного тела стала 16/8.По сравнению с предыдущей ситуацией, вычислительные инструкции составляют гораздо больше улучшен. Достаточно большой коэффициент доступа к вычислительной памяти может улучшить использование вычислительных блоков и скрыть задержку доступа к памяти. Мы можем дополнительно улучшить коэффициент доступа к вычислительной памяти, чтобы производительность ядра была близка к теоретическому пику.

Разделение матрицы и распределение ресурсов

Очевидно, что мы не можем использовать только один блок для вычисления очень большой матрицы, что приведет к большому количеству простоев SM (Streaming Multiprocessor), который требует, чтобы матрица вычислялась блоками, как показано на следующем рисунке:

Различные размеры блоков имеют свои преимущества и недостатки в приложениях для умножения матриц различной формы.В этой статье в качестве примера выбраны блоки размером 128x128.

Из предыдущего раздела мы видим, что улучшение коэффициента доступа к вычислительной памяти дает большие преимущества. Можно ли бесконечно улучшать коэффициент доступа к вычислительной памяти? Ответ — нет. Потому что для улучшения коэффициента доступа к вычислениям одному потоку необходимо вычислять больший блок, для чего требуется больше регистров, но количество регистров ограничено. Взяв в качестве примера GPU архитектуры Turing, общее количество регистров одного SM составляет 65536. Из-за ограничения кодирования инструкций максимальное количество регистров, которые может использовать один поток, равно 255, и чем больше регистров, тем меньше используется, тем лучше. Здесь необходимо ввести понятие занятости (коэффициент занятости).Занятость относится к отношению количества активных деформаций (варпов) в каждом SM к максимальному количеству одновременных деформаций.Высокая занятость не обязательно означает высокую производительность, но она можно выполнить, переключив Warp, чтобы он играл роль в сокрытии задержки. Количество Active Warps в каждом SM зависит от количества ресурсов, используемых блоком, в частности от количества регистров, используемых каждым потоком, и объема общей памяти. Занятость доступна с помощью инструмента CUDA_Occupancy_Calculator.xls, входящего в состав CUDA Toolkit.

Рассмотрим блок для вычисления блоков 128x128, если каждый поток вычисляет 128 результатов, требуемый размер блока равен 128, а одному потоку требуется 128 регистров для хранения результатов расчета, плюс требуемый Gmem для Smem, Smem для Reg и т. д. требуется не менее 180 регистров.Рассчитайте Occupany, и мы увидим, что количество Active Warp в это время составляет всего 8, а Occupany составляет 25%, если размер блока установлен на 256, каждый поток должен вычислить только 64 результата, настроить регистры и Использование Shared Memory и наблюдение Occupany показывает, что если каждый поток использует только 128 регистров, использование Shared Memory в блоке ограничено 32 КБ, а количество Active Warp может достигать 16, что является лучшим выбором. :

А коэффициент доступа к памяти для расчета конфигурации в это время может достигать 64/4 (с использованием векторного чтения), что достаточно, чтобы скрыть задержку доступа к памяти.

Экстремальная оптимизация доступа к памяти

В нормальных условиях после выбора соответствующей конфигурации блочных ресурсов, использования общей памяти для уменьшения задержки доступа к памяти и выполнения развертывания цикла производительность ядра SGEMM достигла хорошего уровня (80% cublas), но это не наш путь. конечная точка. Во-первых, мы можем использовать инструкцию чтения вектораLDS.128Оптимизировать доступ к общей памяти (соответствует типу данных float4), что может значительно сократить количество инструкций доступа к памяти и еще больше улучшить коэффициент доступа к памяти вычислений.Поэтому нам нужно сделать транспонирование перед сохранением матрицы A в smemA:

При этом наше ядро ​​вычисляет блоки 128х128 для потоков 256. Чтобы можно было объединить доступ к Shared Memory, разделим 256 потоков на два измерения, пусть:

    int tx = threadIdx.x % 16;
    int ty = threadIdx.x / 16;

И прочитайте данные в векторе общей памяти следующим образом:

Наконец, один поток вычисляет результаты 2x2 4x4, и схема результатов показана на рисунке:

И с помощью микротеста можно обнаружить, что глобальная задержка доступа к памяти Turing (Tesla T4) составляет около 300 циклов, а задержка доступа к памяти Shared Memory около 30 циклов.Необходимо полностью использовать идею ** Предварительная выборка **, чтобы скрыть чтение глобальной памяти.Задержка доступа к памяти при входе в промежуточный регистр, записи блока данных из глобальной памяти в общую память и чтении блока данных из общей памяти, чтобы предотвратить вычислительного блока из-за слишком долгого бездействия из-за остановки, окончательный псевдокод выглядит следующим образом:

    #define TILE_K 16
    __shared__ float4 smemA[2][TILE_K * 128 / 4];
    __shared__ float4 smemB[2][TILE_K * 128 / 4];
    float4 c[8][2] = {{make_float4(0.f, 0.f, 0.f, 0.f)}};
    float4 ldg_a_reg[2];
    float4 ldg_b_reg[2];
    float4 a_reg[2][2];
    float4 b_reg[2][2];
    
    // transfer first tile from global mem to shared mem
    load_gmem_tile_to_reg(A, 0, ldg_a_reg);
    load_gmem_tile_to_reg(B, 0, ldg_b_reg);
    
    store_reg_to_smem_tile_transpose(ldg_a_reg, 0, smemA[0]);
    store_reg_to_smem_tile(ldg_b_reg, 0, smemB[0]);
    __syncthreads();

    // load first tile from shared mem to register 
    load_smem_tile_to_reg(smemA[0], 0, a_reg[0]);
    load_smem_tile_to_reg(smemB[0], 0, b_reg[0]);

    int write_stage_idx = 1; //ping pong switch
    do {
        i += TILE_K;
        // load next tile from global mem
        load_gmem_tile_to_reg(A, i, ldg_a_reg);
        load_gmem_tile_to_reg(B, i, ldg_b_reg);
        
        int load_stage_idx = write_stage_idx ^ 1;
        
    #pragma unroll
        for(int j = 0; j < TILE_K - 1; ++j) {
            // load next tile from shared mem to register 
            load_smem_tile_to_reg(smemA[load_stage_idx], j + 1, a_reg[(j + 1) % 2]);
            load_smem_tile_to_reg(smemB[load_stage_idx], j + 1, b_reg[(j + 1) % 2]);
            // compute matrix multiply accumulate 8x8
            mma8x8(a_reg[j % 2], b_reg[j % 2], c)；
        }
        
        if(i < K) {
            // store next tile to shared mem
            store_reg_to_smem_tile_transpose(ldg_a_reg, 0, smemA[write_stage_idx]);
            store_reg_to_smem_tile(ldg_b_reg, 0, smemB[write_stage_idx]);
            // use double buffer, only need one sync
            __syncthreads();
            // switch
            write_stage_idx ^= 1;
        }
        
        // load first tile from shared mem to register of next iter
        load_smem_tile_to_reg(smemA[load_stage_idx ^ 1], 0, a_reg[0]);
        load_smem_tile_to_reg(smemB[load_stage_idx ^ 1], 0, b_reg[0]);
        // compute last tile mma 8x8
        mma8x8(a_reg[1], b_reg[1], c)；
    } while (i < K);
    
    store_c(c, C);

Примечание: лениво предполагать, что M, N и K кратны 4. Если они не кратны 4, глобальная память не может использовать float4 для чтения, и результат не может быть записан обратно с float4, и для того, чтобы слияние и обратная запись, необходимо передать разделяемую память подкачки результатов в варпах, чтобы гарантировать, что каждый варп выполняет инструкцию Store для обратной записи в непрерывное пространство памяти.

На данный момент у нас есть полностью оптимизированное ядро ​​SGEMM. Кроме того, добавлен графический процессор Ampere.LDGSTSИнструкции, процесс блоков данных из глобальной памяти в общую память не должен проходить через промежуточные регистры, что может дополнительно оптимизировать производительность SGEMM.

Сравнение производительности

Чтобы cublas не выбирал ядро ​​разделения K, мы зафиксировали K равным 1024, взяли M, N = 2048, 4096, 8192 и 16384 в качестве тестовых случаев и сравнили производительность вышеупомянутого SGEMM Kernel и cublas (тестовый GPU это Tesla T4, заблокированная частота ядра 1100):

Видно, что реализованное ядро ​​SGEMM достигает средней производительности 97,5% cublas.

Помимо cublas: настройка ядра с помощью SASS

На этом этапе у некоторых студентов может возникнуть вопрос.Мы вроде бы использовали все методы оптимизации, которые только могли придумать.Почему написанное ядро ​​CUDA C все еще далеко от cublas?Ответ в том, что используется большое количество ядер Частично это не ядро ​​CUDA, скомпилированное nvcc, а глубоко настроенная версия, написанная на языке ассемблера (Shader Assembly, SASS) графического процессора NVIDIA. Хотя компилятор nvcc постоянно совершенствуется, особенно nvcc в CUDA 11, разрыв между скомпилированным ядром и оптимизированной версией, собранной вручную, значительно сократился, но по-прежнему невозможно полностью избежать влияния конфликта банка регистров и Повторное использование, которое полностью использует регистры Cache (эти две концепции будут подробно описаны ниже), так что пробел все же существует. Даже такой язык псевдоассемблера, как PTX, не может точно контролировать распределение регистров и сталкивается с той же дилеммой, что и CUDA C. Следовательно, чтобы полностью использовать предел производительности графического процессора, требуется точный контроль над инструкциями и регистрами графического процессора, который должен выполняться с помощью SASS на собственном языке ассемблера графического процессора. В этой области было проведено много исследований, таких как углубленное изучение архитектуры NV GPU компанией Citadel.Dissecting the NVidia XXX GPU architecture via microbenchmarkingСерия документов, в которых систематически тестируются, анализируются и подводятся итоги базовой архитектуры.Хотя некоторые выводы могут быть неточными, они, как правило, имеют высокую справочную ценность. В то же время многие ассемблеры с открытым исходным кодом, такие какKeplerAs,maxas(наиболее зрелые и далеко идущие),turingasиCuAssemblerИ ряд ассемблеров SASS с открытым исходным кодом, позволяющих использовать SASS для написания высокопроизводительного ядра.

Регистрация Конфликт банков

Мы знаем, что в Shared Memory есть конфликт банков, и конфликт регистров в банках — аналогичная концепция. Каждый SM графического процессора NVIDIA имеет независимый файл регистров, и файл регистров разделен на несколько банков.В качестве примера возьмем Максвелла, если из одного и того же банка поступает более двух исходных регистров, требуемых инструкцией, будет генерироваться конфликт, и инструкция будет эквивалентна повторному запуску, потере цикла. Номер банка регистрового файла Maxwell/Pascal равен 4, аid%4То есть банк, которому принадлежит регистр (например, R0 принадлежит Банку 0, R5 принадлежит Банку 1),FFMA R1, R0, R4, R1Такая инструкция вызовет конфликт банка регистров. Архитектура Тьюринга была улучшена.Регистровый файл разделен на 2 банка, и каждый банк имеет 2 порта.Если три исходных идентификатора регистра не совпадают с четностью, конфликта не будет, что значительно облегчает конфликт регистрового банка. .

Чтобы смягчить конфликт банков регистров, ядро ​​Maxwell SGEMM SASS в maxas делает деликатное распределение регистров, участвующих в расчете FFMA (см.Документация SGEMM для maxas),Как показано ниже:

После разумного расположения C конфликт банка регистров значительно уменьшается, но его все же нельзя полностью избежать (например, часть, отмеченная черным прямоугольником на приведенном выше рисунке, регистры, используемые A/B, будут генерировать конфликт банков), эта часть конфликта требует использования регистров Reuse для устранения.

Register Reuse

Повторное использование регистров — это механизм, введенный NVIDIA в начале Maxwell для решения проблемы конфликта банков регистров.Nvidia добавляет кэш повторного использования регистра в модуль коллектора, который считывает операнд инструкции. Кэш повторного использования доступен только для чтения, и инструкция получает информацию о том, проходит ли операнд через этот кеш, с помощью управляющего кода инструкции (вики управляющего кода для maxasСуществует подробное введение) как указано, используя cuobjdump для дизассемблирования некоторых ядер, можно обнаружить, что некоторые регистры имеют.reuseФлаг , означающий, что регистр берет значение из кэша повторного использования вместо файла регистров, тем самым устраняя конфликт банка регистров:

# Maxwell GPU
FFMA R2, R64.reuse, R73, R2; # R64 进入 Reuse Cache
FFMA R3, R64.reuse, R72, R3; # R64 从 Reuse Cache 中获取，避免与 R72 冲突

но использовать.reuseДолжны быть соблюдены определенные условия (регистр не будет установлен до перезаписи)..reuse), установка флага повторного использования без разбора может привести к получению исторических значений, что приведет к ошибкам в расчетах.По мнению автора,.reuseБольше похоже на флаг для хранения значения этого регистра в кэше повторного использования. Компиляция ядра CUDA с nvcc также будет использовать кэш повторного использования, чтобы избежать некоторых конфликтов банка регистров, но из-за распределения регистров и расположения инструкций коэффициент использования повторного использования невелик, разберите ядро ​​SGEMM, которое мы только что написали, все FFMA в основном цикле. к статистике инструкций можно обнаружить, что кэш повторного использования достигает только около 20%, в то время как ядро ​​SASS maxas спроектировано так, что коэффициент использования повторного использования может достигать 49%.

Наконец, производительность ядра SGEMM, точно настроенного с помощью SASS, может полностью превзойти cublas Заинтересованные студенты могут скомпилировать ядро ​​SGEMM в maxas и протестировать его на графическом процессоре Maxwell или Pascal. Наконец, хотя использование SASS может полностью использовать производительность GPU, оно сталкивается с тремя основными проблемами: 1. Сторонние сборщики NV GPU полагаются на обратное исследование архитектуры GPU, и могут быть неизвестные ошибки, поскольку они не изучили все детали базового оборудования; 2. Трудно разрабатывать и отлаживать сборочное ядро ​​3. ISA (набор инструкций) каждого поколения NV GPU отличается, и соответствующий ассемблер и сборочное ядро ​​необходимо постоянно разрабатывать. Из-за наличия этих серьезных проблем использование SASS для написания ядра занимает много времени и сил, и если нет необходимости добиваться максимальной производительности, не рекомендуется пробовать его легко.

Расширения GEMM: оптимизация операций свертки

Все мы знаем, что оптимизированная операция свертки может быть реализована путем сопоставления свертки с умножением матриц через im2col.Для приведенного выше ядра SGEMM слегка изменен только процесс перемещения данных из глобальной памяти в общую память, а отображение соответствующая позиция становится отображением im2col, ядро ​​SGEMM преобразуется в ядро, которое вычисляет Conv, что является неявным алгоритмом Gemm операции свертки cudnn. В процессе im2col, если смещение указателя вычисляется напрямую, будет введено большое количество операций целочисленного деления и остатка, что требует больших накладных расходов, поэтому смещение адреса может быть предварительно рассчитано на стороне хоста. , переданный ядру в качестве параметра, вы можете читать из постоянной памяти при необходимости, избегать целочисленного деления и остатка, а также реализовывать Implicit Precomp Gemm.Подробности см. в нашей предыдущей статье.Принцип реализации оператора свертки MegEngine TensorCore. Руководствуясь этой идеей, мы реализуем эффективные операции свертки int8/int4 на основе cutlass (MegEngine cutlass), и в непрерывной итерации вы можете попробовать его.

Суммировать

В этой статье подробно рассказывается о том, как написать высокоэффективное ядро ​​CUDA SGEMM, и рассказывается об использовании программирования SASS, который является экстремальным методом оптимизации производительности, и немного расширяет идею оптимизации операций свертки с помощью Implicit Gemm, надеясь дать тем, кто интересуется экстремальным майнингом Студенты производительности оборудования должны быть вдохновлены. Наконец, вы можете присоединиться к команде MegEngine, чтобы оптимизировать обучение и вывод глубокого обучения.

Бот сказал: Я закончил читать статью~ Вы заинтересованы в том, чтобы присоединиться к разработке фреймворка глубокого обучения? Команда Megengine объявляет набор! С нетерпением жду вашего присоединения~ Возобновление доставки или подробную информацию можете добавить в WeChat: duoduo715495

[Инженер по разработке фреймворков (C++)]

описание работы:

  1. Отвечает за проектирование, эволюцию, внедрение, обслуживание и оптимизацию MegEngine, основной структуры глубины Megvii.
  2. Оптимизация производительности MegEngine на различных вычислительных платформах (CUDA/Arm/x86 и т.д.)
  3. Продолжайте изучать передовые методы оптимизации для улучшения фреймворков глубокого обучения (такие как оптимизация графов, автоматическая генерация кода, сверхнизкобитное квантование, разрежение и т. д.).

Требования к навыкам:

  1. Опыт оптимизации производительности от 1 до 3 лет (X86, CUDA, ARM, OpenCL и т. д.)
  2. Иметь хорошие навыки работы со структурами данных и алгоритмами, а также уметь использовать язык C++ для написания более сложных алгоритмов.
  3. Базовое понимание глубокого обучения и фреймворков глубокого обучения (Caffe, Tensorflow, PyTorch и т. д.) приветствуется.
  4. Приветствуется опыт проектирования сложных систем
Категории