Inline GPU kernel state into prime tester leases

This commit is contained in:
Marek Lemańczyk committed 2025-10-29 18:13:51 +01:00
1 parent 23fa56b356
commit a4ac82e74d
3 files changed
+300 -280

No files matched your search

@@ -213,29 +213,29 @@ public sealed class HeuristicCombinedPrimeTester
ulong[] divisorArray = divisorPool.Rent(batchCapacity); ulong[] divisorArray = divisorPool.Rent(batchCapacity);
var limiter = GpuPrimeWorkLimiter.Acquire(); var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity); var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity);
var accelerator = gpu.Accelerator; var accelerator = gpu.Accelerator;
var state = gpu.State; var kernel = gpu.HeuristicTrialDivisionKernel;
bool compositeDetected = false; bool compositeDetected = false;
int count = 0; int count = 0;
var inputView = gpu.Input.View; var inputView = gpu.Input.View;
var outputView = gpu.Output.View; var outputView = gpu.Output.View;
bool ProcessBatch(int length) bool ProcessBatch(int length)
{ {
inputView.CopyFromCPU(ref divisorArray[0], length); inputView.CopyFromCPU(ref divisorArray[0], length);
byte compositeFlag = 0; byte compositeFlag = 0;
outputView.CopyFromCPU(ref compositeFlag, 1); outputView.CopyFromCPU(ref compositeFlag, 1);
state.HeuristicTrialDivisionKernel(length, inputView, n, outputView); kernel(length, inputView, n, outputView);
accelerator.Synchronize(); accelerator.Synchronize();
outputView.CopyToCPU(ref compositeFlag, 1); outputView.CopyToCPU(ref compositeFlag, 1);
return compositeFlag != 0; return compositeFlag != 0;
} }
ReadOnlySpan<uint> combinedDivisors = GetCombinedDivisors(nMod10); ReadOnlySpan<uint> combinedDivisors = GetCombinedDivisors(nMod10);
ReadOnlySpan<ulong> combinedDivisorSquares = GetCombinedDivisorSquares(nMod10); ReadOnlySpan<ulong> combinedDivisorSquares = GetCombinedDivisorSquares(nMod10);
int length = combinedDivisors.Length; int length = combinedDivisors.Length;
+2 -2
View File
@@ -269,7 +269,7 @@ public sealed class HeuristicPrimeTester
var limiter = GpuPrimeWorkLimiter.Acquire(); var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity); var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity);
var accelerator = gpu.Accelerator; var accelerator = gpu.Accelerator;
var state = gpu.State; var kernel = gpu.HeuristicTrialDivisionKernel;
bool compositeDetected = false; bool compositeDetected = false;
int count = 0; int count = 0;
@@ -282,7 +282,7 @@ public sealed class HeuristicPrimeTester
inputView.CopyFromCPU(ref divisorArray[0], length); inputView.CopyFromCPU(ref divisorArray[0], length);
byte compositeFlag = 0; byte compositeFlag = 0;
outputView.CopyFromCPU(ref compositeFlag, 1); outputView.CopyFromCPU(ref compositeFlag, 1);
state.HeuristicTrialDivisionKernel(length, inputView, n, outputView); kernel(length, inputView, n, outputView);
accelerator.Synchronize(); accelerator.Synchronize();
outputView.CopyToCPU(ref compositeFlag, 1); outputView.CopyToCPU(ref compositeFlag, 1);
+280 -260
View File
@@ -1,6 +1,7 @@
using System; using System;
using System.Buffers; using System.Buffers;
using System.Collections.Concurrent; using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Numerics; using System.Numerics;
using System.Runtime.CompilerServices; using System.Runtime.CompilerServices;
using ILGPU; using ILGPU;
@@ -97,51 +98,51 @@ public sealed class PrimeTester
[MethodImpl(MethodImplOptions.AggressiveInlining)] [MethodImpl(MethodImplOptions.AggressiveInlining)]
public bool IsPrimeGpu(ulong n, CancellationToken ct) public bool IsPrimeGpu(ulong n, CancellationToken ct)
{ {
bool forceCpu = GpuContextPool.ForceCpu; bool forceCpu = GpuContextPool.ForceCpu;
bool belowGpuRange = n < 31UL; bool belowGpuRange = n < 31UL;
if (forceCpu || belowGpuRange) if (forceCpu || belowGpuRange)
{ {
return IsPrimeCpu(n, ct); return IsPrimeCpu(n, ct);
} }
var limiter = GpuPrimeWorkLimiter.Acquire(); var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTesterGpuContextPool.Rent(1); var gpu = PrimeTesterGpuContextPool.Rent(1);
var state = gpu.State; var accelerator = gpu.Accelerator;
var accelerator = gpu.Accelerator; var kernel = gpu.Kernel;
ulong value = n; ulong value = n;
byte flag = 0; byte flag = 0;
var input = gpu.Input; var input = gpu.Input;
var output = gpu.Output; var output = gpu.Output;
input.View.CopyFromCPU(ref value, 1); input.View.CopyFromCPU(ref value, 1);
state.Kernel( kernel(
1, 1,
input.View, input.View,
state.DevicePrimesDefault.View, gpu.DevicePrimesDefault.View,
state.DevicePrimesLastOne.View, gpu.DevicePrimesLastOne.View,
state.DevicePrimesLastSeven.View, gpu.DevicePrimesLastSeven.View,
state.DevicePrimesLastThree.View, gpu.DevicePrimesLastThree.View,
state.DevicePrimesLastNine.View, gpu.DevicePrimesLastNine.View,
state.DevicePrimesPow2Default.View, gpu.DevicePrimesPow2Default.View,
state.DevicePrimesPow2LastOne.View, gpu.DevicePrimesPow2LastOne.View,
state.DevicePrimesPow2LastSeven.View, gpu.DevicePrimesPow2LastSeven.View,
state.DevicePrimesPow2LastThree.View, gpu.DevicePrimesPow2LastThree.View,
state.DevicePrimesPow2LastNine.View, gpu.DevicePrimesPow2LastNine.View,
output.View); output.View);
accelerator.Synchronize(); accelerator.Synchronize();
output.View.CopyToCPU(ref flag, 1); output.View.CopyToCPU(ref flag, 1);
gpu.Dispose(); gpu.Dispose();
limiter.Dispose(); limiter.Dispose();
return flag != 0; return flag != 0;
} }
public static bool IsPrimeCpu(ulong n, CancellationToken ct) public static bool IsPrimeCpu(ulong n, CancellationToken ct)
{ {
// Preserve the legacy entry point for callers that bypass the thread-local caches. // Preserve the legacy entry point for callers that bypass the thread-local caches.
return IsPrime(n, ct); return IsPrime(n, ct);
@@ -152,276 +153,295 @@ public sealed class PrimeTester
[MethodImpl(MethodImplOptions.AggressiveInlining)] [MethodImpl(MethodImplOptions.AggressiveInlining)]
public static void IsPrimeBatchGpu(ReadOnlySpan<ulong> values, Span<byte> results) public static void IsPrimeBatchGpu(ReadOnlySpan<ulong> values, Span<byte> results)
{ {
var limiter = GpuPrimeWorkLimiter.Acquire(); var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize); var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
var state = gpu.State; var accelerator = gpu.Accelerator;
var accelerator = gpu.Accelerator; var kernel = gpu.Kernel;
int totalLength = values.Length; int totalLength = values.Length;
int batchSize = GpuBatchSize; int batchSize = GpuBatchSize;
var input = gpu.Input; var input = gpu.Input;
var output = gpu.Output; var output = gpu.Output;
ArrayPool<ulong> pool = ThreadStaticPools.UlongPool; ArrayPool<ulong> pool = ThreadStaticPools.UlongPool;
ulong[] temp = pool.Rent(batchSize); ulong[] temp = pool.Rent(batchSize);
int pos = 0; int pos = 0;
while (pos < totalLength) while (pos < totalLength)
{ {
int remaining = totalLength - pos; int remaining = totalLength - pos;
int count = remaining > batchSize ? batchSize : remaining; int count = remaining > batchSize ? batchSize : remaining;
values.Slice(pos, count).CopyTo(temp); values.Slice(pos, count).CopyTo(temp);
input.View.CopyFromCPU(ref temp[0], count); input.View.CopyFromCPU(ref temp[0], count);
state.Kernel( kernel(
count, count,
input.View, input.View,
state.DevicePrimesDefault.View, gpu.DevicePrimesDefault.View,
state.DevicePrimesLastOne.View, gpu.DevicePrimesLastOne.View,
state.DevicePrimesLastSeven.View, gpu.DevicePrimesLastSeven.View,
state.DevicePrimesLastThree.View, gpu.DevicePrimesLastThree.View,
state.DevicePrimesLastNine.View, gpu.DevicePrimesLastNine.View,
state.DevicePrimesPow2Default.View, gpu.DevicePrimesPow2Default.View,
state.DevicePrimesPow2LastOne.View, gpu.DevicePrimesPow2LastOne.View,
state.DevicePrimesPow2LastSeven.View, gpu.DevicePrimesPow2LastSeven.View,
state.DevicePrimesPow2LastThree.View, gpu.DevicePrimesPow2LastThree.View,
state.DevicePrimesPow2LastNine.View, gpu.DevicePrimesPow2LastNine.View,
output.View); output.View);
accelerator.Synchronize(); accelerator.Synchronize();
output.View.CopyToCPU(ref results[pos], count); output.View.CopyToCPU(ref results[pos], count);
pos += count; pos += count;
} }
pool.Return(temp, clearArray: false); pool.Return(temp, clearArray: false);
gpu.Dispose(); gpu.Dispose();
limiter.Dispose(); limiter.Dispose();
} }
internal static class PrimeTesterGpuContextPool internal static class PrimeTesterGpuContextPool
{ {
internal sealed class PooledContext internal sealed class PooledContext
{ {
public Context Context { get; } public Context Context { get; }
public Accelerator Accelerator { get; } public Accelerator Accelerator { get; }
public PooledContext() public Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> Kernel { get; }
{
Context = Context.CreateDefault();
Accelerator = Context.GetPreferredDevice(false).CreateAccelerator(Context);
}
public void Dispose() public Action<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>> HeuristicTrialDivisionKernel { get; }
{
ClearGpuCaches(Accelerator);
Accelerator.Dispose();
Context.Dispose();
}
}
private static readonly ConcurrentQueue<PooledContext> Pool = new(); public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesDefault { get; }
internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1) public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastOne { get; }
{
if (Pool.TryDequeue(out var ctx))
{
return new PrimeTesterGpuContextLease(ctx, minBufferCapacity);
}
return new PrimeTesterGpuContextLease(new PooledContext(), minBufferCapacity); public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastSeven { get; }
}
internal static void DisposeAll() public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastThree { get; }
{
while (Pool.TryDequeue(out var ctx))
{
ctx.Dispose();
}
}
private static void Return(PooledContext ctx) public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastNine { get; }
{
ctx.Accelerator.Synchronize();
Pool.Enqueue(ctx);
}
internal sealed class PrimeTesterGpuContextLease : IDisposable public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2Default { get; }
{
private readonly PooledContext _ctx;
public readonly Accelerator Accelerator; public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastOne { get; }
public readonly KernelState State; public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastSeven { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> Input; public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastThree { get; }
public MemoryBuffer1D<byte, Stride1D.Dense> Output; public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastNine { get; }
public int BufferCapacity; public MemoryBuffer1D<ulong, Stride1D.Dense> Input { get; private set; }
internal PrimeTesterGpuContextLease(PooledContext ctx, int minBufferCapacity) public MemoryBuffer1D<byte, Stride1D.Dense> Output { get; private set; }
{
// The pool never requests a zero-length buffer, so the old guard stays commented out to document the invariant.
// if (minBufferCapacity < 1)
// {
// minBufferCapacity = 1;
// }
_ctx = ctx; public int BufferCapacity { get; private set; }
Accelerator = ctx.Accelerator;
State = GpuKernelState.GetOrCreate(Accelerator);
Input = Accelerator.Allocate1D<ulong>(minBufferCapacity);
Output = Accelerator.Allocate1D<byte>(minBufferCapacity);
BufferCapacity = minBufferCapacity;
}
public void EnsureCapacity(int minCapacity) public PooledContext(int minBufferCapacity)
{ {
if (minCapacity <= BufferCapacity) Context = Context.CreateDefault();
{ Accelerator = Context.GetPreferredDevice(false).CreateAccelerator(Context);
return; Kernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
} HeuristicTrialDivisionKernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
Input.Dispose(); var primesDefault = DivisorGenerator.SmallPrimes;
Output.Dispose(); DevicePrimesDefault = Accelerator.Allocate1D<uint>(primesDefault.Length);
DevicePrimesDefault.View.CopyFromCPU(primesDefault);
Input = Accelerator.Allocate1D<ulong>(minCapacity); var primesLastOne = DivisorGenerator.SmallPrimesLastOne;
Output = Accelerator.Allocate1D<byte>(minCapacity); DevicePrimesLastOne = Accelerator.Allocate1D<uint>(primesLastOne.Length);
BufferCapacity = minCapacity; DevicePrimesLastOne.View.CopyFromCPU(primesLastOne);
}
public void Dispose() var primesLastSeven = DivisorGenerator.SmallPrimesLastSeven;
{ DevicePrimesLastSeven = Accelerator.Allocate1D<uint>(primesLastSeven.Length);
Input.Dispose(); DevicePrimesLastSeven.View.CopyFromCPU(primesLastSeven);
Output.Dispose();
BufferCapacity = 0;
Return(_ctx);
}
}
}
// Per-accelerator GPU state for prime sieve (kernel + uploaded primes).
internal sealed class KernelState
{
public Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> Kernel;
public Action<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>> HeuristicTrialDivisionKernel;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesDefault;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastOne;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastSeven;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastThree;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastNine;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2Default;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastOne;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastSeven;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastThree;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastNine;
public bool IsDisposed;
public KernelState(Accelerator accelerator) var primesLastThree = DivisorGenerator.SmallPrimesLastThree;
{ DevicePrimesLastThree = Accelerator.Allocate1D<uint>(primesLastThree.Length);
IsDisposed = false; DevicePrimesLastThree.View.CopyFromCPU(primesLastThree);
// Compile once per accelerator and upload primes once.
Kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
HeuristicTrialDivisionKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
var primesDefault = DivisorGenerator.SmallPrimes; var primesLastNine = DivisorGenerator.SmallPrimesLastNine;
DevicePrimesDefault = accelerator.Allocate1D<uint>(primesDefault.Length); DevicePrimesLastNine = Accelerator.Allocate1D<uint>(primesLastNine.Length);
DevicePrimesDefault.View.CopyFromCPU(primesDefault); DevicePrimesLastNine.View.CopyFromCPU(primesLastNine);
var primesLastOne = DivisorGenerator.SmallPrimesLastOne; var primesPow2Default = DivisorGenerator.SmallPrimesPow2;
DevicePrimesLastOne = accelerator.Allocate1D<uint>(primesLastOne.Length); DevicePrimesPow2Default = Accelerator.Allocate1D<ulong>(primesPow2Default.Length);
DevicePrimesLastOne.View.CopyFromCPU(primesLastOne); DevicePrimesPow2Default.View.CopyFromCPU(primesPow2Default);
var primesLastSeven = DivisorGenerator.SmallPrimesLastSeven; var primesPow2LastOne = DivisorGenerator.SmallPrimesPow2LastOne;
DevicePrimesLastSeven = accelerator.Allocate1D<uint>(primesLastSeven.Length); DevicePrimesPow2LastOne = Accelerator.Allocate1D<ulong>(primesPow2LastOne.Length);
DevicePrimesLastSeven.View.CopyFromCPU(primesLastSeven); DevicePrimesPow2LastOne.View.CopyFromCPU(primesPow2LastOne);
var primesLastThree = DivisorGenerator.SmallPrimesLastThree; var primesPow2LastSeven = DivisorGenerator.SmallPrimesPow2LastSeven;
DevicePrimesLastThree = accelerator.Allocate1D<uint>(primesLastThree.Length); DevicePrimesPow2LastSeven = Accelerator.Allocate1D<ulong>(primesPow2LastSeven.Length);
DevicePrimesLastThree.View.CopyFromCPU(primesLastThree); DevicePrimesPow2LastSeven.View.CopyFromCPU(primesPow2LastSeven);
var primesLastNine = DivisorGenerator.SmallPrimesLastNine; var primesPow2LastThree = DivisorGenerator.SmallPrimesPow2LastThree;
DevicePrimesLastNine = accelerator.Allocate1D<uint>(primesLastNine.Length); DevicePrimesPow2LastThree = Accelerator.Allocate1D<ulong>(primesPow2LastThree.Length);
DevicePrimesLastNine.View.CopyFromCPU(primesLastNine); DevicePrimesPow2LastThree.View.CopyFromCPU(primesPow2LastThree);
var primesPow2Default = DivisorGenerator.SmallPrimesPow2; var primesPow2LastNine = DivisorGenerator.SmallPrimesPow2LastNine;
DevicePrimesPow2Default = accelerator.Allocate1D<ulong>(primesPow2Default.Length); DevicePrimesPow2LastNine = Accelerator.Allocate1D<ulong>(primesPow2LastNine.Length);
DevicePrimesPow2Default.View.CopyFromCPU(primesPow2Default); DevicePrimesPow2LastNine.View.CopyFromCPU(primesPow2LastNine);
var primesPow2LastOne = DivisorGenerator.SmallPrimesPow2LastOne; BufferCapacity = minBufferCapacity < 1 ? 1 : minBufferCapacity;
DevicePrimesPow2LastOne = accelerator.Allocate1D<ulong>(primesPow2LastOne.Length); Input = Accelerator.Allocate1D<ulong>(BufferCapacity);
DevicePrimesPow2LastOne.View.CopyFromCPU(primesPow2LastOne); Output = Accelerator.Allocate1D<byte>(BufferCapacity);
}
var primesPow2LastSeven = DivisorGenerator.SmallPrimesPow2LastSeven; public void EnsureCapacity(int minCapacity)
DevicePrimesPow2LastSeven = accelerator.Allocate1D<ulong>(primesPow2LastSeven.Length); {
DevicePrimesPow2LastSeven.View.CopyFromCPU(primesPow2LastSeven); if (minCapacity < 1)
{
minCapacity = 1;
}
var primesPow2LastThree = DivisorGenerator.SmallPrimesPow2LastThree; if (minCapacity <= BufferCapacity)
DevicePrimesPow2LastThree = accelerator.Allocate1D<ulong>(primesPow2LastThree.Length); {
DevicePrimesPow2LastThree.View.CopyFromCPU(primesPow2LastThree); return;
}
var primesPow2LastNine = DivisorGenerator.SmallPrimesPow2LastNine; Input.Dispose();
DevicePrimesPow2LastNine = accelerator.Allocate1D<ulong>(primesPow2LastNine.Length); Output.Dispose();
DevicePrimesPow2LastNine.View.CopyFromCPU(primesPow2LastNine);
}
public void Clear() Input = Accelerator.Allocate1D<ulong>(minCapacity);
{ Output = Accelerator.Allocate1D<byte>(minCapacity);
if (IsDisposed) BufferCapacity = minCapacity;
{ }
return;
}
DevicePrimesDefault.Dispose(); public void DisposeResources()
DevicePrimesLastOne.Dispose(); {
DevicePrimesLastSeven.Dispose(); Input.Dispose();
DevicePrimesLastThree.Dispose(); Output.Dispose();
DevicePrimesLastNine.Dispose(); DevicePrimesDefault.Dispose();
DevicePrimesPow2Default.Dispose(); DevicePrimesLastOne.Dispose();
DevicePrimesPow2LastOne.Dispose(); DevicePrimesLastSeven.Dispose();
DevicePrimesPow2LastSeven.Dispose(); DevicePrimesLastThree.Dispose();
DevicePrimesPow2LastThree.Dispose(); DevicePrimesLastNine.Dispose();
DevicePrimesPow2LastNine.Dispose(); DevicePrimesPow2Default.Dispose();
IsDisposed = true; DevicePrimesPow2LastOne.Dispose();
} DevicePrimesPow2LastSeven.Dispose();
} DevicePrimesPow2LastThree.Dispose();
DevicePrimesPow2LastNine.Dispose();
Accelerator.Dispose();
Context.Dispose();
}
}
internal static class GpuKernelState private static readonly ConcurrentQueue<PooledContext> Pool = new();
{
// Map accelerator to cached state; use Lazy to serialize kernel creation
private static readonly System.Collections.Concurrent.ConcurrentDictionary<Accelerator, Lazy<KernelState>> States = new();
// TODO: Prewarm this per-accelerator cache during startup (and reuse a simple array keyed by accelerator index)
// once the kernel pool exposes deterministic ordering; the Lazy wrappers showed measurable overhead in the
// GpuModularArithmeticBenchmarks hot path.
[MethodImpl(MethodImplOptions.AggressiveInlining)] internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
public static KernelState GetOrCreate(Accelerator accelerator) {
{ if (!Pool.TryDequeue(out var ctx))
var lazy = States.GetOrAdd( {
accelerator, ctx = new PooledContext(minBufferCapacity);
acc => new Lazy<KernelState>(() => new KernelState(acc), LazyThreadSafetyMode.ExecutionAndPublication)); }
return lazy.Value; else
} {
ctx.EnsureCapacity(minBufferCapacity);
}
public static void Clear(Accelerator accelerator) return new PrimeTesterGpuContextLease(ctx);
{ }
if (States.TryRemove(accelerator, out var lazy))
{
if (lazy.IsValueCreated)
{
var state = lazy.Value;
state.Clear();
}
}
}
}
// Expose cache clearing for accelerator disposal coordination internal static void DisposeAll()
{
while (Pool.TryDequeue(out var ctx))
{
ctx.DisposeResources();
}
}
internal static void Clear(Accelerator accelerator)
{
if (accelerator is null)
{
return;
}
var retained = new List<PooledContext>();
while (Pool.TryDequeue(out var ctx))
{
if (ctx.Accelerator == accelerator)
{
ctx.DisposeResources();
}
else
{
retained.Add(ctx);
}
}
for (int i = 0; i < retained.Count; i++)
{
Pool.Enqueue(retained[i]);
}
}
private static void Return(PooledContext ctx)
{
ctx.Accelerator.Synchronize();
Pool.Enqueue(ctx);
}
internal sealed class PrimeTesterGpuContextLease : IDisposable
{
private readonly PooledContext _ctx;
internal PrimeTesterGpuContextLease(PooledContext ctx)
{
_ctx = ctx;
}
public Accelerator Accelerator => _ctx.Accelerator;
public Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> Kernel => _ctx.Kernel;
public Action<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>> HeuristicTrialDivisionKernel => _ctx.HeuristicTrialDivisionKernel;
public MemoryBuffer1D<ulong, Stride1D.Dense> Input => _ctx.Input;
public MemoryBuffer1D<byte, Stride1D.Dense> Output => _ctx.Output;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesDefault => _ctx.DevicePrimesDefault;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastOne => _ctx.DevicePrimesLastOne;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastSeven => _ctx.DevicePrimesLastSeven;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastThree => _ctx.DevicePrimesLastThree;
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastNine => _ctx.DevicePrimesLastNine;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2Default => _ctx.DevicePrimesPow2Default;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastOne => _ctx.DevicePrimesPow2LastOne;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastSeven => _ctx.DevicePrimesPow2LastSeven;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastThree => _ctx.DevicePrimesPow2LastThree;
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastNine => _ctx.DevicePrimesPow2LastNine;
public int BufferCapacity => _ctx.BufferCapacity;
public void EnsureCapacity(int minCapacity)
{
_ctx.EnsureCapacity(minCapacity);
}
public void Dispose()
{
Return(_ctx);
}
}
}
// Expose cache clearing for accelerator disposal coordination
public static void ClearGpuCaches(Accelerator accelerator) public static void ClearGpuCaches(Accelerator accelerator)
{ {
GpuKernelState.Clear(accelerator); PrimeTesterGpuContextPool.Clear(accelerator);
} }
internal static void DisposeGpuContexts() internal static void DisposeGpuContexts()