Support and accelerate bitmap operations on more platforms

Directory.Packages.props

@@ -22,10 +22,11 @@
     <PackageVersion Include="Microsoft.NET.Test.Sdk" Version="17.13.0" />
     <PackageVersion Include="Microsoft.IdentityModel.Protocols.OpenIdConnect" Version="8.6.1" />
     <PackageVersion Include="Microsoft.IdentityModel.Validators" Version="8.6.1" />
-    <PackageVersion Include="StackExchange.Redis" Version="2.8.16" />
+    <PackageVersion Include="StackExchange.Redis" Version="2.9.25" />
     <PackageVersion Include="System.IdentityModel.Tokens.Jwt" Version="8.6.1" />
     <PackageVersion Include="System.Interactive.Async" Version="6.0.1" />
     <PackageVersion Include="System.Text.Json" Version="9.0.3" />
+    <PackageVersion Include="System.Numerics.Tensors" Version="9.0.9" />
     <PackageVersion Include="Microsoft.Extensions.Hosting" Version="9.0.8" />
     <PackageVersion Include="Microsoft.Extensions.Hosting.WindowsServices" Version="9.0.8" />
   </ItemGroup>

benchmark/BDN.benchmark/Bitmap/BinaryOperations.cs

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System.Runtime.InteropServices;
+using BenchmarkDotNet.Attributes;
+using Garnet.server;
+
+namespace BDN.benchmark.Bitmap
+{
+    public unsafe class BinaryOperations
+    {
+        private const int Alignment = 64;
+
+        [ParamsSource(nameof(GetBitmapSizes))]
+        public Sizes BitmapSizes { get; set; }
+
+        [Params(BitmapOperation.XOR)]
+        public BitmapOperation Op { get; set; }
+
+        public IEnumerable<Sizes> GetBitmapSizes()
+        {
+            yield return new([1 << 21, 1 << 21]);
+            yield return new([1 << 21, (1 << 21) + 1]);
+
+            yield return new([1 << 21, 1 << 21, 1 << 21]);
+            yield return new([1 << 21, 1 << 21, (1 << 21) + 1]);
+
+            yield return new([256, 6 * 512 + 7, 512, 1024]);
+        }
+
+        private int minBitmapSize;
+        private byte** srcPtrs;
+        private byte** srcEndPtrs;
+
+        private int dstLength;
+        private byte* dstPtr;
+
+        [GlobalSetup]
+        public void GlobalSetup_Binary()
+        {
+            minBitmapSize = BitmapSizes.Values.Min();
+            srcPtrs = (byte**)NativeMemory.AllocZeroed((nuint)BitmapSizes.Values.Length, (nuint)sizeof(byte*));
+            srcEndPtrs = (byte**)NativeMemory.AllocZeroed((nuint)BitmapSizes.Values.Length, (nuint)sizeof(byte*));
+
+            for (var i = 0; i < BitmapSizes.Values.Length; i++)
+            {
+                srcPtrs[i] = (byte*)NativeMemory.AlignedAlloc((nuint)BitmapSizes.Values[i], Alignment);
+                srcEndPtrs[i] = srcPtrs[i] + BitmapSizes.Values[i];
+
+                new Random(i).NextBytes(new Span<byte>(srcPtrs[i], BitmapSizes.Values[i]));
+            }
+
+            dstLength = BitmapSizes.Values.Max();
+            dstPtr = (byte*)NativeMemory.AlignedAlloc((nuint)dstLength, Alignment);
+        }
+
+        [Benchmark]
+        public void BinaryOperation()
+        {
+            BitmapManager.InvokeBitOperationUnsafe(Op, BitmapSizes.Values.Length, srcPtrs, srcEndPtrs, dstPtr, dstLength, minBitmapSize);
+        }
+
+        [GlobalCleanup]
+        public void GlobalCleanup()
+        {
+            for (var i = 0; i < BitmapSizes.Values.Length; i++)
+            {
+                NativeMemory.AlignedFree(srcPtrs[i]);
+            }
+
+            NativeMemory.Free(srcPtrs);
+            NativeMemory.Free(srcEndPtrs);
+            NativeMemory.AlignedFree(dstPtr);
+        }
+
+        public record struct Sizes(int[] Values)
+        {
+            public override string ToString() => string.Join(", ", Values);
+        }
+    }
+}

