Manual accelerator stream creations replaced with pool

EvenPerfectBitScanner/Program.cs

@@ -66,7 +66,6 @@ internal static class Program
 			Console.WriteLine("Warming up GPU kernels");
 			PrimeTester.WarmUpGpuKernels(1024);
 			// PrimeTester.WarmUpGpuKernels(gpuPrimeThreads >> 4);
-			Console.WriteLine("Warming up scratch buffers");
 			Console.WriteLine("Starting up threads...");
 			_ = UnboundedTaskScheduler.Instance;
 			int blockSize = Math.Max(1, _cliArguments.BlockSize);

PerfectNumbers.Core.Tests/Gpu/GpuKernelCompilationTests.cs

@@ -3,6 +3,7 @@ using FluentAssertions;
 using ILGPU;
 using ILGPU.Runtime;
 using PerfectNumbers.Core.Gpu;
+using PerfectNumbers.Core.Gpu.Accelerators;
 using Xunit;
 
 namespace PerfectNumbers.Core.Tests.Gpu;
@@ -348,10 +349,11 @@ public class GpuKernelCompilationTests
 	{
 
 		GpuPrimeWorkLimiter.Acquire();
-		AcceleratorStream stream = accelerator.CreateStream();
+		int acceleratorIndex = AcceleratorPool.Shared.Rent();
+		AcceleratorStream stream = AcceleratorStreamPool.Rent(acceleratorIndex);
 		try
 		{
-			var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, _kernelType);
+			var kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, _kernelType);
 			_ = kernels.Order;
 			_ = kernels.Incremental;
 			_ = kernels.Pow2Mod;
@@ -361,15 +363,15 @@ public class GpuKernelCompilationTests
 		finally
 		{
 			stream.Synchronize();
-			stream.Dispose();
+			AcceleratorStreamPool.Return(acceleratorIndex, stream);
 			// GpuPrimeWorkLimiter.Release();
 		}
 
-		stream = accelerator.CreateStream();
+		stream = AcceleratorStreamPool.Rent(acceleratorIndex);
 		GpuPrimeWorkLimiter.Acquire();
 		try
 		{
-			var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, _kernelType);
+			var kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, _kernelType);
 			_ = kernels.Order;
 			_ = kernels.Incremental;
 			_ = kernels.Pow2Mod;
@@ -379,7 +381,7 @@ public class GpuKernelCompilationTests
 		finally
 		{
 			stream.Synchronize();
-			stream.Dispose();
+			AcceleratorStreamPool.Return(acceleratorIndex, stream);
 			// GpuPrimeWorkLimiter.Release();
 		}
 	}

PerfectNumbers.Core/Gpu/Accelerators/HeuristicGpuTables.cs

@@ -1,4 +1,3 @@
-using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
 
@@ -6,50 +5,65 @@ namespace PerfectNumbers.Core.Gpu.Accelerators;
 
 internal sealed class HeuristicGpuTables
 {
-	private static readonly ConcurrentDictionary<Accelerator, HeuristicGpuTables> _sharedTables = new(20_480, 1);
+	private static readonly Accelerator[] _accelerators = AcceleratorPool.Shared.Accelerators;
 
-	internal static HeuristicGpuTables EnsureStaticTables(Accelerator accelerator, AcceleratorStream stream)
+	private static readonly HeuristicGpuTables[] _sharedTables = new HeuristicGpuTables[PerfectNumberConstants.RollingAccelerators];
-		=> _sharedTables.GetOrAdd(accelerator, _ => new HeuristicGpuTables(accelerator, stream));
+
+	private static readonly SemaphoreSlim[] _locks = [..Enumerable.Range(1, PerfectNumberConstants.RollingAccelerators).Select(_ => new SemaphoreSlim(1))];
+
+	internal static HeuristicGpuTables EnsureStaticTables(int acceleratorIndex, AcceleratorStream stream)
+	{
+		var @lock = _locks[acceleratorIndex];
+		@lock.Wait();
+		if (_sharedTables[acceleratorIndex] is {} existing)
+		{
+			@lock.Release();
+			return existing;
+		}
+
+		var accelerator = _accelerators[acceleratorIndex];
+		existing = new(accelerator, stream);
+		_sharedTables[acceleratorIndex] = existing;
+		@lock.Release();
+		return existing;
+	}
 
 	internal HeuristicGpuTables(Accelerator accelerator, AcceleratorStream stream)
 	{
-		// lock (accelerator)
+		var heuristicGroupA = HeuristicPrimeSieves.GroupADivisorsStorage;
-		{
+		HeuristicGroupADivisors = CopySpanToDevice(accelerator, stream, heuristicGroupA);
-			var heuristicGroupA = HeuristicPrimeSieves.GroupADivisorsStorage;
-			HeuristicGroupADivisors = CopySpanToDevice(accelerator, stream, heuristicGroupA);
 
-			var heuristicGroupASquares = HeuristicPrimeSieves.GroupADivisorSquaresStorage;
+		var heuristicGroupASquares = HeuristicPrimeSieves.GroupADivisorSquaresStorage;
-			HeuristicGroupADivisorSquares = CopySpanToDevice(accelerator, stream, heuristicGroupASquares);
+		HeuristicGroupADivisorSquares = CopySpanToDevice(accelerator, stream, heuristicGroupASquares);
 
-			HeuristicGroupBDivisorsEnding1 = CopyUintSpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesLastOneWithoutLastThree);
+		HeuristicGroupBDivisorsEnding1 = CopyUintSpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesLastOneWithoutLastThree);
-			HeuristicGroupBDivisorSquaresEnding1 = CopySpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesPow2LastOneWithoutLastThree);
+		HeuristicGroupBDivisorSquaresEnding1 = CopySpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesPow2LastOneWithoutLastThree);
 
