From 38ebced4bab0fb35d9cd0fd1f16aa523c3b64cc6 Mon Sep 17 00:00:00 2001 From: hadashiA Date: Sat, 8 Aug 2026 19:21:40 +0900 Subject: [PATCH] Add SIMD window to key digest search --- sandbox/DryDB.Benchmark/ReadBenchmark.cs | 19 +++ src/DryDB/BTree/DigestSearch.cs | 109 +++++++++++++++++ src/DryDB/BTree/InternalNodeReader.cs | 69 ++++++----- src/DryDB/BTree/LeafNodeReader.cs | 94 +++++++++++++++ tests/DryDB.Tests/KeyDigestSearchTest.cs | 143 +++++++++++++++++++++++ 5 files changed, 408 insertions(+), 26 deletions(-) create mode 100644 src/DryDB/BTree/DigestSearch.cs create mode 100644 tests/DryDB.Tests/KeyDigestSearchTest.cs diff --git a/sandbox/DryDB.Benchmark/ReadBenchmark.cs b/sandbox/DryDB.Benchmark/ReadBenchmark.cs index df6520a..0275164 100644 --- a/sandbox/DryDB.Benchmark/ReadBenchmark.cs +++ b/sandbox/DryDB.Benchmark/ReadBenchmark.cs @@ -57,6 +57,25 @@ public void DryDB_FindByKey_RandomKeys() } } + uint noRepeatSeed = 123456789u; + + // The variant above re-seeds every op, so the same 1000-key sequence repeats and a + // large branch predictor gradually memorizes its branch history across invocations. + // Carrying the seed across ops makes the sequence genuinely non-repeating — the + // closest model of real random access. + [Benchmark] + public void DryDB_FindByKey_RandomKeys_NoRepeat() + { + var seed = noRepeatSeed; + for (var i = 0; i < Iterations; i++) + { + seed = seed * 1664525u + 1013904223u; + var table = database.GetTable("items"); + using var _ = table.Get((long)(seed % N)); + } + noRepeatSeed = seed; + } + const int ThreadCount = 8; [Benchmark] diff --git a/src/DryDB/BTree/DigestSearch.cs b/src/DryDB/BTree/DigestSearch.cs new file mode 100644 index 0000000..374179e --- /dev/null +++ b/src/DryDB/BTree/DigestSearch.cs @@ -0,0 +1,109 @@ +using System.Runtime.CompilerServices; +#if NET8_0_OR_GREATER +using System.Runtime.Intrinsics; +#endif + +namespace DryDB.BTree; + +/// +/// Lower-bound search over the contiguous key digest array of a node. +/// +/// +/// The tail levels of a binary search carry the unpredictable branches (a ~50% +/// mispredict per level on non-repeating random keys). This kernel runs the branchy +/// binary search only down to a 32-element window and finishes with a branch-free SIMD +/// count, which removes those mispredicts. A full SIMD scan would touch every cache +/// line of the digest array and loses when the node is out of cache; the hybrid wins +/// in both regimes on unpredictable keys. +/// +static class DigestSearch +{ + /// + /// Whether the SIMD window path is available. The lower-bound restructure only + /// pays off together with the SIMD window (losing the classic search's early + /// digest-equality exit costs more than the restructure alone gains), so callers + /// must keep using the classic mixed digest binary search when this is false — + /// notably on netstandard targets, which have no Vector128. + /// +#if NET8_0_OR_GREATER + public static bool IsAccelerated + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get => Vector128.IsHardwareAccelerated; + } +#else + public const bool IsAccelerated = false; +#endif + + /// + /// Returns the number of digests strictly less than , + /// i.e. the index of the first digest >= (or + /// if there is none). + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static int LowerBound(ref byte pageReference, int digestBase, int count, ulong keyDigest) + { + var min = 0; + var max = count; + while (max - min > 32) + { + var mid = min + ((max - min) >> 1); + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, digestBase + mid * sizeof(ulong))); + if (digest < keyDigest) + { + min = mid + 1; + } + else + { + max = mid; + } + } + +#if NET8_0_OR_GREATER + if (Vector128.IsHardwareAccelerated && max >= 32) + { + // Anchor the 32-wide window at [max-32, max): it stays inside the array, + // and every element below `min` it may cover is < keyDigest by the binary + // search invariant, so counting them keeps the result exact — no sentinels + // or masking needed. + var start = max - 32; + var keyVec = Vector128.Create(keyDigest); + // Four independent accumulators: a single accumulator would serialize the + // sixteen subtracts into one long dependency chain, whose latency exceeds + // the mispredict cost it replaces. + var acc0 = Vector128.Zero; + var acc1 = Vector128.Zero; + var acc2 = Vector128.Zero; + var acc3 = Vector128.Zero; + ref var window = ref Unsafe.Add(ref pageReference, digestBase + start * sizeof(ulong)); + for (var i = 0; i < 32 * sizeof(ulong); i += 8 * sizeof(ulong)) + { + // a true lane is all-ones (= -1): accumulate counts by subtraction + acc0 -= Vector128.LessThan(Unsafe.ReadUnaligned>(ref Unsafe.Add(ref window, i)), keyVec); + acc1 -= Vector128.LessThan(Unsafe.ReadUnaligned>(ref Unsafe.Add(ref window, i + 2 * sizeof(ulong))), keyVec); + acc2 -= Vector128.LessThan(Unsafe.ReadUnaligned>(ref Unsafe.Add(ref window, i + 4 * sizeof(ulong))), keyVec); + acc3 -= Vector128.LessThan(Unsafe.ReadUnaligned>(ref Unsafe.Add(ref window, i + 6 * sizeof(ulong))), keyVec); + } + var acc = (acc0 + acc1) + (acc2 + acc3); + return start + (int)(acc.GetElement(0) + acc.GetElement(1)); + } +#endif + + while (min < max) + { + var mid = min + ((max - min) >> 1); + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, digestBase + mid * sizeof(ulong))); + if (digest < keyDigest) + { + min = mid + 1; + } + else + { + max = mid; + } + } + return min; + } +} diff --git a/src/DryDB/BTree/InternalNodeReader.cs b/src/DryDB/BTree/InternalNodeReader.cs index 6ac0eac..6202d78 100644 --- a/src/DryDB/BTree/InternalNodeReader.cs +++ b/src/DryDB/BTree/InternalNodeReader.cs @@ -61,38 +61,55 @@ public bool TrySearch( ref var pageReference = ref MemoryMarshal.GetReference(page); #endif var useDigest = hasKeyDigests && hasKeyDigest; - - var min = 0; - var max = entryCount; - + int min; NodeEntryMeta meta; - while (min < max) + if (useDigest && DigestSearch.IsAccelerated) { - var mid = min + ((max - min) >> 1); - - int cmp; - if (useDigest) + // Branch-free lower bound over the digest array, then advance through the + // run of equal digests with full comparisons to reach the upper bound + // (first entry > key). Entries with a greater digest are already > key. + min = DigestSearch.LowerBound(ref pageReference, DigestBase, entryCount, keyDigest); + while (min < entryCount) { - // One contiguous load instead of dereferencing the variable-length key; - // only digest ties fall back to the full comparison. var digest = Unsafe.ReadUnaligned( - ref Unsafe.Add(ref pageReference, DigestBase + mid * sizeof(ulong))); - cmp = digest != keyDigest - ? (digest < keyDigest ? -1 : 1) - : CompareFull(ref pageReference, mid, key, comparer); - } - else - { - cmp = CompareFull(ref pageReference, mid, key, comparer); + ref Unsafe.Add(ref pageReference, DigestBase + min * sizeof(ulong))); + if (digest != keyDigest) break; + if (CompareFull(ref pageReference, min, key, comparer) > 0) break; + min++; } - - if (cmp <= 0) // upper bounds - { - min = mid + 1; - } - else + } + else + { + min = 0; + var max = entryCount; + while (min < max) { - max = mid; + var mid = min + ((max - min) >> 1); + + int cmp; + if (useDigest) + { + // One contiguous load instead of dereferencing the variable-length + // key; only digest ties fall back to the full comparison. + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, DigestBase + mid * sizeof(ulong))); + cmp = digest != keyDigest + ? (digest < keyDigest ? -1 : 1) + : CompareFull(ref pageReference, mid, key, comparer); + } + else + { + cmp = CompareFull(ref pageReference, mid, key, comparer); + } + + if (cmp <= 0) // upper bounds + { + min = mid + 1; + } + else + { + max = mid; + } } } diff --git a/src/DryDB/BTree/LeafNodeReader.cs b/src/DryDB/BTree/LeafNodeReader.cs index 883e952..17765a7 100644 --- a/src/DryDB/BTree/LeafNodeReader.cs +++ b/src/DryDB/BTree/LeafNodeReader.cs @@ -83,6 +83,36 @@ public bool TryFindValue( ref var pageReference = ref MemoryMarshal.GetReference(page); #endif var useDigest = hasKeyDigests && hasKeyDigest; + if (useDigest && DigestSearch.IsAccelerated) + { + // Branch-free lower bound over the digest array; a match can only sit in + // the run of digests equal to keyDigest, which starts at the bound. + // (Order preservation: digest < keyDigest implies entry < key, and + // digest > keyDigest implies entry > key.) + var i = DigestSearch.LowerBound(ref pageReference, DigestBase, entryCount, keyDigest); + for (; i < entryCount; i++) + { + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, DigestBase + i * sizeof(ulong))); + if (digest != keyDigest) break; + + var compared = CompareFull(ref pageReference, i, key, comparer); + if (compared == 0) + { + var meta = GetMeta(i); + index = i; + valueOffset = meta.PageOffset + meta.KeyLength; + valueLength = meta.ValueLength; + return true; + } + if (compared > 0) break; + } + + index = default; + valueOffset = default; + valueLength = default; + return false; + } var min = 0; var max = entryCount; @@ -144,6 +174,70 @@ public bool TrySearch( ref var pageReference = ref MemoryMarshal.GetReference(page); #endif var useDigest = hasKeyDigests && hasKeyDigest; + if (useDigest && DigestSearch.IsAccelerated) + { + // Branch-free lower bound over the digest array, then resolve the bound + // inside the run of equal digests with full comparisons (run length is + // almost always 0 or 1; exact digests such as Int64 never exceed 1). + // Entries before the bound are < key; the first entry with a greater + // digest is > key, which already satisfies both bound operators. + var i = DigestSearch.LowerBound(ref pageReference, DigestBase, entryCount, keyDigest); + switch (op) + { + case SearchOperator.Equal: + for (; i < entryCount; i++) + { + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, DigestBase + i * sizeof(ulong))); + if (digest != keyDigest) break; + + var compared = CompareFull(ref pageReference, i, key, comparer); + if (compared == 0) + { + index = i; + return true; + } + if (compared > 0) break; + } + index = default; + return false; + + case SearchOperator.LowerBound: + // first entry >= key + while (i < entryCount) + { + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, DigestBase + i * sizeof(ulong))); + if (digest != keyDigest) break; + if (CompareFull(ref pageReference, i, key, comparer) >= 0) break; + i++; + } + break; + + case SearchOperator.UpperBound: + // first entry > key + while (i < entryCount) + { + var digest = Unsafe.ReadUnaligned( + ref Unsafe.Add(ref pageReference, DigestBase + i * sizeof(ulong))); + if (digest != keyDigest) break; + if (CompareFull(ref pageReference, i, key, comparer) > 0) break; + i++; + } + break; + + default: + throw new ArgumentOutOfRangeException(nameof(op), op, null); + } + + if (i >= entryCount) + { + index = default; + return false; + } + index = i; + return true; + } var min = 0; var max = entryCount; diff --git a/tests/DryDB.Tests/KeyDigestSearchTest.cs b/tests/DryDB.Tests/KeyDigestSearchTest.cs new file mode 100644 index 0000000..fd0443e --- /dev/null +++ b/tests/DryDB.Tests/KeyDigestSearchTest.cs @@ -0,0 +1,143 @@ +using System; +using System.Text; +using System.Threading.Tasks; + +namespace DryDB.Tests; + +/// +/// Exercises the digest-array search paths: the SIMD window kernel (nodes with more +/// than 32 entries) and digest-collision runs (ascii keys sharing an 8-byte prefix, +/// where every digest in the node ties and the search must resolve bounds with full +/// key comparisons). +/// +[TestFixture] +public class KeyDigestSearchTest +{ + [Test] + public async Task Get_AsciiKeys_SharedEightBytePrefix() + { + // All keys share the first 8 bytes, so every digest in the tree collides and + // the whole node forms a single digest run. + var table = await TestHelper.BuildTableAsync( + KeyEncoding.Ascii, + tableConfigure: builder => + { + for (var i = 0; i < 300; i++) + { + builder.Append( + Encoding.ASCII.GetBytes($"AAAAAAAA{i:D4}"), + Encoding.ASCII.GetBytes($"value{i:D4}")); + } + }); + + for (var i = 0; i < 300; i++) + { + using var result = table.Get(Encoding.ASCII.GetBytes($"AAAAAAAA{i:D4}")); + Assert.That(result.HasValue, Is.True, $"key {i}"); + Assert.That( + result.Value.Span.SequenceEqual(Encoding.ASCII.GetBytes($"value{i:D4}")), + Is.True, + $"key {i}"); + } + + // Same digest as every stored key, but no exact match. + using var missing = table.Get("AAAAAAAA9999"u8); + Assert.That(missing.HasValue, Is.False); + + using var missingShort = table.Get("AAAAAAAA"u8); + Assert.That(missingShort.HasValue, Is.False); + } + + [Test] + public async Task GetRange_AsciiKeys_SharedEightBytePrefix() + { + var table = await TestHelper.BuildTableAsync( + KeyEncoding.Ascii, + tableConfigure: builder => + { + for (var i = 0; i < 300; i++) + { + builder.Append( + Encoding.ASCII.GetBytes($"AAAAAAAA{i:D4}"), + Encoding.ASCII.GetBytes($"value{i:D4}")); + } + }); + + using var range = table.GetRange("AAAAAAAA0010"u8, "AAAAAAAA0019"u8); + Assert.That(range.Count, Is.EqualTo(10)); + + using var exclusive = table.GetRange( + "AAAAAAAA0010"u8, + "AAAAAAAA0019"u8, + startKeyExclusive: true, + endKeyExclusive: true); + Assert.That(exclusive.Count, Is.EqualTo(8)); + + // Bounds that fall between keys (same digest run, no exact match). + using var between = table.GetRange("AAAAAAAA0010x"u8, "AAAAAAAA0019x"u8); + Assert.That(between.Count, Is.EqualTo(9)); + + Assert.That(table.CountRange("AAAAAAAA0000"u8, "AAAAAAAA0299"u8), Is.EqualTo(300)); + Assert.That(table.CountRange("AAAAAAAA0290"u8, "AAAAAAAA9999"u8), Is.EqualTo(10)); + } + + [Test] + public async Task Get_Int64Keys_LargeNodes() + { + // Default 4KB pages hold >100 entries per node, which drives the search + // through the 32-wide SIMD window path. + var table = await TestHelper.BuildTableAsync( + KeyEncoding.Int64LittleEndian, + tableConfigure: builder => + { + for (var i = 0L; i < 10_000; i++) + { + builder.Append(i * 2, Encoding.ASCII.GetBytes($"value{i:D6}")); + } + }); + + for (var i = 0L; i < 10_000; i++) + { + using var result = table.Get(i * 2); + Assert.That(result.HasValue, Is.True, $"key {i * 2}"); + Assert.That( + result.Value.Span.SequenceEqual(Encoding.ASCII.GetBytes($"value{i:D6}")), + Is.True, + $"key {i * 2}"); + } + + // Odd keys are absent; the digest lower bound lands between entries. + for (var i = 1L; i < 2000; i += 2) + { + using var result = table.Get(i); + Assert.That(result.HasValue, Is.False, $"key {i}"); + } + } + + [Test] + public async Task GetRange_Int64Keys_LargeNodes() + { + var table = await TestHelper.BuildTableAsync( + KeyEncoding.Int64LittleEndian, + tableConfigure: builder => + { + for (var i = 0L; i < 10_000; i++) + { + builder.Append(i * 2, Encoding.ASCII.GetBytes($"value{i:D6}")); + } + }); + + // Bounds on existing keys. + Assert.That(table.CountRange(200L, 400L, false, false), Is.EqualTo(101)); + + // Bounds on absent (odd) keys: exercises the lower/upper-bound miss paths. + Assert.That(table.CountRange(199L, 401L, false, false), Is.EqualTo(101)); + Assert.That(table.CountRange(201L, 399L, false, false), Is.EqualTo(99)); + + using var range = table.GetRange(9_000L, 9_100L); + Assert.That(range.Count, Is.EqualTo(51)); + + using var descending = table.GetRange(100L, 200L, sortOrder: SortOrder.Descending); + Assert.That(descending.Count, Is.EqualTo(51)); + } +}