Merge pull request #42 from mlemanczyk/codex/analyze-program-execution-path-issues

Serialize GPU leases with per-accelerator locks

PerfectNumbers.Core/Gpu/GpuContextPool.cs

@@ -17,6 +17,7 @@ public static class GpuContextPool
                 public Context Context { get; }
                 public Accelerator Accelerator { get; }
                 public bool IsCpu { get; }
+                public object ExecutionLock { get; } = new();
 
 		public PooledContext(bool preferCpu)
 		{
@@ -140,6 +141,8 @@ public static class GpuContextPool
 
                 public object KernelInitLock { get; }
 
+                public object ExecutionLock => _ctx.ExecutionLock;
+
 
                 public void Dispose()
                 {

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -1,5 +1,6 @@
 using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using ILGPU;
 using ILGPU.Runtime;
 using static PerfectNumbers.Core.Gpu.GpuContextPool;
@@ -26,16 +27,16 @@ public readonly struct ResidueAutomatonArgs
 public sealed class KernelContainer
 {
 	// Serializes first-time initialization of kernels/buffers per accelerator.
-	public Action<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>? Order;
+        public Action<AcceleratorStream, Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>? Order;
-    public Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
+    public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
         ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>>? Incremental;
-    public Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
+    public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
         ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>,
         ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>,
         ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>? Pow2Mod;
-	public Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
+        public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
                 ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>? IncrementalOrder;
-    public Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
+    public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
         ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>? Pow2ModOrder;
 
     // Optional device buffer with small divisor cycles (<= 4M). Index = divisor, value = cycle length.
@@ -81,65 +82,112 @@ public ref struct GpuKernelLease(IDisposable limiter, GpuContextLease gpu, Kerne
         private IDisposable? _limiter = limiter;
         private GpuContextLease? _gpu = gpu;
         private KernelContainer _kernels = kernels;
+        private AcceleratorStream? _stream;
+        private object? _executionLock = gpu.ExecutionLock;
 
         public readonly Accelerator Accelerator => _gpu?.Accelerator ?? throw new NullReferenceException("GPU context is already released");
 
-	public readonly Action<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>> OrderKernel
+        public AcceleratorStream Stream
+        {
+                get
+                {
+                        if (_stream is { } stream)
+                        {
+                                return stream;
+                        }
+
+                        stream = Accelerator.CreateStream();
+                        _stream = stream;
+                        return stream;
+                }
+        }
+
+        public readonly ExecutionScope EnterExecutionScope()
+        {
+                return new ExecutionScope(_executionLock);
+        }
+
+        public readonly Action<AcceleratorStream, Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>> OrderKernel
         {
                 get
                 {
                         var accel = Accelerator; // avoid capturing 'this' in lambda
                         lock (_gpu!.Value.KernelInitLock)
                         {
-				return KernelContainer.InitOnce(ref _kernels.Order, () => accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>(OrderKernelScan));
+                                return KernelContainer.InitOnce(ref _kernels.Order, () =>
+                                {
+                                        var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>(OrderKernelScan);
+                                        var kernel = KernelUtil.GetKernel(loaded);
+                                        return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>>();
+                                });
                         }
                 }
         }
 
-    public readonly Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> Pow2ModKernel
+    public readonly Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> Pow2ModKernel
         {
                 get
                 {
                         var accel = Accelerator;
                         lock (_gpu!.Value.KernelInitLock)
                         {
-                                return KernelContainer.InitOnce(ref _kernels.Pow2Mod, () => accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernelScan));
+                                return KernelContainer.InitOnce(ref _kernels.Pow2Mod, () =>
+                                {
+                                        var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernelScan);
+                                        var kernel = KernelUtil.GetKernel(loaded);
+                                        return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>();
+                                });
                         }
                 }
         }
 
-    public readonly Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>> IncrementalKernel
+    public readonly Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>> IncrementalKernel
         {
                 get
                 {
                         var accel = Accelerator;
                         lock (_gpu!.Value.KernelInitLock)
                         {
-				return KernelContainer.InitOnce(ref _kernels.Incremental, () => accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>>(IncrementalKernelScan));
+                                return KernelContainer.InitOnce(ref _kernels.Incremental, () =>
+                                {
+                                        var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>>(IncrementalKernelScan);
+                                        var kernel = KernelUtil.GetKernel(loaded);
+                                        return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<uint, Stride1D.Dense>>>();
+                                });
                         }
                 }
         }
 
-	public readonly Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>> IncrementalOrderKernel
+        public readonly Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>> IncrementalOrderKernel
         {
                 get
                 {
                         var accel = Accelerator;
                         lock (_gpu!.Value.KernelInitLock)
                         {
-				return KernelContainer.InitOnce(ref _kernels.IncrementalOrder, () => accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>(IncrementalOrderKernelScan));
+                                return KernelContainer.InitOnce(ref _kernels.IncrementalOrder, () =>
+                                {
+                                        var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>(IncrementalOrderKernelScan);
+                                        var kernel = KernelUtil.GetKernel(loaded);
+                                        return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>>();
+                                });
                         }
                 }
         }
 
-	public readonly Action<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>> Pow2ModOrderKernel
+        public readonly Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>> Pow2ModOrderKernel
         {
                 get
                 {
                         var accel = Accelerator;
                         lock (_gpu!.Value.KernelInitLock)
                         {
-				return KernelContainer.InitOnce(ref _kernels.Pow2ModOrder, () => accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>(Pow2ModOrderKernelScan));
+                                return KernelContainer.InitOnce(ref _kernels.Pow2ModOrder, () =>
+                                {
+                                        var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>(Pow2ModOrderKernelScan);
+                                        var kernel = KernelUtil.GetKernel(loaded);
+                                        return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<uint, Stride1D.Dense>>>();
+                                });
                         }
                 }
         }
@@ -279,9 +327,14 @@ public ref struct GpuKernelLease(IDisposable limiter, GpuContextLease gpu, Kerne
                 {
                         if (disposing)
                         {
+                                _stream?.Dispose();
+                                _stream = null;
+
                                 _gpu?.Dispose();
                                 _gpu = null;
 
+                                _executionLock = null;
+
                                 _limiter?.Dispose();
                                 _limiter = null;
                         }
@@ -570,6 +623,28 @@ public ref struct GpuKernelLease(IDisposable limiter, GpuContextLease gpu, Kerne
                 // Enabled this when the struct is converted to a class
                 // GC.SuppressFinalize(this);
         }
+
+        public readonly struct ExecutionScope : IDisposable
+        {
+                private readonly object? _lock;
+
+                public ExecutionScope(object? sync)
+                {
+                        _lock = sync;
+                        if (_lock is not null)
+                        {
+                                Monitor.Enter(_lock);
+                        }
+                }
+
+                public void Dispose()
+                {
+                        if (_lock is not null)
+                        {
+                                Monitor.Exit(_lock);
+                        }
+                }
+        }
 }
 
 public class GpuKernelPool
@@ -667,10 +742,12 @@ public class GpuKernelPool
     public static void Run(Action<Accelerator, AcceleratorStream> action)
     {
         var lease = GetKernel(useGpuOrder: true);
+        var execution = lease.EnterExecutionScope();
         var accelerator = lease.Accelerator;
         var stream = accelerator.CreateStream();
         action(accelerator, stream);
         stream.Dispose();
+        execution.Dispose();
         lease.Dispose();
 
     }

PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs

@@ -13,7 +13,9 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
         public void Scan(ulong exponent, UInt128 twoP, bool lastIsSeven, UInt128 maxK, ref bool isPrime)
         {
                 var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+                var execution = gpuLease.EnterExecutionScope();
                 var accelerator = gpuLease.Accelerator;
+                var stream = gpuLease.Stream;
                 int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy
                 UInt128 kStart = 1UL;
                 ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
@@ -54,17 +56,18 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
                                 {
                                         q0.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
                                         var ra = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
-                                        pow2Kernel(currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, divMul,
+                                        pow2Kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, divMul,
-                                                ra, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
+                                                ra, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven,
+                                                primeViews.LastOnePow2, primeViews.LastSevenPow2);
                                 }
                                 else
                                 {
                                         q0.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
-                                        incKernel(currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, divMul,
+                                        incKernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, divMul,
                                                 q0m10, q0m8, q0m3, q0m5, orderBuffer.View, smallCyclesView);
                                 }
 
-                                accelerator.Synchronize();
+                                stream.Synchronize();
                                 orderBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(orders), currentSize);
                                 if (_kernelType == GpuKernelType.Pow2Mod)
                                 {
@@ -95,10 +98,10 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
                 finally
                 {
                         ArrayPool<ulong>.Shared.Return(orderArray);
-                }
-
                         orderBuffer.Dispose();
+                        execution.Dispose();
                         gpuLease.Dispose();
+                }
 
         }
 }

PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs

@@ -10,7 +10,9 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
         public void Scan(ulong exponent, UInt128 twoP, bool lastIsSeven, UInt128 maxK, ref bool isPrime)
         {
                 var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+                var execution = gpuLease.EnterExecutionScope();
                 var accelerator = gpuLease.Accelerator;
+                var stream = gpuLease.Stream;
                 int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy and avoids TDR
                 UInt128 kStart = UInt128.One;
                 ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
@@ -23,6 +25,8 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
         };
 
                 var foundBuffer = accelerator.Allocate1D<int>(1);
+                try
+                {
                         UInt128 remaining;
                         int currentSize;
                         int found = 0;
@@ -46,10 +50,10 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 
                                 // Ensure device has small cycles table for early rejections in order kernels
                                 var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(accelerator);
-            kernel(currentSize, exponent, (GpuUInt128)twoP, (GpuUInt128)kStart, last, divMul,
+                                kernel(stream, currentSize, exponent, (GpuUInt128)twoP, (GpuUInt128)kStart, last, divMul,
                                         ra, foundBuffer.View, smallCyclesView);
 
-                        accelerator.Synchronize();
+                                stream.Synchronize();
                                 foundBuffer.View.CopyToCPU(ref found, 1);
                                 if (found != 0)
                                 {
@@ -59,9 +63,13 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 
                                 kStart += (UInt128)currentSize;
                         }
-
+                }
+                finally
+                {
                         foundBuffer.Dispose();
+                        execution.Dispose();
                         gpuLease.Dispose();
+                }
 
         }
 }

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -28,7 +28,9 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
                 }
 
                 var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+                var execution = gpuLease.EnterExecutionScope();
                 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);
