Как да оптимизирате изключително CUDA HGEMM с Tensor Core?

1 Фон

GEMM (Общо матрично умножение) умножението на матрици е един от най-често използваните и отнемащи много време алгоритми в дълбокото обучение, особено в областите на CNN, RNN, трансформатор и други области. В тези полета голям брой операции за умножение на матрици трябва да бъдат изчислени и обработени бързо. Следователно ефективното внедряване на умножение на матрици е от решаващо значение за изпълнението и точността на задачите за дълбоко обучение.

HGEMM (общо матрично умножение с половин прецизност) с поддръжката на хардуерния модул Tensor Core на Nvidia GPU може значително да увеличи скоростта на изчисление, като същевременно поддържа точност. Получените предимства в производителността могат значително да подобрят дълбокото обучение. Скоростта на изпълнение на задачите за извод и обучение.

Появата на Tensor Core доведе до пробив в оптимизирането на HGEMM. Използването на Tensor Core за оптимизиране на HGEMM на Nvidia GPU се превърна в една от горещите изследователски области в GPU ускорените изчисления през последните години.

2 Резултат

Тази статия използва основно ръкописен WMMA API и MMA PTX CUDA HGEMM ядро ​​за извикване на Tensor Core, след което извършва настройка на производителността и я сравнява с производителността на cublas Tensor Core. Чрез изследване на различни методи за подреждане на матрици и оптимизация, в момента производителността между 256 ~ 16384 измерения не е по-ниска от 95% от производителността на cublas и производителността надвишава cublas в много сценарии. Кодът е с отворен код в cuda_hgemm.

2.1 Условия на теста

  • HGEMM: C (M * N, половина, голям ред) = A (M * K, половина, голям ред) * B (K * N, половина, голям ред)
  • CUDA: 11.3
  • Графичен процесор: RTX3090、RTX A6000

2.2 Спецификации на оборудването

Спецификациите на устройствата на RTX3090 и RTX A6000 са както следва.

2.3 RTX3090

2.4 RTX A6000

3 Matrix Tiling

3.1 Подреждане на блокове

За входна матрица с фиксиран размер, обикновено колкото по-голям е размерът на подреждането на блока, толкова по-голяма е сумата на изчисление в рамките на един блок, но толкова по-малък е броят на блоковете, които трябва да бъдат изчислени паралелно. Това е компромис между сумата за изчисляване на размерите на блока и степента на паралелизъм.

Най-общо казано, за входни матрици с различни размери трябва да се приемат различни стратегии за блокиране, за да се постигне по-добър баланс между обема на изчисление на размерите на блока и паралелизма. Колкото по-малък е размерът на входната матрица, толкова по-малък е размерът на подредбата на блока. Колкото по-голям е размерът на входната матрица, толкова по-голям е размерът на подредбата на блока. От друга страна, в съчетание с ограниченията на хардуерните спецификации, размерът на блоковото подреждане обикновено е комбинация между 32, 64, 128 и 256.

3.2 Изкривяване на плочки

След като определите размера на блоковото подреждане, трябва да продължите да определяте размера на изкривяването. Най-общо казано, колкото по-голям е размерът на изкривяването, толкова по-голяма е сумата за изчисление в рамките на едно изкривяване, но толкова по-малък е броят на изкривяванията, които трябва да бъдат изчислени паралелно. Това е компромисът между количеството на изчислението на размерите на деформацията и паралелизма.

Ако броят на деформациите е твърде малък, заетостта на деформациите ще бъде твърде ниска, т.е. броят на допустимите деформации в планировчика е твърде малък, което може сериозно да повлияе на цикъла на издаване на инструкции вътре в планировчика и в крайна сметка да повлияе на производителността на ядро. От друга страна, в съчетание с ограниченията на хардуерните спецификации, броят на изкривяванията обикновено е 4, 8, 16 и комбиниран с размера на блоковото подреждане, може да се определи оптималният размер на изкривяването на подложките, обикновено комбинация между 8, 16, 32, 64 и 128.

4 Оптимизиране на достъпа до паметта

4.1 Обединен достъп до паметта

Данните на A матрица и B матрица се зареждат от глобалната памет в споделената памет в 16 байта (int4 или float4), доколкото е възможно, за да се намали броят на инструкциите за зареждане, да се подобри коефициентът на достъп до паметта за изчисление на инструкциите и да се намали забавянето на инструкцията издаване.

4.2 Повторно използване на данни