-			HeuristicGroupBDivisorsEnding7 = CopyUintSpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesLastSevenWithoutLastThree);
+		HeuristicGroupBDivisorsEnding7 = CopyUintSpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesLastSevenWithoutLastThree);
-			HeuristicGroupBDivisorSquaresEnding7 = CopySpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesPow2LastSevenWithoutLastThree);
+		HeuristicGroupBDivisorSquaresEnding7 = CopySpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesPow2LastSevenWithoutLastThree);
 
-			HeuristicGroupBDivisorsEnding9 = CopyUintSpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesLastNineWithoutLastThree);
+		HeuristicGroupBDivisorsEnding9 = CopyUintSpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesLastNineWithoutLastThree);
-			HeuristicGroupBDivisorSquaresEnding9 = CopySpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesPow2LastNineWithoutLastThree);
+		HeuristicGroupBDivisorSquaresEnding9 = CopySpanToDevice(accelerator, stream, DivisorGenerator.SmallPrimesPow2LastNineWithoutLastThree);
 
-			var combinedEnding1 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding1;
+		var combinedEnding1 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding1;
-			HeuristicCombinedDivisorsEnding1 = CopySpanToDevice(accelerator, stream, combinedEnding1);
+		HeuristicCombinedDivisorsEnding1 = CopySpanToDevice(accelerator, stream, combinedEnding1);
-			var combinedSquaresEnding1 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding1Squares;
+		var combinedSquaresEnding1 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding1Squares;
-			HeuristicCombinedDivisorSquaresEnding1 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding1);
+		HeuristicCombinedDivisorSquaresEnding1 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding1);
 
-			var combinedEnding3 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding3;
+		var combinedEnding3 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding3;
-			HeuristicCombinedDivisorsEnding3 = CopySpanToDevice(accelerator, stream, combinedEnding3);
+		HeuristicCombinedDivisorsEnding3 = CopySpanToDevice(accelerator, stream, combinedEnding3);
-			var combinedSquaresEnding3 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding3Squares;
+		var combinedSquaresEnding3 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding3Squares;
-			HeuristicCombinedDivisorSquaresEnding3 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding3);
+		HeuristicCombinedDivisorSquaresEnding3 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding3);
 
-			var combinedEnding7 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding7;
+		var combinedEnding7 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding7;
-			HeuristicCombinedDivisorsEnding7 = CopySpanToDevice(accelerator, stream, combinedEnding7);
+		HeuristicCombinedDivisorsEnding7 = CopySpanToDevice(accelerator, stream, combinedEnding7);
-			var combinedSquaresEnding7 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding7Squares;
+		var combinedSquaresEnding7 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding7Squares;
-			HeuristicCombinedDivisorSquaresEnding7 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding7);
+		HeuristicCombinedDivisorSquaresEnding7 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding7);
 
-			var combinedEnding9 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding9;
+		var combinedEnding9 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding9;
-			HeuristicCombinedDivisorsEnding9 = CopySpanToDevice(accelerator, stream, combinedEnding9);
+		HeuristicCombinedDivisorsEnding9 = CopySpanToDevice(accelerator, stream, combinedEnding9);
-			var combinedSquaresEnding9 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding9Squares;
+		var combinedSquaresEnding9 = HeuristicCombinedPrimeTester.CombinedDivisorsEnding9Squares;
-			HeuristicCombinedDivisorSquaresEnding9 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding9);
+		HeuristicCombinedDivisorSquaresEnding9 = CopySpanToDevice(accelerator, stream, combinedSquaresEnding9);
-		}
 	}
 
 	internal readonly MemoryBuffer1D<ulong, Stride1D.Dense> HeuristicGroupADivisors;
