c++ - AVX2, ¿cuál es la forma más eficiente de empacar en función de una máscara?

vectorization sse

Question

Si tiene una matriz de entrada y una matriz de salida, pero solo desea escribir aquellos elementos que pasan por cierta condición, ¿cuál sería la forma más eficiente de hacerlo en AVX2?

He visto en SSE donde se hizo así: (De: https://deplinenoise.files.wordpress.com/2015/03/gdc2015_afredriksson_simd.pdf )

__m128i LeftPack_SSSE3(__m128 mask, __m128 val)
{
 // Move 4 sign bits of mask to 4-bit integer value.
 int mask = _mm_movemask_ps(mask);
 // Select shuffle control data
 __m128i shuf_ctrl = _mm_load_si128(&shufmasks[mask]);
 // Permute to move valid values to front of SIMD register
 __m128i packed = _mm_shuffle_epi8(_mm_castps_si128(val), shuf_ctrl);
 return packed;
}

Esto parece correcto para SSE que tiene un ancho de 4, y por lo tanto solo necesita una LUT de 16 entradas, pero para AVX que tiene un ancho de 8, la LUT se vuelve bastante grande (256 entradas, cada una de 32 bytes, u 8k).

Me sorprende que AVX no parezca tener instrucciones para simplificar este proceso, como una tienda enmascarada con empaque.

Creo que con un poco de orden aleatorio para contar el número de bits de signo establecido a la izquierda, podría generar la tabla de permutación necesaria y luego llamar a _mm256_permutevar8x32_ps. Pero esto también es un buen número de instrucciones, creo.

¿Alguien sabe de algún truco para hacer esto con AVX2? ¿O cuál es el método más eficiente?

Aquí hay una ilustración del problema de embalaje izquierdo del documento anterior:

Gracias

Answer 1

AVX2 + BMI2. Vea mi otra respuesta para AVX512. (Actualización: guardó un pdep en pdep de 64 bits).

Podemos usar AVX2 vpermps ( _mm256_permutevar8x32_ps ) (o el equivalente entero, vpermd ) para hacer una mezcla aleatoria de variables de cruce de carril.

Podemos generar máscaras sobre la marcha , ya que BMI2 pext (Parallel Bits Extract) nos proporciona una versión bitwise de la operación que necesitamos.

Para vectores enteros con elementos de 32 bits o más anchos : 1) _mm256_movemask_ps(_mm256_castsi256_ps(compare_mask)) .
O 2) use _mm256_movemask_epi8 y luego cambie la primera constante PDEP de 0x0101010101010101 a 0x0F0F0F0F0F0F0F para dispersar bloques de 4 bits contiguos. Cambie la multiplicación por 0xFFU en expanded_mask |= expanded_mask<<4; o expanded_mask *= 0x11; (No probado). De cualquier manera, use la máscara aleatoria con VPERMD en lugar de VPERMPS.

Para los elementos enteros o double 64 bits, todo sigue funcionando ; Sucede que la comparación de máscara siempre tiene pares de elementos de 32 bits que son iguales, por lo que la ordenación resultante coloca ambas mitades de cada elemento de 64 bits en el lugar correcto. (Por lo tanto, sigue utilizando VPERMPS o VPERMD, porque VPERMPD y VPERMQ solo están disponibles con operandos de control inmediato).

El algoritmo:

Comience con una constante de índices de 3 bits empaquetados, con cada posición manteniendo su propio índice. es decir, [ 7 6 5 4 3 2 1 0 ] donde cada elemento tiene 3 bits de ancho. 0b111''110''101''...''010''001''000 .

Use pext para extraer los índices que deseamos en una secuencia contigua en la parte inferior de un registro entero. por ejemplo, si queremos los índices 0 y 2, nuestra máscara de control para pext debe ser 0b000''...''111''000''111 . pext tomará los grupos de índice 010 y 000 que se alinean con los 1 bits en el selector. Los grupos seleccionados se empaquetan en los bits bajos de la salida, por lo que la salida será 0b000''...''010''000 . (es decir, [ ... 2 0 ] )

Consulte el código comentado para 0b111000111 cómo generar la entrada pext para pext desde la máscara de vector de entrada.

