Skip to content

Instantly share code, notes, and snippets.

@zhuker
Created March 25, 2021 03:42
Population Count With Java Vector API
package zhuker;
import jdk.incubator.vector.*;
import org.junit.Test;
import java.util.stream.LongStream;
import static jdk.incubator.vector.VectorOperators.*;
import static org.junit.Assert.assertEquals;
public class BitCountTest {
private static final ByteVector lookup128 = ByteVector.fromArray(ByteVector.SPECIES_128, new byte[]{
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4,
}, 0);
private static final ByteVector lookup256 = ByteVector.fromArray(ByteVector.SPECIES_256, new byte[]{
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4,
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4
}, 0);
private static final ByteVector lookup512 = ByteVector.fromArray(ByteVector.SPECIES_512, new byte[]{
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4,
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4,
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4,
0, 1, 1, 2, 1, 2, 2, 3,
1, 2, 2, 3, 2, 3, 3, 4
}, 0);
static int bitCount256(ByteVector v) {
var low_mask = ByteVector.broadcast(ByteVector.SPECIES_256, 0x0f);
var lo = v.and(low_mask);
var hi = v.lanewise(LSHR, 4).and(low_mask);
var popcnt1 = lookup256.rearrange(lo.toShuffle());
var popcnt2 = lookup256.rearrange(hi.toShuffle());
var total = popcnt1.add(popcnt2);
var shortVector = total.castShape(ShortVector.SPECIES_256, 0);
var shortVector1 = total.castShape(ShortVector.SPECIES_256, 1);
return (int) (shortVector.reduceLanesToLong(ADD) + shortVector1.reduceLanesToLong(ADD));
}
static int bitCount128(ByteVector v) {
var low_mask = ByteVector.broadcast(ByteVector.SPECIES_128, 0x0f);
var lo = v.and(low_mask);
var hi = v.lanewise(LSHR, 4).and(low_mask);
var popcnt1 = lookup128.rearrange(lo.toShuffle());
var popcnt2 = lookup128.rearrange(hi.toShuffle());
var total = popcnt1.add(popcnt2);
return total.reduceLanes(ADD) & 0xff;
}
static int bitCount512(ByteVector v) {
var low_mask = ByteVector.broadcast(ByteVector.SPECIES_512, 0x0f);
var lo = v.and(low_mask);
var hi = v.lanewise(LSHR, 4).and(low_mask);
var popcnt1 = lookup512.rearrange(lo.toShuffle());
var popcnt2 = lookup512.rearrange(hi.toShuffle());
var total = popcnt1.add(popcnt2);
var sad0 = total.castShape(ShortVector.SPECIES_512, 0);
var sad1 = total.castShape(ShortVector.SPECIES_512, 1);
return (int) (sad0.reduceLanesToLong(ADD) + sad1.reduceLanesToLong(ADD));
}
static int bitCount(ByteVector v) {
return switch (v.length()) {
case 16 -> bitCount128(v);
case 32 -> bitCount256(v);
case 64 -> bitCount512(v);
default -> -1;
};
}
@Test
public void testBitCount() {
var random = new long[]{
0x1fff9244deadbeefL,
0x2fff9244deadbeefL,
0x3fff9244deadbeefL,
0x4fff9244deadbeefL,
};
assertEquals(169, bitCount(random));
ByteVector v = LongVector.fromArray(LongVector.SPECIES_256, random, 0).reinterpretAsBytes();
int i1 = bitCount(v);
assertEquals(169, i1);
var vmax = ByteVector.broadcast(ByteVector.SPECIES_256, (byte) 0xff);
assertEquals(256, bitCount(vmax));
}
@Test
public void testBitCount128() {
var random = new long[]{
0x1fff9244deadbeefL,
0x3fff9244deadbeefL,
};
assertEquals(85, bitCount(random));
var v = LongVector.fromArray(LongVector.SPECIES_128, random, 0).reinterpretAsBytes();
assertEquals(85, bitCount(v));
var vmax = ByteVector.broadcast(ByteVector.SPECIES_128, (byte) 0xff);
assertEquals(128, bitCount(vmax));
}
int bitCount(long[] array) {
return LongStream.of(array).mapToInt(x -> Long.bitCount(x)).sum();
}
@Test
public void testBitCount512() {
var random = new long[]{
0x1fff9244deadbeefL,
0x2fff9244deadbeefL,
0x3fff9244deadbeefL,
0x4fff9244deadbeefL,
0x5fff9244deadbeefL,
0x6fff9244deadbeefL,
0x7fff9244deadbeefL,
0x8fff9244deadbeefL};
assertEquals(341, bitCount(random));
var v = LongVector.fromArray(LongVector.SPECIES_512, random, 0).reinterpretAsBytes();
int i1 = bitCount(v);
assertEquals(341, i1);
var vmax = ByteVector.broadcast(ByteVector.SPECIES_512, (byte) 0xff);
assertEquals(512, bitCount(vmax));
}
}
@zhuker
Copy link
Author

zhuker commented Mar 25, 2021

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment