Add a few benchmarks and some quick tests with AVX2 lookups

benchmarks/NexusMods.MnemonicDB.Benchmarks/Benchmarks/IndexSegmentABenchmarks.cs

@@ -0,0 +1,48 @@
+using System;
+using BenchmarkDotNet.Attributes;
+using NexusMods.MnemonicDB.Abstractions;
+using NexusMods.MnemonicDB.Abstractions.DatomIterators;
+using NexusMods.MnemonicDB.Abstractions.ElementComparers;
+using NexusMods.MnemonicDB.Abstractions.IndexSegments;
+using NexusMods.MnemonicDB.Abstractions.Internals;
+using NexusMods.MnemonicDB.Storage;
+
+namespace NexusMods.MnemonicDB.Benchmarks.Benchmarks;
+
+public class IndexSegmentABenchmarks
+{
+    private readonly IndexSegment _index;
+
+    public IndexSegmentABenchmarks()
+    {
+        var registry = new AttributeRegistry([]);
+        using var builder = new IndexSegmentBuilder(registry);
+
+        for (int a = 1; a < 100; a++)
+        {
+            var prefix = new KeyPrefix(EntityId.From(42), AttributeId.From((ushort)a), TxId.From(42), false,
+                ValueTags.Null);
+            builder.Add(new Datom(prefix, ReadOnlyMemory<byte>.Empty, registry));
+        }
+
+        _index = builder.Build();
+    }
+
+    [Params(1, 2, 3, 4, 5, 20, 30, 50, 99)]
+    public int ToFind { get; set; }
+
+    [Benchmark]
+    public int FindLinear()
+    {
+        var find = AttributeId.From((ushort)ToFind);
+        for (var i = 0; i < _index.Count; i++)
+        {
+            var datom = _index[i];
+            if (datom.A == find)
+                return i;
+        }
+        return -1;
+    }
+
+
+}

benchmarks/NexusMods.MnemonicDB.Benchmarks/Benchmarks/IndexSegmentBenchmarks.cs → benchmarks/NexusMods.MnemonicDB.Benchmarks/Benchmarks/IndexSegmentEBenchmarks.cs

@@ -10,11 +10,11 @@
 namespace NexusMods.MnemonicDB.Benchmarks.Benchmarks;
 
 [MemoryDiagnoser]
-public class IndexSegmentBenchmarks
+public class IndexSegmentEBenchmarks
 {
     private readonly IndexSegment _index;
 
-    public IndexSegmentBenchmarks()
+    public IndexSegmentEBenchmarks()
     {
         var registry = new AttributeRegistry([]);
         using var builder = new IndexSegmentBuilder(registry);
@@ -31,7 +31,7 @@ public IndexSegmentBenchmarks()
         _index = builder.Build();
     }
 
-    [Params(1, 10, 71, 99)] public int ToFind { get; set; }
+    [Params(0, 1, 2, 3, 10, 71, 99)] public int ToFind { get; set; }
 
 
     [Benchmark]
@@ -84,5 +84,12 @@ public int FindBinarySearchRework()
         return _index.FindFirst(find); // Return the first occurrence found, or -1 if not found
     }
 
+    [Benchmark]
+    public int FindBinarySearchReworkAVX2()
+    {
+        var find = EntityId.From((ulong)ToFind);
+        return _index.FindFirstAVX2(find.Value); // Return the first occurrence found, or -1 if not found
+    }
+
 
 }

benchmarks/NexusMods.MnemonicDB.Benchmarks/Program.cs

@@ -32,7 +32,7 @@
 
 #else
 
-BenchmarkRunner.Run<IndexSegmentBenchmarks>(config: DefaultConfig.Instance.WithOption(ConfigOptions.DisableOptimizationsValidator, true));
+BenchmarkRunner.Run<IndexSegmentEBenchmarks>(config: DefaultConfig.Instance.WithOption(ConfigOptions.DisableOptimizationsValidator, true));
 #endif
 
 

src/NexusMods.MnemonicDB.Abstractions/IndexSegments/IndexSegment.cs

@@ -1,7 +1,10 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Numerics;
 using System.Runtime.InteropServices;
+using System.Runtime.Intrinsics;
+using System.Runtime.Intrinsics.X86;
 using DynamicData;
 using NexusMods.MnemonicDB.Abstractions.DatomIterators;
 using NexusMods.MnemonicDB.Abstractions.Internals;
@@ -203,4 +206,43 @@ public int FindFirst(EntityId find)
         }
         return result; // Return the first occurrence found, or -1 if not found
     }
+
+    public int FindFirstAVX2(ulong find)
+    {
+        var targetVector = Vector256<ulong>.Zero.WithElement(0, find)
+            .WithElement(1, find)
+            .WithElement(2, find)
+            .WithElement(3, find);
+
+        var casted = MemoryMarshal.Cast<ulong, Vector256<ulong>>(_lowers);
+
+        for (var idx = 0; idx < casted.Length; idx += 1)
+        {
+            var lowers = casted[idx];
+
+            var maskHigh = Avx2.And(lowers, Vector256.Create(0x7F00000000000000UL));
+            var maskLow = Avx2.And(Avx2.ShiftRightLogical(lowers, 8), Vector256.Create(0x0000FFFFFFFFFFFFUL));
+            var ored = Avx2.Or(maskHigh, maskLow);
+
+            var comparison = Avx2.CompareEqual(ored, targetVector);
+
+            var mask = Avx2.MoveMask(comparison.As<ulong, byte>());
+            if (mask != 0)
+            {
+                var index = BitOperations.TrailingZeroCount(mask) / 8;
+                return idx * Vector256<ulong>.Count + index;
+            }
+        }
+
+        // Handle remaining elements
+        for (int i = (casted.Length * Vector256<ulong>.Count); i < _lowers.Length; i++)
+        {
+            if (_lowers[i] >= find)
+            {
+                return i;
+            }
+        }
+
+        return -1; // No value found that is greater than or equal to the target
+    }
 }