Найти наибольший элемент в матрице и его индексы столбцов и строк с использованием SSE и AVX

Question

Мне нужно найти самый большой элемент в матрице 1d и его индексы столбцов и строк.

Я использую 1d матрицу, поэтому нужно сначала найти индекс максимального элемента, а затем легко получить строку и столбец.

Моя проблема в том, что я не могу получить этот индекс.

У меня есть рабочая функция, которая находит наибольший элемент и использует SSE, здесь:

float find_largest_element_in_matrix_SSE(float* m, unsigned const int dims) 
{ 
    size_t i; 
    int index = -1; 
    __m128 max_el = _mm_loadu_ps(m); 
    __m128 curr; 

    for (i = 4; i < dims * dims; i += 4) 
    { 
     curr = _mm_loadu_ps(m + i); 
     max_el = _mm_max_ps(max_el, curr); 
    } 

    __declspec(align(16))float max_v[4] = { 0 }; 
    _mm_store_ps(max_v, max_el); 

    return max(max(max(max_v[0], max_v[1]), max_v[2]), max_v[3]); 
} 

, а также у меня есть нерабочая функция, которая использует AVX:

float find_largest_element_in_matrix_AVX(float* m, unsigned const int dims) 
{ 
    size_t i; 
    int index = -1; 
    __m256 max_el = _mm256_loadu_ps(m); 
    __m256 curr; 

    for (i = 8; i < dims * dims; i += 8) 
    { 
     curr = _mm256_loadu_ps(m + i); 
     max_el = _mm256_max_ps(max_el, curr); 
    } 

    __declspec(align(32))float max_v[8] = { 0 }; 
    _mm256_store_ps(max_v, max_el); 

    __m256 y = _mm256_permute2f128_ps(max_el, max_el, 1); 
    __m256 m1 = _mm256_max_ps(max_el, y);m1[1] = max(max_el[1], max_el[3]) 
    __m256 m2 = _mm256_permute_ps(m1, 5); 
    __m256 m_res = _mm256_max_ps(m1, m2); 

    return m[0]; 
} 

---

Может ли кто-нибудь помочь мне с нахождением индекса элемента max и сделать мою версию AVX?

Answer 1

Вот работает SSE (SSE 4) реализация, которая возвращает максимальные Вэла и соответствующий индекс, наряду с опорной скалярной реализации и тестом упряжью:

#include <stdio.h> 
#include <stdint.h> 
#include <stdlib.h> 
#include <time.h> 
#include <smmintrin.h> // SSE 4.1 

float find_largest_element_in_matrix_ref(const float* m, int dims, int *maxIndex) 
{ 
    float maxVal = m[0]; 
    int i; 

    *maxIndex = 0; 

    for (i = 1; i < dims * dims; ++i) 
    { 
     if (m[i] > maxVal) 
     { 
      maxVal = m[i]; 
      *maxIndex = i; 
     } 
    } 
    return maxVal; 
} 

float find_largest_element_in_matrix_SSE(const float* m, int dims, int *maxIndex) 
{ 
    float maxVal = m[0]; 
    float aMaxVal[4]; 
    int32_t aMaxIndex[4]; 
    int i; 

    *maxIndex = 0; 

    const __m128i vIndexInc = _mm_set1_epi32(4); 
    __m128i vMaxIndex = _mm_setr_epi32(0, 1, 2, 3); 
    __m128i vIndex = vMaxIndex; 
    __m128 vMaxVal = _mm_loadu_ps(m); 

    for (i = 4; i < dims * dims; i += 4) 
    { 
     __m128 v = _mm_loadu_ps(&m[i]); 
     __m128 vcmp = _mm_cmpgt_ps(v, vMaxVal); 
     vIndex = _mm_add_epi32(vIndex, vIndexInc); 
     vMaxVal = _mm_blendv_ps(vMaxVal, v, vcmp); 
     vMaxIndex = _mm_blendv_epi8(vMaxIndex, vIndex, _mm_castps_si128(vcmp)); 
    } 
    _mm_storeu_ps(aMaxVal, vMaxVal); 
    _mm_storeu_si128((__m128i *)aMaxIndex, vMaxIndex); 
    maxVal = aMaxVal[0]; 
    *maxIndex = aMaxIndex[0]; 
    for (i = 1; i < 4; ++i) 
    { 
     if (aMaxVal[i] > maxVal) 
     { 
      maxVal = aMaxVal[i]; 
      *maxIndex = aMaxIndex[i]; 
     } 
    } 
    return maxVal; 
} 