benchmark/BDN.benchmark/Bitmap/UnaryOperations.cs

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System.Runtime.InteropServices;
+using BenchmarkDotNet.Attributes;
+using Garnet.server;
+
+namespace BDN.benchmark.Bitmap
+{
+    public unsafe partial class UnaryOperations
+    {
+        private const int Alignment = 64;
+
+        [ParamsSource(nameof(GetBitmapSize))]
+        public int BitmapSize { get; set; }
+
+        public IEnumerable<int> GetBitmapSize()
+        {
+            yield return 256;
+            yield return 1 << 21;
+        }
+
+        private const int Keys = 1;
+        private byte** srcPtrs;
+        private byte** srcEndPtrs;
+
+        private byte* dstPtr;
+
+        [GlobalSetup]
+        public void GlobalSetup_Unary()
+        {
+            srcPtrs = (byte**)NativeMemory.AllocZeroed(Keys, (nuint)sizeof(byte*));
+            srcEndPtrs = (byte**)NativeMemory.AllocZeroed(Keys, (nuint)sizeof(byte*));
+
+            srcPtrs[0] = (byte*)NativeMemory.AlignedAlloc((uint)BitmapSize, Alignment);
+            srcEndPtrs[0] = srcPtrs[0] + (uint)BitmapSize;
+
+            new Random(0).NextBytes(new Span<byte>(srcPtrs[0], BitmapSize));
+
+            dstPtr = (byte*)NativeMemory.AlignedAlloc((nuint)BitmapSize, Alignment);
+        }
+
+        [Benchmark]
+        public void BitOperation_NOT()
+        {
+            BitmapManager.InvokeBitOperationUnsafe(BitmapOperation.NOT, Keys, srcPtrs, srcEndPtrs, dstPtr, BitmapSize, BitmapSize);
+        }
+
+        [GlobalCleanup]
+        public void GlobalCleanup()
+        {
+            NativeMemory.AlignedFree(srcPtrs[0]);
+
+            NativeMemory.Free(srcPtrs);
+            NativeMemory.Free(srcEndPtrs);
+            NativeMemory.AlignedFree(dstPtr);
+        }
+    }
+}

libs/common/Numerics/IBinaryOperator.cs

@@ -0,0 +1,80 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System.Numerics;
+using System.Runtime.Intrinsics;
+
+namespace Garnet.common.Numerics
+{
+    /// <summary>Operator that takes two input values and returns a single value.</summary>
+    public interface IBinaryOperator
+    {
+        /// <summary>
+        /// Computes the binary operation of two scalar values.
+        /// </summary>
+        static abstract T Invoke<T>(T x, T y) where T : IBinaryInteger<T>;
+
+        /// <summary>
+        /// Computes the binary operation of two vectors.
+        /// </summary>
+        static abstract Vector128<byte> Invoke(Vector128<byte> x, Vector128<byte> y);
+
+        /// <inheritdoc cref="Invoke(Vector128{byte}, Vector128{byte})"/>
+        static abstract Vector256<byte> Invoke(Vector256<byte> x, Vector256<byte> y);
+
+        /// <inheritdoc cref="Invoke(Vector128{byte}, Vector128{byte})"/>
+        static abstract Vector512<byte> Invoke(Vector512<byte> x, Vector512<byte> y);
+    }
+
+    /// <summary><c>x &amp; y</c></summary>
+    public readonly struct BitwiseAndOperator : IBinaryOperator
+    {
+        /// <inheritdoc/>
+        public static T Invoke<T>(T x, T y) where T : IBinaryInteger<T> => x & y;
+        /// <inheritdoc/>
+        public static Vector128<byte> Invoke(Vector128<byte> x, Vector128<byte> y) => x & y;
+        /// <inheritdoc/>
+        public static Vector256<byte> Invoke(Vector256<byte> x, Vector256<byte> y) => x & y;
+        /// <inheritdoc/>
+        public static Vector512<byte> Invoke(Vector512<byte> x, Vector512<byte> y) => x & y;
+    }
+
+    /// <summary><c>x | y</c></summary>
+    public readonly struct BitwiseOrOperator : IBinaryOperator
+    {
+        /// <inheritdoc/>
+        public static T Invoke<T>(T x, T y) where T : IBinaryInteger<T> => x | y;
+        /// <inheritdoc/>
+        public static Vector128<byte> Invoke(Vector128<byte> x, Vector128<byte> y) => x | y;
+        /// <inheritdoc/>
+        public static Vector256<byte> Invoke(Vector256<byte> x, Vector256<byte> y) => x | y;
+        /// <inheritdoc/>
+        public static Vector512<byte> Invoke(Vector512<byte> x, Vector512<byte> y) => x | y;
+    }
+
+    /// <summary><c>x ^ y</c></summary>
+    public readonly struct BitwiseXorOperator : IBinaryOperator
+    {
+        /// <inheritdoc/>
+        public static T Invoke<T>(T x, T y) where T : IBinaryInteger<T> => x ^ y;
+        /// <inheritdoc/>
+        public static Vector128<byte> Invoke(Vector128<byte> x, Vector128<byte> y) => x ^ y;
+        /// <inheritdoc/>
+        public static Vector256<byte> Invoke(Vector256<byte> x, Vector256<byte> y) => x ^ y;
+        /// <inheritdoc/>
+        public static Vector512<byte> Invoke(Vector512<byte> x, Vector512<byte> y) => x ^ y;
+    }
+
+    /// <summary><c>x &amp; ~y</c></summary>
+    public readonly struct BitwiseAndNotOperator : IBinaryOperator
+    {
+        /// <inheritdoc/>
+        public static T Invoke<T>(T x, T y) where T : IBinaryInteger<T> => x & ~y;
+        /// <inheritdoc/>
+        public static Vector128<byte> Invoke(Vector128<byte> x, Vector128<byte> y) => x & ~y;
+        /// <inheritdoc/>
+        public static Vector256<byte> Invoke(Vector256<byte> x, Vector256<byte> y) => x & ~y;
+        /// <inheritdoc/>
+        public static Vector512<byte> Invoke(Vector512<byte> x, Vector512<byte> y) => x & ~y;
+    }
+}

