Fast Keccak cache (wait-free and lockless)

src/Nethermind/Nethermind.Core.Test/KeccakCacheTests.cs

@@ -0,0 +1,96 @@
+// SPDX-FileCopyrightText: 2022 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Buffers.Binary;
+using System.Linq;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Nethermind.Core.Crypto;
+using Nethermind.Core.Extensions;
+using NUnit.Framework;
+
+namespace Nethermind.Core.Test
+{
+    [TestFixture]
+    public class KeccakCacheTests
+    {
+        [Test]
+        public void Multiple()
+        {
+            const int spins = 10;
+
+            var random = new Random(13);
+            var bytes = new byte[31]; // misaligned length
+            random.NextBytes(bytes);
+
+            ValueHash256 expected = ValueKeccak.Compute(bytes);
+
+            for (int i = 0; i < spins; i++)
+            {
+                ValueHash256 actual = KeccakCache.Compute(bytes);
+                actual.Equals(expected).Should().BeTrue();
+            }
+        }
+
+        [Test]
+        public void Empty()
+        {
+            ReadOnlySpan<byte> span = ReadOnlySpan<byte>.Empty;
+            KeccakCache.Compute(span).Should().Be(ValueKeccak.Compute(span));
+        }
+
+        [Test]
+        public void Very_long()
+        {
+            ReadOnlySpan<byte> span = new byte[192];
+            KeccakCache.Compute(span).Should().Be(ValueKeccak.Compute(span));
+        }
+
+        [Test]
+        public void Collision()
+        {
+            var colliding = new[]
+            {
+                "f8ae910727f29363002d948385ff15bc6a9bacbef13bc2afc0aa8d02749668",
+                "baec2065df3da176cee21714b7bfb00d0c57f37a21daf2b2d4056f67270290",
+                "924cc47a10ad801c74a491b19492563d2351b285ff1679e5b5264e57b13bbb",
+                "4fb39f7800b3a43e4e722dc6fed03b126e0125d7ca713b0558564a29903ea9",
+            };
+
+            var collisions = colliding.Length;
+            var array = colliding.Select(c => Bytes.FromHexString(c)).ToArray();
+            var values = array.Select(a => ValueKeccak.Compute(a)).ToArray();
+
+            var bucket = KeccakCache.GetBucket(array[0]);
+
+            for (int i = 1; i < collisions; i++)
+            {
+                var input = array[i];
+                bucket.Should().Be(KeccakCache.GetBucket(input));
+                KeccakCache.Compute(input).Should().Be(values[i]);
+            }
+
+            Parallel.ForEach(array, (a, state, index) =>
+            {
+                ValueHash256 v = values[index];
+
+                for (int i = 0; i < 100_000; i++)
+                {
+                    KeccakCache.Compute(a).Should().Be(v);
+                }
+            });
+        }
+
+        [Test]
+        public void Spin_through_all()
+        {
+            Span<byte> span = stackalloc byte[4];
+            for (int i = 0; i < KeccakCache.Count; i++)
+            {
+                BinaryPrimitives.WriteInt32LittleEndian(span, i);
+                KeccakCache.Compute(span).Should().Be(ValueKeccak.Compute(span));
+            }
+        }
+    }
+}

src/Nethermind/Nethermind.Core/Crypto/KeccakCache.cs

