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,42 +97,41 @@ 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);
-        {
+
-            composite = CheckDivisors(
+        var kernel = GetKernel(accelerator);
-            prime,
+        var resources = RentBatchResources(accelerator, batchCapacity);
-            allowedMax,
+
-            accelerator,
+        composite = CheckDivisors(
-            kernel,
+        prime,
-            resources.DivisorsBuffer,
+        allowedMax,
-            resources.ExponentBuffer,
+        accelerator,
-            resources.HitsBuffer,
+        kernel,
-            resources.DivisorDeltaBuffer,
+        resources.DivisorsBuffer,
-            resources.RemainderDeltaBuffer,
+        resources.ExponentBuffer,
-            resources.RemainderBuffer,
+        resources.HitsBuffer,
-            resources.RemainderDeltaKernel,
+        resources.DivisorDeltaBuffer,
-            resources.RemainderScanKernel,
+        resources.RemainderDeltaBuffer,
-            resources.CandidateMaskKernel,
+        resources.RemainderBuffer,
-            resources.Divisors,
+        resources.RemainderDeltaKernel,
-            resources.Exponents,
+        resources.RemainderScanKernel,
-            resources.Hits,
+        resources.CandidateMaskKernel,
-            resources.DivisorData,
+        resources.Divisors,
-            resources.CycleCandidates,
+        resources.Exponents,
-            resources.CycleLengths,
+        resources.Hits,
-            resources.CycleCapacity,
+        resources.DivisorData,
-            out lastProcessed,
+        resources.CycleCandidates,
-            out coveredRange,
+        resources.CycleLengths,
-            out processedCount
+        resources.CycleCapacity,
-            );
+        out lastProcessed,
-        }
+        out coveredRange,
-        finally
+        out processedCount
-        {
+        );
-            ReturnBatchResources(accelerator, resources);
+
-            gpuLease.Dispose();
+        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,92 +17,104 @@ 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 accelerator = gpuLease.Accelerator;
+        var execution = gpuLease.EnterExecutionScope();
-        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;
-        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;
-
-        var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);
-        // 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);
-        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
         {
-            while (kStart < limit && Volatile.Read(ref isPrime))
+            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;
+            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;
+
+            var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);
+            // 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 remaining = limit - kStart;
+                UInt128 batchSize128 = (UInt128)batchSize;
-                int currentSize = remaining > batchSize128 ? batchSize : (int)remaining;
+                UInt128 fullBatchStep = twoP * batchSize128;
-                Span<ulong> orders = orderArray.AsSpan(0, currentSize);
+                UInt128 q = twoP * kStart + UInt128.One;
-                ref ulong ordersRef = ref MemoryMarshal.GetReference(orders);
+                // 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.
 
-                q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
+                while (kStart < limit && Volatile.Read(ref isPrime))
-                var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
+                {
+                    UInt128 remaining = limit - kStart;
+                    int currentSize = remaining > batchSize128 ? batchSize : (int)remaining;
+                    Span<ulong> orders = orderArray.AsSpan(0, currentSize);
+                    ref ulong ordersRef = ref MemoryMarshal.GetReference(orders);
 
-                kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
+                    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();
+                    accelerator.Synchronize();
-                orderBuffer.View.CopyToCPU(ref ordersRef, currentSize);
+                    orderBuffer.View.CopyToCPU(ref ordersRef, currentSize);
-                if (!Volatile.Read(ref isPrime))
+                    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)
+                    bool compositeFound = false;
-                {
+                    for (int i = 0; i < currentSize; i++)
-                    break;
-                }
-
-                UInt128 processed = (UInt128)currentSize;
-                kStart += processed;
-                if (kStart < limit)
-                {
-                    if (processed == batchSize128)
                     {
-                        q += fullBatchStep;
+                        if (orders[i] != 0UL)
+                        {
+                            Volatile.Write(ref isPrime, false);
+                            compositeFound = true;
+                            break;
+                        }
                     }
-                    else
+
+                    if (compositeFound)
                     {
-                        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.
+                        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
+            {
+                ArrayPool<ulong>.Shared.Return(orderArray, clearArray: false);
+                orderBuffer.Dispose();
+            }
         }
         finally
         {
-            ArrayPool<ulong>.Shared.Return(orderArray);
+            execution.Dispose();
-            orderBuffer.Dispose();
             gpuLease.Dispose();
         }
     }