@@ -90,10 +104,7 @@ internal sealed class HeuristicGpuTables
 	private static MemoryBuffer1D<ulong, Stride1D.Dense> CopySpanToDevice(Accelerator accelerator, AcceleratorStream stream, in ulong[] span)
 	{
 		MemoryBuffer1D<ulong, Stride1D.Dense>? buffer;
-		// lock (accelerator)
+		buffer = accelerator.Allocate1D(stream, span);
-		{
-			buffer = accelerator.Allocate1D(stream, span);
-		}
 
 		return buffer;
 	}
@@ -101,10 +112,7 @@ internal sealed class HeuristicGpuTables
 	private static MemoryBuffer1D<ulong, Stride1D.Dense> CopyUintSpanToDevice(Accelerator accelerator, AcceleratorStream stream, ReadOnlySpan<uint> span)
 	{
 		MemoryBuffer1D<ulong, Stride1D.Dense>? buffer;
-		// lock (accelerator)
+		buffer = accelerator.Allocate1D<ulong>(span.Length);
-		{
-			buffer = accelerator.Allocate1D<ulong>(span.Length);
-		}
 
 		if (!span.IsEmpty)
 		{

PerfectNumbers.Core/Gpu/Accelerators/HeuristicPrimeTesterAccelerator.cs

@@ -15,24 +15,18 @@ internal sealed class HeuristicPrimeTesterAccelerator
 	{
 		Accelerator[]? accelerators = AcceleratorPool.Shared.Accelerators;
 		int acceleratorCount = accelerators.Length;
-		Span<AcceleratorStream> streams = new(new AcceleratorStream[acceleratorCount]);
 		for (var i = 0; i < acceleratorCount; i++)
 		{
 			Console.WriteLine($"Preparing heuristic accelerator {i}...");
 			var accelerator = accelerators[i];
 			AcceleratorStream stream = accelerator.CreateStream();
-			Accelerators.HeuristicGpuTables.EnsureStaticTables(accelerator, stream);
+			Accelerators.HeuristicGpuTables.EnsureStaticTables(i, stream);
 			// _ = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.None);
 			// KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream);
 			// GpuStaticTableInitializer.EnsureStaticTables(accelerator, kernels, stream);
-			streams[i] = stream;
-		}
-
-		for (var i = 0; i < acceleratorCount; i++)
-		{
-			AcceleratorStream stream = streams[i];
 			stream.Synchronize();
 			stream.Dispose();
+			accelerator.Synchronize();
 			Console.WriteLine($"Heuristic accelerator {i} is ready");
 		}
 	}
@@ -148,7 +142,7 @@ internal sealed class HeuristicPrimeTesterAccelerator
 		
 		// GpuStaticTableInitializer.EnsureStaticTables(accelerator, kernels, stream);
 
-		_sharedTables = Accelerators.HeuristicGpuTables.EnsureStaticTables(accelerator, stream);
+		_sharedTables = Accelerators.HeuristicGpuTables.EnsureStaticTables(acceleratorIndex, stream);
 		_heuristicDivisorTables = _sharedTables.CreateHeuristicDivisorTables();
 
 		stream.Synchronize();

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -1,12 +1,13 @@
 using ILGPU;
 using ILGPU.Runtime;
+using PerfectNumbers.Core.Gpu.Accelerators;
 
 namespace PerfectNumbers.Core.Gpu;
 
 public class GpuKernelPool
 {
 	[ThreadStatic]
-	private static Dictionary<Accelerator, KernelContainer>? _pool;
+	private static KernelContainer[]? _pool;
 
 	public static T InitOnce<T>(ref T? target, Func<T> valueFactory) where T : class
 	{
@@ -19,14 +20,18 @@ public class GpuKernelPool
 		return Interlocked.CompareExchange(ref target, newValue, null) ?? newValue;
 	}
 
-	internal static KernelContainer GetOrAddKernels(Accelerator accelerator, AcceleratorStream stream, KernelType kernelType)
+	private static readonly Accelerator[] _accelerators = AcceleratorPool.Shared.Accelerators;
+
+
+	internal static KernelContainer GetOrAddKernels(int acceleratorIndex, AcceleratorStream stream, KernelType kernelType)
 	{
 		var pool = _pool ??= [];
-		if (!pool.TryGetValue(accelerator, out var kernels))
+		if (pool[acceleratorIndex] is not {} kernels)
 		{
-			kernels = new KernelContainer();			
+			kernels = new();
 		}
 
+		Accelerator accelerator = _accelerators[acceleratorIndex];
 		if (kernelType.HasFlag(KernelType.OrderKernelScan))
 		{
 			InitOrderKernelScan(accelerator, kernels);
@@ -64,7 +69,7 @@ public class GpuKernelPool
 
 		PreloadStaticTables(accelerator, kernels, stream, kernelType);
 
-		pool[accelerator] = kernels;
+		pool[acceleratorIndex] = kernels;
 		return kernels;
 
 		static void InitOrderKernelScan(Accelerator accelerator, KernelContainer kernels)
@@ -281,8 +286,9 @@ public class GpuKernelPool
 	public static void Run(Action<Accelerator, AcceleratorStream> action)
 	{
 		// GpuPrimeWorkLimiter.Acquire();
-		var accelerator = SharedGpuContext.CreateAccelerator();
+		var acceleratorIndex = AcceleratorPool.Shared.Rent();
-		var stream = accelerator.CreateStream();
+		var accelerator = _accelerators[acceleratorIndex];
+		var stream = AcceleratorStreamPool.Rent(acceleratorIndex);
 		action(accelerator, stream);
 		stream.Synchronize();
 		stream.Dispose();

PerfectNumbers.Core/Gpu/InputOutputScratchBuffers.cs

@@ -1,126 +0,0 @@
-using System.Collections.Concurrent;
-using ILGPU;
-using ILGPU.Runtime;
-
-namespace PerfectNumbers.Core.Gpu;
-
-public sealed class GpuInputOutputScratchBuffer<TInput, TOutput>
-	where TInput : unmanaged
-	where TOutput: unmanaged
-{
-	private readonly Accelerator _accelerator;
-	public MemoryBuffer1D<TInput, Stride1D.Dense> Input { get; private set; }
-	public MemoryBuffer1D<TOutput, Stride1D.Dense> Output { get; private set; }
-
-	public GpuInputOutputScratchBuffer(Accelerator accelerator,int inputCapacity, int outputCapacity)
-	{
-		_accelerator = accelerator;
-		// lock (accelerator)
-		{
-			Input = accelerator.Allocate1D<TInput>(inputCapacity);
-			Output = accelerator.Allocate1D<TOutput>(outputCapacity);
-		}
-	}
-
-	public void EnsureCapacity(int inputCapacity, int outputCapacity)
-	{
-		if (Input.Length >= inputCapacity && Output.Length >= outputCapacity)
-		{
-			return;
-		}
-
-		Accelerator accelerator = _accelerator;
-		if (Input.Length >= inputCapacity && Output.Length >= outputCapacity)
-		{
-			return;
-		}
-
-		// lock (accelerator)
-		{
-			if (Input.Length < inputCapacity)
-			{
-				Input.Dispose();
-				Input = accelerator.Allocate1D<TInput>(inputCapacity);
-			}
-
-			if (Output.Length < outputCapacity)
-			{
-				Output.Dispose();
-				Output = accelerator.Allocate1D<TOutput>(outputCapacity);
-			}
-		}
-	}
-	
-	public void Dispose()
-	{
-		Input.Dispose();
-		Output.Dispose();
-	}
-}
-
-public static class GpuInputOutputScratchBufferPool<TInput, TOutput>
-	where TInput : unmanaged
-	where TOutput : unmanaged
-{
-	private const int WarmUpCapacity = 20480;
-
-	[ThreadStatic]
-	private static Dictionary<Accelerator, Queue<GpuInputOutputScratchBuffer<TInput, TOutput>>>? _pool;
-
-	private static Queue<GpuInputOutputScratchBuffer<TInput, TOutput>> GetPools(Accelerator accelerator)
-	{
-		var pool = _pool ??= [];
-
-		if (!pool.TryGetValue(accelerator, out var bufferQueue))
-		{
-			bufferQueue = [];
-			pool[accelerator] = bufferQueue;
-		}
-		
-		return bufferQueue;
-	}
-
-	// private static readonly ConcurrentDictionary<Accelerator, ConcurrentQueue<InputOutputScratchBuffer>> _pools = new();
-
-	private const int DefaultSmallPrimeFactorSlotCount = 64; // From PrimeOrderCalculator.Gpu.cs
-	private const int DefaultSpecialMaxFactorCapacity = 1024; // A reasonable default, will be resized if needed
-
-	public static GpuInputOutputScratchBuffer<TInput, TOutput> Rent(Accelerator accelerator, int inputCapacity, int outputCapacity)
-	{
-		var bufferQueue = GetPools(accelerator);
-		if (bufferQueue.TryDequeue(out var buffer))
-		{
-			if (buffer.Input.Length < inputCapacity || buffer.Output.Length < outputCapacity)
-			{
-				// lock (accelerator)
-				// {
-				// Console.WriteLine($"Resizing GPU scratch buffer from pool ({buffer.SmallPrimeFactorPrimeSlots.Length} / {smallPrimeFactorSlotCount}), ({buffer.SpecialMaxFactors.Length}/{specialMaxFactorCapacity})");
-				buffer.EnsureCapacity(inputCapacity, outputCapacity);
-				// }
-			}
-
-			return buffer;
-		}
-
-		return new GpuInputOutputScratchBuffer<TInput, TOutput>(accelerator, inputCapacity, outputCapacity);
-	}
-
-	public static void Return(Accelerator accelerator, GpuInputOutputScratchBuffer<TInput, TOutput> buffer) => GetPools(accelerator).Enqueue(buffer);
-
-	public static void WarmUp(int count, int inputCapacity, int outputCapacity)
-	{
-		// List<GpuInputOutputScratchBuffer<TInput, TOutput>> rented = new(count);
-		// for (int i = 0; i < count; i++)
-		// {
-		// 	rented.Add(Rent(inputCapacity, outputCapacity));
-		// }
-
-		// for (int i = 0; i < count; i++)
-		// {
-		// 	Return(rented[i]);
-		// }
-
-		// rented.Clear();
-		// rented.Capacity = 1;
-	}
-}

PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs

@@ -27,7 +27,7 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
         ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
         byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters
 
-		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.Pow2ModKernelScan | KernelType.IncrementalKernelScan);
+		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.Pow2ModKernelScan | KernelType.IncrementalKernelScan);
 		var pow2Kernel = kernels.Pow2Mod!;
         var incKernel = kernels.Incremental!;
         ulong step10 = (exponent.Mod10() << 1).Mod10();

PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs

@@ -21,7 +21,7 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 		ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
 		byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters
 
-		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.Pow2ModOrderKernelScan | KernelType.IncrementalOrderKernelScan);
+		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.Pow2ModOrderKernelScan | KernelType.IncrementalOrderKernelScan);
 		var kernel = _kernelType switch
 		{
 			GpuKernelType.Pow2Mod => kernels.Pow2ModOrder!,

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -1,4 +1,3 @@
-using System.Collections.Concurrent;
 using System.Runtime.InteropServices;
 using System.Runtime.CompilerServices;
 using ILGPU;
@@ -36,7 +35,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
 		var stream = accelerator.CreateStream();
 		var kernel = GetPow2ModKernel(accelerator);
 		var resources = RentResources(accelerator, GpuConstants.ScanBatchSize);
-		var gpuKernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.SmallCycles | KernelType.SmallPrimes);
+		var gpuKernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.SmallCycles | KernelType.SmallPrimes);
 
 		var orderBuffer = resources.OrderBuffer;
 		ulong[] orderArray = resources.OrderArray;

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -25,16 +25,16 @@ internal static partial class PrimeOrderGpuHeuristics
 	private static readonly ConcurrentDictionary<UInt128, byte> OverflowedPrimesWide = new();
 
 	[ThreadStatic]
-	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>? _pow2ModKernel;
+	private static Kernel[]? _pow2ModKernel;
 
 	[ThreadStatic]
-	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>>? _partialFactorKernel;
+	private static Kernel[]? _partialFactorKernel;
 
 	[ThreadStatic]
-	private static Dictionary<Accelerator, OrderKernelLauncher>? _orderKernel;
+	private static Kernel[]? _orderKernel;
 
 	[ThreadStatic]
-	private static Dictionary<Accelerator, SmallPrimeDeviceCache>? _smallPrimeDeviceCache;
+	private static SmallPrimeDeviceCache[]? _smallPrimeDeviceCache;
 
 	public readonly struct OrderKernelConfig(ulong previousOrder, byte hasPreviousOrder, uint smallFactorLimit, int maxPowChecks, int mode)
 	{
@@ -69,17 +69,6 @@ internal static partial class PrimeOrderGpuHeuristics
 		public readonly ArrayView1D<byte, Stride1D.Dense> Status = status;
 	}
 
-	private delegate void OrderKernelLauncher(
-		AcceleratorStream stream,
-		Index1D extent,
-		ulong prime,
-		OrderKernelConfig config,
-		MontgomeryDivisorData divisor,
-		ArrayView1D<uint, Stride1D.Dense> primes,
-		ArrayView1D<ulong, Stride1D.Dense> squares,
-		int primeCount,
-		OrderKernelBuffers buffers);
-
 	private sealed class SmallPrimeDeviceCache
 	{
 		public MemoryBuffer1D<uint, Stride1D.Dense>? Primes;
@@ -111,10 +100,10 @@ internal static partial class PrimeOrderGpuHeuristics
 		s_capability = PrimeOrderGpuCapability.Default;
 	}
 
-	private static SmallPrimeDeviceCache GetSmallPrimeDeviceCache(Accelerator accelerator, AcceleratorStream stream)
+	private static SmallPrimeDeviceCache GetSmallPrimeDeviceCache(int acceleratorIndex, AcceleratorStream stream)
 	{
 		var pool = _smallPrimeDeviceCache ??= [];
-		if (pool.TryGetValue(accelerator, out var cached))
+		if (pool[acceleratorIndex] is {} cached)
 		{
 			return cached;
 		}
@@ -123,11 +112,10 @@ internal static partial class PrimeOrderGpuHeuristics
 		ulong[] squares = PrimesGenerator.SmallPrimesPow2;
 		MemoryBuffer1D<uint, Stride1D.Dense>? primeBuffer;
 		MemoryBuffer1D<ulong, Stride1D.Dense>? squareBuffer;
-		// lock (accelerator)
+
-		{
+		var accelerator = _accelerators[acceleratorIndex];
-			primeBuffer = accelerator.Allocate1D<uint>(primes.Length);
+		primeBuffer = accelerator.Allocate1D<uint>(primes.Length);
-			squareBuffer = accelerator.Allocate1D<ulong>(squares.Length);
+		squareBuffer = accelerator.Allocate1D<ulong>(squares.Length);
-		}
 
 		primeBuffer.View.CopyFromCPU(stream, primes);
 		squareBuffer.View.CopyFromCPU(stream, squares);
@@ -138,25 +126,23 @@ internal static partial class PrimeOrderGpuHeuristics
 			Count = primes.Length,
 		};
 
-		pool[accelerator] = cached;
+		pool[acceleratorIndex] = cached;
 		return cached;
 	}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> GetPartialFactorKernel(Accelerator accelerator)
