From 3e7606025d32698fd731bdba85f43ea47d6c2d86 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Lema=C5=84czyk?= <marek@lemanczyk.eu>
Date: Tue, 21 Oct 2025 22:34:30 +0200
Subject: [PATCH] Add dedicated GPU context pool for PrimeTester

---
 PerfectNumbers.Core/Gpu/GpuContextPool.cs |  2 +
 PerfectNumbers.Core/PrimeTester.cs        | 96 ++++++++++++++++++++++-
 2 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/PerfectNumbers.Core/Gpu/GpuContextPool.cs b/PerfectNumbers.Core/Gpu/GpuContextPool.cs
index 52b6f6d..bbf8074 100644
--- a/PerfectNumbers.Core/Gpu/GpuContextPool.cs
+++ b/PerfectNumbers.Core/Gpu/GpuContextPool.cs
@@ -100,6 +100,8 @@ public static class GpuContextPool
 		{
 			gpu.Dispose();
 		}
+
+		PrimeTester.DisposeGpuContexts();
 	}
 
 	private static void Return(PooledContext ctx)
diff --git a/PerfectNumbers.Core/PrimeTester.cs b/PerfectNumbers.Core/PrimeTester.cs
index 2a177ff..ea5464a 100644
--- a/PerfectNumbers.Core/PrimeTester.cs
+++ b/PerfectNumbers.Core/PrimeTester.cs
@@ -1,4 +1,5 @@
 using System.Buffers;
+using System.Collections.Concurrent;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using ILGPU;
@@ -103,7 +104,7 @@ public sealed class PrimeTester(bool useInternal = false)
     {
         // Limit concurrency and declare variables outside loops for performance and reuse
         var limiter = GpuPrimeWorkLimiter.Acquire();
-        var gpu = GpuContextPool.Rent();
+        var gpu = PrimeTesterGpuContextPool.Rent();
 
         try
         {
@@ -156,6 +157,92 @@ public sealed class PrimeTester(bool useInternal = false)
         }
     }
 
+    private static class PrimeTesterGpuContextPool
+    {
+        internal sealed class PooledContext : IDisposable
+        {
+            private bool _disposed;
+
+            public Context Context { get; }
+
+            public Accelerator Accelerator { get; }
+
+            public object ExecutionLock { get; } = new();
+
+            public PooledContext()
+            {
+                Context = Context.CreateDefault();
+                Accelerator = Context.GetPreferredDevice(false).CreateAccelerator(Context);
+            }
+
+            public void Dispose()
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+
+                PrimeTester.ClearGpuCaches(Accelerator);
+                Accelerator.Dispose();
+                Context.Dispose();
+                _disposed = true;
+            }
+        }
+
+        private static readonly ConcurrentQueue<PooledContext> Pool = new();
+        private static readonly object CreationLock = new();
+
+        internal static PrimeTesterGpuContextLease Rent()
+        {
+            if (Pool.TryDequeue(out var ctx))
+            {
+                return new PrimeTesterGpuContextLease(ctx);
+            }
+
+            lock (CreationLock)
+            {
+                return new PrimeTesterGpuContextLease(new PooledContext());
+            }
+        }
+
+        internal static void DisposeAll()
+        {
+            while (Pool.TryDequeue(out var ctx))
+            {
+                ctx.Dispose();
+            }
+        }
+
+        private static void Return(PooledContext ctx)
+        {
+            ctx.Accelerator.Synchronize();
+            Pool.Enqueue(ctx);
+        }
+
+        internal struct PrimeTesterGpuContextLease : IDisposable
+        {
+            private PooledContext? _ctx;
+
+            internal PrimeTesterGpuContextLease(PooledContext ctx)
+            {
+                _ctx = ctx;
+            }
+
+            public Accelerator Accelerator => _ctx!.Accelerator;
+
+            public object ExecutionLock => _ctx!.ExecutionLock;
+
+            public void Dispose()
+            {
+                if (_ctx is { } ctx)
+                {
+                    Return(ctx);
+                    _ctx = null;
+                }
+            }
+        }
+    }
+
     // Per-accelerator GPU state for prime sieve (kernel + uploaded primes).
     private sealed class KernelState
     {
@@ -268,6 +355,11 @@ public sealed class PrimeTester(bool useInternal = false)
         GpuKernelState.Clear(accelerator);
     }
 
+    internal static void DisposeGpuContexts()
+    {
+        PrimeTesterGpuContextPool.DisposeAll();
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     internal static bool SharesFactorWithMaxExponent(ulong n)
     {
@@ -283,7 +375,7 @@ public sealed class PrimeTester(bool useInternal = false)
         // TODO: Route this batch helper through the shared GPU kernel pool from
         // GpuUInt128BinaryGcdBenchmarks so we reuse cached kernels, pinned host buffers,
         // and divisor-cycle staging instead of allocating new device buffers per call.
-        var gpu = GpuContextPool.Rent();
+        var gpu = PrimeTesterGpuContextPool.Rent();
         var accelerator = gpu.Accelerator;
         var kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SharesFactorKernel);