Прийоми використання масочных регістрів в AVX512 коді

В процесорах компанії Intel на зміну AVX2 приходить новий набір інструкцій AVX512, в якому з'явилася концепція масочных регістрів. Автор цієї статті вже кілька років займається розробкою версії бібліотеки Intel Integrated Performance Primitives, оптимізованої для AVX512, і накопичив досить великий досвід використання AVX512 інструкцій з масками, який було вирішено об'єднати в одну окрему статтю, оскільки саме використання таких інструкцій з масками дозволяє спростити і прискорити код на додаток до прискорення від дворазового збільшення ширини регістрів.

Приклад 1. Застосовуємо інструкції з масками «до і після» основного циклу

Розглянемо функцію додавання двох зображень.
Лист 1.1
void add_ref( const float* pSrc1, int src1Step, const float* pSrc2, int src2Step,
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
for( i = 0; i < width; i++ ) {
pDst[i] = pSrc1[i] + pSrc2[i];
}
pSrc1 = pSrc1 + src1Step;
pSrc2 = pSrc2 + src2Step;
pDst = pDst + dstStep;
}
}


Код функції дуже простий і в принципі компілятор icl цілком здатний векторизовать даний код. Однак наша мета – продемонструвати використання масочных avx512 регістрів, тому напишемо першу і очевидну версію avx512 коду.
Лист 1.2
void add_avx512( const float* pSrc1, int src1Step, const float* pSrc2, int src2Step,
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h<height; h++ ) {
i = 0;
__m512 zmm0, zmm1;
for(i=i; i < (width&(~15)); i+=16 ) {
zmm0 = _mm512_loadu_ps(pSrc1+i );
zmm1 = _mm512_loadu_ps(pSrc2+i );
zmm0 = _mm512_add_ps(zmm0, zmm1);
_mm512_storeu_ps (pDst+i, zmm0 );
}
for(i=i; i < width; i++ ) {
pDst[i] = pSrc1[i] + pSrc2[i];
}
pSrc1 = pSrc1 + src1Step;
pSrc2 = pSrc2 + src2Step;
pDst = pDst + dstStep;
}
}


В даному коді додано цикл, який складав відразу по 16 елементів за ітерацію. Оскільки ширина зображення може бути довільною і не кратній 16, то в наступному циклі складаються елементи, що залишилися до кінця рядка. Число таких елементів може досягати 15-ти, і вплив цього циклу на загальну продуктивність може бути досить велика. Використання ж масочных регістрів дозволяє істотно прискорити обробку «хвоста».
Лист 1.3
__mmask16 len2mask[] = { 0x0000, 0x0001, 0x0003, 0x0007,
0x000F, 0x001F, 0x003F, 0x007F,
0x00FF, 0x01FF, 0x03FF, 0x07FF,
0x0FFF, 0x1FFF, 0x3FFF, 0x7FFF,
0xFFFF };
void add_avx512_2( float* pSrc1, int src1Step, const float* pSrc2, int src2Step,
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
i = 0;
__m512 zmm0, zmm1;
__mmask16 msk;
for(i=i; i < (width&(~15)); i+=16 ) {
zmm0 = _mm512_loadu_ps(pSrc1+i );
zmm1 = _mm512_loadu_ps(pSrc2+i );
zmm0 = _mm512_add_ps (zmm0, zmm1);
_mm512_storeu_ps (pDst+i, zmm0 );
}
msk = len2mask[width - i];
if(msk){
zmm0 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc1+i );
zmm1 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc2+i );
zmm0 = _mm512_add_ps (zmm0, zmm1);
_mm512_mask_storeu_ps (pDst+i, msk, zmm0 );
}
pSrc1 = pSrc1 + src1Step;
pSrc2 = pSrc2 + src2Step;
pDst = pDst + dstStep;
}
}


