Merge pull request #152 from mlemanczyk/codex/optimize-accelerator-allocation-for-mersennedivisorcycles

Remove GPU fallback from context pool

PerfectNumbers.Core/Gpu/GpuContextPool.cs

@@ -6,156 +6,159 @@ namespace PerfectNumbers.Core.Gpu;
 
 public static class GpuContextPool
 {
-        private static readonly bool PoolingEnabled = true;
+    private static readonly bool PoolingEnabled = true;
-        // TODO: Introduce accelerator-specific warmup so pooled contexts precompile the ProcessEightBitWindows kernels and load
+    // TODO: Introduce accelerator-specific warmup so pooled contexts precompile the ProcessEightBitWindows kernels and load
-        // divisor-cycle data before the first scan begins.
+    // divisor-cycle data before the first scan begins.
-        // Default device preference for generic GPU kernels (prime scans, NTT, etc.)
+    // Default device preference for generic GPU kernels (prime scans, NTT, etc.)
-        public static bool ForceCpu { get; set; } = false;
+    public static bool ForceCpu { get; set; } = false;
 
-        internal sealed class PooledContext
+    internal sealed class PooledContext
+    {
+        private bool disposed;
+
+        public Context Context { get; }
+        public Accelerator Accelerator { get; }
+        public bool IsCpu { get; }
+        public object ExecutionLock { get; } = new();
+
+        public PooledContext(bool preferCpu)
         {
-                private bool disposed;
+            Context = Context.CreateDefault();
-
+            Accelerator = Context.GetPreferredDevice(preferCpu).CreateAccelerator(Context);
-                public Context Context { get; }
+            IsCpu = Accelerator.AcceleratorType == AcceleratorType.CPU;
-                public Accelerator Accelerator { get; }
+            // NOTE: Avoid loading/compiling any kernel here to prevent implicit
-                public bool IsCpu { get; }
+            // CL stream/queue creation during accelerator construction.
-                public object ExecutionLock { get; } = new();
+            // Some OpenCL drivers are fragile when a queue is created immediately
-
+            // after device init. Real kernels will be JIT-loaded at first use.
-		public PooledContext(bool preferCpu)
-		{
-			Context = Context.CreateDefault();
-			Accelerator = Context.GetPreferredDevice(preferCpu).CreateAccelerator(Context);
-			IsCpu = Accelerator.AcceleratorType == AcceleratorType.CPU;
-			// NOTE: Avoid loading/compiling any kernel here to prevent implicit
-			// CL stream/queue creation during accelerator construction.
-			// Some OpenCL drivers are fragile when a queue is created immediately
-			// after device init. Real kernels will be JIT-loaded at first use.
-		}
-
-		public void Dispose()
-		{
-			if (disposed)
-			{
-				return;
-			}
-
-			// Ensure all cached GPU buffers for this accelerator are released.
-			NttGpuMath.ClearCaches(Accelerator);
-			// Release PrimeTester GPU state for this accelerator (kernel/device primes)
-			// Clear any per-accelerator cached resources to avoid releasing
-			// after the accelerator is destroyed.
-			PrimeTester.ClearGpuCaches(Accelerator);
-
-			Accelerator.Dispose();
-			Context.Dispose();
-			disposed = true;
-		}
-	}
-
-	private static readonly ConcurrentQueue<PooledContext> CpuPool = new();
-	private static readonly ConcurrentQueue<PooledContext> GpuPool = new();
-	private static readonly object CreationLock = new();
-
-	public static GpuContextLease Rent()
-	{
-		return RentPreferred(ForceCpu);
-	}
-
-	// Allows callers to choose CPU/GPU per use-case, decoupled from ForceCpu.
-	public static GpuContextLease RentPreferred(bool preferCpu)
-	{
-		if (PoolingEnabled)
-		{
-			if (preferCpu && CpuPool.TryDequeue(out var cpu))
-			{
-				return new GpuContextLease(cpu, CreationLock);
-			}
-
-			if (!preferCpu && GpuPool.TryDequeue(out var gpu))
-			{
-				return new GpuContextLease(gpu, CreationLock);
-			}
-		}
-
-		// Serialize accelerator creation to avoid concurrent CL device/stream setup across threads.
-		lock (CreationLock)
-		{
-			return new GpuContextLease(new PooledContext(preferCpu), CreationLock);
-		}
-	}
-
-	public static void DisposeAll()
-	{
-		if (!PoolingEnabled)
-		{
-			return;
-		}
-
-		while (CpuPool.TryDequeue(out var cpu))
-		{
-			cpu.Dispose();
-		}
-
-		while (GpuPool.TryDequeue(out var gpu))
-		{
-			gpu.Dispose();
-		}
-	}
-
-	private static void Return(PooledContext ctx)
-	{
-		if (PoolingEnabled)
-		{
-			ctx.Accelerator.Synchronize();
-			if (ctx.IsCpu)
-			{
-				CpuPool.Enqueue(ctx);
-			}
-			else
-			{
-				GpuPool.Enqueue(ctx);
-			}
-		}
-
-                else
-                {
-                        lock (CreationLock)
-                        {
-                                ctx.Dispose();
-                        }
-                }
         }
 
+        public void Dispose()
+        {
+            if (disposed)
+            {
+                return;
+            }
+
+            // Ensure all cached GPU buffers for this accelerator are released.
+            NttGpuMath.ClearCaches(Accelerator);
+            // Release PrimeTester GPU state for this accelerator (kernel/device primes)
+            // Clear any per-accelerator cached resources to avoid releasing
+            // after the accelerator is destroyed.
+            PrimeTester.ClearGpuCaches(Accelerator);
+
+            Accelerator.Dispose();
+            Context.Dispose();
+            disposed = true;
+        }
+    }
+
+    private static readonly ConcurrentQueue<PooledContext> CpuPool = new();
+    private static readonly ConcurrentQueue<PooledContext> GpuPool = new();
+    private static readonly object CreationLock = new();
+
+    public static GpuContextLease Rent()
+    {
+        return RentPreferred(ForceCpu);
+    }
+
+    // Allows callers to choose CPU/GPU per use-case, decoupled from ForceCpu.
+    public static GpuContextLease RentPreferred(bool preferCpu)
+    {
+        if (PoolingEnabled)
+        {
+            if (preferCpu && CpuPool.TryDequeue(out var cpu))
+            {
+                return new GpuContextLease(cpu, CreationLock);
+            }
+
+            if (!preferCpu && GpuPool.TryDequeue(out var gpu))
+            {
+                return new GpuContextLease(gpu, CreationLock);
+            }
+        }
+
+        // Serialize accelerator creation to avoid concurrent CL device/stream setup across threads.
+        lock (CreationLock)
+        {
+            return new GpuContextLease(new PooledContext(preferCpu), CreationLock);
+        }
+    }
+
+    public static void DisposeAll()
+    {
+        if (!PoolingEnabled)
+        {
+            return;
+        }
+
+        while (CpuPool.TryDequeue(out var cpu))
+        {
+            cpu.Dispose();
+        }
+
+        while (GpuPool.TryDequeue(out var gpu))
+        {
+            gpu.Dispose();
+        }
+    }
+
+    private static void Return(PooledContext ctx)
+    {
+        if (PoolingEnabled)
+        {
+            ctx.Accelerator.Synchronize();
+            if (ctx.IsCpu)
+            {
+                CpuPool.Enqueue(ctx);
+            }
+            else
+            {
+                GpuPool.Enqueue(ctx);
+            }
+        }
+
+        else
+        {
+            lock (CreationLock)
+            {
+                ctx.Dispose();
+            }
+        }
+    }
+
     public struct GpuContextLease
     {
         private readonly PooledContext _ctx;
-                private bool disposed;
+        private bool disposed;
 
-                internal GpuContextLease(PooledContext ctx, object kernelInitLock)
+        internal GpuContextLease(PooledContext ctx, object kernelInitLock)
         {
             _ctx = ctx;
-                        KernelInitLock = kernelInitLock;
+            KernelInitLock = kernelInitLock;
-                }
+        }
 
         public Context Context => _ctx.Context;
 
         public Accelerator Accelerator => _ctx.Accelerator;
 
-                public object KernelInitLock { get; }
+        public object KernelInitLock { get; }
 
-                public object ExecutionLock => _ctx.ExecutionLock;
+        public object ExecutionLock => _ctx.ExecutionLock;
 
+        internal void ReturnToPool()
+        {
+            if (disposed)
+            {
+                return;
+            }
 
-                public void Dispose()
+            Return(_ctx);
-                {
+            disposed = true;
-                        if (disposed)
+        }
-			{
-				return;
-			}
 
-			Return(_ctx);
+        public void Dispose()
-			disposed = true;
+        {
+            ReturnToPool();
         }
     }
 }
