Implement LFU sketch using arm64 intrinsics (redux) (#648)

bitfaster · web-flow · commit 5669d38f9b90 · 2024-11-30T16:45:18.000-08:00
* basic impl

* run tests

* fix

* table lookup

* opt

* opt

* temp

* cleanup

* endif

* fix return

* cleanup

---------
diff --git a/BitFaster.Caching.Benchmarks/BitFaster.Caching.Benchmarks.csproj b/BitFaster.Caching.Benchmarks/BitFaster.Caching.Benchmarks.csproj
@@ -3,12 +3,14 @@
   <PropertyGroup>
     <OutputType>Exe</OutputType>
     <LangVersion>latest</LangVersion>
-    <TargetFrameworks>net48;net6.0;net8.0</TargetFrameworks>
+    <TargetFrameworks>net6.0;net8.0</TargetFrameworks>
     <AllowUnsafeBlocks>True</AllowUnsafeBlocks>
     <!-- https://stackoverflow.com/a/59916801/131345 -->
     <IsWindows Condition="'$([System.Runtime.InteropServices.RuntimeInformation]::IsOSPlatform($([System.Runtime.InteropServices.OSPlatform]::Windows)))' == 'true'">true</IsWindows> 
     <IsLinux Condition="'$([System.Runtime.InteropServices.RuntimeInformation]::IsOSPlatform($([System.Runtime.InteropServices.OSPlatform]::Linux)))' == 'true'">true</IsLinux>
     <IsMacOS Condition="'$([System.Runtime.InteropServices.RuntimeInformation]::IsOSPlatform($([System.Runtime.InteropServices.OSPlatform]::OSX)))' == 'true'">true</IsMacOS>
+    <IsArm64 Condition="$([System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture) == Arm64">true</IsArm64>
+    <IsX64 Condition="$([System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture) == X64">true</IsX64>
   </PropertyGroup>
 
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|AnyCPU'">
@@ -41,5 +43,11 @@
     <PropertyGroup Condition="'$(IsMacOS)'=='true'">
         <DefineConstants>MacOS</DefineConstants>
     </PropertyGroup>
+    <PropertyGroup Condition="'$(IsArm64)'=='true'">
+        <DefineConstants>Arm64</DefineConstants>
+    </PropertyGroup>
+    <PropertyGroup Condition="'$(IsX64)'=='true'">
+        <DefineConstants>X64</DefineConstants>
+    </PropertyGroup>
 
-</Project>
+</Project>
diff --git a/BitFaster.Caching.Benchmarks/Lfu/CmSketchNoPin.cs b/BitFaster.Caching.Benchmarks/Lfu/CmSketchNoPin.cs
@@ -1,9 +1,12 @@
 ﻿using System;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
 
 #if NET6_0_OR_GREATER
 using System.Runtime.Intrinsics;
+using System.Runtime.Intrinsics.Arm;
 using System.Runtime.Intrinsics.X86;
 #endif
 
@@ -61,6 +64,12 @@ public int EstimateFrequency(T value)
             {
                 return EstimateFrequencyAvx(value);
             }
+#if NET6_0_OR_GREATER
+            else if (isa.IsArm64Supported)
+            {
+                return EstimateFrequencyArm(value);
+            }
+#endif
             else
             {
                 return EstimateFrequencyStd(value);
@@ -84,6 +93,12 @@ public void Increment(T value)
             {
                 IncrementAvx(value);
             }
+#if NET6_0_OR_GREATER
+            else if (isa.IsArm64Supported)
+            {
+                IncrementArm(value);
+            }
+#endif
             else
             {
                 IncrementStd(value);
@@ -314,5 +329,94 @@ private unsafe void IncrementAvx(T value)
             }
         }
 #endif
+
+#if NET6_0_OR_GREATER
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private unsafe void IncrementArm(T value)
+        {
+            int blockHash = Spread(comparer.GetHashCode(value));
+            int counterHash = Rehash(blockHash);
+            int block = (blockHash & blockMask) << 3;
+
+            Vector128<int> h = AdvSimd.ShiftArithmetic(Vector128.Create(counterHash), Vector128.Create(0, -8, -16, -24));
+            Vector128<int> index = AdvSimd.And(AdvSimd.ShiftRightLogical(h, 1), Vector128.Create(0xf));
+            Vector128<int> blockOffset = AdvSimd.Add(AdvSimd.Add(Vector128.Create(block), AdvSimd.And(h, Vector128.Create(1))), Vector128.Create(0, 2, 4, 6));
+
+            fixed (long* tablePtr = table)
+            {
+                int t0 = AdvSimd.Extract(blockOffset, 0);
+                int t1 = AdvSimd.Extract(blockOffset, 1);
+                int t2 = AdvSimd.Extract(blockOffset, 2);
+                int t3 = AdvSimd.Extract(blockOffset, 3);
+
+                Vector128<long> tableVectorA = Vector128.Create(AdvSimd.LoadVector64(tablePtr + t0), AdvSimd.LoadVector64(tablePtr + t1));
+                Vector128<long> tableVectorB = Vector128.Create(AdvSimd.LoadVector64(tablePtr + t2), AdvSimd.LoadVector64(tablePtr + t3));
+
+                index = AdvSimd.ShiftLeftLogicalSaturate(index, 2);
+
+                Vector128<int> longOffA = AdvSimd.Arm64.InsertSelectedScalar(AdvSimd.Arm64.InsertSelectedScalar(Vector128<int>.Zero, 0, index, 0), 2, index, 1);
+                Vector128<int> longOffB = AdvSimd.Arm64.InsertSelectedScalar(AdvSimd.Arm64.InsertSelectedScalar(Vector128<int>.Zero, 0, index, 2), 2, index, 3);
+
+                Vector128<long> fifteen = Vector128.Create(0xfL);
+                Vector128<long> maskA = AdvSimd.ShiftArithmetic(fifteen, longOffA.AsInt64());
+                Vector128<long> maskB = AdvSimd.ShiftArithmetic(fifteen, longOffB.AsInt64());
+
+                Vector128<long> maskedA = AdvSimd.Not(AdvSimd.Arm64.CompareEqual(AdvSimd.And(tableVectorA, maskA), maskA));
+                Vector128<long> maskedB = AdvSimd.Not(AdvSimd.Arm64.CompareEqual(AdvSimd.And(tableVectorB, maskB), maskB));
+
+                var one = Vector128.Create(1L);
+                Vector128<long> incA = AdvSimd.And(maskedA, AdvSimd.ShiftArithmetic(one, longOffA.AsInt64()));
+                Vector128<long> incB = AdvSimd.And(maskedB, AdvSimd.ShiftArithmetic(one, longOffB.AsInt64()));
+
+                tablePtr[t0] += AdvSimd.Extract(incA, 0);
+                tablePtr[t1] += AdvSimd.Extract(incA, 1);
+                tablePtr[t2] += AdvSimd.Extract(incB, 0);
+                tablePtr[t3] += AdvSimd.Extract(incB, 1);
+
+                var max = AdvSimd.Arm64.MaxAcross(AdvSimd.Arm64.InsertSelectedScalar(AdvSimd.Arm64.MaxAcross(incA.AsInt32()), 1, AdvSimd.Arm64.MaxAcross(incB.AsInt32()), 0).AsInt16());
+
+                if (max.ToScalar() != 0 && (++size == sampleSize))
+                {
+                    Reset();
+                }
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private unsafe int EstimateFrequencyArm(T value)
+        {
+            int blockHash = Spread(comparer.GetHashCode(value));
+            int counterHash = Rehash(blockHash);
+            int block = (blockHash & blockMask) << 3;
+
+            Vector128<int> h = AdvSimd.ShiftArithmetic(Vector128.Create(counterHash), Vector128.Create(0, -8, -16, -24));
+            Vector128<int> index = AdvSimd.And(AdvSimd.ShiftRightLogical(h, 1), Vector128.Create(0xf));
+            Vector128<int> blockOffset = AdvSimd.Add(AdvSimd.Add(Vector128.Create(block), AdvSimd.And(h, Vector128.Create(1))), Vector128.Create(0, 2, 4, 6));
+
+            fixed (long* tablePtr = table)
+            {
+                Vector128<long> tableVectorA = Vector128.Create(AdvSimd.LoadVector64(tablePtr + AdvSimd.Extract(blockOffset, 0)), AdvSimd.LoadVector64(tablePtr + AdvSimd.Extract(blockOffset, 1)));
+                Vector128<long> tableVectorB = Vector128.Create(AdvSimd.LoadVector64(tablePtr + AdvSimd.Extract(blockOffset, 2)), AdvSimd.LoadVector64(tablePtr + AdvSimd.Extract(blockOffset, 3)));
+
+                index = AdvSimd.ShiftLeftLogicalSaturate(index, 2);
+
+                Vector128<int> indexA = AdvSimd.Negate(AdvSimd.Arm64.InsertSelectedScalar(AdvSimd.Arm64.InsertSelectedScalar(Vector128<int>.Zero, 0, index, 0), 2, index, 1));
+                Vector128<int> indexB = AdvSimd.Negate(AdvSimd.Arm64.InsertSelectedScalar(AdvSimd.Arm64.InsertSelectedScalar(Vector128<int>.Zero, 0, index, 2), 2, index, 3));
+
+                var fifteen = Vector128.Create(0xfL);
+                Vector128<long> a = AdvSimd.And(AdvSimd.ShiftArithmetic(tableVectorA, indexA.AsInt64()), fifteen);
+                Vector128<long> b = AdvSimd.And(AdvSimd.ShiftArithmetic(tableVectorB, indexB.AsInt64()), fifteen);
+
+                // Before: < 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, A, B, C, D, E, F >
+                // After:  < 0, 1, 2, 3, 8, 9, A, B, 4, 5, 6, 7, C, D, E, F >
+                var min = AdvSimd.Arm64.VectorTableLookup(a.AsByte(), Vector128.Create(0x0B0A090803020100, 0xFFFFFFFFFFFFFFFF).AsByte());
+                min = AdvSimd.Arm64.VectorTableLookupExtension(min, b.AsByte(), Vector128.Create(0xFFFFFFFFFFFFFFFF, 0x0B0A090803020100).AsByte());
+
+                var min32 = AdvSimd.Arm64.MinAcross(min.AsInt32());
+
+                return min32.ToScalar();
+            }
+        }
+#endif
     }
 }
diff --git a/BitFaster.Caching.Benchmarks/Lfu/SketchFrequency.cs b/BitFaster.Caching.Benchmarks/Lfu/SketchFrequency.cs
@@ -51,7 +51,7 @@ public int FrequencyFlat()
 
             return count;
         }
-
+#if X64
         [Benchmark(OperationsPerInvoke = iterations)]
         public int FrequencyFlatAvx()
         {
@@ -61,7 +61,7 @@ public int FrequencyFlatAvx()
 
             return count;
         }
-
+#endif
         [Benchmark(OperationsPerInvoke = iterations)]
         public int FrequencyBlock()
         {
@@ -73,7 +73,11 @@ public int FrequencyBlock()
         }
 
         [Benchmark(OperationsPerInvoke = iterations)]
+#if Arm64
+        public int FrequencyBlockNeonNotPinned()
+#else
         public int FrequencyBlockAvxNotPinned()
+#endif
         {
             int count = 0;
             for (int i = 0; i < iterations; i++)
@@ -83,7 +87,12 @@ public int FrequencyBlockAvxNotPinned()
         }
 
         [Benchmark(OperationsPerInvoke = iterations)]
+
+#if Arm64
+        public int FrequencyBlockNeonPinned()
+#else
         public int FrequencyBlockAvxPinned()
+#endif
         {
             int count = 0;
             for (int i = 0; i < iterations; i++)
diff --git a/BitFaster.Caching.Benchmarks/Lfu/SketchIncrement.cs b/BitFaster.Caching.Benchmarks/Lfu/SketchIncrement.cs
@@ -27,6 +27,7 @@ public class SketchIncrement
         private CmSketchNoPin<int, DetectIsa> blockAvxNoPin;
         private CmSketchCore<int, DetectIsa> blockAvx;
 
+
         [Params(32_768, 524_288, 8_388_608, 134_217_728)]
         public int Size { get; set; }
 
@@ -49,7 +50,7 @@ public void IncFlat()
                 flatStd.Increment(i);
             }
         }