@@ -84,10 +86,10 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
                                         q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
                                         var ra = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
 
-                                        kernel(currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
+                                        kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
                                                 ra, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
 
-                                        accelerator.Synchronize();
+                                        stream.Synchronize();
                                         orderBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(orders), currentSize);
                                         if (!Volatile.Read(ref isPrime))
                                         {
@@ -111,6 +113,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
                 {
                         ArrayPool<ulong>.Shared.Return(orderArray);
                         orderBuffer.Dispose();
+                        execution.Dispose();
                         gpuLease.Dispose();
                 }
 

PerfectNumbers.Core/MersenneNumberTester.cs

@@ -161,9 +161,13 @@ public sealed class MersenneNumberTester(
 		UInt128[] qs = qsBuffer;
 
                 var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+                var execution = gpuLease.EnterExecutionScope();
                 var accelerator = gpuLease.Accelerator;
+                var stream = gpuLease.Stream;
                 var orderKernel = gpuLease.OrderKernel!;
 
+                try
+                {
                         // Guard long-running kernels by chunking the warm-up across batches.
                         // Reuse the same ScanBatchSize knob used for scanning.
                         int batchSize = GpuConstants.ScanBatchSize;
@@ -183,8 +187,8 @@ public sealed class MersenneNumberTester(
                                 // Copy only the populated elements to the device buffer to avoid overrun on pooled arrays.
                                 qBuffer.View.CopyFromCPU(ref tmp[0], count);
                                 System.Buffers.ArrayPool<PerfectNumbers.Core.Gpu.GpuUInt128>.Shared.Return(tmp);
-			orderKernel(count, exponent, divMul, qBuffer.View, orderBuffer.View);
+                                orderKernel(stream, count, exponent, divMul, qBuffer.View, orderBuffer.View);
-			accelerator.Synchronize();
+                                stream.Synchronize();
 
                                 ulong[] orders = orderBuffer.GetAsArray1D();
                                 for (int i = 0; i < count; i++)
@@ -205,10 +209,14 @@ public sealed class MersenneNumberTester(
                                 orderBuffer.Dispose();
                                 qBuffer.Dispose();
                         }
-
+                }
+                finally
+                {
+                        execution.Dispose();
                         gpuLease.Dispose();
                         ArrayPool<UInt128>.Shared.Return(qs);
                 }
+        }
 
 	public bool IsMersennePrime(ulong exponent) => IsMersennePrime(exponent, out _);