Масив len2mask використовується для перетворення числа у відповідне число бітів поспіль. І замість скалярного циклу, ми отримуємо лише один if, який в-принципі теж не обов'язковий, оскільки у разі маски складається з одних нулів, читання і запис не будуть здійснюватися.
Для досягнення максимальної продуктивності рекомендується вирівнювати дані на ширину кеш лінії, а завантаження здійснювати по вирівняній на ширину регістра адресою. У Skylake ширина кеш лінії по колишньому 64 байта, тому в нашому коді ми можемо додати таке вирівнювання по pDst знову ж таки з допомогою масочных операцій, але тільки до основного циклу.
Лист 1.4
void add_avx512_3(float* pSrc1, int src1Step, const float* pSrc2, int src2Step,
float* pDst, int dstStep, int width, int height)
{
int h, i, t; 
for( h = 0; h < height; h++ ) {
i = 0;
__m512 zmm0, zmm1;
__mmask16 msk;
t = (((( int)pDst) & (63)) >> 2);
msk = len2mask[t];
if(msk){
zmm0 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc1+i );
zmm1 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc2+i );
zmm0 = _mm512_add_ps (zmm0, zmm1);
_mm512_mask_storeu_ps (pDst+i, msk, zmm0 );
}
i += t;
for(i=i; i < (width&(~15)); i+=16 ) {
zmm0 = _mm512_loadu_ps(pSrc1+i );
zmm1 = _mm512_loadu_ps(pSrc2+i );
zmm0 = _mm512_add_ps (zmm0, zmm1);
_mm512_storeu_ps (pDst+i, zmm0 );
}
msk = len2mask[width - i];
if(msk){
zmm0 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc1+i );
zmm1 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc2+i );
zmm0 = _mm512_add_ps (zmm0, zmm1);
_mm512_mask_storeu_ps (pDst+i, msk, zmm0 );
}
pSrc1 = pSrc1 + src1Step;
pSrc2 = pSrc2 + src2Step;
pDst = pDst + dstStep;
}
}
<source>
</spoiler>
Розглянута функція add досить проста і тут цілком допустимо зробити завантаження по масці до і після основного циклу. Але в більш складних, наприклад цілочисельних з округленням до найближчого парним, алгоритмах кількість таких операцій може досягати десятків і така методика маски для вирівнювання і обробки хвоста може призвести до суттєвого збільшення розміру коду. А правити помилки доведеться в трьох різних місцях цього досить довгого коду. Використання ж масочных регістрів і дозволяє знизити розмір коду і спростити його подальшу підтримку. Аркуш. 1.5 демонструє цю ідею

<spoiler title=" Лист 1.5">
<source lang="C++">
#define min(a,b) ((a)<(b)?(a):(b))
void add_avx512_4( const float* pSrc1, int src1Step, const float* pSrc2, int src2Step,
float* pDst, int dstStep,
int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
i = 0;
__m512 zmm0, zmm1;
__mmask16 msk;
for(i=i; i < width; i+=16 ) {
msk = len2mask[min(16, width - i)];
zmm0 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc1+i );
zmm1 = _mm512_mask_loadu_ps(_mm512_setzero_ps(),msk,pSrc2+i );
zmm0 = _mm512_add_ps (zmm0, zmm1);
/*THE LONG PIPELINE HERE*/
_mm512_mask_storeu_ps (pDst+i, msk, zmm0 );
}
pSrc1 = pSrc1 + src1Step;
pSrc2 = pSrc2 + src2Step;
pDst = pDst + dstStep;
}
}


Код став виглядати коротше, і при достатньо довгій послідовності обчислень, витрати на обчислення маски в кожній ітерації на такому тлі можуть виявитися незначними.

Приклад 2. Маска як безпосередня частина реалізованого алгоритму.
В різних алгоритмах обробки зображень одним з вхідних параметрів може бути бінарна маска. Приміром, така маска використовується в операціях морфології зображень.
Лист 2.1
void morph_3x3_ref( const float* pSrc1, int src1Step, char* pMask, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
for(i=0; i < width; i++ ){
int x, y;
char* pm = pMask;
/*we assume that pMask is not zero total*/
float m = 3.402823466 e+38f;
float val = 0.0;
for (y = 0; y < 3;y++){
for (x = 0; x < 3; x++){
if (*pm){
val = pSrc1[i + src1Step*y + x];
if (val < m) {
m = val;
}
}
pm++;
}
}
pDst[i] = m;
}
pSrc1 = pSrc1 + src1Step;
pDst = pDst + dstStep;
}
}