+	private static Kernel GetPartialFactorKernel(int acceleratorIndex)
 	{
 		var pool = _partialFactorKernel ??= [];
-		if (pool.TryGetValue(accelerator, out var cached))
+		if (pool[acceleratorIndex] is {} cached)
 		{
 			return cached;
 		}
 
+		var accelerator = _accelerators[acceleratorIndex];
 		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>(PartialFactorKernel);
 
-		var kernel = KernelUtil.GetKernel(loaded);
+		cached = KernelUtil.GetKernel(loaded);
-
+		pool[acceleratorIndex] = cached;
-		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>>();
-
-		pool[accelerator] = cached;
 		return cached;
 	}
 
@@ -226,7 +212,7 @@ internal static partial class PrimeOrderGpuHeuristics
 			Accelerator accelerator = _accelerators[acceleratorIndex];
 			AcceleratorStream stream = accelerator.CreateStream();
 
-			SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator, stream);
+			SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(acceleratorIndex, stream);
 
 			// TODO: We should create / reallocate these buffer only once or if the new required length is bigger than capacity.
 			MemoryBuffer1D<ulong, Stride1D.Dense>? factorBuffer = accelerator.Allocate1D<ulong>(primeTargets.Length);
@@ -242,11 +228,12 @@ internal static partial class PrimeOrderGpuHeuristics
 			// remainingBuffer.MemSetToZero(stream);
 			// fullyFactoredBuffer.MemSetToZero(stream);
 
