Isolate residue GPU leases and drop try/finally

PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs

@@ -26,6 +26,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
     private readonly ConcurrentDictionary<Accelerator, Action<KernelConfig, Index1D, ArrayView<byte>, ArrayView<byte>, byte, byte>> _remainderScanKernelCache = new(AcceleratorReferenceComparer.Instance);
     private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ArrayView<byte>, ArrayView<byte>, ArrayView<byte>, ArrayView<byte>, byte, ArrayView<byte>>> _candidateMaskKernelCache = new(AcceleratorReferenceComparer.Instance);
     private readonly ConcurrentDictionary<Accelerator, ConcurrentBag<BatchResources>> _resourcePools = new(AcceleratorReferenceComparer.Instance);
+    private readonly ConcurrentBag<GpuContextPool.GpuContextLease> _acceleratorPool = new();
     private readonly ConcurrentBag<DivisorScanSession> _sessionPool = new();
 
     private Action<Index1D, ArrayView<MontgomeryDivisorData>, ArrayView<ulong>, ArrayView<byte>> GetKernel(Accelerator accelerator) =>
@@ -95,7 +96,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         ulong processedCount;
         ulong lastProcessed;
 
-        var gpuLease = GpuContextPool.RentPreferred(preferCpu: false);
+        var gpuLease = RentAccelerator();
         var accelerator = gpuLease.Accelerator;
 
         Monitor.Enter(gpuLease.ExecutionLock);
@@ -104,34 +105,33 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         var resources = RentBatchResources(accelerator, batchCapacity);
 
         composite = CheckDivisors(
-        prime,
+            prime,
-        allowedMax,
+            allowedMax,
-        accelerator,
+            accelerator,
-        kernel,
+            kernel,
-        resources.DivisorsBuffer,
+            resources.DivisorsBuffer,
-        resources.ExponentBuffer,
+            resources.ExponentBuffer,
-        resources.HitsBuffer,
+            resources.HitsBuffer,
-        resources.DivisorDeltaBuffer,
+            resources.DivisorDeltaBuffer,
-        resources.RemainderDeltaBuffer,
+            resources.RemainderDeltaBuffer,
-        resources.RemainderBuffer,
+            resources.RemainderBuffer,
-        resources.RemainderDeltaKernel,
+            resources.RemainderDeltaKernel,
-        resources.RemainderScanKernel,
+            resources.RemainderScanKernel,
-        resources.CandidateMaskKernel,
+            resources.CandidateMaskKernel,
-        resources.Divisors,
+            resources.Divisors,
-        resources.Exponents,
+            resources.Exponents,
-        resources.Hits,
+            resources.Hits,
-        resources.DivisorData,
+            resources.DivisorData,
-        resources.CycleCandidates,
+            resources.CycleCandidates,
-        resources.CycleLengths,
+            resources.CycleLengths,
-        resources.CycleCapacity,
+            resources.CycleCapacity,
-        out lastProcessed,
+            out lastProcessed,
-        out coveredRange,
+            out coveredRange,
-        out processedCount
+            out processedCount);
-        );
 
         ReturnBatchResources(accelerator, resources);
         Monitor.Exit(gpuLease.ExecutionLock);
-        gpuLease.Dispose();
+        ReturnAccelerator(gpuLease);
 
         if (composite)
         {
@@ -652,6 +652,18 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         }
     }
 
