diff --git a/src/main/java/dev/morling/onebrc/CalculateAverage_royvanrijn.java b/src/main/java/dev/morling/onebrc/CalculateAverage_royvanrijn.java index c74415e..aa22bef 100644 --- a/src/main/java/dev/morling/onebrc/CalculateAverage_royvanrijn.java +++ b/src/main/java/dev/morling/onebrc/CalculateAverage_royvanrijn.java @@ -15,24 +15,22 @@ */ package dev.morling.onebrc; -import java.io.File; import java.io.IOException; -import java.io.RandomAccessFile; -import java.nio.ByteBuffer; +import java.lang.foreign.Arena; +import java.lang.reflect.Field; import java.nio.ByteOrder; import java.nio.channels.FileChannel; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.StandardOpenOption; -import java.util.ArrayList; import java.util.Arrays; -import java.util.List; import java.util.Objects; import java.util.TreeMap; import java.util.stream.Collectors; +import java.util.stream.IntStream; import java.util.stream.Stream; +import sun.misc.Unsafe; + /** * Changelog: * @@ -50,6 +48,7 @@ import java.util.stream.Stream; * Not using SWAR for EOL: 2850 ms * Inlining hash calculation: 2450 ms * Replacing branchless code: 2200 ms (sometimes we need to kill the things we love) + * Added unsafe memory access: 1900 ms (keeping the long[] small and local) * * Best performing JVM on MacBook M2 Pro: 21.0.1-graal * `sdk use java 21.0.1-graal` @@ -58,7 +57,140 @@ import java.util.stream.Stream; public class CalculateAverage_royvanrijn { private static final String FILE = "./measurements.txt"; - // private static final String FILE = "./src/test/resources/samples/measurements-10000-unique-keys.txt"; + + private static final Unsafe UNSAFE = initUnsafe(); + private static final boolean isBigEndian = ByteOrder.nativeOrder().equals(ByteOrder.BIG_ENDIAN); + + private static Unsafe initUnsafe() { + try { + Field theUnsafe = Unsafe.class.getDeclaredField("theUnsafe"); + theUnsafe.setAccessible(true); + return (Unsafe) theUnsafe.get(Unsafe.class); + } + catch (NoSuchFieldException | IllegalAccessException e) { + throw new RuntimeException(e); + } + } + + public static void main(String[] args) throws Exception { + new CalculateAverage_royvanrijn().run(); + } + + public void run() throws Exception { + + // Calculate input segments. + int numberOfChunks = Runtime.getRuntime().availableProcessors(); + long[] chunks = getSegments(numberOfChunks); + + // Parallel processing of segments. + TreeMap results = IntStream.range(0, chunks.length - 1) + .mapToObj(chunkIndex -> process(chunks[chunkIndex], chunks[chunkIndex + 1])).parallel() + .flatMap(MeasurementRepository::get) + .collect(Collectors.toMap(e -> e.city, MeasurementRepository.Entry::measurement, Measurement::updateWith, TreeMap::new)); + + System.out.println(results); + } + + private static long[] getSegments(int numberOfChunks) throws IOException { + try (var fileChannel = FileChannel.open(Path.of(FILE), StandardOpenOption.READ)) { + long fileSize = fileChannel.size(); + long segmentSize = (fileSize + numberOfChunks - 1) / numberOfChunks; + long[] chunks = new long[numberOfChunks + 1]; + long mappedAddress = fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileSize, Arena.global()).address(); + chunks[0] = mappedAddress; + long endAddress = mappedAddress + fileSize; + for (int i = 1; i < numberOfChunks; ++i) { + long chunkAddress = mappedAddress + i * segmentSize; + // Align to first row start. + while (chunkAddress < endAddress && UNSAFE.getByte(chunkAddress++) != '\n') { + // nop + } + chunks[i] = Math.min(chunkAddress, endAddress); + } + chunks[numberOfChunks] = endAddress; + return chunks; + } + } + + private MeasurementRepository process(long fromAddress, long toAddress) { + + MeasurementRepository repository = new MeasurementRepository(); + long ptr = fromAddress; + long[] dataBuffer = new long[16]; + while ((ptr = processEntity(dataBuffer, ptr, toAddress, repository)) < toAddress) + ; + + return repository; + } + + private static final long SEPARATOR_PATTERN = compilePattern((byte) ';'); + + /** + * Already looping the longs here, lets shoehorn in making a hash + */ + private long processEntity(final long[] data, final long start, final long limit, final MeasurementRepository measurementRepository) { + int hash = 1; + long i; + int dataPtr = 0; + for (i = start; i <= limit - 8; i += 8) { + long word = UNSAFE.getLong(i); + if (isBigEndian) { + word = Long.reverseBytes(word); // Reversing the bytes is the cheapest way to do this + } + final long match = word ^ SEPARATOR_PATTERN; + long mask = ((match - 0x0101010101010101L) & ~match) & 0x8080808080808080L; + + if (mask != 0) { + + final long partialWord = word & ((mask >> 7) - 1); + hash = longHashStep(hash, partialWord); + data[dataPtr] = partialWord; + + final int index = Long.numberOfTrailingZeros(mask) >> 3; + return process(start, i + index, hash, data, measurementRepository); + } + data[dataPtr++] = word; + hash = longHashStep(hash, word); + } + // Handle remaining bytes near the limit of the buffer: + long partialWord = 0; + int len = 0; + for (; i < limit; i++) { + byte read; + if ((read = UNSAFE.getByte(i)) == ';') { + hash = longHashStep(hash, partialWord); + data[dataPtr] = partialWord; + return process(start, i, hash, data, measurementRepository); + } + partialWord = partialWord | ((long) read << (len << 3)); + len++; + } + return limit; + } + + private static final long DOT_BITS = 0x10101000; + private static final long MAGIC_MULTIPLIER = (100 * 0x1000000 + 10 * 0x10000 + 1); + + private long process(final long startAddress, final long delimiterAddress, final int hash, final long[] data, final MeasurementRepository measurementRepository) { + + long word = UNSAFE.getLong(delimiterAddress + 1); + if (isBigEndian) { + word = Long.reverseBytes(word); + } + final long invWord = ~word; + final int decimalSepPos = Long.numberOfTrailingZeros(invWord & DOT_BITS); + final long signed = (invWord << 59) >> 63; + final long designMask = ~(signed & 0xFF); + final long digits = ((word & designMask) << (28 - decimalSepPos)) & 0x0F000F0F00L; + final long absValue = ((digits * MAGIC_MULTIPLIER) >>> 32) & 0x3FF; + final int measurement = (int) ((absValue ^ signed) - signed); + + // Store: + measurementRepository.update(startAddress, data, (int) (delimiterAddress - startAddress), hash, measurement); + + return delimiterAddress + (decimalSepPos >> 3) + 4; // Determine next start: + // return nextAddress; + } static final class Measurement { int min, max, count; @@ -94,176 +226,6 @@ public class CalculateAverage_royvanrijn { } } - public static void main(String[] args) throws Exception { - new CalculateAverage_royvanrijn().run(); - // new CalculateAverage_royvanrijn().runTests(); - } - - private void run() throws Exception { - - var results = getFileSegments(new File(FILE)).stream().map(segment -> { - - long segmentEnd = segment.end(); - try (var fileChannel = (FileChannel) Files.newByteChannel(Path.of(FILE), StandardOpenOption.READ)) { - var bb = fileChannel.map(FileChannel.MapMode.READ_ONLY, segment.start(), segmentEnd - segment.start()); - - // Work with any UTF-8 city name, up to 100 in length: - var cityNameAsLongArray = new long[16]; - var delimiterPointerAndHash = new int[2]; - - // Calculate using native ordering (fastest?): - bb.order(ByteOrder.nativeOrder()); - - // Record the order it is and calculate accordingly: - final boolean bufferIsBigEndian = bb.order().equals(ByteOrder.BIG_ENDIAN); - MeasurementRepository measurements = new MeasurementRepository(); - - int startPointer; - int limit = bb.limit(); - while ((startPointer = bb.position()) < limit) { - - int delimiterPointer, endPointer; - - // SWAR method to find delimiter *and* record the cityname as long[] *and* calculate a hash: - findNextDelimiterAndCalculateHash(bb, SEPARATOR_PATTERN, startPointer, limit, delimiterPointerAndHash, cityNameAsLongArray, bufferIsBigEndian); - delimiterPointer = delimiterPointerAndHash[0]; - - // Simple lookup is faster for '\n' (just three options) - if (delimiterPointer >= limit) { - return measurements; - } - // Extract the measurement value (10x): - final int cityNameLength = delimiterPointer - startPointer; - - int measuredValue; - int neg = 1; - if (bb.get(++delimiterPointer) == '-') { - neg = -1; - delimiterPointer++; - } - byte dot; - if ((dot = (bb.get(delimiterPointer + 1))) == '.') { - measuredValue = neg * ((bb.get(delimiterPointer)) * 10 + (bb.get(delimiterPointer + 2)) - 528); - endPointer = delimiterPointer + 3; - } - else { - measuredValue = neg * (bb.get(delimiterPointer) * 100 + dot * 10 + bb.get(delimiterPointer + 3) - 5328); - endPointer = delimiterPointer + 4; - } - - // Store everything in a custom hashtable: - measurements.update(cityNameAsLongArray, bb, cityNameLength, delimiterPointerAndHash[1]).updateWith(measuredValue); - - bb.position(endPointer + 1); // skip to next line. - } - return measurements; - } - catch (IOException e) { - throw new RuntimeException(e); - } - }).parallel() - .flatMap(v -> v.get()) - .collect(Collectors.toMap(e -> e.cityName, MeasurementRepository.Entry::measurement, Measurement::updateWith, TreeMap::new)); - - System.out.println(results); - - // System.out.println("Processed: " + results.entrySet().stream().mapToLong(e -> e.getValue().count).sum()); - } - - /** - * -------- This section contains SWAR code (SIMD Within A Register) which processes a bytebuffer as longs to find values: - */ - private static final long SEPARATOR_PATTERN = compilePattern((byte) ';'); - - /** - * Already looping the longs here, lets shoehorn in making a hash - */ - private void findNextDelimiterAndCalculateHash(final ByteBuffer bb, final long pattern, final int start, final int limit, final int[] output, - final long[] asLong, final boolean bufferBigEndian) { - int hash = 1; - int i; - int lCnt = 0; - for (i = start; i <= limit - 8; i += 8) { - long word = bb.getLong(i); - if (bufferBigEndian) { - word = Long.reverseBytes(word); // Reversing the bytes is the cheapest way to do this - } - final long match = word ^ pattern; - long mask = ((match - 0x0101010101010101L) & ~match) & 0x8080808080808080L; - - if (mask != 0) { - final int index = Long.numberOfTrailingZeros(mask) >> 3; - output[0] = (i + index); - - final long partialHash = word & ((mask >> 7) - 1); - asLong[lCnt] = partialHash; - output[1] = longHashStep(hash, partialHash); - return; - } - asLong[lCnt++] = word; - hash = longHashStep(hash, word); - } - // Handle remaining bytes near the limit of the buffer: - long partialHash = 0; - int len = 0; - for (; i < limit; i++) { - byte read; - if ((read = bb.get(i)) == (byte) pattern) { - asLong[lCnt] = partialHash; - output[0] = i; - output[1] = longHashStep(hash, partialHash); - return; - } - partialHash = partialHash | ((long) read << (len << 3)); - len++; - } - output[0] = limit; // delimiter not found - } - - private static int longHashStep(final int hash, final long word) { - return 31 * hash + (int) (word ^ (word >>> 32)); - } - - private static long compilePattern(final byte value) { - return ((long) value << 56) | ((long) value << 48) | ((long) value << 40) | ((long) value << 32) | - ((long) value << 24) | ((long) value << 16) | ((long) value << 8) | (long) value; - } - - record FileSegment(long start, long end) { - } - - private static List getFileSegments(final File file) throws IOException { - final int numberOfSegments = Runtime.getRuntime().availableProcessors(); - final long fileSize = file.length(); - final long segmentSize = fileSize / numberOfSegments; - final List segments = new ArrayList<>(); - if (segmentSize < 1000) { - segments.add(new FileSegment(0, fileSize)); - return segments; - } - try (RandomAccessFile randomAccessFile = new RandomAccessFile(file, "r")) { - long segStart = 0; - long segEnd = segmentSize; - while (segStart < fileSize) { - segEnd = findSegment(randomAccessFile, segEnd, fileSize); - segments.add(new FileSegment(segStart, segEnd)); - segStart = segEnd; // Just re-use the end and go from there. - segEnd = Math.min(fileSize, segEnd + segmentSize); - } - } - return segments; - } - - private static long findSegment(RandomAccessFile raf, long location, final long fileSize) throws IOException { - raf.seek(location); - while (location < fileSize) { - location++; - if (raf.read() == '\n') - return location; - } - return location; - } - // branchless max (unprecise for large numbers, but good enough) static int max(final int a, final int b) { final int diff = a - b; @@ -278,85 +240,69 @@ public class CalculateAverage_royvanrijn { return b + (diff & dsgn); } + private static int longHashStep(final int hash, final long word) { + return 31 * hash + (int) (word ^ (word >>> 32)); + } + + private static long compilePattern(final byte value) { + return ((long) value << 56) | ((long) value << 48) | ((long) value << 40) | ((long) value << 32) | + ((long) value << 24) | ((long) value << 16) | ((long) value << 8) | (long) value; + } + /** * A normal Java HashMap does all these safety things like boundary checks... we don't need that, we need speeeed. * * So I've written an extremely simple linear probing hashmap that should work well enough. */ class MeasurementRepository { - private int tableSize = 1 << 20; // can grow in theory, made large enough not to (this is faster) + private int tableSize = 1 << 20; // large enough for the contest. private int tableMask = (tableSize - 1); - private int tableLimit = (int) (tableSize * LOAD_FACTOR); - private int tableFilled = 0; - private static final float LOAD_FACTOR = 0.8f; - private Entry[] table = new Entry[tableSize]; + private MeasurementRepository.Entry[] table = new MeasurementRepository.Entry[tableSize]; + + record Entry(long address, long[] data, int length, int hash, String city, Measurement measurement) { - record Entry(int hash, long[] nameBytesInLong, String cityName, Measurement measurement) { @Override public String toString() { - return cityName + "=" + measurement; + return city + "=" + measurement; } } - public Measurement update(long[] nameBytesInLong, ByteBuffer bb, int length, int calculatedHash) { + public void update(long address, long[] data, int length, int hash, int temperature) { - final int nameBytesInLongLength = 1 + (length >>> 3); - - int index = calculatedHash & tableMask; - Entry tableEntry; + int dataLength = length >> 3; + int index = hash & tableMask; + MeasurementRepository.Entry tableEntry; while ((tableEntry = table[index]) != null - && (tableEntry.hash != calculatedHash || !arrayEquals(tableEntry.nameBytesInLong, nameBytesInLong, nameBytesInLongLength))) { // search for the right spot + && (tableEntry.hash != hash || tableEntry.length != length || !arrayEquals(tableEntry.data, data, dataLength))) { // search for the right spot index = (index + 1) & tableMask; } if (tableEntry != null) { - return tableEntry.measurement; + tableEntry.measurement.updateWith(temperature); + return; } // --- This is a brand new entry, insert into the hashtable and do the extra calculations (once!) do slower calculations here. Measurement measurement = new Measurement(); - // Now create a string: - byte[] buffer = new byte[length]; - bb.get(buffer, 0, length); - String cityName = new String(buffer, 0, length); + byte[] bytes = new byte[length]; + for (int i = 0; i < length; i++) { + bytes[i] = UNSAFE.getByte(address + i); + } + String city = new String(bytes); - // Store the long[] for faster equals: - long[] nameBytesInLongCopy = new long[nameBytesInLongLength]; - System.arraycopy(nameBytesInLong, 0, nameBytesInLongCopy, 0, nameBytesInLongLength); + long[] dataCopy = new long[dataLength]; + System.arraycopy(data, 0, dataCopy, 0, dataLength); // And add entry: - Entry toAdd = new Entry(calculatedHash, nameBytesInLongCopy, cityName, measurement); + MeasurementRepository.Entry toAdd = new MeasurementRepository.Entry(address, dataCopy, length, hash, city, measurement); table[index] = toAdd; - // Resize the table if filled too much: - if (++tableFilled > tableLimit) { - resizeTable(); - } - - return toAdd.measurement; + toAdd.measurement.updateWith(temperature); } - private void resizeTable() { - // Resize the table: - Entry[] oldEntries = table; - table = new Entry[tableSize <<= 2]; // x2 - tableMask = (tableSize - 1); - tableLimit = (int) (tableSize * LOAD_FACTOR); - - for (Entry entry : oldEntries) { - if (entry != null) { - int updatedTableIndex = entry.hash & tableMask; - while (table[updatedTableIndex] != null) { - updatedTableIndex = (updatedTableIndex + 1) & tableMask; - } - table[updatedTableIndex] = entry; - } - } - } - - public Stream get() { + public Stream get() { return Arrays.stream(table).filter(Objects::nonNull); } } @@ -372,49 +318,4 @@ public class CalculateAverage_royvanrijn { return true; } - public void runTests() { - // Method used for debugging purposes, easy to make mistakes with all the bit hacking. - - // These all have the same hashes: - testInput("Delft;-12.4", 0, true, new int[]{ 5, 1718384401 }, new long[]{ 499934586180L }); - testInput("aDelft;-12.4", 1, true, new int[]{ 6, 1718384401 }, new long[]{ 499934586180L }); - - testInput("Delft;-12.4", 0, false, new int[]{ 5, 1718384401 }, new long[]{ 499934586180L }); - testInput("aDelft;-12.4", 1, false, new int[]{ 6, 1718384401 }, new long[]{ 499934586180L }); - - testInput("Rotterdam;-12.4", 0, true, new int[]{ 9, -784321989 }, new long[]{ 7017859899421126482L, 109L }); - testInput("abcdefghijklmnpoqrstuvwxyzRotterdam;-12.4", 26, true, new int[]{ 35, -784321989 }, new long[]{ 7017859899421126482L, 109L }); - testInput("abcdefghijklmnpoqrstuvwxyzARotterdam;-12.4", 27, true, new int[]{ 36, -784321989 }, new long[]{ 7017859899421126482L, 109L }); - - testInput("Rotterdam;-12.4", 0, false, new int[]{ 9, -784321989 }, new long[]{ 7017859899421126482L, 109L }); - testInput("abcdefghijklmnpoqrstuvwxyzRotterdam;-12.4", 26, false, new int[]{ 35, -784321989 }, new long[]{ 7017859899421126482L, 109L }); - testInput("abcdefghijklmnpoqrstuvwxyzARotterdam;-12.4", 27, false, new int[]{ 36, -784321989 }, new long[]{ 7017859899421126482L, 109L }); - - // These have different hashes from the strings above: - testInput("abcdefghijklmnpoqrstuvwxyzAROtterdam;-12.4", 27, true, new int[]{ 36, -792194501 }, new long[]{ 7017859899421118290L, 109L }); - testInput("abcdefghijklmnpoqrstuvwxyzAROtterdam;-12.4", 27, false, new int[]{ 36, -792194501 }, new long[]{ 7017859899421118290L, 109L }); - } - - private void testInput(final String inputString, final int start, final boolean bigEndian, final int[] expectedDelimiterAndHash, final long[] expectedCityNameLong) { - - byte[] input = inputString.getBytes(StandardCharsets.UTF_8); - - ByteBuffer buffer = ByteBuffer.wrap(input).order(bigEndian ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN); - - int[] output = new int[2]; - long[] cityName = new long[128]; - findNextDelimiterAndCalculateHash(buffer, SEPARATOR_PATTERN, start, buffer.limit(), output, cityName, bigEndian); - - if (!Arrays.equals(output, expectedDelimiterAndHash)) { - System.out.println("Error in delimiter or hash"); - System.out.println("Expected: " + Arrays.toString(expectedDelimiterAndHash)); - System.out.println("Received: " + Arrays.toString(output)); - } - int amountLong = 1 + ((output[0] - start) >>> 3); - if (!Arrays.equals(cityName, 0, amountLong, expectedCityNameLong, 0, amountLong)) { - System.out.println("Error in long array"); - System.out.println("Expected: " + Arrays.toString(expectedCityNameLong)); - System.out.println("Received: " + Arrays.toString(cityName)); - } - } }