-			var partialFactorKernel = GetPartialFactorKernel(accelerator);
+			var partialFactorKernel = GetPartialFactorKernel(acceleratorIndex);
 
-			partialFactorKernel(
+			partialFactorKernel.Launch(
 				stream,
 				1,
+				1,
 				cache.Primes!.View,
 				cache.Squares!.View,
 				cache.Count,
@@ -388,7 +375,7 @@ internal static partial class PrimeOrderGpuHeuristics
 		int acceleratorIndex = AcceleratorPool.Shared.Rent();
 		Accelerator accelerator = _accelerators[acceleratorIndex];
 		AcceleratorStream stream = AcceleratorStreamPool.Rent(acceleratorIndex);
-		SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator, stream);
+		SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(acceleratorIndex, stream);
 
 		// TODO: These buffers should be created reallocated once and assigned to an accelerator. Callers should use their own
 		// thread static cache to prevent other threads from using taking accelerators with pre-allocated buffers if they don't
@@ -447,10 +434,11 @@ internal static partial class PrimeOrderGpuHeuristics
 			resultBuffer.View,
 			statusBuffer.View);
 
-		var orderKernel = GetOrderKernel(accelerator);
+		var orderKernel = GetOrderKernel(acceleratorIndex);
-		orderKernel(
+		orderKernel.Launch(
 			stream,
 			1,
+			1,
 			prime,
 			kernelConfig,
 			divisorData,
@@ -497,19 +485,20 @@ internal static partial class PrimeOrderGpuHeuristics
 		return order != 0UL;
 	}
 
-	private static OrderKernelLauncher GetOrderKernel(Accelerator accelerator)
+	private static Kernel GetOrderKernel(int acceleratorIndex)
 	{
 		var pool = _orderKernel ??= [];
-		if (pool.TryGetValue(accelerator, out var cached))
+		if (pool[acceleratorIndex] is {} cached)
 		{
 			return cached;
 		}
 
+		var accelerator = _accelerators[acceleratorIndex];
 		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, OrderKernelConfig, MontgomeryDivisorData, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, OrderKernelBuffers>(CalculateOrderKernel);
+
 		var kernel = KernelUtil.GetKernel(loaded);
-		cached = kernel.CreateLauncherDelegate<OrderKernelLauncher>();
+		pool[acceleratorIndex] = kernel;
-		pool[accelerator] = cached;
+		return kernel;
-		return cached;
 	}
 
 	private static GpuPow2ModStatus TryPow2ModBatchInternal(ReadOnlySpan<UInt128> exponents, UInt128 prime, Span<UInt128> remainders)
@@ -576,10 +565,10 @@ internal static partial class PrimeOrderGpuHeuristics
 		// Pow2Mod kernel always assigns the required output elements so we don't need to worry about clearing these.
 		// remainderBuffer.MemSetToZero(stream);
 
-		var pow2ModKernel = GetPow2ModKernel(accelerator);
+		var pow2ModKernel = GetPow2ModKernel(acceleratorIndex);
 		try
 		{
-			pow2ModKernel(stream, exponents.Length, exponentBuffer.View, divisorData, remainderBuffer.View);
+			pow2ModKernel.Launch(stream, 1, exponents.Length, exponentBuffer.View, divisorData, remainderBuffer.View);
 			remainderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(results), exponents.Length);
 			stream.Synchronize();
 
@@ -684,21 +673,19 @@ internal static partial class PrimeOrderGpuHeuristics
 		return ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisorData, exponent);
 	}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> GetPow2ModKernel(Accelerator accelerator)