-
+#if X64
         [Benchmark(OperationsPerInvoke = iterations)]
         public void IncFlatAvx()
         {
@@ -58,7 +59,7 @@ public void IncFlatAvx()
                 flatAvx.Increment(i);
             }
         }
-
+#endif
         [Benchmark(OperationsPerInvoke = iterations)]
         public void IncBlock()
         {
@@ -69,7 +70,11 @@ public void IncBlock()
         }
 
         [Benchmark(OperationsPerInvoke = iterations)]
+#if Arm64
+        public void IncBlockNeonNotPinned()
+#else
         public void IncBlockAvxNotPinned()
+#endif
         {
             for (int i = 0; i < iterations; i++)
             {
@@ -78,7 +83,11 @@ public void IncBlockAvxNotPinned()
         }
 
         [Benchmark(OperationsPerInvoke = iterations)]
+#if Arm64
+        public void IncBlockNeonPinned()
+#else
         public void IncBlockAvxPinned()
+#endif
         {
             for (int i = 0; i < iterations; i++)
             {
diff --git a/BitFaster.Caching.UnitTests/Intrinsics.cs b/BitFaster.Caching.UnitTests/Intrinsics.cs
@@ -1,6 +1,10 @@
 ﻿#if NETCOREAPP3_1_OR_GREATER
 using System.Runtime.Intrinsics.X86;
 #endif
+#if NET6_0_OR_GREATER
+using System.Runtime.Intrinsics.Arm;
+#endif
+
 using Xunit;
 
 namespace BitFaster.Caching.UnitTests
@@ -10,8 +14,14 @@ public static class Intrinsics
         public static void SkipAvxIfNotSupported<I>()
         {
 #if NETCOREAPP3_1_OR_GREATER
+    #if NET6_0_OR_GREATER
+            // when we are trying to test Avx2/Arm64, skip the test if it's not supported
+            Skip.If(typeof(I) == typeof(DetectIsa) && !(Avx2.IsSupported || AdvSimd.Arm64.IsSupported));
+    #else
             // when we are trying to test Avx2, skip the test if it's not supported
             Skip.If(typeof(I) == typeof(DetectIsa) && !Avx2.IsSupported);
+    #endif
+
 #else
             Skip.If(true);
 #endif
diff --git a/BitFaster.Caching.UnitTests/Lfu/CmSketchTests.cs b/BitFaster.Caching.UnitTests/Lfu/CmSketchTests.cs
@@ -6,12 +6,12 @@
 
 namespace BitFaster.Caching.UnitTests.Lfu
 {
-    // Test with AVX2 if it is supported
-    public class CMSketchAvx2Tests : CmSketchTestBase<DetectIsa>
+    // Test with AVX2/ARM64 if it is supported
+    public class CMSketchIntrinsicsTests : CmSketchTestBase<DetectIsa>
     {
     }
 
-    // Test with AVX2 disabled
+    // Test with AVX2/ARM64 disabled
     public class CmSketchTests : CmSketchTestBase<DisableHardwareIntrinsics>
     {
     }
@@ -29,14 +29,23 @@ public CmSketchTestBase()
         public void Repro()
         {
             sketch = new CmSketchCore<int, I>(1_048_576, EqualityComparer<int>.Default);
+            var baseline = new CmSketchCore<int, DisableHardwareIntrinsics>(1_048_576, EqualityComparer<int>.Default);
 
             for (int i = 0; i < 1_048_576; i++)
             {
                 if (i % 3 == 0)
                 {
                     sketch.Increment(i);
+                    baseline.Increment(i);
                 }
             }
+
+            baseline.Size.Should().Be(sketch.Size);
+
+            for (int i = 0; i < 1_048_576; i++)
+            {
+                sketch.EstimateFrequency(i).Should().Be(baseline.EstimateFrequency(i));
+            }
         }
 
 
diff --git a/BitFaster.Caching/Intrinsics.cs b/BitFaster.Caching/Intrinsics.cs
@@ -2,6 +2,10 @@
 using System.Runtime.Intrinsics.X86;
 #endif
 
+#if NET6_0
+using System.Runtime.Intrinsics.Arm;
+#endif
+
 namespace BitFaster.Caching
 {
     /// <summary>
@@ -12,7 +16,14 @@ public interface IsaProbe
         /// <summary>
         /// Gets a value indicating whether AVX2 is supported.
         /// </summary>
-        bool IsAvx2Supported { get; }
+        bool IsAvx2Supported { get; }
+
+#if NET6_0_OR_GREATER
+        /// <summary>
+        /// Gets a value indicating whether Arm64 is supported.
+        /// </summary>
+        bool IsArm64Supported { get => false; }
+#endif
     }
 
     /// <summary>
@@ -25,7 +36,15 @@ public interface IsaProbe
         public bool IsAvx2Supported => false;
 #else
         /// <inheritdoc/>
-        public bool IsAvx2Supported => Avx2.IsSupported;
+        public bool IsAvx2Supported => Avx2.IsSupported;
+#endif
+
+#if NET6_0_OR_GREATER
+        /// <inheritdoc/>
+        public bool IsArm64Supported => AdvSimd.Arm64.IsSupported;
+#else
+        /// <inheritdoc/>
+        public bool IsArm64Supported => false;
 #endif
     }
 
@@ -35,6 +54,9 @@ public interface IsaProbe
     public readonly struct DisableHardwareIntrinsics : IsaProbe
     {
         /// <inheritdoc/>
-        public bool IsAvx2Supported => false;
+        public bool IsAvx2Supported => false;
+
+        /// <inheritdoc/>
+        public bool IsArm64Supported => false;
     }
 }
diff --git a/BitFaster.Caching/Lfu/CmSketchCore.cs b/BitFaster.Caching/Lfu/CmSketchCore.cs

Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ public int FrequencyFlat()`
`51`	`51`
`52`	`52`	`return count;`
`53`	`53`	`}`
`54`		`-`
	`54`	`+#if X64`
`55`	`55`	`[Benchmark(OperationsPerInvoke = iterations)]`
`56`	`56`	`public int FrequencyFlatAvx()`
`57`	`57`	`{`
`@@ -61,7 +61,7 @@ public int FrequencyFlatAvx()`
`61`	`61`
`62`	`62`	`return count;`
`63`	`63`	`}`
`64`		`-`
	`64`	`+#endif`
`65`	`65`	`[Benchmark(OperationsPerInvoke = iterations)]`
`66`	`66`	`public int FrequencyBlock()`
`67`	`67`	`{`
`@@ -73,7 +73,11 @@ public int FrequencyBlock()`
`73`	`73`	`}`
`74`	`74`
`75`	`75`	`[Benchmark(OperationsPerInvoke = iterations)]`
	`76`	`+#if Arm64`
	`77`	`+ public int FrequencyBlockNeonNotPinned()`
	`78`	`+#else`
`76`	`79`	`public int FrequencyBlockAvxNotPinned()`
	`80`	`+#endif`
`77`	`81`	`{`
`78`	`82`	`int count = 0;`
`79`	`83`	`for (int i = 0; i < iterations; i++)`
`@@ -83,7 +87,12 @@ public int FrequencyBlockAvxNotPinned()`
`83`	`87`	`}`
`84`	`88`
`85`	`89`	`[Benchmark(OperationsPerInvoke = iterations)]`
	`90`	`+`
	`91`	`+#if Arm64`
	`92`	`+ public int FrequencyBlockNeonPinned()`
	`93`	`+#else`
`86`	`94`	`public int FrequencyBlockAvxPinned()`
	`95`	`+#endif`
`87`	`96`	`{`
`88`	`97`	`int count = 0;`
`89`	`98`	`for (int i = 0; i < iterations; i++)`
Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@ public class SketchIncrement`
`27`	`27`	`private CmSketchNoPin<int, DetectIsa> blockAvxNoPin;`
`28`	`28`	`private CmSketchCore<int, DetectIsa> blockAvx;`
`29`	`29`
	`30`	`+`
`30`	`31`	`[Params(32_768, 524_288, 8_388_608, 134_217_728)]`
`31`	`32`	`public int Size { get; set; }`
`32`	`33`
`@@ -49,7 +50,7 @@ public void IncFlat()`
`49`	`50`	`flatStd.Increment(i);`
`50`	`51`	`}`
`51`	`52`	`}`
`52`		`-`
	`53`	`+#if X64`
