c++ - Encuentre el índice del elemento máximo en el vector SIMD x86

sse avx

Question

Estoy pensando en implementar 8-ary heapsort para uint32_t''s. Para hacer esto, necesito una función que seleccione el índice del elemento máximo en un vector de 8 elementos para que pueda compararlo con el elemento padre y realice condicionalmente swap y otros pasos siftDown.

(8 uint32_ts se pueden cambiar, por ejemplo, a 16 uint32_ts u 8 uint64_t o lo que sea que x86 SIMD pueda soportar de manera eficiente).

Tengo algunas ideas sobre cómo hacerlo, pero estoy buscando algo más rápido que el código no vectorizado, especialmente estoy buscando algo que me permita hacer heapsort rápido.

Tengo clang ++ 3.3 y Core i7-4670, así que probablemente debería ser capaz de usar incluso las más recientes características de SIMD x86.

(Por cierto, eso es parte de un proyecto más grande: https://github.com/tarsa/SortingAlgorithmsBenchmark y hay, por ejemplo, heatsort cuaternario, así que después de implementar heapsort SIMD pude compararlos al instante)

Para repetir, la pregunta es : ¿cuál es la forma más eficiente de calcular el índice del elemento máximo en el vector x86 SIMD?

PD: No es un duplicado de las preguntas vinculadas: observe que estoy pidiendo un índice de un elemento máximo, no solo el valor del elemento.

Answer 1

La forma más eficiente de hacer una operación horizontal (producto punto, suma, índice máximo, lo que sea) en un vector SIMD de n vías es hacer n de una vez, transponiéndolos y usando operaciones verticales en su lugar. Algunas arquitecturas SIMD tienen mejor soporte para operaciones horizontales, pero en general el enfoque de transposición por bloques será más flexible y eficiente.

Answer 2

Usando la biblioteca Vc escribiría:

size_t maximumIndex(Vc::uint_v vec) {
  const unsigned int max = vec.max();
  return (max == vec).firstOne();
}

Con los intrínsecos debería ser algo similar a esto (esto es AVX sin AVX2, con AVX2 se vuelve un poco más fácil):

size_t maximumIndex(_mm256i vec) {
  __m128i lo = _mm256_castsi256_si128(vec);
  __m128i hi = _mm256_extractf128_si256(vec, 1);
  __m128i tmp = _mm_max_epu32(lo, hi);
  tmp = _mm_max_epu32(tmp, _mm_shuffle_epi32(tmp, _MM_SHUFFLE(1, 0, 3, 2)));
  tmp = _mm_max_epu32(tmp, _mm_shufflelo_epi16(tmp, _MM_SHUFFLE(1, 0, 3, 2))); // using lo_epi16 for speed here
  const int max = _mm_cvtsi128_si32(tmp);
  tmp = _mm_packs_epi16(_mm_packs_epi32(_mm_cmpeq_epi32(_mm_set1_epi32(max), lo),
                                        _mm_cmpeq_epi32(_mm_set1_epi32(max), hi)),
                        _mm_setzero_si128());
  return _bit_scan_forward(_mm_movemask_epi8(tmp));
}

Por cierto, si quieres inspiración de SIMDized merge-sort mira aquí: http://code.compeng.uni-frankfurt.de/projects/vc/repository/revisions/master/entry/src/avx_sorthelper.cpp

Answer 3

Las operaciones horizontales son malas noticias con SIMD, y particularmente con AVX, donde la mayoría de las instrucciones de 256 bits se dividen en dos operaciones separadas de 128 bits. Habiendo dicho eso, si realmente debe hacer un máximo horizontal de 32 bits en 8 elementos, entonces creo que el enfoque general debería ser:

• encuentre el valor máximo (típicamente varias operaciones de cambio / permutación y máximo)
• splat el valor máximo en los 8 elementos de un segundo vector (puede combinarse con la operación anterior)
• hacer una comparación entre el vector original y el vector máximo ( _mm256_cmpeq_epi32 )
• extraer la máscara escalar ( _mm256_movemask_epi8 )
• convertir la máscara escalar al índice

Aquí hay una primera pasada en una implementación de AVX2 que acabo de armar: lo probé y lo evalué en un Haswell a 2.6 GHz y funciona a aproximadamente 1.7 ns / vector (incluyendo la carga del vector y el almacenamiento del índice resultante):

uint8_t _mm256_hmax_index(const __m256i v)
{
    __m256i vmax = v;

    vmax = _mm256_max_epu32(vmax, _mm256_alignr_epi8(vmax, vmax, 4));
    vmax = _mm256_max_epu32(vmax, _mm256_alignr_epi8(vmax, vmax, 8));
    vmax = _mm256_max_epu32(vmax, _mm256_permute2x128_si256(vmax, vmax, 0x01));

    __m256i vcmp = _mm256_cmpeq_epi32(v, vmax);

    uint32_t mask = _mm256_movemask_epi8(vcmp);

    return __builtin_ctz(mask) >> 2;
}