Ahora estamos en el mismo barco que la LUT comprimida: desempaquetamos hasta 8 índices empaquetados.

En el momento en que pones todas las piezas juntas, hay tres pext / pdep total. Trabajé hacia atrás desde lo que quería, por lo que probablemente también es más fácil entenderlo en esa dirección. (es decir, comenzar con la línea de reproducción aleatoria y trabajar hacia atrás desde allí).

Podemos simplificar el desempaquetado si trabajamos con índices uno por byte en lugar de en grupos de 3 bits empaquetados . Como tenemos 8 índices, esto solo es posible con un código de 64 bits.

Vea esta y una versión de solo 32 bits en el Explorador del compilador de Godbolt . Utilicé #ifdef s para que #ifdef manera óptima con -m64 o -m32 . gcc desperdicia algunas instrucciones, pero clang hace un código realmente bueno.

#include <stdint.h>
#include <immintrin.h>

// Uses 64bit pdep / pext to save a step in unpacking.
__m256 compress256(__m256 src, unsigned int mask /* from movmskps */)
{
  uint64_t expanded_mask = _pdep_u64(mask, 0x0101010101010101);  // unpack each bit to a byte
  expanded_mask *= 0xFF;    // mask |= mask<<1 | mask<<2 | ... | mask<<7;
  // ABC... -> AAAAAAAABBBBBBBBCCCCCCCC...: replicate each bit to fill its byte

  const uint64_t identity_indices = 0x0706050403020100;    // the identity shuffle for vpermps, packed to one index per byte
  uint64_t wanted_indices = _pext_u64(identity_indices, expanded_mask);

  __m128i bytevec = _mm_cvtsi64_si128(wanted_indices);
  __m256i shufmask = _mm256_cvtepu8_epi32(bytevec);

  return _mm256_permutevar8x32_ps(src, shufmask);
}

Esto compila a código sin cargas de memoria, solo constantes inmediatas. (Vea el enlace de Godbolt para esta y la versión de 32 bits).

    # clang 3.7.1 -std=gnu++14 -O3 -march=haswell
    mov     eax, edi                   # just to zero extend: goes away when inlining
    movabs  rcx, 72340172838076673     # The constants are hoisted after inlining into a loop
    pdep    rax, rax, rcx              # ABC       -> 0000000A0000000B....
    imul    rax, rax, 255              # 0000000A0000000B.. -> AAAAAAAABBBBBBBB..
    movabs  rcx, 506097522914230528
    pext    rax, rcx, rax
    vmovq   xmm1, rax
    vpmovzxbd       ymm1, xmm1         # 3c latency since this is lane-crossing
    vpermps ymm0, ymm1, ymm0
    ret

Entonces, de acuerdo con los números de Agner Fog , esto es 6 puntos (sin contar las constantes, o el movimiento de extensión cero que desaparece cuando está en línea). En Intel Haswell, es 16c de latencia (1 para vmovq, 3 para cada pdep / imul / pext / vpmovzx / vpermps). No hay paralelismo a nivel de instrucción. Sin embargo, en un bucle en el que esto no forma parte de una dependencia transmitida por un bucle (como la que incluí en el enlace de Godbolt), el cuello de botella es, simplemente, el rendimiento, manteniendo múltiples iteraciones de esto a la vez.

Esto puede administrar un rendimiento de uno por 3 ciclos, con un cuello de botella en el puerto 1 para pdep / pext / imul. Por supuesto, con las cargas / tiendas y la sobrecarga del bucle (incluyendo la comparación, movmsk y popcnt), el rendimiento total de uop puede ser un problema fácilmente. (por ejemplo, el bucle de filtro en mi enlace de godbolt es 14 uops con clang, con -fno-unroll-loops para que sea más fácil de leer. Podría mantener una iteración por 4c, mantenerse al día con el front-end, si tenemos suerte , pero creo que clang no tuvo en cuenta la falsa dependencia de popcnt en su salida, por lo que se producirá un cuello de botella en 3/5 de la latencia de la función compress256 .)