libs/server/Garnet.server.csproj

@@ -20,6 +20,7 @@
     <PackageReference Include="Microsoft.Extensions.Logging" />
     <PackageReference Include="Microsoft.IdentityModel.Protocols.OpenIdConnect" />
     <PackageReference Include="System.IdentityModel.Tokens.Jwt" />
+    <PackageReference Include="System.Numerics.Tensors" />
     <PackageReference Include="KeraLua" />
   </ItemGroup>
 

libs/server/Lua/LuaRunner.Functions.cs

@@ -2857,6 +2857,7 @@ internal int AclCheckCommand(nint luaStatePtr)
                         case RespCommand.BITOP_OR: state.PushConstantString(constStrs.OR); break;
                         case RespCommand.BITOP_XOR: state.PushConstantString(constStrs.XOR); break;
                         case RespCommand.BITOP_NOT: state.PushConstantString(constStrs.NOT); break;
+                        case RespCommand.BITOP_DIFF: state.PushConstantString(constStrs.DIFF); break;
 
                         default: throw new InvalidOperationException($"Unexpected BITOP sub command: {subCommand}");
                     }

libs/server/Lua/LuaRunner.Strings.cs

@@ -161,6 +161,8 @@ private readonly struct ConstantStringRegistryIndexes
             internal int XOR { get; }
             /// <see cref="CmdStrings.LUA_NOT"/>
             internal int NOT { get; }
+            /// <see cref="CmdStrings.LUA_DIFF"/>
+            internal int DIFF { get; }
             /// <see cref="CmdStrings.LUA_KEYS"/>
             internal int KEYS { get; }
             /// <see cref="CmdStrings.LUA_ARGV"/>
@@ -246,6 +248,7 @@ internal ConstantStringRegistryIndexes(ref LuaStateWrapper state)
                 OR = ConstantStringToRegistry(ref state, CmdStrings.LUA_OR);
                 XOR = ConstantStringToRegistry(ref state, CmdStrings.LUA_XOR);
                 NOT = ConstantStringToRegistry(ref state, CmdStrings.LUA_NOT);
+                DIFF = ConstantStringToRegistry(ref state, CmdStrings.LUA_DIFF);
                 KEYS = ConstantStringToRegistry(ref state, CmdStrings.LUA_KEYS);
                 ARGV = ConstantStringToRegistry(ref state, CmdStrings.LUA_ARGV);
             }

libs/server/Resp/Bitmap/BitmapCommands.cs

@@ -40,7 +40,12 @@ public enum BitmapOperation : byte
         /// <summary>
         /// NOT
         /// </summary>
-        NOT
+        NOT,
+
+        /// <summary>
+        /// DIFF
+        /// </summary>
+        DIFF
     }
 
     internal enum BitFieldOverflow : byte
@@ -317,6 +322,11 @@ private bool NetworkStringBitOperation<TGarnetApi>(BitmapOperation bitOp, ref TG
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_NUMBER_OF_ARGUMENTS);
             }
 
+            if (bitOp == BitmapOperation.DIFF && parseState.Count < 3)
+            {
+                return AbortWithErrorMessage(CmdStrings.RESP_ERR_BITOP_DIFF_TWO_SOURCE_KEYS_REQUIRED);
+            }
+
             if (parseState.Count > 64)
             {
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_BITOP_KEY_LIMIT);