Функція шукає мінімальний за величиною піксель всередині квадрата 3x3. Вхідна маска pMask являє собою масив з 3x3=9 байтів. Якщо значення байта не дорівнює нулю, то вхідний піксель з квадрата 3x3 бере участь в пошуку, якщо одно, то не бере участь.
Пишемо avx512 код.
Лист 2.2
void morph_3x3_avx512( const float* pSrc1, int src1Step, char* pMask, float* pDst, int dstStep, int width, int height)
{
int h, i; 
__mmask msk[9], tail_msk;
for ( i = 0; i < 9; i++){
msk[i] = (!pMask[i]) ? 0 : 0xFFFF;//create load mask
}
tail_msk = len2mask[width & 15]; //tail mask
for( h = 0; h < height; h++ ) {
__m512 zmm0, zmmM;
int x, y;

i = 0;

for(i=i; i < (width&(~15)); i+=16 ) {
zmmM = _mm512_set1_ps(3.402823466 e+38f);
for (y = 0; y < 3;y++){
for (x = 0; x < 3; x++){
zmm0 = _mm512_mask_loadu_ps(zmmM, msk[3*y+x], &pSrc1[i + src1Step*y + x]);
zmmM = _mm512_min_ps(zmm0, zmmM);
}
}
_mm512_storeu_ps (pDst+i, zmmM );
}

if(tail_msk) {
zmmM = _mm512_set1_ps(3.402823466 e+38f);
for (y = 0; y < 3;y++){
for (x = 0; x < 3; x++){
zmm0=_mm512_mask_loadu_ps(zmmM,msk[3*y+x]&tail_msk,&pSrc1[i+src1Step*y+x]);
zmmM=_mm512_min_ps(zmm0, zmmM);
}
}
_mm512_mask_storeu_ps (pDst+i, tail_msk, zmmM );
}

pSrc1 = pSrc1 + src1Step;
pDst = pDst + dstStep;
}
}


Спочатку формується масив масок __mmask msk[9]. Кожна маска виходить заміною байта з pMask на біт (0 на 0, всі інші значення на 1) і цей біт розмножується на 16 елементів. Основний цикл завантажує по цій масці 16 елементів. Причому, якщо елемент не бере участь в пошуку, то він і не буде завантажений. В даному коді також за допомогою масок обробляється і хвіст, ми просто виконуємо операцію & над маскою операції морфології і маскою хвоста.
Звичайно, продемонстрированый код не зовсім оптимальний, однак він демонструє саму ідею, а ускладнення коду призвело б до втрати наочності.

Приклад 3. Інструкції сімейства expand/compress
Ще в першому 256 bit наборі інструкцій AVX регістри були розділені бар'єром на 2 частини, звані lane. Більшість векторних інструкції незалежно обробляють ці частини. Виглядає це як 2 паралельні інструкції SSE. У avx512 регістр поділяється вже на чотири 128-бітні частини по чотири float/int елемента. А в багатьох алгоритмах обробки зображень використовуються три канали і об'єднання їх по 4 досить проблематично. В цьому випадку можна розглянути можливість використання інструкцій виду expand/compress

__m512 _mm512_mask_expandloadu_ps (__m512 src, __mmask16 k, const void* mem_addr)
void _mm512_mask_compressstoreu_ps (void* base_addr, __mmask16 k, __m512 a)

Інструкція _mm512_mask_expandloadu_ps завантажує з пам'яті безперервний блок float даних довжиною рівною кількістю одиничних біт у масці. Таким чином може бути завантажений блок довжиною від 0 до 16 елементів. У регістр-приймач дані поміщаються наступним чином. Починаючи з самого молодшого біта маски перевіряється, якщо біт дорівнює 1, то елемент з пам'яті записується в регістр, якщо 0, то переходимо до розгляду наступного біта і того ж самого елемента, рис. 3.1
Рис. 3.1 Демонстрація роботи _mm512_mask_expandloadu_ps

Видно, що область пам'яті як би «розтягується» (expand) по всьому 512 бітного регістру. Інструкція _mm512_mask_compressstoreu_ps працює в протилежну сторону — «стискає»(compress) регістр по масці і записує в безперервну область пам'яті.
Отже, припустимо нам необхідно перейти колірного простору RGB у XYZ.
Лист 3.1

void rgb_ref( const float* pSrc1, int src1Step, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
for(i=0; i < width; i++ ) {
pDst[3*i+0]=(0.412 f*pSrc1[3*i]+0.357 f*pSrc1[3*i+1]+0.180 f*pSrc1[3*i+2]);
pDst[3*i+1]=(0.212 f*pSrc1[3*i]+0.715 f*pSrc1[3*i+1]+0.072 f*pSrc1[3*i+2]);
pDst[3*i+2]=(0.019 f*pSrc1[3*i]+0.119 f*pSrc1[3*i+1]+0.950 f*pSrc1[3*i+2]);
if (pDst[3 * i + 2] < 0.0){
pDst[3 * i + 2] = 0.0;
}
if (pDst[3 * i + 2] > 1.0){
pDst[3 * i + 2] = 1.0;
}
}
pSrc1 = pSrc1 + src1Step;
pDst = pDst + dstStep;
}
}