int main() 
{ 
    const int dims = 1024; 
    float m[dims * dims]; 
    float maxVal_ref, maxVal_SSE; 
    int maxIndex_ref = -1, maxIndex_SSE = -1; 
    int i; 

    srand(time(NULL)); 

    for (i = 0; i < dims * dims; ++i) 
    { 
     m[i] = (float)rand()/RAND_MAX; 
    } 

    maxVal_ref = find_largest_element_in_matrix_ref(m, dims, &maxIndex_ref); 
    maxVal_SSE = find_largest_element_in_matrix_SSE(m, dims, &maxIndex_SSE); 

    if (maxVal_ref == maxVal_SSE && maxIndex_ref == maxIndex_SSE) 
    { 
     printf("PASS: maxVal = %f, maxIndex = %d\n", 
         maxVal_ref, maxIndex_ref); 
    } 
    else 
    { 
     printf("FAIL: maxVal_ref = %f, maxVal_SSE = %f, maxIndex_ref = %d, maxIndex_SSE = %d\n", 
         maxVal_ref, maxVal_SSE, maxIndex_ref, maxIndex_SSE); 
    } 
    return 0; 
} 

компилировать и запускать:

$ gcc -Wall -msse4 Yakovenko.c && ./a.out 
PASS: maxVal = 0.999999, maxIndex = 120409 

Очевидно, вы можете получить индексы строк и столбцов, если необходимо:

int rowIndex = maxIndex/dims; 
int colIndex = maxIndex % dims; 

Отсюда должно быть довольно просто написать реализацию AVX2.

Answer 2

Одним из подходов было бы вычисление максимума в первом проходе и поиск индекса по линейному поиску во втором проходе. Вот пример реализации в SSE2:

#define anybit __builtin_ctz //or lookup table with 16 entries... 
float find_largest_element_in_matrix_SSE(const float* m, int dims, int *maxIndex) { 
    //first pass: calculate maximum as usual 
    __m128 vMaxVal = _mm_loadu_ps(m); 
    for (int i = 4; i < dims * dims; i += 4) 
     vMaxVal = _mm_max_ps(vMaxVal, _mm_loadu_ps(&m[i])); 
    //perform in-register reduction 
    vMaxVal = _mm_max_ps(vMaxVal, _mm_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(2, 3, 0, 1))); 
    vMaxVal = _mm_max_ps(vMaxVal, _mm_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(1, 0, 3, 2))); 
    //second pass: search for maximal value 
    for (int i = 0; i < dims * dims; i += 4) { 
     __m128 vIsMax = _mm_cmpeq_ps(vMaxVal, _mm_loadu_ps(&m[i])); 
     if (int mask = _mm_movemask_ps(vIsMax)) { 
      *maxIndex = i + anybit(mask); 
      return _mm_cvtss_f32(vMaxVal); 
     } 
    } 
} 

Обратите внимание, что ветвь во втором цикле должно быть почти идеально предсказывал, если ваш ввод данных очень мало.

Решение страдает от ряда проблем, в частности:

1. Он может работать неправильно при наличии точечных значений с плавающей странно, например с NaNs.

