Manage GPU lease buffers without kernel scratch pool

PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs

@@ -213,7 +213,7 @@ public sealed class HeuristicCombinedPrimeTester
 	        ulong[] divisorArray = divisorPool.Rent(batchCapacity);
 
 	        var limiter = GpuPrimeWorkLimiter.Acquire();
-	        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity, 1);
+	        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity);
 	        var accelerator = gpu.Accelerator;
 	        var state = gpu.State;
 

PerfectNumbers.Core/HeuristicPrimeTester.cs

@@ -267,7 +267,7 @@ public sealed class HeuristicPrimeTester
         ulong[] divisorArray = divisorPool.Rent(batchCapacity);
 
         var limiter = GpuPrimeWorkLimiter.Acquire();
-        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity, 1);
+        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity);
         var accelerator = gpu.Accelerator;
         var state = gpu.State;
 

PerfectNumbers.Core/PrimeTester.cs

@@ -183,7 +183,6 @@ public sealed class PrimeTester
 			public Accelerator Accelerator { get; }
 
 			private KernelState? _kernelState;
-			private KernelState.ScratchBuffers? _scratch;
 
 			public PooledContext()
 			{
@@ -207,46 +206,8 @@ public sealed class PrimeTester
 				}
 			}
 
-			public KernelState.ScratchBuffers GetScratch(int minCapacity, int outputClearLength)
-			{
-				var scratch = _scratch;
-				var state = KernelState;
-				if (scratch is null)
-				{
-					scratch = state.RentScratch(minCapacity, Accelerator);
-					_scratch = scratch;
-				}
-				else if (scratch.Capacity < minCapacity)
-				{
-					scratch.Dispose();
-					scratch = state.RentScratch(minCapacity, Accelerator);
-					_scratch = scratch;
-				}
-
-				int clearLength = outputClearLength <= 0 ? 0 : outputClearLength;
-				if (clearLength > 0)
-				{
-					int outputLength = (int)scratch.Output.Length;
-					if (clearLength > outputLength)
-					{
-						clearLength = outputLength;
-					}
-
-					scratch.Output.View.SubView(0, clearLength).MemSetToZero();
-					Accelerator.Synchronize();
-				}
-
-				return scratch;
-			}
-
 			public void Dispose()
 			{
-				if (_scratch is not null)
-				{
-					_scratch.Dispose();
-					_scratch = null;
-				}
-
 				ClearGpuCaches(Accelerator);
 				_kernelState = null;
 				Accelerator.Dispose();
@@ -254,16 +215,62 @@ public sealed class PrimeTester
 			}
 		}
 
+		private sealed class DeviceBuffers : IDisposable
+		{
+			public DeviceBuffers(Accelerator accelerator, int capacity)
+			{
+				Accelerator = accelerator;
+				Capacity = Math.Max(1, capacity);
+				Input = accelerator.Allocate1D<ulong>(Capacity);
+				Output = accelerator.Allocate1D<byte>(Capacity);
+			}
+
+			public Accelerator Accelerator { get; }
+
+			public MemoryBuffer1D<ulong, Stride1D.Dense> Input { get; private set; }
+
+			public MemoryBuffer1D<byte, Stride1D.Dense> Output { get; private set; }
+
+			public int Capacity { get; private set; }
+
+			public void EnsureCapacity(int minCapacity)
+			{
+				int required = Math.Max(1, minCapacity);
+				if (Capacity >= required)
+				{
+					return;
+				}
+
+				DisposeBuffers();
+				Capacity = required;
+				Input = Accelerator.Allocate1D<ulong>(Capacity);
+				Output = Accelerator.Allocate1D<byte>(Capacity);
+			}
+
+			private void DisposeBuffers()
+			{
+				Output.Dispose();
+				Input.Dispose();
+			}
+
+			public void Dispose()
+			{
+				DisposeBuffers();
+			}
+		}
+
 		private static readonly ConcurrentQueue<PooledContext> Pool = new();
 
-		internal static PrimeTesterGpuContextLease Rent(int minScratchCapacity = 0, int outputClearLength = 0)
+		private static readonly ConcurrentQueue<DeviceBuffers> BufferPool = new();
+
+		internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 0)
 		{
 			if (Pool.TryDequeue(out var ctx))
 			{
-				return new PrimeTesterGpuContextLease(ctx, minScratchCapacity, outputClearLength);
+				return new PrimeTesterGpuContextLease(ctx, minBufferCapacity);
 			}
 
-			return new PrimeTesterGpuContextLease(new PooledContext(), minScratchCapacity, outputClearLength);
+			return new PrimeTesterGpuContextLease(new PooledContext(), minBufferCapacity);
 		}
 
 		internal static void DisposeAll()
@@ -272,6 +279,37 @@ public sealed class PrimeTester
 			{
 				ctx.Dispose();
 			}
+
+			while (BufferPool.TryDequeue(out var buffers))
+			{
+				buffers.Dispose();
+			}
+		}
+
+		private static DeviceBuffers? AcquireBuffers(Accelerator accelerator, int minCapacity)
+		{
+			if (minCapacity <= 0)
+			{
+				return null;
+			}
+
+			while (BufferPool.TryDequeue(out var buffers))
+			{
+				if (buffers.Accelerator == accelerator)
+				{
+					buffers.EnsureCapacity(minCapacity);
+					return buffers;
+				}
+
+				buffers.Dispose();
+			}
+
+			return new DeviceBuffers(accelerator, minCapacity);
+		}
+
+		private static void ReturnBuffers(DeviceBuffers buffers)
+		{
+			BufferPool.Enqueue(buffers);
 		}
 
 		private static void Return(PooledContext ctx)