Використовуючи expand/compress avx512 код може виглядати наступним чином.
Лист 3.2
void rgb_avx512( const float* pSrc1, int src1Step, float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
__m512 zmm0, zmm1, zmm2, zmm3;
i = 0;
for(i=i; i < (width&(~3)); i+=4 ) {
zmm0 = _mm512_mask_expandloadu_ps(_mm512_setzero_ps(), 0x7777, &pSrc1[3*i ]);
zmm1 = _mm512_mul_ps(_mm512_set4_ps(0.0 f, 0.019 f, 0.212 f, 0.412 f), _mm512_shuffle_ps(zmm0, zmm0, 0x00));
zmm2 = _mm512_mul_ps(_mm512_set4_ps(0.0 f, 0.119 f, 0.715 f, 0.357 f), _mm512_shuffle_ps(zmm0, zmm0, 0x55));
zmm3 = _mm512_mul_ps(_mm512_set4_ps(0.0 f, 0.950 f, 0.072 f, 0.180 f), _mm512_shuffle_ps(zmm0, zmm0, 0xAA));
zmm0 = _mm512_add_ps(zmm1, zmm2);
zmm0 = _mm512_add_ps(zmm0, zmm3);
zmm0 = _mm512_mask_max_ps(zmm0, 0x4444,_mm512_set1_ps(0.0 f), zmm0);
zmm0 = _mm512_mask_min_ps(zmm0, 0x4444,_mm512_set1_ps(1.0 f), zmm0);
_mm512_mask_compressstoreu_ps (pDst+3*i, 0x7777, zmm0 );
}

pSrc1 = pSrc1 + src1Step;
pDst = pDst + dstStep;
}
}


З допомогою _mm512_mask_expandloadu_ps ми поміщаємо 4 пікселя в різні lane, після чого послідовно формуємо r3r3r3r3… r0r0r0, ...g0g0g0, b0b0b0 і множимо на коефіцієнти перетворення. Для перевірки переповнення також використовуються операції з масками _mm512_mask_max/min_ps. Запис перетворених даних назад в пам'ять виконується командою _mm512_mask_compressstoreu_ps. У цій функції для обробки хвоста також можна використовувати маски, в залежності від довжини хвоста.

Приклад 4. Векторизація розгалуження
Якщо Ви дочитали до цього місця, то вже підготовлені до найбільш, на мій погляд, цікавою області застосування масочных регістрів. Це векторизація циклів з умовами. Мова йде про якійсь подобі предикатних регістрів, наявних в процесорах сімейства Itanium. Розглянемо просту функцію, у якої є if всередині циклу for.
Лист 4.1

#define ABS(A) (A)>=0.0 f?(A):(-(A))
void avg_ref( float* pSrc1, int src1Step, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
for(i=0; i < width; i++ ) {
float dv, dh;
float values = pSrc1[src1Step*(h - 1) + i ];
float valD = pSrc1[src1Step*(h + 1) + i ];
float valL = pSrc1[src1Step*( h ) + (i-1)];
float valR = pSrc1[src1Step*( h ) + (i+1)];

dv = ABS(values - valD);
dh = ABS(valL - valR);
if(dv<=dh){
pDst[i] = (values + valD) * 0.5 f; //A branch
} else {
pDst[i] = (valL + valR) * 0.5 f; //B branch
}
}
pDst = pDst + dstStep;
}
}


