otimizando máximo por coluna com SIMD

Aug 15 2020

Eu tenho essa função em que gastei uma quantidade significativa de tempo em meu código e gostaria de otimizá-la por intrínsecos do compilador SIMD de vetorização, se possível.

Ele essencialmente encontra o valor e a localização do máximo em uma matriz sobre colunas e os armazena:

val_ptr: matriz de entrada: coluna principal (estilo Fortran) n_rows-by-n_cols (onde normalmente n_rows>>n_cols)
opt_pos_ptr : vetor int de comprimento n_rows onde armazenar a posição do máximo. Na entrada preenchida com zeros.
max_ptr: vetor float de comprimento n_rows onde armazenar o máximo. Na entrada preenchida com cópias da primeira coluna de val_ptr
A função será chamada em um loop paralelo
A região da memória é garantida para não se sobrepor
Eu realmente não preciso que o max_ptr seja preenchido, atualmente ele é usado apenas para contabilidade e para evitar alocação de memória
Eu uso MSVC, C ++ 17 no Windows 10. Destina-se a executar CPUs Intel modernas

O código, onde o tipo de modelo deve ser float ou double:

template <typename eT>
find_max(const int n_cols, 
         const int n_rows, 
         const eT* val_ptr,
         int* opt_pos_ptr,
         eT* max_ptr){
    for (int col = 1; col < n_cols; ++col)
    {
        //Getting the pointer to the beginning of the column
        const auto* value_col = val_ptr + col * n_rows;
        //Looping over the rows
        for (int row = 0; row < n_rows; ++row)
        {
            //If the value is larger than the current maximum, we replace and we store its positions
            if (value_col[row] > max_ptr[row])
            {
                max_ptr[row] = value_col[row];
                opt_pos_ptr[row] = col;
            }
        }
    }
}

O que tentei até agora:

Tentei usar o OpenMP paralelo no loop interno, mas traz algo apenas em linhas muito grandes, um pouco maiores do que meu uso atual.
O if in inner loop impede que #pragma omp simd funcione, e não consegui reescrevê-lo sem ele.

Respostas

3 AndreySemashev Aug 15 2020 at 21:55

Com base no exemplo de código que você postou, parece que você deseja calcular um valor máximo vertical, o que significa que, no seu caso, as "colunas" são horizontais. Em C/C++, sequências horizontais de elementos (ou seja, onde dois elementos adjacentes têm distância de um elemento na memória) são normalmente chamadas de linhas e verticais (onde dois elementos adjacentes têm distância do tamanho da linha na memória) - colunas. Na minha resposta abaixo, usarei a terminologia tradicional, onde as linhas são horizontais e as colunas são verticais.

Além disso, por brevidade, focarei em um tipo possível de elemento de matriz - float. A ideia básica é a mesma para double, com a principal diferença sendo o número de elementos por vetor e a seleção dos intrínsecos _ps/ . _pdVou fornecer uma versão para doubleno final.

A ideia é que você possa calcular o máximo vertical para várias colunas em paralelo usando _mm_max_ps/ _mm_max_pd. Para registrar também a posição do máximo encontrado, você pode comparar o máximo anterior com os elementos atuais. O resultado da comparação é uma máscara, onde os elementos são todos uns onde o máximo é atualizado. Essa máscara também pode ser usada para selecionar qual posição precisa ser atualizada.

Devo observar que o algoritmo abaixo assume que não é importante qual posição do elemento máximo é registrada, se houver vários elementos máximos iguais em uma coluna. Além disso, presumo que a matriz não contenha valores NaN, o que afetaria as comparações. Mais sobre isso mais tarde.