Както е показано на фигурата по-долу, когато block изчислява зеления блок на матрицата C, той се нуждае от целия светлосин блок от матрицата A и целия светложълт блок от матрицата B. За различни изкривявания в блока, когато различни изкривявания изчисляват един и същ ред или колона, данните от матрица A или матрица B, които трябва да бъдат заредени, са еднакви, т.е. може да има многократно зареждане на данни от матрица A или матрица B вътре в блока.

Решението е първо да се заредят данните от A матрицата или B матрицата в споделената памет, за да се постигне споделяне в рамките на блока, а след това деформацията в блока зарежда данните от споделената памет в регистъра за изчисление. Тъй като латентността на достъпа на споделената памет е много по-малка от тази на глобалната памет, тя може значително да облекчи пречките в честотната лента на ядрото.

4.3 Асинхронно копие

cp.async е неблокираща PTX инструкция за асинхронно копиране, която може да копира данни от глобалната памет в споделената памет. Той не само има определена стратегия за контрол на кеша, но също така може независимо да контролира множество групи инструкции за излъчване. До известна степен обхватът на оптимизация на миграцията на данни и стратегията за паралелна оптимизация на тръбопровода са разширени.

cp.async.ca.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], cp-size{, src-size}{, cache-policy} ;
cp.async.cg.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], 16{, src-size}{, cache-policy} ;
cp.async.ca.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], cp-size{, ignore-src}{, cache-policy} ;
cp.async.cg.shared{::cta}.global{.level::cache_hint}{.level::prefetch_size}
                         [dst], [src], 16{, ignore-src}{, cache-policy} ;

.level::cache_hint =     { .L2::cache_hint }
.level::prefetch_size =  { .L2::64B, .L2::128B, .L2::256B }
cp-size =                { 4, 8, 16 }

4.4 Премахване на банков конфликт

Най-общо казано, споделената памет на графичния процесор на Nvidia е разделена на 32 банки, всяка банка по подразбиране е 4 байта, общо 128 байта. Обикновено, когато нишка в един варп има достъп до различни банки споделена памет, се изисква само една заявка за видео памет. Ако нишка в едно изкривяване има достъп до различни полета в една и съща банка споделена памет, ще възникне конфликт на банката и заявката за памет ще стане серийно изпълнение, значително увеличавайки забавянето на достъпа до паметта.

Следователно елиминирането на банковия конфликт на споделената памет е почти основната операция на цялата оптимизация на ядрото на CUDA, обикновено използвайки подложки или пермутирани методи.

(1) Подложка

Както е показано на фигурата по-долу, ако приемем, че синята зона е споделената област на паметта, до която има достъп от един варп, ако се осъществи директен достъп, ще възникне очевиден банков конфликт.

Ако при кандидатстване за споделена памет, всеки ред важи за 4 допълнителни жълти области на банката. По това време споделената област на паметта, до която има достъп от един варп, се оказва в различна банка, като по този начин се избягва възникването на банков конфликт.

Методът за подпълване може ефективно да реши проблема с банковия конфликт на споделената памет, но ще се прилага за допълнителна споделена памет и ще увеличи използването на споделена памет вътре в блока.

(2) Пермутирани

Може ли проблемът с банковия конфликт да бъде разрешен без кандидатстване за допълнителна споделена памет? Отговорът е да.

Както е показано на фигурата по-долу, ако приемем, че матрица А първо се зарежда от глобалната памет в споделената памет, една и съща цветна област представлява транзакция за достъп до паметта на единична деформация. Ако съответното местоположение е съхранено и след това инструкцията ldmatrix PTX се използва за зареждането му от споделената памет в регистъра, ще възникне очевиден банков конфликт. Ако извършите пермутирана операция при зареждане в споделена памет и след това използвате инструкцията ldmatrix PTX за зареждане от споделена памет в регистър, банковият конфликт ще бъде избегнат.

4.5 Подобряване на скоростта на попадение в L2 кеша

Проблемът с блоковото разделяне беше споменат по-рано. Ако блоковото разделяне е завършено, как трябва да се проектира изчислителната последователност на тези блокови разделения по време на действителните изчисления. Най-прекият начин е да се смята по ред. Този метод на изчисление ще създаде проблем и ще намали значително скоростта на попадение в L2 кеша.

Най-общо казано, за блока на матрицата C в същия ред, необходимите блокови данни на матрицата A са същите. По подобен начин, за блока на C матрицата на същата колона, необходимите блокови данни на B матрицата са същите. Ако се изчисли по редове, за матрица A, данните са едни и същи и процентът на попадения в кеша L2 е много висок. Но за матрица B данните са различни и коефициентът на попадение в L2 кеша е много нисък. Взети заедно, като се има предвид капацитетът на L2 кеша, процентът на попадения в L2 кеша няма да бъде твърде висок.