Функція здійснює інтерполяцію за сусіднім горизонтальним або вертикальнам елементів, між якими різниця мінімальна. Головне тут те, що всередині циклу йде поділу на дві гілки, коли if-умова истенно, і коли ні. Але, ми можемо обчислити на avx512 регістрах значення в тому і іншому випадку, а потім об'єднати їх по масці.
Лист 4.2
void avg_avx512( const float* pSrc1, int src1Step, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
i = 0;
for(i=i; i < (width&(~15)); i+=16 ) {
__m512 zvU, zvD, zvL, zvR;
__m512 zdV, zdH;
__m512 zavgV, zavgH, zavg;
__mmask16 mskV;

zvU = _mm512_loadu_ps(pSrc1+src1Step*(h - 1) + i );
zvD = _mm512_loadu_ps(pSrc1+src1Step*(h + 1) + i );
zvL = _mm512_loadu_ps(pSrc1+src1Step*( h ) + (i-1));
zvR = _mm512_loadu_ps(pSrc1+src1Step*( h ) + (i+1));

zdV = _mm512_sub_ps(zvU, zvD);
zdH = _mm512_sub_ps(zvL, zvR);
zdV = _mm512_abs_ps(zdV);
zdH = _mm512_abs_ps(zdH);

mskV = _mm512_cmp_ps_mask(zdV, zdH, _CMP_LE_OS);
zavgV = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvU, zvD));
zavgH = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvL, zvR));
zavg = _mm512_mask_or_ps(zavgH, mskV, zavgV, zavgV);
_mm512_storeu_ps(pDst + i, zavg);
}
//remainder skipped
pDst = pDst + dstStep;
}
}


Таких if реалізується в алгоритмі може бути кілька, для них можна завести нові маски до тих пір поки код буде швидше скалярного.
Насправді, компілятор icc цілком здатний векторизовать код 4.1 Для цього достатньо всього лише додати ключове слово restrict до вказівників pSrc1 і pDst і ключ –Qrestrict.
void avg_ref( float* restrict pSrc1, int src1Step, float* restrict pDst, int dstStep, int width, int height)

Нагадаю, що модифікатор restrict вказує компілятору, що доступ до об'єкта здійснюється тільки через цей покажчик і таким чином вектора pSrc1 і pDst не перетинаються, що і робить можливою векторизацію.
Вимірювання на внутрішньому симуляторі CPU з поддержкоий avx512 показують, що продуктивність практично дорівнює продуктивності нашого avx512 коду. Т. е компілятор теж вміє ефективно використовувати маски

Здавалося б, на цьому все. Але подивимося, що буде якщо ми злегка модицифируем нашу функцію і внесемо в неї залежність між ітераціями.
Лист 4.3
void avg_ref( float* restrict pSrc1, int src1Step, 
float* restrict pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
int t = 1;
for(i=0; i < width; i++ ) {
float dv, dh;
float values = pSrc1[src1Step*(h - 1) + i ];
float valD = pSrc1[src1Step*(h + 1) + i ];
float valL = pSrc1[src1Step*( h ) + (i-1)];
float valR = pSrc1[src1Step*( h ) + (i+1)];

dv = ABS(values - valD);
dh = ABS(valL - valR);
if(dv<dh){
pDst[i] = (values + valD) * 0.5 f;
t = 1;
} else if (dv>dh){
pDst[i] = (valL + valR) * 0.5 f;
t = 0;
} else if (t == 1) {
pDst[i] = (values + valD) * 0.5 f;
} else {
pDst[i] = (valL + valR) * 0.5 f;
}
}
pDst = pDst + dstStep;
}
}


Функція також здійснює інтерполяцію по сусідніх пікселях, і якщо різниця по вертикалі і горизонталі однакова, то використовується той напрямок інтерполяції, яке було в попередній ітерації. На алгоритмах такого роду знижується ефект від механізму out of order, реалізованих у сучасних cpu з-за того, що утворюється довга послідовність залежних операцій. Компілятор також тепер не може векторизовать цикл і продуктивність стала в 17 разів повільніше. Тобто як раз приблизно на ширину 16 float елементів в AVX512 регістрі. Тепер спробуємо якось модифікувати і наш avx512 код, щоб отримати хоч якийсь прискорення. Введемо в розгляд такі бінарні маски-змінні

mskV(n) – для n-го елемента V різниця мінімальна

mskE(n) – для n-го елемента V і H різниці однакові

mskD(n) – для n-го елемента використана V інтерполяція.

Тепер побудуємо таблицю істинності, як формується mskD(n) залежно від mskV(n),mskE(n) і попередньої використання маски – mskD(n-1)