`53`	`54`	`[Benchmark(OperationsPerInvoke = iterations)]`
`54`	`55`	`public void IncFlatAvx()`
`55`	`56`	`{`
`@@ -58,7 +59,7 @@ public void IncFlatAvx()`
`58`	`59`	`flatAvx.Increment(i);`
`59`	`60`	`}`
`60`	`61`	`}`
`61`		`-`
	`62`	`+#endif`
`62`	`63`	`[Benchmark(OperationsPerInvoke = iterations)]`
`63`	`64`	`public void IncBlock()`
`64`	`65`	`{`
`@@ -69,7 +70,11 @@ public void IncBlock()`
`69`	`70`	`}`
`70`	`71`
`71`	`72`	`[Benchmark(OperationsPerInvoke = iterations)]`
	`73`	`+#if Arm64`
	`74`	`+ public void IncBlockNeonNotPinned()`
	`75`	`+#else`
`72`	`76`	`public void IncBlockAvxNotPinned()`
	`77`	`+#endif`
`73`	`78`	`{`
`74`	`79`	`for (int i = 0; i < iterations; i++)`
`75`	`80`	`{`
`@@ -78,7 +83,11 @@ public void IncBlockAvxNotPinned()`
`78`	`83`	`}`
`79`	`84`
`80`	`85`	`[Benchmark(OperationsPerInvoke = iterations)]`
	`86`	`+#if Arm64`
	`87`	`+ public void IncBlockNeonPinned()`
	`88`	`+#else`