+	private static Kernel GetPow2ModKernel(int acceleratorIndex)
 	{
 		var pool = _pow2ModKernel ??= [];
-		if (pool.TryGetValue(accelerator, out var cached))
+		if (pool[acceleratorIndex] is {} cached)
 		{
 			return cached;
 		}
 
-		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernel);
+		var accelerator = _accelerators[acceleratorIndex];
+		var loaded = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernel));
 
-		var kernel = KernelUtil.GetKernel(loaded);
+		pool[acceleratorIndex] = loaded;
-
+		return loaded;
-		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
-		pool[accelerator] = cached;
-		return cached;
 	}
 
 	internal readonly record struct PrimeOrderGpuCapability(int ModulusBits, int ExponentBits)

PerfectNumbers.Core/MersenneDivisorCycles.cs

@@ -9,6 +9,7 @@ using PerfectNumbers.Core.Gpu;
 using ILGPU;
 using ILGPU.Algorithms;
 using ILGPU.Runtime;
+using PerfectNumbers.Core.Gpu.Accelerators;
 
 namespace PerfectNumbers.Core;
 
@@ -789,7 +790,7 @@ public class MersenneDivisorCycles
 	}
 
 	[ThreadStatic]
-	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>? _divisorCycleKernel;
+	private static Kernel[]? _divisorCycleKernel;
 
 	private static readonly Accelerator[] _accelerators = AcceleratorPool.Shared.Accelerators;
 