З таблиці випливає, що
mskD(n) = mskV(n) | (mskE(n)&mskD(n-1)),
що загалом-то й так було очевидно. Отже наш avx512 код буде виглядати наступним чином.
Лист 4.4
void avg_avx512( const float* pSrc1, int src1Step, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
__mmask16 mskD = 0xFFFF;
i = 0;
for(i=i; i < (width&(~15)); i+=16 ) {
__m512 zvU, zvD, zvL, zvR;
__m512 zdV, zdH;
__m512 zavgV, zavgH, zavg;
__mmask16 mskV, mskE;

zvU = _mm512_loadu_ps(pSrc1+src1Step*(h - 1) + i );
zvD = _mm512_loadu_ps(pSrc1+src1Step*(h + 1) + i );
zvL = _mm512_loadu_ps(pSrc1+src1Step*( h ) + (i-1));
zvR = _mm512_loadu_ps(pSrc1+src1Step*( h ) + (i+1));

zdV = _mm512_sub_ps(zvU, zvD);
zdH = _mm512_sub_ps(zvL, zvR);
zdV = _mm512_abs_ps(zdV);
zdH = _mm512_abs_ps(zdH);

mskV = _mm512_cmp_ps_mask(zdV, zdH, _CMP_LT_OS);
mskE = _mm512_cmp_ps_mask(zdV, zdH, _CMP_EQ_OS);

mskD = mskV | (mskE & (mskD >>15) & (1<<0)); // 0 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<1)); // 1 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<2)); // 2 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<3)); // 3 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<4)); // 4 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<5)); // 5 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<6)); // 6 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<7)); // 7 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<8)); // 8 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<9)); // 9 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<10)); // 10 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<11)); // 11 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<12)); // 12 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<13)); // 13 bit
mskD = mskD | (mskE & (mskD << 1) & (1<<14)); // 14 bit
mskD = mskD | (mskE & (mskD << 1) & (1 << 15)); // 15 bit
zavgV = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvU, zvD));
zavgH = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvL, zvR));
zavg = _mm512_mask_or_ps(zavgH, mskD, zavgV, zavgV);
_mm512_storeu_ps(pDst + i, zavg);
}
pDst = pDst + dstStep;
}
}


В ньому послідовно перебираються всі 16 біт маски. На розмірі 64x64 він працює в ~1.7 X разів швидше ніж З-шний. Ну хоча б щось вдалося отримати. Наступна можлива оптимізація полягає в тому, що можна заздалегідь прорахувати всі комбінації масок. На кожній ітерації по 16 елементів у нас є 16 біт mskV, 16 біт mskE, і 1 біт від попередньої ітерації. Разом 2^33 ступеня варіантів для mskD. Це багато. А що якщо обробляти не 16, а з 8 елементів за ітерацію? Отримуємо 2^(8+8+1)=128кбайт таблицю. А це цілком нормальний розмір. Створюємо функцію ініціалізації.
Лист 4.5
unsigned char table[2*256 * 256];
extern init_table_mask()
{
int mskV, mskE, mskD=0;
for (mskD = 0; mskD < 2; mskD++){
for (mskV = 0; mskV < 256; mskV++){
for (mskE = 0; mskE < 256; mskE++){
int msk;
msk = mskV | (mskE & (mskD ) & (1 << 0)); // 0 bit
msk = msk | (mskE & (msk << 1) & (1 << 1)); // 1 bit
msk = msk | (mskE & (msk << 1) & (1 << 2)); // 2 bit
msk = msk | (mskE & (msk << 1) & (1 << 3)); // 3 bit
msk = msk | (mskE & (msk << 1) & (1 << 4)); // 4 bit
msk = msk | (mskE & (msk << 1) & (1 << 5)); // 5 bit
msk = msk | (mskE & (msk << 1) & (1 << 6)); // 6 bit
msk = msk | (mskE & (msk << 1) & (1 << 7)); // 7 bit
table[256*256*mskD+256 * mskE + mskV] = (unsigned char)msk;
}
}
}
}