@@ -0,0 +1,208 @@
+// SPDX-FileCopyrightText: 2024 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Diagnostics;
+using System.Numerics;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using System.Threading;
+
+namespace Nethermind.Core.Crypto;
+
+/// <summary>
+/// This is a minimalistic one-way set associative cache for Keccak values.
+///
+/// It allocates only 8MB of memory to store 64k of entries.
+/// No misaligned reads. Everything is aligned to both cache lines as well as to boundaries so no torn reads.
+/// Requires a single CAS to lock and <see cref="Volatile.Write(ref int,int)"/> to unlock.
+/// On lock failure, it just moves on with execution.
+/// Uses copying on the stack to get the entry, have it copied and release the lock ASAP. This is 128 bytes to copy that quite likely will be the hit.
+/// </summary>
+public static unsafe class KeccakCache
+{
+    private static readonly uint s_instanceRandom = (uint)System.Security.Cryptography.RandomNumberGenerator.GetInt32(int.MinValue, int.MaxValue);
+    /// <summary>
+    /// Count is defined as a +1 over bucket mask. In the future, just change the mask as the main parameter.
+    /// </summary>
+    public const int Count = BucketMask + 1;
+    private const int BucketMask = 0x0000_FFFF;
+    private const uint HashMask = unchecked((uint)~BucketMask);
+
+    private static readonly Entry* Memory;
+
+    static KeccakCache()
+    {
+        const UIntPtr size = Count * Entry.Size;
+
+        // Aligned, so that no torn reads if fields of Entry are properly aligned.
+        Memory = (Entry*)NativeMemory.AlignedAlloc(size, Entry.Size);
+        NativeMemory.Clear(Memory, size);
+    }
+
+    [SkipLocalsInit]
+    public static ValueHash256 Compute(ReadOnlySpan<byte> input)
+    {
+        Unsafe.SkipInit(out ValueHash256 hash);
+
+        // Special cases first
+        if (input.Length == 0)
+        {
+            hash = ValueKeccak.OfAnEmptyString;
+            goto Return;
+        }
+
+        if (input.Length > Entry.MaxPayloadLength)
+        {
+            hash = ValueKeccak.Compute(input);
+            goto Return;
+        }
+
+        var fast = FastHash(input);
+        var index = fast & BucketMask;
+
+        Debug.Assert(index < Count);
+
+        uint hashAndLength = (fast & HashMask) | (ushort)input.Length;
+
+        ref Entry e = ref Unsafe.Add(ref Unsafe.AsRef<Entry>(Memory), index);
+
+        // Read aligned, volatile, won't be torn, check with computed
+        if (Volatile.Read(ref e.HashAndLength) == hashAndLength)
+        {
+            // There's a possibility of a hit, try lock.
+            if (Interlocked.CompareExchange(ref e.Lock, Entry.Locked, Entry.Unlocked) == Entry.Unlocked)
+            {
+                if (e.HashAndLength != hashAndLength)
+                {
+                    // The value has been changed between reading and taking a lock.
+                    // Release the lock and compute.
+                    Volatile.Write(ref e.Lock, Entry.Unlocked);
+                    goto Compute;
+                }
+
+                // Local copy of 128 bytes, to release the lock as soon as possible and make a key comparison without holding it.
+                Entry copy = e;
+
+                // Release the lock
+                Volatile.Write(ref e.Lock, Entry.Unlocked);
+
+                // Lengths are equal, the input length can be used without any additional operation.
+                if (MemoryMarshal.CreateReadOnlySpan(ref copy.Payload, input.Length).SequenceEqual(input))
+                {
+                    hash = copy.Value;
+                    goto Return;
+                }
+            }
+        }
+
+    Compute:
+        hash = ValueKeccak.Compute(input);
+
+        // Try lock and memoize
+        if (Interlocked.CompareExchange(ref e.Lock, Entry.Locked, Entry.Unlocked) == Entry.Unlocked)
+        {
+            e.HashAndLength = hashAndLength;
+            e.Value = hash;
+
+            input.CopyTo(MemoryMarshal.CreateSpan(ref e.Payload, input.Length));
+
+            // Release the lock
+            Volatile.Write(ref e.Lock, Entry.Unlocked);
+        }
+
+    Return:
+        return hash;
+    }
+
+    /// <summary>
+    /// Gets the bucket for tests.
+    /// </summary>
+    public static uint GetBucket(ReadOnlySpan<byte> input) => FastHash(input) & BucketMask;
+
+    [SkipLocalsInit]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static uint FastHash(ReadOnlySpan<byte> input)
+    {
+        Debug.Assert(input.Length >= 1, "Cannot hash empty");
+
+        var length = input.Length;
+
+        ref var b = ref MemoryMarshal.GetReference(input);
+
+        // Start with first
+        uint hash = s_instanceRandom ^ b;
+
+        // This is done below, without branches
+        // if ((length & 1) == 1)
+        // {
+        //     hash = b;
+        //     b = ref Unsafe.Add(ref b, 1);
+        //     length -= 1;
+        // }
+
+        var bit = length & 1;
+        b = ref Unsafe.Add(ref b, bit);
+        length -= bit;
+
+        if ((length & 2) == 2)
+        {
+            hash = BitOperations.Crc32C(hash, Unsafe.ReadUnaligned<ushort>(ref b));
+            b = ref Unsafe.Add(ref b, 2);
+            length -= 2;
+        }
+        if ((length & 4) == 4)
+        {
+            hash = BitOperations.Crc32C(hash, Unsafe.ReadUnaligned<uint>(ref b));
+            b = ref Unsafe.Add(ref b, 4);
+            length -= 4;
+        }
+
+        while (length > 0)
+        {
+            hash = BitOperations.Crc32C(hash, Unsafe.ReadUnaligned<ulong>(ref b));
+            b = ref Unsafe.Add(ref b, 8);
+            length -= 8;
+        }
+
+        return hash;
+    }
+
+    /// <summary>
+    /// An entry to cache keccak
+    /// </summary>
+    [StructLayout(LayoutKind.Explicit, Size = Size)]
+    private struct Entry
+    {
+        public const int Unlocked = 0;
+        public const int Locked = 1;
+
+        /// <summary>
+        /// Should work for both ARM and x64 and be aligned.
+        /// </summary>
+        public const int Size = 128;
+
+        private const int PayloadStart = 8;
+        private const int ValueStart = Size - ValueHash256.MemorySize;
+        public const int MaxPayloadLength = ValueStart - PayloadStart;
+
+        [FieldOffset(0)]
+        public int Lock;
+
+        /// <summary>
+        /// The mix of hash and length allows for a fast comparison and a single volatile read.
+        /// The length is encoded as the low part, while the hash as the high part of uint.
+        /// </summary>
+        [FieldOffset(4)]
+        public uint HashAndLength;
+
+        [FieldOffset(PayloadStart)]
+        public byte Payload;
+
+        /// <summary>
+        /// The actual value
+        /// </summary>
+        [FieldOffset(ValueStart)]
+        public ValueHash256 Value;
+    }
+}