@@ -809,7 +810,7 @@ public class MersenneDivisorCycles
 		// GpuPrimeWorkLimiter.Acquire();
 		var acceleratorIndex = AcceleratorPool.Shared.Rent();
 		var accelerator = _accelerators[acceleratorIndex];
-		var stream = accelerator.CreateStream();
+		var stream = AcceleratorStreamPool.Rent(acceleratorIndex);
 
 		using Stream outputStream = new FileStream(path, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.Read, BufferSize10M, useAsync: true);
 		if (nextPosition > 0L)
@@ -838,7 +839,7 @@ public class MersenneDivisorCycles
 			skipCount--;
 		}
 
-		var divisorCycleKernel = GetDivisorCycleKernel(accelerator);
+		var divisorCycleKernel = GetDivisorCycleKernel(acceleratorIndex);
 		while (start <= maxDivisor)
 		{
 			end = Math.Min(start + batchSizeUL - 1UL, maxDivisor);
@@ -865,8 +866,9 @@ public class MersenneDivisorCycles
 				bufferCycle = accelerator.Allocate1D<ulong>(idx);				
 			}
 
-			divisorCycleKernel(
+			divisorCycleKernel.Launch(
 					stream,
+					1,
 					idx,
 					bufferDiv.View,
 					bufferCycle.View);
@@ -896,31 +898,28 @@ public class MersenneDivisorCycles
 			start = end + 1UL;
 		}
 
