Merge pull request #219 from mlemanczyk/codex/refactor-gpu-kernel-launch-process

Inline GPU divisor locking adjustments

PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs

@@ -2,6 +2,7 @@ using System.Buffers;
 using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using System.Threading;
 using ILGPU;
 using ILGPU.Runtime;
 
@@ -96,11 +97,12 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
 
         var gpuLease = GpuContextPool.RentPreferred(preferCpu: false);
         var accelerator = gpuLease.Accelerator;
-        var kernel = GetKernel(accelerator);
-        BatchResources resources = RentBatchResources(accelerator, batchCapacity);
 
-        try
+        Monitor.Enter(gpuLease.ExecutionLock);
-        {
+
+        var kernel = GetKernel(accelerator);
+        var resources = RentBatchResources(accelerator, batchCapacity);
+
         composite = CheckDivisors(
         prime,
         allowedMax,
@@ -126,12 +128,10 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         out coveredRange,
         out processedCount
         );
-        }
+
-        finally
-        {
         ReturnBatchResources(accelerator, resources);
+        Monitor.Exit(gpuLease.ExecutionLock);
         gpuLease.Dispose();
-        }
 
         if (composite)
         {
@@ -540,7 +540,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         private readonly MersenneNumberDivisorByDivisorGpuTester _owner;
         private readonly GpuContextPool.GpuContextLease _lease;
         private readonly Accelerator _accelerator;
-        private readonly Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> _kernel;
+        private Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> _kernel;
         private MemoryBuffer1D<ulong, Stride1D.Dense> _exponentsBuffer;
         private MemoryBuffer1D<ulong, Stride1D.Dense> _resultsBuffer;
         private ulong[] _hostBuffer;
@@ -552,11 +552,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
             _owner = owner;
             _lease = GpuContextPool.RentPreferred(preferCpu: false);
             _accelerator = _lease.Accelerator;
-            _kernel = owner.GetKernelByPrimeExponent(_accelerator);
             _capacity = Math.Max(1, owner._gpuBatchSize);
-            _exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
-            _resultsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
-            _hostBuffer = ArrayPool<ulong>.Shared.Rent(_capacity);
         }
 
         internal void Reset()
@@ -564,6 +560,14 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
             _disposed = false;
         }
 
+        private void EnsureExecutionResourcesLocked()
+        {
+            _kernel = _owner.GetKernelByPrimeExponent(_accelerator);
+            _exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
+            _resultsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
+            _hostBuffer = ArrayPool<ulong>.Shared.Rent(_capacity);
+        }
+
         public void CheckDivisor(ulong divisor, in MontgomeryDivisorData divisorData, ulong divisorCycle, in ReadOnlySpan<ulong> primes, Span<byte> hits)
         {
             if (_disposed)
@@ -583,6 +587,13 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
                 return;
             }
 
+            Monitor.Enter(_lease.ExecutionLock);
+
+            if (_kernel == null)
+            {
+                EnsureExecutionResourcesLocked();
+            }
+
             ArrayView1D<ulong, Stride1D.Dense> exponentsView = _exponentsBuffer.View;
             ArrayView1D<ulong, Stride1D.Dense> resultsView = _resultsBuffer.View;
             Span<ulong> hostSpan = _hostBuffer.AsSpan(0, _capacity);
@@ -609,6 +620,8 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
 
                 offset += batchSize;
             }
+
+            Monitor.Exit(_lease.ExecutionLock);
         }
 
         public void Dispose()

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -17,11 +17,16 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
     public void Scan(ulong exponent, UInt128 twoP, bool lastIsSeven, UInt128 maxK, ref bool isPrime)
     {
         var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+        var execution = gpuLease.EnterExecutionScope();
+        try
+        {
             var accelerator = gpuLease.Accelerator;
             var stream = gpuLease.Stream;
+
             // Ensure device has small cycles and primes tables for in-kernel lookup
             var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(accelerator);
             ResiduePrimeViews primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(accelerator);
+
             int batchSize = GpuConstants.ScanBatchSize;
             UInt128 kStart = 1UL;
             UInt128 limit = maxK + UInt128.One;
@@ -41,13 +46,15 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
             // testers such as the order kernels.
             orderBuffer.MemSetToZero();
             ulong[] orderArray = ArrayPool<ulong>.Shared.Rent(batchSize);
-        UInt128 batchSize128 = (UInt128)batchSize;
-        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.
 
             try
             {
+                UInt128 batchSize128 = (UInt128)batchSize;
+                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))
                 {
                     UInt128 remaining = limit - kStart;
@@ -101,8 +108,13 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
             }
             finally
             {
-            ArrayPool<ulong>.Shared.Return(orderArray);
+                ArrayPool<ulong>.Shared.Return(orderArray, clearArray: false);
                 orderBuffer.Dispose();
+            }
+        }
+        finally
+        {
+            execution.Dispose();
             gpuLease.Dispose();
         }
     }