void find_max(const int n_cols, 
         const int n_rows, 
         const float* val_ptr,
         int* opt_pos_ptr,
         float* max_ptr){
    const __m128i mm_one = _mm_set1_epi32(1);

    // Pre-compute the number of rows that can be processed in full vector width.
    // In a 128-bit vector there are 4 floats or 2 doubles
    int tail_size = n_rows & 3;
    int n_rows_aligned = n_rows - tail_size;
    int row = 0;
    for (; row < n_rows_aligned; row += 4)
    {
        const auto* col_ptr = val_ptr + row;
        __m128 mm_max = _mm_loadu_ps(col_ptr);
        __m128i mm_max_pos = _mm_setzero_si128();
        __m128i mm_pos = mm_one;
        col_ptr += n_rows;
        for (int col = 1; col < n_cols; ++col)
        {
            __m128 mm_value = _mm_loadu_ps(col_ptr);

            // See if this value is greater than the old maximum
            __m128 mm_mask = _mm_cmplt_ps(mm_max, mm_value);
            // If it is, save its position
            mm_max_pos = _mm_blendv_epi8(mm_max_pos, mm_pos, _mm_castps_si128(mm_mask));

            // Compute the maximum
            mm_max = _mm_max_ps(mm_value, mm_max);

            mm_pos = _mm_add_epi32(mm_pos, mm_one);
            col_ptr += n_rows;
        }

        // Store the results
        _mm_storeu_ps(max_ptr + row, mm_max);
        _mm_storeu_si128(reinterpret_cast< __m128i* >(opt_pos_ptr + row), mm_max_pos);
    }

    // Process tail serially
    for (; row < n_rows; ++row)
    {
        const auto* col_ptr = val_ptr + row;
        auto max = *col_ptr;
        int max_pos = 0;
        col_ptr += n_rows;
        for (int col = 1; col < n_cols; ++col)
        {
            auto value = *col_ptr;
            if (value > max)
            {
                max = value;
                max_pos = col;
            }

            col_ptr += n_rows;
        }

        max_ptr[row] = max;
        opt_pos_ptr[row] = max_pos;
    }
}

O código acima requer SSE4.1 devido aos intrínsecos da combinação. Você pode substituí-los por uma combinação de _mm_and_si128/ _ps, _mm_andnot_si128/ _pse _mm_or_si128/ _ps, caso em que os requisitos serão reduzidos para SSE2. Consulte o Guia Intel Intrinsics para obter mais detalhes sobre os intrínsecos específicos, incluindo quais extensões de conjunto de instruções eles exigem.

Uma observação sobre os valores NaN. Se sua matriz puder ter NaNs, o _mm_cmplt_psteste sempre retornará falso. Quanto a _mm_max_ps, geralmente não se sabe o que ele retornará. A maxpsinstrução para a qual o intrínseco traduz retorna seu segundo operando (origem) se qualquer um dos operandos for um NaN, portanto, ao organizar os operandos da instrução, você pode obter qualquer um dos comportamentos. No entanto, não está documentado qual argumento do _mm_max_psintrínseco representa qual operando da instrução, e é até possível que o compilador use associações diferentes em casos diferentes. Veja esta resposta para mais detalhes.

A fim de garantir o comportamento correto wrt. NaNs, você pode usar o montador embutido para forçar a ordem correta dos maxpsoperandos. Infelizmente, essa não é uma opção com MSVC para destino x86-64, que você disse que está usando, então, em vez disso, você pode reutilizar o _mm_cmplt_psresultado para uma segunda combinação como esta:

// Compute the maximum
mm_max = _mm_blendv_ps(mm_max, mm_value, mm_mask);

Isso suprimirá NaNs nos valores máximos resultantes. Se você quiser manter NaNs, você pode usar uma segunda comparação para detectar NaNs:

// Detect NaNs
__m128 mm_nan_mask = _mm_cmpunord_ps(mm_value, mm_value);

// Compute the maximum
mm_max = _mm_blendv_ps(mm_max, mm_value, _mm_or_ps(mm_mask, mm_nan_mask));

Provavelmente, você poderia melhorar ainda mais o desempenho do algoritmo acima se usar vetores mais largos ( __m256ou __m512) e desenrolar o loop externo por um pequeno fator, de modo que pelo menos uma linha de cache de dados de linha seja carregada em cada iteração do loop interno.

Aqui está um exemplo de implementação para double. O ponto importante a ser observado aqui é que, como há apenas dois doubleelementos por vetor e ainda há quatro posições por vetor, temos que desenrolar o loop externo para processar dois vetores de doublepor vez e, em seguida, comprimir as duas máscaras de comparações com o máximos anteriores para combinar as posições de 32 bits.