gcc realiza la multiplicación por 0xFF con varias instrucciones, utilizando un desplazamiento a la izquierda por 8 y un sub . Esto requiere instrucciones de mov adicionales, pero el resultado final es una multiplicación con una latencia de 2. (Haswell maneja mov en la etapa de cambio de nombre de registro con latencia cero).

Dado que todo el hardware que admite AVX2 también admite BMI2, es probable que no tenga sentido proporcionar una versión para AVX2 sin BMI2.

Si necesita hacer esto en un bucle muy largo, la LUT probablemente valga la pena si las faltas de caché iniciales se amortizan en iteraciones suficientes con la menor sobrecarga de simplemente desempaquetar la entrada de la LUT. Aún necesita movmskps , por lo que puede colocar la máscara y usarla como un índice LUT, pero guarda un pdep / imul / pexp.

Puede desempaquetar las entradas de LUT con la misma secuencia de enteros que usé, pero el set1() / vpsrlvd / vpand vpsrlvd vpand es probablemente mejor cuando la entrada de la LUT comienza en la memoria y no necesita ingresar en los registros de enteros en primer lugar. (Una carga de difusión de 32 bits no necesita una ALU uop en las CPU Intel). Sin embargo, un cambio variable es 3 uops en Haswell (pero solo 1 en Skylake).

Answer 2

En caso de que a alguien le interese, aquí hay una solución para SSE2 que utiliza una instrucción LUT en lugar de una tabla de datos LUT o tabla de salto. Sin embargo, con AVX necesitaría 256 casos.

Cada vez que llama LeftPack_SSE2 continuación, utiliza esencialmente tres instrucciones: jmp, shufps, jmp. Cinco de los dieciséis casos no necesitan modificar el vector.

static inline __m128 LeftPack_SSE2(__m128 val, int mask)  {
  switch(mask) {
  case  0:
  case  1: return val;
  case  2: return _mm_shuffle_ps(val,val,0x01);
  case  3: return val;
  case  4: return _mm_shuffle_ps(val,val,0x02);
  case  5: return _mm_shuffle_ps(val,val,0x08);
  case  6: return _mm_shuffle_ps(val,val,0x09);
  case  7: return val;
  case  8: return _mm_shuffle_ps(val,val,0x03);
  case  9: return _mm_shuffle_ps(val,val,0x0c);
  case 10: return _mm_shuffle_ps(val,val,0x0d);
  case 11: return _mm_shuffle_ps(val,val,0x34);
  case 12: return _mm_shuffle_ps(val,val,0x0e);
  case 13: return _mm_shuffle_ps(val,val,0x38);
  case 14: return _mm_shuffle_ps(val,val,0x39);
  case 15: return val;
  }
}

__m128 foo(__m128 val, __m128 maskv) {
  int mask = _mm_movemask_ps(maskv);
  return LeftPack_SSE2(val, mask);
}

Answer 3

Se me ocurrió este método, que utiliza una LUT comprimida, que es de 768 bytes (+1 relleno), en lugar de 8k. Requiere una transmisión de un solo valor escalar, que luego se desplaza una cantidad diferente en cada carril, luego se enmascara a los 3 bits más bajos, lo que proporciona una LUT de 0-7.

Aquí está la versión intrínseca, junto con el código para construir LUT.

//Generate Move mask via: _mm256_movemask_ps(_mm256_castsi256_ps(mask)); etc
__m256i MoveMaskToIndices(int moveMask) {
    u8 *adr = g_pack_left_table_u8x3 + moveMask * 3;
    __m256i indices = _mm256_set1_epi32(*reinterpret_cast<u32*>(adr));//lower 24 bits has our LUT

    __m256i m = _mm256_sllv_epi32(indices, _mm256_setr_epi32(29, 26, 23, 20, 17, 14, 11, 8));

    //now shift it right to get 3 bits at bottom
    __m256i shufmask = _mm256_srli_epi32(m, 29);
    return shufmask;
}

u32 get_nth_bits(int a) {
    u32 out = 0;
    int c = 0;
    for (int i = 0; i < 8; ++i) {
        auto set = (a >> i) & 1;
        if (set) {
            out |= (i << (c * 3));
            c++;
        }
    }
    return out;
}
u8 g_pack_left_table_u8x3[256 * 3 + 1];

