Offload node score buffers

2023-07-07 15:45:42 +02:00
parent 1386f9150c
commit e4d9e3a52e
10 changed files with 188 additions and 97 deletions
@@ -0,0 +1,90 @@
+using System;
+using System.ComponentModel;
+using System.Diagnostics.Contracts;
+using System.Numerics;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+
+namespace Craftimizer.Solver.Crafty;
+
+// Adapted from https://github.com/dtao/ConcurrentList/blob/4fcf1c76e93021a41af5abb2d61a63caeba2adad/ConcurrentList/ConcurrentList.cs
+public struct NodeScoresBuffer
+{
+    public sealed class ScoresBatch
+    {
+        public Memory<float> ScoreSum;
+        public Memory<float> MaxScore;
+        public Memory<int> Visits;
+
+        public ScoresBatch()
+        {
+            ScoreSum = new float[BatchSize];
+            MaxScore = new float[BatchSize];
+            Visits = new int[BatchSize];
+        }
+    }
+
+    // Technically 25, but it's very unlikely to actually get to there.
+    // The benchmark reaches 20 at most, but here we have a little leeway just in case.
+    private const int MaxSize = 24;
+
+    private static int BatchSize = Vector<float>.Count;
+    private static int BatchSizeBits = int.Log2(BatchSize);
+    private static int BatchSizeMask = BatchSize - 1;
+
+    private static int BatchCount = MaxSize / BatchSize;
+
+    public ScoresBatch[] Data;
+    private int index;
+    private int count;
+
+    public readonly int Count => count;
+
+    public void AddConcurrent()
+    {
+        if (Data == null)
+            Interlocked.CompareExchange(ref Data, new ScoresBatch[BatchCount], null);
+
+        var idx = Interlocked.Increment(ref index) - 1;
+
+        var (arrayIdx, _) = GetArrayIndex(idx);
+
+        if (Data[arrayIdx] == null)
+            Interlocked.CompareExchange(ref Data[arrayIdx], new ScoresBatch(), null);
+
+        Interlocked.Increment(ref count);
+    }
+
+    public void Add()
+    {
+        Data ??= new ScoresBatch[BatchCount];
+
+        var idx = count++;
+
+        var (arrayIdx, _) = GetArrayIndex(idx);
+
+        Data[arrayIdx] ??= new();
+    }
+
+    public readonly void VisitConcurrent((int arrayIdx, int subIdx) at, float score)
+    {
+        Intrinsics.CASAdd(ref Data[at.arrayIdx].ScoreSum.Span[at.subIdx], score);
+        Intrinsics.CASMax(ref Data[at.arrayIdx].MaxScore.Span[at.subIdx], score);
+        Interlocked.Increment(ref Data[at.arrayIdx].Visits.Span[at.subIdx]);
+    }
+
+    public readonly void Visit((int arrayIdx, int subIdx) at, float score)
+    {
+        Data[at.arrayIdx].ScoreSum.Span[at.subIdx] += score;
+        Data[at.arrayIdx].MaxScore.Span[at.subIdx] = Math.Max(Data[at.arrayIdx].MaxScore.Span[at.subIdx], score);
+        Data[at.arrayIdx].Visits.Span[at.subIdx]++;
+    }
+
+    public readonly int GetVisits((int arrayIdx, int subIdx) at) =>
+        Data[at.arrayIdx].Visits.Span[at.subIdx];
+
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static (int arrayIdx, int subIdx) GetArrayIndex(int idx) =>
+        (idx >> BatchSizeBits, idx & BatchSizeMask);
+}