-

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -396,7 +396,10 @@ public ref struct GpuKernelLease(IDisposable limiter, GpuContextLease gpu, Kerne
                 _stream?.Dispose();
                 _stream = null;
 
-                _gpu?.Dispose();
+                if (_gpu.HasValue)
+                {
+                    _gpu.Value.ReturnToPool();
+                }
                 _gpu = null;
 
                 _executionLock = null;

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -1715,10 +1715,10 @@ internal static partial class PrimeOrderGpuHeuristics
         GpuUInt128[]? rentedExponents = null;
         GpuUInt128[]? rentedResults = null;
         var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+        var execution = lease.EnterExecutionScope();
 
         try
         {
-            var execution = lease.EnterExecutionScope();
             Accelerator accelerator = lease.Accelerator;
             AcceleratorStream stream = lease.Stream;
             var kernel = GetPow2ModWideKernel(accelerator);
@@ -1755,7 +1755,6 @@ internal static partial class PrimeOrderGpuHeuristics
 
             exponentBuffer.Dispose();
             remainderBuffer.Dispose();
-            execution.Dispose();
             return true;
         }
         catch (Exception)
@@ -1775,6 +1774,7 @@ internal static partial class PrimeOrderGpuHeuristics
                 ArrayPool<GpuUInt128>.Shared.Return(rentedResults, clearArray: false);
             }
 
+            execution.Dispose();
             lease.Dispose();
         }
     }