+    private GpuContextPool.GpuContextLease RentAccelerator()
+    {
+        if (_acceleratorPool.TryTake(out var lease))
+        {
+            return lease;
+        }
+
+        return GpuContextPool.RentPreferred(preferCpu: false);
+    }
+
+    private void ReturnAccelerator(GpuContextPool.GpuContextLease lease) => _acceleratorPool.Add(lease);
+
     private BatchResources RentBatchResources(Accelerator accelerator, int capacity)
     {
         var bag = _resourcePools.GetOrAdd(accelerator, static _ => []);

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -1,5 +1,11 @@
+using System;
 using System.Buffers;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Runtime.InteropServices;
+using System.Threading;
+using System.Runtime.CompilerServices;
+using ILGPU;
 using ILGPU.Runtime;
 
 namespace PerfectNumbers.Core.Gpu;
@@ -7,6 +13,10 @@ namespace PerfectNumbers.Core.Gpu;
 public class MersenneNumberResidueGpuTester(bool useGpuOrder)
 {
     private readonly bool _useGpuOrder = useGpuOrder;
+    private readonly ConcurrentBag<GpuContextPool.GpuContextLease> _acceleratorPool = new();
+    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>> _pow2ModKernelCache = new(AcceleratorReferenceComparer.Instance);
+    private readonly ConcurrentDictionary<Accelerator, ConcurrentBag<ResidueResources>> _resourcePools = new(AcceleratorReferenceComparer.Instance);
+
     // TODO: Integrate MersenneDivisorCycles.Shared to consult cycle lengths for small q (<= 4M) and fast-reject
     // candidates without launching heavy kernels. Consider a device-visible constant buffer per-accelerator
     // via GpuKernelPool to avoid host round-trips.
@@ -16,106 +26,169 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
     // GPU residue variant: check 2^p % q == 1 for q = 2*p*k + 1.
     public void Scan(ulong exponent, UInt128 twoP, bool lastIsSeven, UInt128 maxK, ref bool isPrime)
     {
-        var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+        var limiter = GpuPrimeWorkLimiter.Acquire();
-        var execution = gpuLease.EnterExecutionScope();
+        var gpuLease = RentAccelerator();
-        try
+        var accelerator = gpuLease.Accelerator;
+
+        Monitor.Enter(gpuLease.ExecutionLock);
+
+        var stream = accelerator.CreateStream();
+        var kernel = GetPow2ModKernel(accelerator);
+        var resources = RentResources(accelerator, GpuConstants.ScanBatchSize);
+
+        var orderBuffer = resources.OrderBuffer;
+        ulong[] orderArray = resources.OrderArray;
+        int batchSize = resources.Capacity;
+
+        // Ensure device has small cycles and primes tables for in-kernel lookup
+        var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(accelerator);
+        ResiduePrimeViews primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(accelerator);
+
+        UInt128 kStart = 1UL;
+        UInt128 limit = maxK + UInt128.One;
+        byte last = lastIsSeven ? (byte)1 : (byte)0;
+        twoP.Mod10_8_5_3(out ulong step10, out ulong step8, out ulong step5, out ulong step3);
+        step10 = step10.Mod10();
+
+        GpuUInt128 twoPGpu = (GpuUInt128)twoP;
+        UInt128 batchSize128 = (UInt128)batchSize;
+        UInt128 fullBatchStep = twoP * batchSize128;
+        UInt128 q = twoP * kStart + UInt128.One;
+
+        while (kStart < limit && Volatile.Read(ref isPrime))
         {
-            var accelerator = gpuLease.Accelerator;
+            UInt128 remaining = limit - kStart;
-            var stream = gpuLease.Stream;
+            int currentSize = remaining > batchSize128 ? batchSize : (int)remaining;
+            Span<ulong> orders = orderArray.AsSpan(0, currentSize);
+            ref ulong ordersRef = ref MemoryMarshal.GetReference(orders);
 
-            // Ensure device has small cycles and primes tables for in-kernel lookup
+            q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
-            var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(accelerator);
+            var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
-            ResiduePrimeViews primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(accelerator);
 
-            int batchSize = GpuConstants.ScanBatchSize;
+            kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
-            UInt128 kStart = 1UL;
+                kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
-            UInt128 limit = maxK + UInt128.One;
-            byte last = lastIsSeven ? (byte)1 : (byte)0;
-            var kernel = gpuLease.Pow2ModKernel;
-            twoP.Mod10_8_5_3(out ulong step10, out ulong step8, out ulong step5, out ulong step3);
-            step10 = step10.Mod10();
 
-            GpuUInt128 twoPGpu = (GpuUInt128)twoP;
+            stream.Synchronize();
-
+            orderBuffer.View.CopyToCPU(ref ordersRef, currentSize);
-            var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);
+            if (!Volatile.Read(ref isPrime))
-            // Device memory returned by Allocate1D is not guaranteed to be zeroed.
-            // Ensure the buffer starts with a known state so that lanes skipped by the
-            // kernel (for example due to early rejections) do not leave stale garbage
-            // values that would be interpreted as composite witnesses when copied back
-            // to the host. This mirrors the explicit zeroing performed by other GPU
-            // testers such as the order kernels.
-            orderBuffer.MemSetToZero();
-            ulong[] orderArray = ArrayPool<ulong>.Shared.Rent(batchSize);
-
-            try
             {
-                UInt128 batchSize128 = (UInt128)batchSize;
+                break;
-                UInt128 fullBatchStep = twoP * batchSize128;
+            }
-                UInt128 q = twoP * kStart + UInt128.One;
-                // TODO: Replace the direct UInt128 multiply/add updates below with the residue-cycle stepping helper so that q advances
-                // reuse cached cycle increments rather than recomputing twoP multiples on every iteration.
 
-                while (kStart < limit && Volatile.Read(ref isPrime))
+            bool compositeFound = false;
+            for (int i = 0; i < currentSize; i++)
+            {
+                if (orders[i] != 0UL)
                 {
-                    UInt128 remaining = limit - kStart;
+                    Volatile.Write(ref isPrime, false);
-                    int currentSize = remaining > batchSize128 ? batchSize : (int)remaining;
+                    compositeFound = true;
-                    Span<ulong> orders = orderArray.AsSpan(0, currentSize);
+                    break;
-                    ref ulong ordersRef = ref MemoryMarshal.GetReference(orders);
-
-                    q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
-                    var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
-
-                    kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
-                        kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
-
-                    accelerator.Synchronize();
-                    orderBuffer.View.CopyToCPU(ref ordersRef, currentSize);
-                    if (!Volatile.Read(ref isPrime))
-                    {
-                        break;
-                    }
-
-                    bool compositeFound = false;
-                    for (int i = 0; i < currentSize; i++)
-                    {
-                        if (orders[i] != 0UL)
-                        {
-                            Volatile.Write(ref isPrime, false);
-                            compositeFound = true;
-                            break;
-                        }
-                    }
-
-                    if (compositeFound)
-                    {
-                        break;
-                    }
-
-                    UInt128 processed = (UInt128)currentSize;
-                    kStart += processed;
-                    if (kStart < limit)
-                    {
-                        if (processed == batchSize128)
-                        {
-                            q += fullBatchStep;
-                        }
-                        else
-                        {
-                            q += twoP * processed; // TODO: Swap this multiplication for the shared cycle stepping helper once residue cycles are mandatory so we can reuse the cached remainder ladder instead of performing UInt128 multiply operations.
-                        }
-                    }
                 }
             }
-            finally
+
+            if (compositeFound)
             {
-                ArrayPool<ulong>.Shared.Return(orderArray, clearArray: false);
+                break;
-                orderBuffer.Dispose();
+            }
+
+            UInt128 processed = (UInt128)currentSize;
+            kStart += processed;
+            if (kStart < limit)
+            {
+                if (processed == batchSize128)
+                {
+                    q += fullBatchStep;
+                }
+                else
+                {
+                    q += twoP * processed; // TODO: Swap this multiplication for the shared cycle stepping helper once residue cycles are mandatory so we can reuse the cached remainder ladder instead of performing UInt128 multiply operations.
+                }
             }
         }
-        finally
+
+        ReturnResources(accelerator, resources);
+        stream.Dispose();
+        Monitor.Exit(gpuLease.ExecutionLock);
+        ReturnAccelerator(gpuLease);
+        limiter.Dispose();
+    }
+
+    private Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> GetPow2ModKernel(Accelerator accelerator) =>
+        _pow2ModKernelCache.GetOrAdd(
+            accelerator,
+            static acc =>
+            {
+                var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModKernelScan);
+                return (Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>)Delegate.CreateDelegate(
+                    typeof(Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>),
+                    loaded.Target,
+                    loaded.Method);
+            });
+
+    private GpuContextPool.GpuContextLease RentAccelerator()
+    {
+        if (_acceleratorPool.TryTake(out var lease))
         {
-            execution.Dispose();
+            return lease;
-            gpuLease.Dispose();
+        }
+
+        return GpuContextPool.RentPreferred(preferCpu: !_useGpuOrder);
+    }
+
+    private void ReturnAccelerator(GpuContextPool.GpuContextLease lease) => _acceleratorPool.Add(lease);
+
+    private ResidueResources RentResources(Accelerator accelerator, int capacity)
+    {
+        var bag = _resourcePools.GetOrAdd(accelerator, static _ => new ConcurrentBag<ResidueResources>());
+        while (bag.TryTake(out var resources))
+        {
+            if (resources.Capacity >= capacity)
+            {
+                resources.OrderBuffer.MemSetToZero();
+                return resources;
+            }
+
+            resources.Dispose();
+        }
+
+        var created = new ResidueResources(accelerator, capacity);
+        created.OrderBuffer.MemSetToZero();
+        return created;
+    }
+
+    private void ReturnResources(Accelerator accelerator, ResidueResources resources) =>
+        _resourcePools.GetOrAdd(accelerator, static _ => new ConcurrentBag<ResidueResources>()).Add(resources);
+
+    private sealed class ResidueResources : IDisposable
+    {
+        internal ResidueResources(Accelerator accelerator, int capacity)
+        {
+            Capacity = Math.Max(1, capacity);
+            OrderBuffer = accelerator.Allocate1D<ulong>(Capacity);
+            OrderArray = ArrayPool<ulong>.Shared.Rent(Capacity);
+        }
+
+        internal MemoryBuffer1D<ulong, Stride1D.Dense> OrderBuffer { get; }
+
+        internal ulong[] OrderArray { get; }
+
+        internal int Capacity { get; }
+
+        public void Dispose()
+        {
+            OrderBuffer.Dispose();
+            ArrayPool<ulong>.Shared.Return(OrderArray, clearArray: false);
         }
     }
+
+    private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>
+    {
+        public static AcceleratorReferenceComparer Instance { get; } = new();
+
+        public bool Equals(Accelerator? x, Accelerator? y) => ReferenceEquals(x, y);
+
+        public int GetHashCode(Accelerator obj) => RuntimeHelpers.GetHashCode(obj);
+    }
+
+
 }