void BuildPackMask() {
    for (int i = 0; i < 256; ++i) {
        *reinterpret_cast<u32*>(&g_pack_left_table_u8x3[i * 3]) = get_nth_bits(i);
    }
}

Aquí está el ensamblado generado por VS2015:

lea eax, DWORD PTR [rcx+rcx*2]
movsxd  rcx, eax
lea rax, OFFSET FLAT:?g_pack_left_table_u8x3@@3PAEA ; g_pack_left_table_u8x3
vpbroadcastd ymm0, DWORD PTR [rcx+rax]
vpsllvd ymm0, ymm0, YMMWORD PTR __ymm@000000080000000b0000000e0000001100000014000000170000001a0000001d
vpsrld  ymm0, ymm0, 29

Answer 4

Vea mi otra respuesta para AVX2 + BMI2 sin LUT.

Como mencionó una inquietud sobre la escalabilidad en AVX512: no se preocupe, hay una instrucción AVX512F para exactamente esto :

VCOMPRESSPS - Almacene los valores de punto flotante de precisión simple de VCOMPRESSPS en una memoria densa. (También hay versiones para elementos enteros dobles y de 32 o 64 bits ( vpcompressq ), pero no bytes ni palabras (16 bits)). Es como BMI2 pdep / pext , pero para vectores en lugar de bits en un registro entero.

El destino puede ser un registro vectorial o un operando de memoria, mientras que la fuente es un vector y un registro de máscara. Con un registro dest, puede fusionar o poner a cero los bits superiores. Con un destino de memoria, "Sólo el vector contiguo se escribe en la ubicación de la memoria de destino".

Para averiguar hasta dónde puede avanzar su puntero para el siguiente vector, coloque la máscara.

Digamos que quiere filtrar todo menos los valores> = 0 de una matriz:

#include <stdint.h>
#include <immintrin.h>
size_t filter_non_negative(float *__restrict__ dst, const float *__restrict__ src, size_t len) {
    const float *endp = src+len;
    float *dst_start = dst;
    do {
        __m512      sv  = _mm512_loadu_ps(src);
        __mmask16 keep = _mm512_cmp_ps_mask(sv, _mm512_setzero_ps(), _CMP_GE_OQ);  // true for src >= 0.0, false for unordered and src < 0.0
        _mm512_mask_compressstoreu_ps(dst, keep, sv);   // clang is missing this intrinsic, which can''t be emulated with a separate store

        src += 16;
        dst += _mm_popcnt_u64(keep);   // popcnt_u64 instead of u32 helps gcc avoid a wasted movsx, but is potentially slower on some CPUs
    } while (src < endp);
    return dst - dst_start;
}

Esto compila (con gcc4.9 o posterior) a ( Godbolt Compiler Explorer ):

 # Output from gcc6.1, with -O3 -march=haswell -mavx512f.  Same with other gcc versions
    lea     rcx, [rsi+rdx*4]             # endp
    mov     rax, rdi
    vpxord  zmm1, zmm1, zmm1             # vpxor  xmm1, xmm1,xmm1 would save a byte, using VEX instead of EVEX
.L2:
    vmovups zmm0, ZMMWORD PTR [rsi]
    add     rsi, 64
    vcmpps  k1, zmm0, zmm1, 29           # AVX512 compares have mask regs as a destination
    kmovw   edx, k1                      # There are some insns to add/or/and mask regs, but not popcnt
    movzx   edx, dx                      # gcc is dumb and doesn''t know that kmovw already zero-extends to fill the destination.
    vcompressps     ZMMWORD PTR [rax]{k1}, zmm0
    popcnt  rdx, rdx
    ## movsx   rdx, edx         # with _popcnt_u32, gcc is dumb.  No casting can get gcc to do anything but sign-extend.  You''d expect (unsigned) would mov to zero-extend, but no.
    lea     rax, [rax+rdx*4]             # dst += ...
    cmp     rcx, rsi
    ja      .L2

    sub     rax, rdi
    sar     rax, 2                       # address math -> element count
    ret