src/Nethermind/Nethermind.Evm/VirtualMachine.cs

@@ -1132,7 +1132,8 @@ private CallResult ExecuteCode<TTracingInstructions, TTracingRefunds, TTracingSt
                         if (!UpdateMemoryCost(vmState, ref gasAvailable, in a, b)) goto OutOfGas;
 
                         bytes = vmState.Memory.LoadSpan(in a, b);
-                        stack.PushBytes(ValueKeccak.Compute(bytes).BytesAsSpan);
+
+                        stack.PushBytes(KeccakCache.Compute(bytes).BytesAsSpan);
                         break;
                     }
                 case Instruction.ADDRESS:

src/Nethermind/Nethermind.State/StateTree.cs

@@ -42,14 +42,14 @@ public StateTree(ITrieStore? store, ILogManager? logManager)
         [DebuggerStepThrough]
         public Account? Get(Address address, Hash256? rootHash = null)
         {
-            ReadOnlySpan<byte> bytes = Get(ValueKeccak.Compute(address.Bytes).BytesAsSpan, rootHash);
+            ReadOnlySpan<byte> bytes = Get(KeccakCache.Compute(address.Bytes).BytesAsSpan, rootHash);
             return bytes.IsEmpty ? null : _decoder.Decode(bytes);
         }
 
         [DebuggerStepThrough]
         public bool TryGetStruct(Address address, out AccountStruct account, Hash256? rootHash = null)
         {
-            ReadOnlySpan<byte> bytes = Get(ValueKeccak.Compute(address.Bytes).BytesAsSpan, rootHash);
+            ReadOnlySpan<byte> bytes = Get(KeccakCache.Compute(address.Bytes).BytesAsSpan, rootHash);
             Rlp.ValueDecoderContext valueDecoderContext = new Rlp.ValueDecoderContext(bytes);
             if (bytes.IsEmpty)
             {
@@ -69,7 +69,7 @@ public bool TryGetStruct(Address address, out AccountStruct account, Hash256? ro
 
         public void Set(Address address, Account? account)
         {
-            ValueHash256 keccak = ValueKeccak.Compute(address.Bytes);
+            ValueHash256 keccak = KeccakCache.Compute(address.Bytes);
             Set(keccak.BytesAsSpan, account is null ? null : account.IsTotallyEmpty ? EmptyAccountRlp : Rlp.Encode(account));
         }
 

src/Nethermind/Nethermind.State/StorageTree.cs

@@ -5,6 +5,8 @@
 using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using System.Runtime.Intrinsics;
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Extensions;
 using Nethermind.Logging;
@@ -46,12 +48,15 @@ public StorageTree(IScopedTrieStore? trieStore, Hash256 rootHash, ILogManager? l
             TrieType = TrieType.Storage;
         }
 
-        private static void ComputeKey(in UInt256 index, ref Span<byte> key)
+        private static void ComputeKey(in UInt256 index, in Span<byte> key)
         {
             index.ToBigEndian(key);
 
-            // in situ calculation
-            KeccakHash.ComputeHashBytesToSpan(key, key);
+            ValueHash256 keyHash = KeccakCache.Compute(key);
+
+            // Assign to update the argument
+            Unsafe.As<byte, Vector256<byte>>(ref MemoryMarshal.GetReference(key))
+                = Unsafe.As<ValueHash256, Vector256<byte>>(ref keyHash);
         }
 
         [SkipLocalsInit]
@@ -63,7 +68,7 @@ public byte[] Get(in UInt256 index, Hash256? storageRoot = null)
             }
 
             Span<byte> key = stackalloc byte[32];
-            ComputeKey(index, ref key);
+            ComputeKey(index, key);
             return GetArray(key, storageRoot);
         }
 
@@ -92,7 +97,7 @@ public void Set(in UInt256 index, byte[] value)
             else
             {
                 Span<byte> key = stackalloc byte[32];
-                ComputeKey(index, ref key);
+                ComputeKey(index, in key);
                 SetInternal(key, value);
             }
         }