Переписуємо код.
Лист 4.6
void avg_avx512( const float* pSrc1, int src1Step, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
__mmask16 mskD = 0xFFFF;
i = 0;
mskD = 0x00FF;
for(i=i; i < (width&(~7)); i+=8 ) {
__m512 Z = _mm512_setzero_ps();
__m512 zvU, zvD, zvL, zvR;
__m512 zdV, zdH;
__m512 zavgV, zavgH, zavg;
__mmask16 mskV, mskE;

zvU = _mm512_mask_loadu_ps(Z,0xFF, pSrc1+src1Step*(h - 1) + i );
zvD = _mm512_mask_loadu_ps(Z,0xFF, pSrc1+src1Step*(h + 1) + i );
zvL = _mm512_mask_loadu_ps(Z,0xFF, pSrc1+src1Step*( h ) + (i-1));
zvR = _mm512_mask_loadu_ps(Z,0xFF, pSrc1+src1Step*( h ) + (i+1));

zdV = _mm512_sub_ps(zvU, zvD);
zdH = _mm512_sub_ps(zvL, zvR);
zdV = _mm512_abs_ps(zdV);
zdH = _mm512_abs_ps(zdH);

mskV = _mm512_cmp_ps_mask(zdV, zdH, _CMP_LT_OS)&0xFF;
mskE = _mm512_cmp_ps_mask(zdV, zdH, _CMP_EQ_OS)&0xFF;
mskD = table[256 * 256 * (mskD >> 7) + 256 * mskE + mskV];
zavgV = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvU, zvD));
zavgH = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvL, zvR));
zavg = _mm512_mask_or_ps(zavgH, mskD, zavgV, zavgV);
_mm512_mask_storeu_ps(pDst + i, 0xFF, zavg);
}
pDst = pDst + dstStep;
}
}


Код став працювати у ~4.1 X разів швидше. Можна добававить обробку по 16 елементів, але до таблиці доведеться звернутися двічі за ітерацію.
Лист 4.7
void avg_avx512( const float* pSrc1, int src1Step, 
float* pDst, int dstStep, int width, int height)
{
int h, i; 
for( h = 0; h < height; h++ ) {
__mmask16 mskD = 0xFFFF;
i = 0;
for(i=i; i < (width&(~15)); i+=16 ) {
__m512 zvU, zvD, zvL, zvR;
__m512 zdV, zdH;
__m512 zavgV, zavgH, zavg;
__mmask16 mskV, mskE;
ushort mskD_0_7=0, mskD_8_15=0;

zvU = _mm512_loadu_ps(pSrc1+src1Step*(h - 1) + i );
zvD = _mm512_loadu_ps(pSrc1+src1Step*(h + 1) + i );
zvL = _mm512_loadu_ps(pSrc1+src1Step*( h ) + (i-1));
zvR = _mm512_loadu_ps(pSrc1+src1Step*( h ) + (i+1));

zdV = _mm512_sub_ps(zvU, zvD);
zdH = _mm512_sub_ps(zvL, zvR);
zdV = _mm512_abs_ps(zdV);
zdH = _mm512_abs_ps(zdH);

mskV = _mm512_cmp_ps_mask(zdV, zdH, _CMP_LT_OS);
mskE = _mm512_cmp_ps_mask(zdV, zdH, _CMP_EQ_OS);

mskD_0_7 = table[256 * 256 * (((ushort)mskD) >> 15) +
256 * (((ushort)mskE) & 0xFF) + (((ushort)mskV) & 0xFF)];
mskD_8_15 = table[256 * 256 * (((ushort)mskD_0_7) >> 7) +
256 * (((ushort)mskE)>> 8 ) + (((ushort)mskV) >>8 )];
mskD = (mskD_8_15 << 8) | mskD_0_7;

zavgV = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvU, zvD));
zavgH = _mm512_mul_ps(_mm512_set1_ps(0.5 f), _mm512_add_ps(zvL, zvR));
zavg = _mm512_mask_or_ps(zavgH, mskD, zavgV, zavgV);
_mm512_storeu_ps(pDst + i, zavg);
}
pDst = pDst + dstStep;
}
}


Тепер прискорення становить ~5.6 X у порівнянні з початковим З кодом, що дуже навіть непогано для алгоритму з зворотним зв'язком. Таким чином комбінація масочного регістра і попередньо обчисленої таблиці дозволяє отримати значний приріст продуктивності. Код такого роду зовсім не є спеціально підібраним для статті і зустрічається, наприклад, в алгоритмах обробки фотографій, при конвертації зображення з RAW формату. Правда там використовуються цілочисельні значення, але такий табличний метод цілком можна застосувати і для них.

Підсумок
У цій статті показано лише деякі способи застосування масочных інструкцій AVX 512. Ви можете застосовувати їх в своїх додатках або придумувати свої прийоми. А можете скористатися бібліотекою IPP, в якій вже є код, оптимізований спеціально для процесорів Intel з підтримкою avx512, важливою частиною якого є інструкції з масочными регістрами.

Джерело: Хабрахабр

0 коментарів

Тільки зареєстровані та авторизовані користувачі можуть залишати коментарі.