Malé řazení ve vektoru - pcmpistrm

20. 9. 2021 #řazení #SIMD #D

Znalost řadících algoritmů je zbytečná, zní univerzální pravda z branže. Já nejsem svázán profesionálními okovy a tak můžu vesele experimentovat a netrápit se u toho, jestli zapadám do předschváleného mustru.

Nedávno tu padlo pár vět o bitonic sortu pro velice specifické případy malého řazení, pro nějž jsem měl uplatnění. Už nějakou dobu mi v hlavě sedí myšlenka, jestli by nešla k řazení zneužít i moje oblíbená nejciscovatější x86 instrukce pcmpXstrX. Ta umí mnoho: porovnávat elementy paralelně jako běžná vektorová operace, hledat jestli je jeden vektor podřetězcem druhého nebo pro každý element jednoho vektoru zjistit, jestli se nachází kdekoli ve vektoru druhém. Všechno tohle cílí primárně na práci se stringy (z toho důvodu pracuje jen s 1B nebo 2B elementy), i když to najde využítí i jinde. Na řazení to ale nestačí, rovnost je málo, je potřeba porovnat nerovnosti.

pcmpXstrX je švýcarský nůž a zvládá i tohle. Jedna varianta instrukce interpretuje první vstupní vektor jako 1 až 4 intervaly. Jeho 8 položek jsou až 4 páry horních a dolních mezí. Výsledkem je pak maska, jeden bit pro korespondující prvek druhého vektoru, nastavená na 1, pokud se element nachází v některém z rozpětí. Tohle pro řazení stačí.

Řazení je proces permutace, který přemístí prvek na index 𝑖 takový, že mu vždy předchází 𝑖 menších prvků (pokud počítáme od nuly a ve vstupu nejsou duplicity)¹ . Tohle lze verbatim převést do podoby algoritmu: Pro každý prvek vektoru spočítám, kolik ostatních je menších, tj. v rozmezí mezi nulou a daným prvkem, a to bude výsledný index 𝑖. Hotovo. Technicky vzato jde o n² operaci, ale n je malá konstanta a procesor bez námahy odpálí jeden pcmpistrm každých 3-5 taktů.

alias pcmpistrm = __builtin_ia32_pcmpistrm128;

void pcmpstrSort(ref ushort[8] arr) {
  ushort[8] res;
  ushort8 vec = *(cast(ushort8*) arr.ptr);

  ushort8 range = 0;
  range.array[0] = 1;

  static foreach (i; 0 .. 8) {{
    range.array[1] = arr[i];
    auto mask = pcmpistrm(cast(ubyte16) range, cast(ubyte16) vec, 0b0_00_01_01);
    ulong pos = popcnt((cast(ulong2)mask).array[0]);
    res.ptr[pos-1] = arr[i];
  }}

  arr = res;
}

Na Sandy Bridge laptopu (i5-2520M) to trvá 13.3 ns.

Nefunguje to pro nuly, ty se berou jako ukončovače stringu ve stylu jazyka C, pro ty je třeba o něco pomalejší pcmpestrm, kde hsou délky svou vektorů předány ve dvojici registrů. Ale jako průzkum terénu, co se dá udělat pro tenhle velice specifický případ, musí říct, že za mě dobré.

Když jsem výše uvedené napsal, byl jsem se sebou spokojený asi tak třináct vteřin, než mi došlo, že to samé se dá udělat s normálním SSE/AVX bez šílenosti pcmpXstrX. Pochopitelně.

void vecPrefixSort(ref ushort[8] arr) {
  ushort[8] res;
  ushort8 vec = *(cast(ushort8*) arr.ptr);

  static foreach (i; 0 .. 8) {{
    ushort8 x = vec.array[i];
    auto byteMask = __builtin_ia32_pcmpgtw128(cast(short8)x, cast(short8)vec);
    auto mask = __builtin_ia32_pmovmskb128(cast(ubyte16)byteMask);
    ulong pos = ulong(popcnt(mask));
    *(cast(ushort*) ((cast(ubyte*) res.ptr) + pos)) = arr.ptr[i];
  }}

  arr = res;
}

A navíc je to rychlejší. Pochopitelně. (Kompletní zdroják tady.)

Kompilováno přes gdc-10 -02 -march=native:

Sandy Bridge i5-2520M
pcmpstrSort   13.3 ns / 42.6 taktů, 5.3 taktů/short)
vecPrefixSort 12.0 ns / 38.4 taktů, 4.8 taktů/short)

Haswell i5-4570
pcmpstrSort   10.4 ns / 37.4 taktů (4.7 taktů/short)
vecPrefixSort  8.9 ns / 32.0 taktů (4.0 taktů/short)

Samo o sobě jde o nepříliš užitečný kus programu, ale dovedu si představit, že by se dal použít jako dílčí komponenta jiných algoritmů.

Přesně tato formulace je základ některých radix sortů a paralelního řazení.

0xDEADBEEF

Malé řazení ve vektoru - pcmpistrm