merykitty's second attempt

This commit is contained in:
Quan Anh Mai 2024-01-11 02:24:19 +07:00 committed by GitHub
parent 786a52034c
commit 97b1f014ad
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 144 additions and 197 deletions

View File

@ -16,5 +16,5 @@
# #
JAVA_OPTS="--enable-preview --add-modules=jdk.incubator.vector" # -XX:+UnlockDiagnosticVMOptions -XX:PrintAssemblyOptions=intel -XX:CompileCommand=print,*.CalculateAverage_merykitty::iterate" JAVA_OPTS="--enable-preview --add-modules=jdk.incubator.vector -XX:-TieredCompilation" # -XX:+UnlockDiagnosticVMOptions -XX:PrintAssemblyOptions=intel -XX:CompileCommand=print,*.CalculateAverage_merykitty::iterate"
java $JAVA_OPTS --class-path target/average-1.0.0-SNAPSHOT.jar dev.morling.onebrc.CalculateAverage_merykitty java $JAVA_OPTS --class-path target/average-1.0.0-SNAPSHOT.jar dev.morling.onebrc.CalculateAverage_merykitty

View File

@ -25,8 +25,6 @@ import java.nio.channels.FileChannel.MapMode;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.nio.file.Path; import java.nio.file.Path;
import java.nio.file.StandardOpenOption; import java.nio.file.StandardOpenOption;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map; import java.util.Map;
import java.util.TreeMap; import java.util.TreeMap;
import jdk.incubator.vector.ByteVector; import jdk.incubator.vector.ByteVector;
@ -35,13 +33,21 @@ import jdk.incubator.vector.VectorSpecies;
public class CalculateAverage_merykitty { public class CalculateAverage_merykitty {
private static final String FILE = "./measurements.txt"; private static final String FILE = "./measurements.txt";
private static final VectorSpecies<Byte> BYTE_SPECIES = ByteVector.SPECIES_PREFERRED; private static final VectorSpecies<Byte> BYTE_SPECIES = ByteVector.SPECIES_PREFERRED.length() >= 32
? ByteVector.SPECIES_256
: ByteVector.SPECIES_128;
private static final ValueLayout.OfLong JAVA_LONG_LT = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN); private static final ValueLayout.OfLong JAVA_LONG_LT = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
private static final long KEY_MAX_SIZE = 100; private static final long KEY_MAX_SIZE = 100;
private record ResultRow(double min, double mean, double max) { private static class Aggregator {
private int keySize;
private long min = Integer.MAX_VALUE;
private long max = Integer.MIN_VALUE;
private long sum;
private long count;
public String toString() { public String toString() {
return round(min) + "/" + round(mean) + "/" + round(max); return round(min / 10.) + "/" + round(sum / (double) (10 * count)) + "/" + round(max / 10.);
} }
private double round(double value) { private double round(double value) {
@ -49,96 +55,100 @@ public class CalculateAverage_merykitty {
} }
} }
private static class Aggregator {
private long min = Integer.MAX_VALUE;
private long max = Integer.MIN_VALUE;
private long sum;
private long count;
}
// An open-address map that is specialized for this task // An open-address map that is specialized for this task
private static class PoorManMap { private static class PoorManMap {
static final int R_LOAD_FACTOR = 2;
private static class PoorManMapNode { // 100-byte key + 4-byte hash + 4-byte size +
byte[] data; // 2-byte min + 2-byte max + 8-byte sum + 8-byte count
long size; private static final int KEY_SIZE = 128;
int hash;
Aggregator aggr;
PoorManMapNode(MemorySegment data, long offset, long size, int hash) { // There is an assumption that map size <= 10000;
this.hash = hash; private static final int CAPACITY = 1 << 17;
this.size = size; private static final int BUCKET_MASK = CAPACITY - 1;
this.data = new byte[BYTE_SPECIES.vectorByteSize() + (int) KEY_MAX_SIZE];
this.aggr = new Aggregator(); byte[] keyData;
MemorySegment.copy(data, offset, MemorySegment.ofArray(this.data), BYTE_SPECIES.vectorByteSize(), size); Aggregator[] nodes;
}
PoorManMap() {
this.keyData = new byte[CAPACITY * KEY_SIZE];
this.nodes = new Aggregator[CAPACITY];
} }
MemorySegment data; void observe(Aggregator node, long value) {
PoorManMapNode[] nodes; node.min = Math.min(node.min, value);
int size; node.max = Math.max(node.max, value);
node.sum += value;
PoorManMap(MemorySegment data) { node.count++;
this.data = data;
this.nodes = new PoorManMapNode[1 << 10];
} }
Aggregator indexSimple(long offset, long size, int hash) { Aggregator indexSimple(MemorySegment data, long offset, int size) {
hash = rehash(hash); int x;
int bucketMask = nodes.length - 1; int y;
int bucket = hash & bucketMask; if (size >= Integer.BYTES) {
for (;; bucket = (bucket + 1) & bucketMask) { x = data.get(ValueLayout.JAVA_INT_UNALIGNED, offset);
PoorManMapNode node = nodes[bucket]; y = data.get(ValueLayout.JAVA_INT_UNALIGNED, offset + size - Integer.BYTES);
}
else {
x = data.get(ValueLayout.JAVA_BYTE, offset);
y = data.get(ValueLayout.JAVA_BYTE, offset + size - Byte.BYTES);
}
int hash = hash(x, y);
int bucket = hash & BUCKET_MASK;
for (;; bucket = (bucket + 1) & BUCKET_MASK) {
var node = this.nodes[bucket];
if (node == null) { if (node == null) {
this.size++; return insertInto(bucket, data, offset, size);
if (this.size * R_LOAD_FACTOR > nodes.length) {
grow();
bucketMask = nodes.length - 1;
for (bucket = hash & bucketMask; nodes[bucket] != null; bucket = (bucket + 1) & bucketMask) {
} }
} else if (keyEqualScalar(bucket, data, offset, size)) {
node = new PoorManMapNode(this.data, offset, size, hash); return node;
nodes[bucket] = node;
return node.aggr;
}
else if (keyEqualScalar(node, offset, size, hash)) {
return node.aggr;
} }
} }
} }
void grow() { Aggregator insertInto(int bucket, MemorySegment data, long offset, int size) {
var oldNodes = this.nodes; var node = new Aggregator();
var newNodes = new PoorManMapNode[oldNodes.length * 2]; node.keySize = size;
int bucketMask = newNodes.length - 1; this.nodes[bucket] = node;
for (var node : oldNodes) { MemorySegment.copy(data, offset, MemorySegment.ofArray(this.keyData), (long) bucket * KEY_SIZE, size);
return node;
}
void mergeInto(Map<String, Aggregator> target) {
for (int i = 0; i < CAPACITY; i++) {
var node = this.nodes[i];
if (node == null) { if (node == null) {
continue; continue;
} }
int bucket = node.hash & bucketMask;
for (; newNodes[bucket] != null; bucket = (bucket + 1) & bucketMask) { String key = new String(this.keyData, i * KEY_SIZE, node.keySize, StandardCharsets.UTF_8);
} target.compute(key, (k, v) -> {
newNodes[bucket] = node; if (v == null) {
} v = new Aggregator();
this.nodes = newNodes;
} }
static int rehash(int x) { v.min = Math.min(v.min, node.min);
x = ((x >>> 16) ^ x) * 0x45d9f3b; v.max = Math.max(v.max, node.max);
x = ((x >>> 16) ^ x) * 0x45d9f3b; v.sum += node.sum;
x = (x >>> 16) ^ x; v.count += node.count;
return x; return v;
});
}
} }
private boolean keyEqualScalar(PoorManMapNode node, long offset, long size, int hash) { static int hash(int x, int y) {
if (node.hash != hash || node.size != size) { int seed = 0x9E3779B9;
int rotate = 5;
return (Integer.rotateLeft(x * seed, rotate) ^ y) * seed; // FxHash
}
private boolean keyEqualScalar(int bucket, MemorySegment data, long offset, int size) {
if (this.nodes[bucket].keySize != size) {
return false; return false;
} }
// Be simple // Be simple
for (int i = 0; i < size; i++) { for (int i = 0; i < size; i++) {
int c1 = node.data[BYTE_SPECIES.vectorByteSize() + i]; int c1 = this.keyData[bucket * KEY_SIZE + i];
int c2 = data.get(ValueLayout.JAVA_BYTE, offset + i); int c2 = data.get(ValueLayout.JAVA_BYTE, offset + i);
if (c1 != c2) { if (c1 != c2) {
return false; return false;
@ -152,7 +162,7 @@ public class CalculateAverage_merykitty {
// 1 - 2 digits to the left and 1 digits to the right of the separator to a // 1 - 2 digits to the left and 1 digits to the right of the separator to a
// fix-precision format. It returns the offset of the next line (presumably followed // fix-precision format. It returns the offset of the next line (presumably followed
// the final digit and a '\n') // the final digit and a '\n')
private static long parseDataPoint(Aggregator aggr, MemorySegment data, long offset) { private static long parseDataPoint(PoorManMap aggrMap, Aggregator node, MemorySegment data, long offset) {
long word = data.get(JAVA_LONG_LT, offset); long word = data.get(JAVA_LONG_LT, offset);
// The 4th binary digit of the ascii of a digit is 1 while // The 4th binary digit of the ascii of a digit is 1 while
// that of the '.' is 0. This finds the decimal separator // that of the '.' is 0. This finds the decimal separator
@ -176,16 +186,13 @@ public class CalculateAverage_merykitty {
// That was close :) // That was close :)
long absValue = ((digits * 0x640a0001) >>> 32) & 0x3FF; long absValue = ((digits * 0x640a0001) >>> 32) & 0x3FF;
long value = (absValue ^ signed) - signed; long value = (absValue ^ signed) - signed;
aggr.min = Math.min(value, aggr.min); aggrMap.observe(node, value);
aggr.max = Math.max(value, aggr.max);
aggr.sum += value;
aggr.count++;
return offset + (decimalSepPos >>> 3) + 3; return offset + (decimalSepPos >>> 3) + 3;
} }
// Tail processing version of the above, do not over-fetch and be simple // Tail processing version of the above, do not over-fetch and be simple
private static long parseDataPointTail(Aggregator aggr, MemorySegment data, long offset) { private static long parseDataPointSimple(PoorManMap aggrMap, Aggregator node, MemorySegment data, long offset) {
int point = 0; int value = 0;
boolean negative = false; boolean negative = false;
if (data.get(ValueLayout.JAVA_BYTE, offset) == '-') { if (data.get(ValueLayout.JAVA_BYTE, offset) == '-') {
negative = true; negative = true;
@ -195,110 +202,80 @@ public class CalculateAverage_merykitty {
int c = data.get(ValueLayout.JAVA_BYTE, offset); int c = data.get(ValueLayout.JAVA_BYTE, offset);
if (c == '.') { if (c == '.') {
c = data.get(ValueLayout.JAVA_BYTE, offset + 1); c = data.get(ValueLayout.JAVA_BYTE, offset + 1);
point = point * 10 + (c - '0'); value = value * 10 + (c - '0');
offset += 3; offset += 3;
break; break;
} }
point = point * 10 + (c - '0'); value = value * 10 + (c - '0');
} }
point = negative ? -point : point; value = negative ? -value : value;
aggr.min = Math.min(point, aggr.min); aggrMap.observe(node, value);
aggr.max = Math.max(point, aggr.max);
aggr.sum += point;
aggr.count++;
return offset; return offset;
} }
// An iteration of the main parse loop, parse some lines starting from offset. // An iteration of the main parse loop, parse a line starting from offset.
// This requires offset to be the start of a line and there is spare space so // This requires offset to be the start of the line and there is spare space so
// that we have relative freedom in processing // that we have relative freedom in processing
// It returns the offset of the next line that it needs to be processed // It returns the offset of the next line that it needs processing
private static long iterate(PoorManMap aggrMap, MemorySegment data, long offset) { private static long iterate(PoorManMap aggrMap, MemorySegment data, long offset) {
// This method fetches a segment of the file starting from offset and returns after
// finishing processing that segment
var line = ByteVector.fromMemorySegment(BYTE_SPECIES, data, offset, ByteOrder.nativeOrder()); var line = ByteVector.fromMemorySegment(BYTE_SPECIES, data, offset, ByteOrder.nativeOrder());
// Find the delimiter ';' // Find the delimiter ';'
long semicolons = line.compare(VectorOperators.EQ, ';').toLong(); int keySize = line.compare(VectorOperators.EQ, ';').firstTrue();
// If we cannot find the delimiter in the current segment, that means the key is // If we cannot find the delimiter in the vector, that means the key is
// longer than the segment, fall back to scalar processing // longer than the vector, fall back to scalar processing
if (semicolons == 0) { if (keySize == BYTE_SPECIES.vectorByteSize()) {
long semicolonPos = BYTE_SPECIES.vectorByteSize(); while (data.get(ValueLayout.JAVA_BYTE, offset + keySize) != ';') {
for (; data.get(ValueLayout.JAVA_BYTE, offset + semicolonPos) != ';'; semicolonPos++) { keySize++;
} }
int hash = line.reinterpretAsInts().lane(0); var node = aggrMap.indexSimple(data, offset, keySize);
var aggr = aggrMap.indexSimple(offset, semicolonPos, hash); return parseDataPoint(aggrMap, node, data, offset + 1 + keySize);
return parseDataPoint(aggr, data, offset + 1 + semicolonPos);
}
long currOffset = offset;
while (true) {
// Process line by line, currOffset is the offset of the current line in
// the file, localOffset is the offset of the current line with respect
// to the start of the iteration segment
int localOffset = (int) (currOffset - offset);
// The key length
long semicolonPos = Long.numberOfTrailingZeros(semicolons) - localOffset;
int hash = data.get(ValueLayout.JAVA_INT_UNALIGNED, currOffset);
if (semicolonPos < Integer.BYTES) {
hash = (byte) hash;
} }
// We inline the searching of the value in the hash map // We inline the searching of the value in the hash map
Aggregator aggr; int x;
hash = PoorManMap.rehash(hash); int y;
int bucketMask = aggrMap.nodes.length - 1; if (keySize >= Integer.BYTES) {
int bucket = hash & bucketMask; x = data.get(ValueLayout.JAVA_INT_UNALIGNED, offset);
for (;; bucket = (bucket + 1) & bucketMask) { y = data.get(ValueLayout.JAVA_INT_UNALIGNED, offset + keySize - Integer.BYTES);
PoorManMap.PoorManMapNode node = aggrMap.nodes[bucket]; }
else {
x = data.get(ValueLayout.JAVA_BYTE, offset);
y = data.get(ValueLayout.JAVA_BYTE, offset + keySize - Byte.BYTES);
}
int hash = PoorManMap.hash(x, y);
int bucket = hash & PoorManMap.BUCKET_MASK;
Aggregator node;
for (;; bucket = (bucket + 1) & PoorManMap.BUCKET_MASK) {
node = aggrMap.nodes[bucket];
if (node == null) { if (node == null) {
aggrMap.size++; node = aggrMap.insertInto(bucket, data, offset, keySize);
if (aggrMap.size * PoorManMap.R_LOAD_FACTOR > aggrMap.nodes.length) {
aggrMap.grow();
bucketMask = aggrMap.nodes.length - 1;
for (bucket = hash & bucketMask; aggrMap.nodes[bucket] != null; bucket = (bucket + 1) & bucketMask) {
}
}
node = new PoorManMap.PoorManMapNode(data, currOffset, semicolonPos, hash);
aggrMap.nodes[bucket] = node;
aggr = node.aggr;
break; break;
} }
if (node.keySize != keySize) {
if (node.hash != hash || node.size != semicolonPos) {
continue; continue;
} }
// The technique here is to align the key in both vectors so that we can do an var nodeKey = ByteVector.fromArray(BYTE_SPECIES, aggrMap.keyData, bucket * PoorManMap.KEY_SIZE);
// element-wise comparison and check if all characters match long eqMask = line.compare(VectorOperators.EQ, nodeKey).toLong();
var nodeKey = ByteVector.fromArray(BYTE_SPECIES, node.data, BYTE_SPECIES.length() - localOffset); long validMask = -1L >>> -keySize;
var eqMask = line.compare(VectorOperators.EQ, nodeKey).toLong();
long validMask = (-1L >>> -semicolonPos) << localOffset;
if ((eqMask & validMask) == validMask) { if ((eqMask & validMask) == validMask) {
aggr = node.aggr;
break; break;
} }
} }
long nextOffset = parseDataPoint(aggr, data, currOffset + 1 + semicolonPos); return parseDataPoint(aggrMap, node, data, offset + keySize + 1);
semicolons &= (semicolons - 1);
if (semicolons == 0) {
return nextOffset;
}
currOffset = nextOffset;
}
} }
// Process all lines that start in [offset, limit) // Process all lines that start in [offset, limit)
private static PoorManMap processFile(MemorySegment data, long offset, long limit) { private static PoorManMap processFile(MemorySegment data, long offset, long limit) {
var aggrMap = new PoorManMap(data); var aggrMap = new PoorManMap();
// Find the start of a new line // Find the start of a new line
if (offset != 0) { if (offset != 0) {
offset--; offset--;
for (; offset < limit;) { while (offset < limit) {
if (data.get(ValueLayout.JAVA_BYTE, offset++) == '\n') { if (data.get(ValueLayout.JAVA_BYTE, offset++) == '\n') {
break; break;
} }
@ -318,18 +295,12 @@ public class CalculateAverage_merykitty {
// Now we are at the tail, just be simple // Now we are at the tail, just be simple
while (offset < limit) { while (offset < limit) {
long semicolonPos = 0; int keySize = 0;
for (; data.get(ValueLayout.JAVA_BYTE, offset + semicolonPos) != ';'; semicolonPos++) { while (data.get(ValueLayout.JAVA_BYTE, offset + keySize) != ';') {
keySize++;
} }
int hash; var node = aggrMap.indexSimple(data, offset, keySize);
if (semicolonPos >= Integer.BYTES) { offset = parseDataPointSimple(aggrMap, node, data, offset + 1 + keySize);
hash = data.get(ValueLayout.JAVA_INT_UNALIGNED, offset);
}
else {
hash = data.get(ValueLayout.JAVA_BYTE, offset);
}
var aggr = aggrMap.indexSimple(offset, semicolonPos, hash);
offset = parseDataPointTail(aggr, data, offset + 1 + semicolonPos);
} }
return aggrMap; return aggrMap;
@ -337,7 +308,7 @@ public class CalculateAverage_merykitty {
public static void main(String[] args) throws InterruptedException, IOException { public static void main(String[] args) throws InterruptedException, IOException {
int processorCnt = Runtime.getRuntime().availableProcessors(); int processorCnt = Runtime.getRuntime().availableProcessors();
var res = HashMap.<String, Aggregator> newHashMap(processorCnt); var res = new TreeMap<String, Aggregator>();
try (var file = FileChannel.open(Path.of(FILE), StandardOpenOption.READ); try (var file = FileChannel.open(Path.of(FILE), StandardOpenOption.READ);
var arena = Arena.ofShared()) { var arena = Arena.ofShared()) {
var data = file.map(MapMode.READ_ONLY, 0, file.size(), arena); var data = file.map(MapMode.READ_ONLY, 0, file.size(), arena);
@ -348,9 +319,7 @@ public class CalculateAverage_merykitty {
int index = i; int index = i;
long offset = i * chunkSize; long offset = i * chunkSize;
long limit = Math.min((i + 1) * chunkSize, data.byteSize()); long limit = Math.min((i + 1) * chunkSize, data.byteSize());
var thread = new Thread(() -> { var thread = new Thread(() -> resultList[index] = processFile(data, offset, limit));
resultList[index] = processFile(data, offset, limit);
});
threadList[index] = thread; threadList[index] = thread;
thread.start(); thread.start();
} }
@ -360,32 +329,10 @@ public class CalculateAverage_merykitty {
// Collect the results // Collect the results
for (var aggrMap : resultList) { for (var aggrMap : resultList) {
for (var node : aggrMap.nodes) { aggrMap.mergeInto(res);
if (node == null) {
continue;
}
byte[] keyData = Arrays.copyOfRange(node.data, BYTE_SPECIES.vectorByteSize(), BYTE_SPECIES.vectorByteSize() + (int) node.size);
String key = new String(keyData, StandardCharsets.UTF_8);
var aggr = node.aggr;
var resAggr = new Aggregator();
var existingAggr = res.putIfAbsent(key, resAggr);
if (existingAggr != null) {
resAggr = existingAggr;
}
resAggr.min = Math.min(resAggr.min, aggr.min);
resAggr.max = Math.max(resAggr.max, aggr.max);
resAggr.sum += aggr.sum;
resAggr.count += aggr.count;
}
} }
} }
Map<String, ResultRow> measurements = new TreeMap<>(); System.out.println(res);
for (var entry : res.entrySet()) {
String key = entry.getKey();
var aggr = entry.getValue();
measurements.put(key, new ResultRow((double) aggr.min / 10, (double) aggr.sum / (aggr.count * 10), (double) aggr.max / 10));
}
System.out.println(measurements);
} }
} }