Vectorized CLZ en CTZ

Ik gebruikte ctz in een FPU-emulator waaraan ik werk, maar ontdekte hoe ik dit kon vermijden met floating-point trucs. Ik realiseerde me dat dit op een indirecte manier kan worden gegeneraliseerd naar een vectoriseerbare ctz-polyfill. Voor de volledigheid heb ik ook clz geïmplementeerd.

CLZ

Laten we beginnen met clz, de makkelijkste van de twee:

fn clz(x: u32) -> u32 {
    let a = 2.0f64.powi(-970);
    32 - ((f64::from_bits(a.to_bits() | x as u64) - a).to_bits() >> 52) as u32
}

Het algemene idee is als volgt: Floating-point exponenten zijn biased logaritmen van hun waarden. Door een 32-bit getal x in de mantisse van een waarde $2^k$ te plaatsen, krijgen we een double die $2^k(1 + 2^{-52}x)$ representeert. We kunnen vervolgens $2^k$ als double aftrekken om $2^{k-52}x$ te krijgen. Het extraheren van de exponent geeft $k - 52 + 31 - \text{clz}(x)$, waaruit clz kan worden berekend met een bitgewijze aftrekking. Op basis daarvan kan $k$ zo worden gekozen dat clz correct werkt voor $x=0$.

We hebben 64-bit doubles nodig om 32-bit inputs te verwerken; dit betekent helaas dat deze truc niet werkt voor willekeurige 64-bit inputs, maar slechts tot 52 bits.

Uitgaande van inputs en outputs die zijn opgeslagen in u64x4, compileert dit naar:

vpbroadcastq ymm1, [rip + bias]
vorpd ymm0, ymm0, [rip + a]
vsubpd ymm0, ymm0, [rip + a]
vpsrlq ymm0, ymm0, 52
vpsubq ymm0, ymm1, ymm0

a:
.quad 0x350000000000000, 0x350000000000000, 0x350000000000000, 0x350000000000000
bias:
.quad 32

Op mijn Haswell draait dit op 0,45 ns/iteratie, vergeleken met 1 ns voor de scalaire versie. Wanneer de latentie bepalend is, stijgen de getallen naar 2 ns tegenover 1 ns (maar als je latentie-gebonden bent bij een gevectoriseerde ctz, doe je waarschijnlijk iets fout).

Ian Qvist testte dit op Alder Lake en behaalde 0,29 ns/iteratie, vergeleken met 0,85 ns voor de scalaire versie, en 1,3 ns tegenover 0,85 ns wanneer de latentie bepalend is. Op moderne Intel CPU's zouden de cijfers hetzelfde of beter moeten zijn.

AMD CPU's maken lzcnt zo goedkoop dat een scalaire versie waarschijnlijk zal winnen. Houd er echter rekening mee dat Zen CPU's AVX-512 ondersteunen, wat vplzcntd heeft, dus dat is ook een optie.

CTZ

Nu voor ctz:

fn ctz(x: u32) -> u32 {
    let a = f64::from_bits((0x340000100000001 ^ x as u64) ^ (x as u64 + u32::MAX as u64))
    - f64::from_bits(0x340000000000000);
    (a.to_bits() >> 52) as u32
}

We beginnen met $x \oplus (x-1)$ om het laagste ingestelde bit te isoleren. ctz is gelijk aan de logaritme van die waarde, die we bepalen door $2^k$ bitgewijs op te tellen en $2^k$ als double af te trekken. Door vervolgens de exponent te inspecteren, bevat deze — bij een goed gekozen $k$ — de onbevooroordeelde (unbiased) ctz. We mengen $2^{32}$ vooraf in de mantisse en gebruiken $x + 2^{32}-1$ in plaats van $x-1$ om $x=0$ correct af te handelen, en mengen 1 vooraf in de mantisse om ervoor te zorgen dat oneven $x$ een a=0 genereren en geen trage subnormaal.

Deze functie compileert naar:

vpxor ymm1, ymm0, [rip + c1]
vpaddq ymm0, ymm0, [rip + u32_max]
vpxor ymm0, ymm1, ymm0
vsubpd ymm0, ymm0, [rip + c2]
vpsrlq ymm0, ymm0, 52

c1:
.qword 0x340000100000001, 0x340000100000001, 0x340000100000001, 0x340000100000001
u32_max:
.qword 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff
c2:
.qword 0x340000000000000, 0x340000000000000, 0x340000000000000, 0x340000000000000