@@ -283,50 +321,41 @@ public sealed class PrimeTester
 		internal readonly struct PrimeTesterGpuContextLease
 		{
 			private readonly PooledContext? _ctx;
-			private readonly KernelState? _state;
-			private readonly KernelState.ScratchBuffers? _scratch;
-			private readonly MemoryBuffer1D<ulong, Stride1D.Dense>? _input;
-			private readonly MemoryBuffer1D<byte, Stride1D.Dense>? _output;
 
-			internal PrimeTesterGpuContextLease(PooledContext ctx, int minScratchCapacity, int outputClearLength)
+			private readonly KernelState? _state;
+
+			private readonly DeviceBuffers? _buffers;
+
+			internal PrimeTesterGpuContextLease(PooledContext ctx, int minBufferCapacity)
 			{
 				_ctx = ctx;
 				_state = ctx.KernelState;
-				if (minScratchCapacity > 0 || outputClearLength > 0)
+				_buffers = AcquireBuffers(ctx.Accelerator, minBufferCapacity);
-				{
-					var scratch = ctx.GetScratch(minScratchCapacity, outputClearLength);
-					_scratch = scratch;
-					_input = scratch.Input;
-					_output = scratch.Output;
-				}
-				else
-				{
-					_scratch = null;
-					_input = null;
-					_output = null;
-				}
 			}
 
 			public Accelerator Accelerator => (_ctx ?? throw new InvalidOperationException("GPU context lease is not initialized.")).Accelerator;
 
 			public KernelState State => _state ?? throw new InvalidOperationException("GPU kernel state is not initialized.");
 
-			public MemoryBuffer1D<ulong, Stride1D.Dense> Input => _input ?? throw new InvalidOperationException("Input buffer was not rented for this lease.");
+			public MemoryBuffer1D<ulong, Stride1D.Dense> Input => (_buffers ?? throw new InvalidOperationException("Input buffer was not rented for this lease.")).Input;
 
-			public MemoryBuffer1D<byte, Stride1D.Dense> Output => _output ?? throw new InvalidOperationException("Output buffer was not rented for this lease.");
+			public MemoryBuffer1D<byte, Stride1D.Dense> Output => (_buffers ?? throw new InvalidOperationException("Output buffer was not rented for this lease.")).Output;
 
-			public int ScratchCapacity => _scratch?.Capacity ?? 0;
+			public int BufferCapacity => _buffers?.Capacity ?? 0;
 
 			public void Dispose()
 			{
 				var ctx = _ctx ?? throw new InvalidOperationException("GPU context lease is not initialized.");
 
+				if (_buffers is not null)
+				{
+					ReturnBuffers(_buffers);
+				}
+
 				Return(ctx);
 			}
 		}
-
 	}
-
 	// Per-accelerator GPU state for prime sieve (kernel + uploaded primes).
 	internal sealed class KernelState
 	{
@@ -343,10 +372,6 @@ public sealed class PrimeTester
 		public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastThree;
 		public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastNine;
 		public bool IsDisposed;
-		private readonly ConcurrentBag<ScratchBuffers> _scratchPool = [];
-		// TODO: Replace this ConcurrentBag with the lock-free ring buffer variant validated in
-		// GpuModularArithmeticBenchmarks so renting scratch buffers stops contending on the bag's internal locks when
-		// thousands of GPU batches execute per second.
 
 		public KernelState(Accelerator accelerator)
 		{
@@ -396,48 +421,6 @@ public sealed class PrimeTester
 			DevicePrimesPow2LastNine.View.CopyFromCPU(primesPow2LastNine);
 		}
 
-		internal sealed class ScratchBuffers
-		{
-			public MemoryBuffer1D<ulong, Stride1D.Dense> Input { get; private set; }
-			public MemoryBuffer1D<byte, Stride1D.Dense> Output { get; private set; }
-			public int Capacity { get; private set; }
-
-			public ScratchBuffers(Accelerator accel, int capacity)
-			{
-				Capacity = Math.Max(1, capacity);
-				Input = accel.Allocate1D<ulong>(Capacity);
-				Output = accel.Allocate1D<byte>(Capacity);
-			}
-
-			public void Dispose()
-			{
-				Output.Dispose();
-				Input.Dispose();
-			}
-		}
-
-		[MethodImpl(MethodImplOptions.AggressiveInlining)]
-		public ScratchBuffers RentScratch(int minCapacity, Accelerator accel)
-		{
-			while (_scratchPool.TryTake(out var sb))
-			{
-				if (sb.Capacity >= minCapacity)
-				{
-					return sb;
-				}
-
-				sb.Dispose();
-			}
-
-			return new ScratchBuffers(accel, minCapacity);
-		}
-
-		[MethodImpl(MethodImplOptions.AggressiveInlining)]
-		public void ReturnScratch(ScratchBuffers scratch)
-		{
-			_scratchPool.Add(scratch);
-		}
-
 		public void Clear()
 		{
 			if (IsDisposed)
@@ -445,11 +428,6 @@ public sealed class PrimeTester
 				return;
 			}
 
-			while (_scratchPool.TryTake(out var sb))
-			{
-				sb.Dispose();
-			}
-
 			DevicePrimesDefault.Dispose();
 			DevicePrimesLastOne.Dispose();
 			DevicePrimesLastSeven.Dispose();