`81`	`89`	`public void IncBlockAvxPinned()`
	`90`	`+#endif`
`82`	`91`	`{`
`83`	`92`	`for (int i = 0; i < iterations; i++)`
`84`	`93`	`{`
Original file line number	Diff line number	Diff line change
`@@ -6,12 +6,12 @@`
`6`	`6`
`7`	`7`	`namespace BitFaster.Caching.UnitTests.Lfu`
`8`	`8`	`{`
`9`		`- // Test with AVX2 if it is supported`
`10`		`- public class CMSketchAvx2Tests : CmSketchTestBase<DetectIsa>`
	`9`	`+ // Test with AVX2/ARM64 if it is supported`
	`10`	`+ public class CMSketchIntrinsicsTests : CmSketchTestBase<DetectIsa>`
`11`	`11`	`{`
`12`	`12`	`}`
`13`	`13`
`14`		`- // Test with AVX2 disabled`
	`14`	`+ // Test with AVX2/ARM64 disabled`
`15`	`15`	`public class CmSketchTests : CmSketchTestBase<DisableHardwareIntrinsics>`
`16`	`16`	`{`
`17`	`17`	`}`
`@@ -29,14 +29,23 @@ public CmSketchTestBase()`
`29`	`29`	`public void Repro()`
`30`	`30`	`{`
`31`	`31`	`sketch = new CmSketchCore<int, I>(1_048_576, EqualityComparer<int>.Default);`
	`32`	`+ var baseline = new CmSketchCore<int, DisableHardwareIntrinsics>(1_048_576, EqualityComparer<int>.Default);`
`32`	`33`
`33`	`34`	`for (int i = 0; i < 1_048_576; i++)`
`34`	`35`	`{`
`35`	`36`	`if (i % 3 == 0)`
`36`	`37`	`{`
`37`	`38`	`sketch.Increment(i);`
	`39`	`+ baseline.Increment(i);`
`38`	`40`	`}`
`39`	`41`	`}`
	`42`	`+`
	`43`	`+ baseline.Size.Should().Be(sketch.Size);`
	`44`	`+`
	`45`	`+ for (int i = 0; i < 1_048_576; i++)`
	`46`	`+ {`
	`47`	`+ sketch.EstimateFrequency(i).Should().Be(baseline.EstimateFrequency(i));`
	`48`	`+ }`
`40`	`49`	`}`
`41`	`50`
`42`	`51`