Следователно може да се приеме методът за изчисление на swizzle, както е показано на фигурата по-долу, т.е. изчислението за „отглеждане на едър рогат добитък“, като се вземе предвид процентът на попадения в L2 кеша на матрицата A и B матрицата и се подобрява общото попадение в L2 кеша процент. В същото време размерът на блока и капацитетът на L2 кеша могат да се комбинират, за да се коригира размерът на стъпката на „отглеждане на едър рогат добитък“, за да се получи най-висок процент на попадения в L2 кеша.

4.6 Подобрете процента на повторно използване на регистъра

Проблемът с изкривяването на плочките беше споменат по-рано. Ако подреждането на основата е завършено, тогава как да проектирате изчислителната последователност на вътрешната плочка на основата (изчислителен размер на Tensor Core) по време на действителното изчисление. Тук изходният код на cutlass дава идея.

За устройства с изчислителни възможности от 8.0 и по-високи, приемете метода за изчисление „Right Left Right Left“, за да подобрите честотата на повторно използване на регистъра на матрицата A.

За устройства с изчислителни възможности под 8.0, приемете метода за изчисление „Надолу нагоре, надолу нагоре“, за да подобрите честотата на повторно използване на регистрите на B матрица.

5 Оптимизация на тръбопровода

Най-общо казано, процесът на изчисление на CUDA HGEMM е първо да копира блоковите данни на A матрица и B матрица от глобалната памет в споделена памет на партиди според K измерения и след това да копира данните в споделената памет в регистъра партида по партида за матрично умножение. Изчислете, докато всички данни бъдат завършени, изчислението на умножението на матрицата, схематичната диаграма е както следва.

5.1 Двоен буфер

Двойният буфер е метод за предварително извличане на данни и също е един от често използваните методи в оптимизацията на GEMM. Прилага се главно за два споделени буфера на паметта и зарежда данни последователно. Когато един от буферите зарежда данни, другият буфер е завършил зареждането на данни и може да извършва изчисления за умножение на матрици. По този начин изчисленията за зареждане на данни и умножение на матрици могат да бъдат паралелизирани и забавянето на зареждането на данни може да бъде скрито. Както е показано на фигурата по-долу, е показан работният процес на двоен буфер. Светлосиният е буферът за зареждане на данни, а тъмносиният е буферът за изчисление на умножението на матрицата. Двата тръбопровода могат да се изпълняват паралелно.

5.2 Етап

Въз основа на двоен буфер, помислете за този проблем. Ако в стъпка 1 зареждането на данни и изчислението на умножението на матрицата се изпълняват паралелно, но изчислението на умножението на матрицата се изпълнява първо и зареждането на данни не е завършено, тогава ще се появи изчислителната единица Tensor Core. Когато чака данни, т.е. SM ще бъде неактивен и общото използване на SM ще бъде ниско. Как да разрешим този проблем е как да попречим на изчислителната единица Tensor Core да чака данни.

Тъй като това е проблем с тясното място на честотната лента, могат да бъдат въведени повече споделени буфери на паметта. Зареждането на данни на множество буфери може да бъде завършено по време на предварително извличане, като по този начин се избягва ситуацията, при която изчислителната единица Tensor Core чака данни. Фигурата по-долу показва работния процес, когато етап е 3. Зареждането на данни на всяка стъпка е успоредно на изчислението на умножението на матрицата на две стъпки, като допълнително скрива забавянето на зареждането на данни. По същия начин, когато споделената памет е достатъчна, етапът може да продължи да се увеличава, докато закъснението в зареждането на данни може да бъде напълно скрито. Следователно изборът на етап зависи от честотната лента на устройството и изчислителната мощност на Tensor Core.

6 Други

6.1 Метод за оптимизация

Тази статия представя главно общите методи за оптимизация и някои специални методи за оптимизация на CUDA HGEMM. По-късно може да говорим за опита на оптимизация в детайли на определена точка на оптимизация. За различните версии на GPU и CUDA оптимизационните стратегии за постигане на оптимална производителност са различни.

6.2 Изходен код

Всички методи за оптимизация, използвани в тази статия, са с отворен код в cuda_hgemm, включително кода за внедряване на WMMA API и MMA PTX. Размерът на блоковата подложка (256*128) и размерът на изкривената подложка (64*64) са фиксирани и могат да бъдат анализирани заедно с изходния код в бъдеще.