Merge pull request #253 from mlemanczyk/codex/optimize-evenperfectbitscanner-for-primetestergpucontextleas

Add GPU lease warm-up and shared limiter

EvenPerfectBitScanner/Program.cs

@@ -58,14 +58,15 @@ internal static class Program
 			bool startPrimeProvided = _cliArguments.StartPrimeProvided;
 			int threadCount = Math.Max(1, _cliArguments.ThreadCount);
 			int gpuPrimeBatch = Math.Max(1, _cliArguments.GpuPrimeBatch);
+			int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads);
 			UnboundedTaskScheduler.ConfigureThreadCount(threadCount);
 			PrimeTester.GpuBatchSize = gpuPrimeBatch;
+			GpuPrimeWorkLimiter.SetLimit(gpuPrimeThreads);
 			GpuContextPool.WarmUpPool(threadCount);
-			PrimeTester.WarmUpGpuKernels(threadCount);
+			PrimeTester.WarmUpGpuKernels(gpuPrimeThreads);
 			Console.WriteLine("Starting up threads...");
 			_ = UnboundedTaskScheduler.Instance;
 			int blockSize = Math.Max(1, _cliArguments.BlockSize);
-			int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads);
 			GpuKernelType kernelType = _cliArguments.KernelType;
 			bool useBitTransform = _cliArguments.UseBitTransform;
 			bool useOrder = _cliArguments.UseOrder;
@@ -117,9 +118,6 @@ internal static class Program
 				return;
 			}
 
-			// Apply GPU prime sieve runtime configuration
-			GpuPrimeWorkLimiter.SetLimit(gpuPrimeThreads);
-			PrimeTester.GpuBatchSize = gpuPrimeBatch;
 
 			MersenneDivisorCycles mersenneDivisorCycles = new();
 			if (!File.Exists(cyclesPath) || continueCyclesGeneration)

PerfectNumbers.Core/Gpu/GpuContextPool.cs

@@ -59,6 +59,7 @@ public static class GpuContextPool
 		return new GpuContextLease(new PooledContext());
 	}
 
+
 	public static void WarmUpPool(int threadCount)
 	{
 		if (threadCount <= WarmedGpuContextCount)
@@ -98,21 +99,30 @@ public static class GpuContextPool
 		GpuPool.Enqueue(ctx);
 	}
 
-	public readonly struct GpuContextLease
+	public struct GpuContextLease
 	{
-		private readonly PooledContext _ctx;
+		private PooledContext? _ctx;
+		private bool _disposed;
 
 		internal GpuContextLease(PooledContext ctx)
 		{
 			_ctx = ctx;
+			_disposed = false;
 		}
 
-		public readonly Context Context => _ctx.Context;
+		public Context Context => _ctx!.Context;
-		public readonly Accelerator Accelerator => _ctx.Accelerator;
+		public Accelerator Accelerator => _ctx!.Accelerator;
 
-		public readonly void Dispose()
+		public void Dispose()
 		{
-			Return(_ctx);
+			if (_disposed)
+			{
+				return;
+			}
+
+			_disposed = true;
+			Return(_ctx!);
+			_ctx = null;
 		}
 	}
 }

PerfectNumbers.Core/Gpu/GpuPrimeWorkLimiter.cs

@@ -1,3 +1,6 @@
+using System;
+using System.Threading;
+
 namespace PerfectNumbers.Core.Gpu;
 
 public static class GpuPrimeWorkLimiter
@@ -10,12 +13,17 @@ public static class GpuPrimeWorkLimiter
         // TODO: Inline the pooled limiter guard from the GpuLimiterThroughputBenchmarks so prime-sieve
         // GPU jobs stop allocating Lease instances on every acquisition.
         var sem = _semaphore;
-        // sem.Wait();
+        sem.Wait();
         return new Lease(sem);
     }
 
     public static void SetLimit(int value)
     {
+        if (value <= 0)
+        {
+            value = 1;
+        }
+
         if (value == _currentLimit)
         {
             return;
@@ -23,11 +31,13 @@ public static class GpuPrimeWorkLimiter
 
         // TODO: Switch to a shared limiter implementation with GpuWorkLimiter so we can coordinate CPU/GPU prime work using
         // the same semaphore pool without reallocating per adjustment.
+        var previous = _semaphore;
         _semaphore = new SemaphoreSlim(value, value);
         _currentLimit = value;
+        previous.Dispose();
     }
 
-    public sealed class Lease
+    public sealed class Lease : IDisposable
     {
         private readonly SemaphoreSlim _sem;
 
@@ -38,8 +48,7 @@ public static class GpuPrimeWorkLimiter
 
         public void Dispose()
         {
-            // _sem.Release();
+            _sem.Release();
         }
     }
 }
-

PerfectNumbers.Core/PrimeTester.cs

@@ -159,18 +159,7 @@ public sealed class PrimeTester
 				return;
 			}
 
-			int toWarm = target - WarmedGpuLeaseCount;
+			PrimeTesterGpuContextPool.WarmUp(target, GpuBatchSize);
-			var leases = new PrimeTesterGpuContextPool.PrimeTesterGpuContextLease[toWarm];
-			for (int i = 0; i < toWarm; i++)
-			{
-				leases[i] = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
-			}
-
-			for (int i = 0; i < toWarm; i++)
-			{
-				leases[i].Dispose();
-			}
-
 			WarmedGpuLeaseCount = target;
 		}
 	}
@@ -229,6 +218,8 @@ public sealed class PrimeTester
 	internal static class PrimeTesterGpuContextPool
 	{
 		private static readonly ConcurrentQueue<PrimeTesterGpuContextLease> Pool = new();
+		private static readonly object WarmUpLock = new();
+		private static int WarmedLeaseCount;
 		private static readonly object SharedTablesLock = new();
 		private static readonly Dictionary<string, SharedHeuristicGpuTables> SharedTables = new();
 
@@ -267,6 +258,30 @@ public sealed class PrimeTester
 			}
 		}
 
+		internal static void WarmUp(int target, int minBufferCapacity)
+		{
+			if (target <= 0)
+			{
+				return;
+			}
+
+			lock (WarmUpLock)
+			{
+				if (target <= WarmedLeaseCount)
+				{
+					return;
+				}
+
+				int toCreate = target - WarmedLeaseCount;
+				for (int i = 0; i < toCreate; i++)
+				{
+					Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity));
+				}
+
+				WarmedLeaseCount = target;
+			}
+		}
+
                 internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
                 {
                         if (!Pool.TryDequeue(out var lease))
@@ -282,6 +297,7 @@ public sealed class PrimeTester
                         return lease;
                 }
 
+
 		internal static void Return(PrimeTesterGpuContextLease lease)
 		{
 			lease.Accelerator.Synchronize();
@@ -304,6 +320,8 @@ public sealed class PrimeTester
 
 				SharedTables.Clear();
 			}
+
+			WarmedLeaseCount = 0;
 		}
 
 		internal static void Clear(Accelerator accelerator)
@@ -313,6 +331,8 @@ public sealed class PrimeTester
 				return;
 			}
 
+			lock (WarmUpLock)
+			{
 				var retained = new List<PrimeTesterGpuContextLease>();
 				while (Pool.TryDequeue(out var lease))
 				{
@@ -331,6 +351,9 @@ public sealed class PrimeTester
 				{
 					Pool.Enqueue(retained[i]);
 				}
+
+				WarmedLeaseCount = Pool.Count;
+			}
 		}
 
 		internal sealed class PrimeTesterGpuContextLease : IDisposable