Op Haswell draait dit op 0,49 ns/iteratie en 2,3 ns wanneer de latentie bepalend is. Op Alder Lake is dit 0,35 ns/iteratie en 1,3 ns bij latentie-gebondenheid. De vertraging ten opzichte van clz komt door het gebruik van één extra instructie. Dit kan worden vermeden door vpternlogq te gebruiken als AVX-512 aanwezig is, maar op dat punt kun je net zo goed vpopcntd uitvoeren op (x - 1) & !x. De scalaire versie gedraagt zich niet anders dan clz.

Update van 10 oktober

Nikolay Malkovsky wees erop dat de Bruijn-sequenties een andere vectoriseerbare aanpak bieden. Na wat testen kwam ik tot de volgende code:

const char table[32] = {
    0, 4, 5, 6, 11, 9, 7, 12, 15, 3, 10, 8, 14, 2, 13, 1,
    0, 4, 5, 6, 11, 9, 7, 12, 15, 3, 10, 8, 14, 2, 13, 1,
};

__m256i bit = _mm256_andnot_si256(x, _mm256_sub_epi32(x, _mm256_set1_epi32(1)));
__m256i high = _mm256_madd_epi16(
    _mm256_cmpeq_epi16(bit, _mm256_set1_epi16(-1)),
    _mm256_set1_epi16(-16)
);
__m256i index = _mm256_srli_epi32(_mm256_mullo_epi32(bit, _mm256_set1_epi32(0xf0a6f0a7)), 28);
__m256i low = _mm256_shuffle_epi8(_mm256_loadu_si256((__m256i*)table), index);
return _mm256_add_epi32(low, high);

We kunnen geen echte 32-byte LUT (Lookup Table) gebruiken omdat vpshufb geen 16-byte lanes kan overschrijden. De aanpak die ik in plaats daarvan gebruik is lastig uit te leggen, maar in essentie gebruiken we een 16-bit de Bruijn-sequentie die twee keer wordt herhaald om bits 0-3 van de ctz te berekenen, en voegen we vervolgens 16 of 32 toe, afhankelijk van welke helften nul zijn. De constante 0xf0a6f0a7 is een van de slechts vier magische constanten die dit laten werken.

Dit kost 1 ns op Haswell (0,7 ns op Alder Lake), maar heeft een dubbele doorvoersnelheid (throughput), waardoor het iets sneller kan zijn dan de FP-gebaseerde aanpak als het helpt om shuffling te vermijden.

Als je geen rekening hoeft te houden met $x=0$ (of wilt dat ctz(0) 0 is en niet 32), brengt het volgende fragment de tijd terug naar 0,82 ns:

__m256i high = _mm256_and_si256(
    _mm256_cmpgt_epi32(bit, _mm256_set1_epi32(0x7fff)),
    _mm256_set1_epi32(16)
);

Update van 11 oktober

Arya Dradjica en ik brainstormden over een 64-bit versie van het algoritme en slaagden erin te vinden wat ik denk dat het optimale algoritme is:

const char table[32] = { // x -> 4 - popcnt(x)
    4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0,
    4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0,
};

__m256i lut = _mm256_loadu_si256((__m256i*)table);
__m256i bit = _mm256_andnot_si256(x, _mm256_sub_epi64(x, _mm256_set1_epi64x(1)));
__m256i low = _mm256_shuffle_epi8(lut, bit);
__m256i high = _mm256_shuffle_epi8(lut, _mm256_srli_epi16(bit, 4));
return _mm256_sad_epu8(_mm256_add_epi8(low, high), _mm256_set1_epi8(8));

Dit geeft correct ctz(0)=64 terug en kost 0,63 ns op mijn Haswell (2,8 ns wanneer de latentie bepalend is). Het berekent de geïnverteerde popcounts van de twee nibbles van elke byte van (x - 1) & !x afzonderlijk (geïnverteerd vanwege enkele eigenaardigheden van pshufb) en gebruikt vervolgens vpsadbw om de horizontale som uit te voeren en de inversie in één keer te corrigeren.

vpsadbw kan slechts exact 8 bytes horizontaal optellen, wat de reden is dat dit algoritme gespecialiseerd is voor 64 bits en niet werkt met 32-bit inputs. De aanpak op basis van de de Bruijn-sequentie is in dat scenario nog steeds de snelste.