-		stream.Dispose();
+		AcceleratorStreamPool.Return(acceleratorIndex, stream);
 		pool.Return(divisors, clearArray: false);
 		pool.Return(outCycles, clearArray: false);
 
 		// GpuPrimeWorkLimiter.Release();
 	}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> GetDivisorCycleKernel(Accelerator accelerator)
+	private static Kernel GetDivisorCycleKernel(int acceleratorIndex)
 	{
 		var pool = _divisorCycleKernel ??= [];
-		if (pool.TryGetValue(accelerator, out var cached))
+		if (pool[acceleratorIndex] is {} cached)
 		{
 			return cached;
 		}
 
-		var loaded = accelerator.LoadAutoGroupedStreamKernel<
+		var loaded = KernelUtil.GetKernel(_accelerators[acceleratorIndex].LoadAutoGroupedStreamKernel<
 				Index1D,
 				ArrayView1D<ulong, Stride1D.Dense>,
-				ArrayView1D<ulong, Stride1D.Dense>>(DivisorCycleKernels.GpuDivisorCycleKernel);
+				ArrayView1D<ulong, Stride1D.Dense>>(DivisorCycleKernels.GpuDivisorCycleKernel));
 
-		var kernel = KernelUtil.GetKernel(loaded);
+		pool[acceleratorIndex] = loaded;
-
+		return loaded;
-		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>();
-		pool[accelerator] = cached;
-		return cached;
 	}
 
 	public static ulong CalculateCycleLengthGpu(ulong divisor) => DivisorCycleKernels.CalculateCycleLengthGpu(divisor);

PerfectNumbers.Core/MersenneNumberTester.cs

@@ -6,6 +6,7 @@ using ILGPU.Runtime;
 using PerfectNumbers.Core.Cpu;
 using PerfectNumbers.Core.Gpu;
 using UInt128 = System.UInt128;
+using PerfectNumbers.Core.Gpu.Accelerators;
 
 namespace PerfectNumbers.Core;
 
@@ -152,8 +153,8 @@ public sealed class MersenneNumberTester(
 		// GpuPrimeWorkLimiter.Acquire();
 		var acceleratorIndex = AcceleratorPool.Shared.Rent();
 		var accelerator = _accelerators[acceleratorIndex];
-		var stream = accelerator.CreateStream();
+		var stream = AcceleratorStreamPool.Rent(acceleratorIndex);
-		var orderKernel = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.OrderKernelScan).Order!;
+		var orderKernel = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.OrderKernelScan).Order!;
 
 		// Guard long-running kernels by chunking the warm-up across batches.
 		// Reuse the same ScanBatchSize knob used for scanning.
@@ -203,7 +204,7 @@ public sealed class MersenneNumberTester(
 			qBuffer.Dispose();
 		}
 
-		stream.Dispose();
+		AcceleratorStreamPool.Return(acceleratorIndex,stream);
 		// GpuPrimeWorkLimiter.Release();
 		uInt128Pool.Return(qs);
 	}

PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs

@@ -70,7 +70,7 @@ internal static partial class PrimeOrderCalculator
 
 		// GpuPrimeWorkLimiter.Acquire();
 		var stream = AcceleratorStreamPool.Rent(acceleratorIndex);
-		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.SmallPrimeFactorKernelScan);
+		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.SmallPrimeFactorKernelScan);
 		SmallPrimeFactorViews tables = GpuKernelPool.GetSmallPrimeFactorTables(kernels);
 
 		var kernel = kernels.SmallPrimeFactor!;
@@ -141,7 +141,7 @@ internal static partial class PrimeOrderCalculator
 		Accelerator accelerator = gpu.Accelerator;
 		// GpuPrimeWorkLimiter.Acquire();
 		var stream = AcceleratorStreamPool.Rent(acceleratorIndex);
-		var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream, KernelType.EvaluateSpecialMaxCandidatesKernel);
+		var kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.EvaluateSpecialMaxCandidatesKernel);
 
 		ArrayView1D<ulong, Stride1D.Dense> specialMaxFactorsView = gpu.SpecialMaxFactors.View;
 		specialMaxFactorsView.SubView(0, factorCount).CopyFromCPU(stream, ref MemoryMarshal.GetReference(factors), factorCount);

PerfectNumbers.Core/PrimeTester.cs

@@ -269,9 +269,6 @@ public sealed class PrimeTester
 		return BinaryGcd(n, m) != 1UL;
 	}
 
-	[ThreadStatic]
-	private static readonly Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<byte>>>? _kernel;
-
 	internal static void SharesFactorWithMaxExponentBatch(ReadOnlySpan<ulong> values, Span<byte> results)
 	{
 		// TODO: Route this batch helper through the shared GPU kernel pool from