2. Если ваша матрица не вписывается в кеш процессора, тогда код будет считывать матрицу дважды из основной памяти, поэтому она будет в два раза медленнее, чем однопроходный. Это можно решить для больших матриц путем блочной обработки.

3. В первом цикле каждая итерация зависит от предыдущей (vMaxVal как модифицирована, так и прочитана), поэтому она будет замедляться с задержкой _mm_max_ps. Возможно, было бы здорово развернуть первый цикл бит (2x или 4x), имея 4 независимых регистра для vMaxVal (фактически, второй цикл также выиграет от разворачивания).

Портирование на AVX должно быть довольно прямо вперед, для уменьшения в регистре, за исключением:

vMaxVal = _mm256_max_ps(vMaxVal, _mm256_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(2, 3, 0, 1))); 
vMaxVal = _mm256_max_ps(vMaxVal, _mm256_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(1, 0, 3, 2))); 
vMaxVal = _mm256_max_ps(vMaxVal, _mm256_permute2f128_ps(vMaxVal, vMaxVal, 1)); 

Answer 3

еще один подход:

void find_largest_element_in_matrix_SSE(float * matrix, size_t n, int * row, int * column, float * v){ 

    __m128 indecies = _mm_setr_ps(0, 1, 2, 3); 
    __m128 update = _mm_setr_ps(4, 4, 4, 4); 
    __m128 max_indecies = _mm_setr_ps(0, 1, 2, 3); 
    __m128 max = _mm_load_ps(matrix); 
    for (int i = 4; i < n * n; i+=4){ 
     indecies = _mm_add_ps(indecies, update); 
     __m128 pm2 = _mm_load_ps(&matrix[i]); 
     __m128 mask = _mm_cmpge_ps(max, pm2); 
     max = _mm_max_ps(max, pm2); 
     max_indecies = _mm_or_ps(_mm_and_ps(max_indecies, mask), _mm_andnot_ps(mask, indecies)); 
    } 
    __declspec (align(16)) int max_ind[4]; 
    __m128i maxi = _mm_cvtps_epi32(max_indecies); 
    _mm_store_si128((__m128i *) max_ind, maxi); 
    int c = max_ind[0]; 
    for (int i = 1; i < 4; i++) 
     if (matrix[max_ind[i]] >= matrix[c] && max_ind[i] < c){ 
      c = max_ind[i]; 
     } 

    *v = matrix[c]; 
    *row = c/n; 
    *column = c % n; 
} 

void find_largest_element_in_matrix_AVX(float * matrix, size_t n, int * row, int * column, float * v){ 
    __m256 indecies = _mm256_setr_ps(0, 1, 2, 3, 4, 5, 6, 7); 
    __m256 update = _mm256_setr_ps(8, 8, 8, 8, 8, 8, 8, 8); 
    __m256 max_indecies = _mm256_setr_ps(0, 1, 2, 3, 4, 5, 6, 7); 
    __m256 max = _mm256_load_ps(matrix); 

    for (int i = 8; i < n * n; i += 8){ 
     indecies = _mm256_add_ps(indecies, update); 
     __m256 pm2 = _mm256_load_ps(&matrix[i]); 
     __m256 mask = _mm256_cmp_ps(max, pm2, _CMP_GE_OQ); 
     max = _mm256_max_ps(max, pm2); 
     max_indecies = _mm256_or_ps(_mm256_and_ps(max_indecies, mask), _mm256_andnot_ps(mask, indecies)); 
    } 
    __declspec (align(32)) int max_ind[8]; 
    __m256i maxi = _mm256_cvtps_epi32(max_indecies); 

    _mm256_store_si256((__m256i *) max_ind, maxi); 

    int c = max_ind[0]; 
    for (int i = 1; i < 8; i++) 
     if (matrix[max_ind[i]] >= matrix[c] && max_ind[i] < c){ 
      c = max_ind[i]; 
     } 

    *v = matrix[c]; 
    *row = c/n; 
    *column = c % n; 
}