Optimize workers

Author: emelrad12

JobScheduler.Benchmarks/Benchmark.cs

@@ -126,7 +126,7 @@ public long End(int jobs, string type)
 
 public class Benchmark
 {
-    private const int jobCount = 20000;
+    private const int jobCount = 2000;
     private const int loopCount = 1000;
     private static JobScheduler jobScheduler = new();
 
@@ -208,14 +208,16 @@ private static long BenchVector(bool useVector)
     {
         var timer = new JobTimer();
         var data = new VectorCalculationJob { a = new float[jobCount], b = new float[jobCount], result = new float[jobCount], Repetitions = 500 };
+        var parentJob = jobScheduler.Schedule();
         for (var sindex = 0; sindex < loopCount; sindex++)
         {
             var job = new ParallelJobProducer<VectorCalculationJob>(0, jobCount, data, jobScheduler, 16, !useVector);
             job.CheckAndSplit();
+            job.GetHandle().SetParent(parentJob);
             jobScheduler.Flush(job.GetHandle());
-            jobScheduler.Wait(job.GetHandle());
         }
-
+        jobScheduler.Flush(parentJob);
+        jobScheduler.Wait(parentJob);
         return timer.End(jobCount * loopCount, $"Use vector: {useVector}");
     }
 
@@ -234,7 +236,7 @@ private static void Main(string[] args)
         // config = config.WithOptions(ConfigOptions.DisableOptimizationsValidator);
         // BenchmarkRunner.Run<JobSchedulerBenchmark>(config);
         // return;
-        for (var i = 0;; i++)
+        for (var i = 0; i < 20; i++)
         {
             // CorrectnessTestJob();
             // BenchB();
@@ -244,6 +246,8 @@ private static void Main(string[] args)
             var nonVectorized = BenchVector(false);
             Console.WriteLine($"Ratio {(double)nonVectorized / vectorized}");
         }
+        jobScheduler.Dispose();
+
         //using var jobScheduler = new JobScheduler();
 
         // Spawn massive jobs and wait for finish

JobScheduler/JobScheduler.cs

@@ -102,14 +102,9 @@ public float CalculateThreadUsage()
 
     private void FlushDirect(JobHandle job)
     {
-        // Round Robin,
-        var workerIndex = NextWorkerIndex;
-        while (!Workers[workerIndex].Enqueue(job))
+        while (!Worker.Enqueue(job))
         {
-            NextWorkerIndex = (NextWorkerIndex + 1) % Workers.Count;
         }
-
-        NextWorkerIndex = (NextWorkerIndex + 1) % Workers.Count;
     }
 
     /// <summary>
@@ -143,24 +138,28 @@ public void Wait(JobHandle job)
     {
         while (!job.IsFinished())
         {
-            var jobWasFound = false;
-            for (var i = 0; i < Workers.Count; i++)
+            var nextJob = Worker.TryStealJobExternal(out var stolenJob);
+            if (!nextJob)
             {
-                var nextJob = Workers[i].Queue.TrySteal(out var stolenJob);
-                if (!nextJob)
+                for (var i = 0; i < Workers.Count; i++)
                 {
-                    continue;
+                    if (nextJob)
+                    {
+                        break;
+                    }
+                    nextJob = Workers[i].Queue.TrySteal(out stolenJob);
                 }
+            }
 
-                jobWasFound = true;
+            if (nextJob)
+            {
                 stolenJob.Job?.Execute();
                 Finish(stolenJob);
             }
 
             // Dont yield if you can find something to execute.
-            if (!jobWasFound)
+            if (!nextJob)
             {
-                // Console.WriteLine($"Cpu usage: {CalculateThreadUsages()}");
                 Thread.Yield();
             }
         }

JobScheduler/Utils/BulkQueue.cs

@@ -0,0 +1,80 @@
+using System.Collections.Concurrent;
+
+namespace Schedulers.Utils;
+
+/// <summary>
+/// Thread-safe bulk queue implementation.
+/// Provides around 10-20% performance in jobs compared to ordinary flush.
+/// </summary>
+public class BulkQueue<T>
+{
+    public BulkQueue(int segmentCount, int segmentSize)
+    {
+        fullSegments = [];
+        emptySegments = [];
+        for (var i = 0; i < segmentCount; i++)
+        {
+            emptySegments.Add(new(segmentSize));
+        }
+    }
+
+    public struct Segment
+    {
+        private T[] items;
+        private int index;
+
+        public Segment(int capacity)
+        {
+            items = new T[capacity];
+        }
+
+        public bool Dequeue(out T item)
+        {
+            if (index == 0)
+            {
+                item = default!;
+                return false;
+            }
+
+            index--;
+            item = items[index];
+            items[index] = default!; // Clear the reference
+            return true;
+        }
+
+        public bool Enqueue(T item)
+        {
+            if (index >= items.Length)
+            {
+                return false; // Segment is full
+            }
+
+            items[index] = item;
+            index++;
+            return true;
+        }
+    }
+
+    private ConcurrentBag<Segment> emptySegments = new();
+    private ConcurrentBag<Segment> fullSegments = new();
+
+    public bool GetEmtySegment(out Segment segment)
+    {
+        return emptySegments.TryTake(out segment);
+    }
+
+    public bool GetFullSegment(out Segment segment)
+    {
+        return fullSegments.TryTake(out segment);
+    }
+
+    public void Return(Segment segment)
+    {
+        emptySegments.Add(segment);
+    }
+
+    public void Enqueue(Segment segment)
+    {
+        fullSegments.Add(segment);
+    }
+}

JobScheduler/Utils/IParallelJobProducer.cs

@@ -75,7 +75,6 @@ private ParallelJobProducer(T producer, JobScheduler scheduler, JobHandle parent
         _loopSize = loopSize;
         _onlySingle = onlySingle;
         _selfHandle = _scheduler.Schedule(this, parent);
-        _scheduler.Flush(_selfHandle);
     }
 
     /// <summary>
@@ -141,6 +140,10 @@ public bool CheckAndSplit()
     private void Split()
     {
         var childrenToSplitInto = CalculateChildrenToSplitInto();
+        BulkQueue<JobHandle>.Segment segment;
+        while (!Worker.GetEmptySegment(out segment))
+        {
+        }
         _selfHandle.Job = null;
         for (var i = 0; i < childrenToSplitInto; i++)
         {
@@ -151,8 +154,11 @@ private void Split()
                 throw new($"Invalid range from {start} to {end}");
             }
 
-            _ = new ParallelJobProducer<T>(_producer, _scheduler, _selfHandle, start, end, _loopSize, _onlySingle);
+            var pjob = new ParallelJobProducer<T>(_producer, _scheduler, _selfHandle, start, end, _loopSize, _onlySingle);
+            segment.Enqueue(pjob._selfHandle);
+            pjob._selfHandle.UnfinishedJobs--;
         }
+        Worker.EnqueueFullSegment(segment);
     }
 
     /// <summary>

JobScheduler/Utils/JobHandlePool.cs

@@ -6,7 +6,6 @@ namespace Schedulers.Utils;
 /// This <see cref="JobHandlePool"/> class
 /// acts as a pool for <see cref="JobHandle"/> ids.
 /// </summary>
-// This class somehow provides a boost of around 4.4m to 6.1m in the benchmark.
 internal class JobHandlePool
 {
     // Handles are returned to the _freeHandles queue.
@@ -36,7 +35,7 @@ private void SwapQueues()
     /// <param name="size">Its maximum size.</param>
     public JobHandlePool(int size)
     {
-        MainThreadHandleCutoff = (ushort)(size / 2);
+        MainThreadHandleCutoff = (ushort)(size * 0.5f); // 80% of the handles are for the main thread.
         _mainThreadFreeHandles = new();
         _freeHandlesOtherThreads = new();
         _handleCache = new(MainThreadHandleCutoff);

JobScheduler/Utils/UnorderedQueue.cs

@@ -5,7 +5,7 @@ namespace Schedulers.Utils;
 public class UnorderedThreadSafeQueue<T>
 {
     private ConcurrentQueue<T> _queue = new();
-    private const int maxItems = 512;
+    private const int maxItems = 512 * 100;
 
     public UnorderedThreadSafeQueue()
     {

JobScheduler/Worker.cs

@@ -1,4 +1,5 @@
-using Schedulers.Utils;
+using System.Runtime.InteropServices;
+using Schedulers.Utils;
 
 namespace Schedulers;
 
@@ -11,15 +12,21 @@ internal class Worker
     private readonly int _workerId;
     private readonly Thread _thread;
 
+    /// <summary>
+    /// Use a single combined <see cref="UnorderedThreadSafeQueue{T}"/> for all <see cref="Worker"/>s to enqueue <see cref="JobHandle"/>s.
+    /// Because there really isn't performance advantage to having multiple queues for each <see cref="Worker"/>,
+    /// </summary>
     private static readonly UnorderedThreadSafeQueue<JobHandle> _incomingQueue = new();
-    private readonly WorkStealingDeque<JobHandle> _queue;
 
+    private readonly WorkStealingDeque<JobHandle> _queue;
     private readonly JobScheduler _jobScheduler;
     private volatile CancellationTokenSource _cancellationToken;
+    private static readonly BulkQueue<JobHandle> _bulkQueue = new(1000, 128);
     public bool IsCurrentlyWorking { get; private set; } = false;
+
     // use a high spin count to avoid sleeping the thread under variable load.
     // 2047 is the maximum value for a ManualResetEventSlim spin count.
-    private readonly ManualResetEventSlim _workAvailable = new(false, 2047);
+    private readonly static ManualResetEventSlim _workAvailable = new(false, 2047);
 
     /// <summary>
     /// Creates a new <see cref="Worker"/>.
@@ -39,16 +46,25 @@ public Worker(JobScheduler jobScheduler, int id)
         _thread.Name = $"Arch Worker #{_workerId}";
     }
 
-    public bool Enqueue(JobHandle handle)
+    public static bool Enqueue(JobHandle handle)
     {
         var result = _incomingQueue.TryEnqueue(handle);
-        _workAvailable.Set();
         return result;
     }
 
     public void ForceEnqueue(JobHandle handle)
     {
-       _incomingQueue.ForceEnqueue(handle);
+        _incomingQueue.ForceEnqueue(handle);
+    }
+
+    public static bool GetEmptySegment(out BulkQueue<JobHandle>.Segment segment)
+    {
+        return _bulkQueue.GetEmtySegment(out segment);
+    }
+
+    public static void EnqueueFullSegment(BulkQueue<JobHandle>.Segment segment)
+    {
+        _bulkQueue.Enqueue(segment);
     }
 
     /// <summary>
@@ -81,7 +97,14 @@ public void Start()
     public void Stop()
     {
         _cancellationToken.Cancel();
-        // _workAvailable.Set();
+        _workAvailable.Set();
+    }
+
+    public static bool TryStealJobExternal(out JobHandle job)
+    {
+        // job = new JobHandle();
+        // return false;
+        return _incomingQueue.TryDequeue(out job);
     }
 
     /// <summary>
@@ -95,11 +118,29 @@ private void Run(CancellationToken token)
         {
             while (!token.IsCancellationRequested)
             {
+                // it is faster to do it in a single threaded manner to lower the queue contention.
+                if (_workerId == 0)
+                {
+                    while (_bulkQueue.GetFullSegment(out var segment))
+                    {
+                        while (segment.Dequeue(out var item))
+                        {
+                            ForceEnqueue(item);
+                        }
+                        _bulkQueue.Return(segment);
+                    }
+                }
+
                 IsCurrentlyWorking = true;
                 var noWorkFound = true;
                 // Pass jobs to the local queue
                 while (_queue.Size() < 32 && _incomingQueue.TryDequeue(out var jobHandle))
                 {
+                    if (_workerId == 0)
+                    {
+                        _workAvailable.Set();
+                    }
+
                     _queue.PushBottom(jobHandle);
                     noWorkFound = false;
                 }
@@ -134,6 +175,7 @@ private void Run(CancellationToken token)
                         break;
                     }
                 }
+
                 if (noWorkFound)
                 {
                     IsCurrentlyWorking = false;
@@ -146,7 +188,7 @@ private void Run(CancellationToken token)
                     else
                     {
                         _workAvailable.Reset();
-                        _workAvailable.Wait(1, token);
+                        _workAvailable.Wait(100, token);
                     }
                 }
             }