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
+204 -184

No files matched your search

@@ -215,7 +215,7 @@ public sealed class HeuristicCombinedPrimeTester
var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity);
var accelerator = gpu.Accelerator;
var state = gpu.State;
var kernel = gpu.HeuristicTrialDivisionKernel;
bool compositeDetected = false;
int count = 0;
@@ -228,7 +228,7 @@ public sealed class HeuristicCombinedPrimeTester
inputView.CopyFromCPU(ref divisorArray[0], length);
byte compositeFlag = 0;
outputView.CopyFromCPU(ref compositeFlag, 1);
state.HeuristicTrialDivisionKernel(length, inputView, n, outputView);
kernel(length, inputView, n, outputView);
accelerator.Synchronize();
outputView.CopyToCPU(ref compositeFlag, 1);
+2 -2
View File
@@ -269,7 +269,7 @@ public sealed class HeuristicPrimeTester
var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity);
var accelerator = gpu.Accelerator;
var state = gpu.State;
var kernel = gpu.HeuristicTrialDivisionKernel;
bool compositeDetected = false;
int count = 0;
@@ -282,7 +282,7 @@ public sealed class HeuristicPrimeTester
inputView.CopyFromCPU(ref divisorArray[0], length);
byte compositeFlag = 0;
outputView.CopyFromCPU(ref compositeFlag, 1);
state.HeuristicTrialDivisionKernel(length, inputView, n, outputView);
kernel(length, inputView, n, outputView);
accelerator.Synchronize();
outputView.CopyToCPU(ref compositeFlag, 1);
+200 -180
View File
@@ -1,6 +1,7 @@
using System;
using System.Buffers;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Numerics;
using System.Runtime.CompilerServices;
using ILGPU;
@@ -108,8 +109,8 @@ public sealed class PrimeTester
var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTesterGpuContextPool.Rent(1);
var state = gpu.State;
var accelerator = gpu.Accelerator;
var kernel = gpu.Kernel;
ulong value = n;
byte flag = 0;
@@ -118,19 +119,19 @@ public sealed class PrimeTester
var output = gpu.Output;
input.View.CopyFromCPU(ref value, 1);
state.Kernel(
kernel(
1,
input.View,
state.DevicePrimesDefault.View,
state.DevicePrimesLastOne.View,
state.DevicePrimesLastSeven.View,
state.DevicePrimesLastThree.View,
state.DevicePrimesLastNine.View,
state.DevicePrimesPow2Default.View,
state.DevicePrimesPow2LastOne.View,
state.DevicePrimesPow2LastSeven.View,
state.DevicePrimesPow2LastThree.View,
state.DevicePrimesPow2LastNine.View,
gpu.DevicePrimesDefault.View,
gpu.DevicePrimesLastOne.View,
gpu.DevicePrimesLastSeven.View,
gpu.DevicePrimesLastThree.View,
gpu.DevicePrimesLastNine.View,
gpu.DevicePrimesPow2Default.View,
gpu.DevicePrimesPow2LastOne.View,
gpu.DevicePrimesPow2LastSeven.View,
gpu.DevicePrimesPow2LastThree.View,
gpu.DevicePrimesPow2LastNine.View,
output.View);
accelerator.Synchronize();
output.View.CopyToCPU(ref flag, 1);
@@ -155,8 +156,8 @@ public sealed class PrimeTester
{
var limiter = GpuPrimeWorkLimiter.Acquire();
var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
var state = gpu.State;
var accelerator = gpu.Accelerator;
var kernel = gpu.Kernel;
int totalLength = values.Length;
int batchSize = GpuBatchSize;
@@ -174,19 +175,19 @@ public sealed class PrimeTester
values.Slice(pos, count).CopyTo(temp);
input.View.CopyFromCPU(ref temp[0], count);
state.Kernel(
kernel(
count,
input.View,
state.DevicePrimesDefault.View,
state.DevicePrimesLastOne.View,
state.DevicePrimesLastSeven.View,
state.DevicePrimesLastThree.View,
state.DevicePrimesLastNine.View,
state.DevicePrimesPow2Default.View,
state.DevicePrimesPow2LastOne.View,
state.DevicePrimesPow2LastSeven.View,
state.DevicePrimesPow2LastThree.View,
state.DevicePrimesPow2LastNine.View,
gpu.DevicePrimesDefault.View,
gpu.DevicePrimesLastOne.View,
gpu.DevicePrimesLastSeven.View,
gpu.DevicePrimesLastThree.View,
gpu.DevicePrimesLastNine.View,
gpu.DevicePrimesPow2Default.View,
gpu.DevicePrimesPow2LastOne.View,
gpu.DevicePrimesPow2LastSeven.View,
gpu.DevicePrimesPow2LastThree.View,
gpu.DevicePrimesPow2LastNine.View,
output.View);
accelerator.Synchronize();
output.View.CopyToCPU(ref results[pos], count);
@@ -208,78 +209,95 @@ public sealed class PrimeTester
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; }
public Action<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>> HeuristicTrialDivisionKernel { get; }
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesDefault { get; }
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastOne { get; }
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastSeven { get; }
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastThree { get; }
public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimesLastNine { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2Default { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastOne { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastSeven { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastThree { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimesPow2LastNine { get; }
public MemoryBuffer1D<ulong, Stride1D.Dense> Input { get; private set; }
public MemoryBuffer1D<byte, Stride1D.Dense> Output { get; private set; }
public int BufferCapacity { get; private set; }
public PooledContext(int minBufferCapacity)
{
Context = Context.CreateDefault();
Accelerator = Context.GetPreferredDevice(false).CreateAccelerator(Context);
}
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);
public void Dispose()
{
ClearGpuCaches(Accelerator);
Accelerator.Dispose();
Context.Dispose();
}
}
var primesDefault = DivisorGenerator.SmallPrimes;
DevicePrimesDefault = Accelerator.Allocate1D<uint>(primesDefault.Length);
DevicePrimesDefault.View.CopyFromCPU(primesDefault);
private static readonly ConcurrentQueue<PooledContext> Pool = new();
var primesLastOne = DivisorGenerator.SmallPrimesLastOne;
DevicePrimesLastOne = Accelerator.Allocate1D<uint>(primesLastOne.Length);
DevicePrimesLastOne.View.CopyFromCPU(primesLastOne);
internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
{
if (Pool.TryDequeue(out var ctx))
{
return new PrimeTesterGpuContextLease(ctx, minBufferCapacity);
}
var primesLastSeven = DivisorGenerator.SmallPrimesLastSeven;
DevicePrimesLastSeven = Accelerator.Allocate1D<uint>(primesLastSeven.Length);
DevicePrimesLastSeven.View.CopyFromCPU(primesLastSeven);
return new PrimeTesterGpuContextLease(new PooledContext(), minBufferCapacity);
}
var primesLastThree = DivisorGenerator.SmallPrimesLastThree;
DevicePrimesLastThree = Accelerator.Allocate1D<uint>(primesLastThree.Length);
DevicePrimesLastThree.View.CopyFromCPU(primesLastThree);
internal static void DisposeAll()
{
while (Pool.TryDequeue(out var ctx))
{
ctx.Dispose();
}
}
var primesLastNine = DivisorGenerator.SmallPrimesLastNine;
DevicePrimesLastNine = Accelerator.Allocate1D<uint>(primesLastNine.Length);
DevicePrimesLastNine.View.CopyFromCPU(primesLastNine);
private static void Return(PooledContext ctx)
{
ctx.Accelerator.Synchronize();
Pool.Enqueue(ctx);
}
var primesPow2Default = DivisorGenerator.SmallPrimesPow2;
DevicePrimesPow2Default = Accelerator.Allocate1D<ulong>(primesPow2Default.Length);
DevicePrimesPow2Default.View.CopyFromCPU(primesPow2Default);
internal sealed class PrimeTesterGpuContextLease : IDisposable
{
private readonly PooledContext _ctx;
var primesPow2LastOne = DivisorGenerator.SmallPrimesPow2LastOne;
DevicePrimesPow2LastOne = Accelerator.Allocate1D<ulong>(primesPow2LastOne.Length);
DevicePrimesPow2LastOne.View.CopyFromCPU(primesPow2LastOne);
public readonly Accelerator Accelerator;
var primesPow2LastSeven = DivisorGenerator.SmallPrimesPow2LastSeven;
DevicePrimesPow2LastSeven = Accelerator.Allocate1D<ulong>(primesPow2LastSeven.Length);
DevicePrimesPow2LastSeven.View.CopyFromCPU(primesPow2LastSeven);
public readonly KernelState State;
var primesPow2LastThree = DivisorGenerator.SmallPrimesPow2LastThree;
DevicePrimesPow2LastThree = Accelerator.Allocate1D<ulong>(primesPow2LastThree.Length);
DevicePrimesPow2LastThree.View.CopyFromCPU(primesPow2LastThree);
public MemoryBuffer1D<ulong, Stride1D.Dense> Input;
var primesPow2LastNine = DivisorGenerator.SmallPrimesPow2LastNine;
DevicePrimesPow2LastNine = Accelerator.Allocate1D<ulong>(primesPow2LastNine.Length);
DevicePrimesPow2LastNine.View.CopyFromCPU(primesPow2LastNine);
public MemoryBuffer1D<byte, Stride1D.Dense> Output;
public int BufferCapacity;
internal PrimeTesterGpuContextLease(PooledContext ctx, int minBufferCapacity)
{
// 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;
Accelerator = ctx.Accelerator;
State = GpuKernelState.GetOrCreate(Accelerator);
Input = Accelerator.Allocate1D<ulong>(minBufferCapacity);
Output = Accelerator.Allocate1D<byte>(minBufferCapacity);
BufferCapacity = minBufferCapacity;
BufferCapacity = minBufferCapacity < 1 ? 1 : minBufferCapacity;
Input = Accelerator.Allocate1D<ulong>(BufferCapacity);
Output = Accelerator.Allocate1D<byte>(BufferCapacity);
}
public void EnsureCapacity(int minCapacity)
{
if (minCapacity < 1)
{
minCapacity = 1;
}
if (minCapacity <= BufferCapacity)
{
return;
@@ -293,87 +311,10 @@ public sealed class PrimeTester
BufferCapacity = minCapacity;
}
public void Dispose()
public void DisposeResources()
{
Input.Dispose();
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)
{
IsDisposed = false;
// 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;
DevicePrimesDefault = accelerator.Allocate1D<uint>(primesDefault.Length);
DevicePrimesDefault.View.CopyFromCPU(primesDefault);
var primesLastOne = DivisorGenerator.SmallPrimesLastOne;
DevicePrimesLastOne = accelerator.Allocate1D<uint>(primesLastOne.Length);
DevicePrimesLastOne.View.CopyFromCPU(primesLastOne);
var primesLastSeven = DivisorGenerator.SmallPrimesLastSeven;
DevicePrimesLastSeven = accelerator.Allocate1D<uint>(primesLastSeven.Length);
DevicePrimesLastSeven.View.CopyFromCPU(primesLastSeven);
var primesLastThree = DivisorGenerator.SmallPrimesLastThree;
DevicePrimesLastThree = accelerator.Allocate1D<uint>(primesLastThree.Length);
DevicePrimesLastThree.View.CopyFromCPU(primesLastThree);
var primesLastNine = DivisorGenerator.SmallPrimesLastNine;
DevicePrimesLastNine = accelerator.Allocate1D<uint>(primesLastNine.Length);
DevicePrimesLastNine.View.CopyFromCPU(primesLastNine);
var primesPow2Default = DivisorGenerator.SmallPrimesPow2;
DevicePrimesPow2Default = accelerator.Allocate1D<ulong>(primesPow2Default.Length);
DevicePrimesPow2Default.View.CopyFromCPU(primesPow2Default);
var primesPow2LastOne = DivisorGenerator.SmallPrimesPow2LastOne;
DevicePrimesPow2LastOne = accelerator.Allocate1D<ulong>(primesPow2LastOne.Length);
DevicePrimesPow2LastOne.View.CopyFromCPU(primesPow2LastOne);
var primesPow2LastSeven = DivisorGenerator.SmallPrimesPow2LastSeven;
DevicePrimesPow2LastSeven = accelerator.Allocate1D<ulong>(primesPow2LastSeven.Length);
DevicePrimesPow2LastSeven.View.CopyFromCPU(primesPow2LastSeven);
var primesPow2LastThree = DivisorGenerator.SmallPrimesPow2LastThree;
DevicePrimesPow2LastThree = accelerator.Allocate1D<ulong>(primesPow2LastThree.Length);
DevicePrimesPow2LastThree.View.CopyFromCPU(primesPow2LastThree);
var primesPow2LastNine = DivisorGenerator.SmallPrimesPow2LastNine;
DevicePrimesPow2LastNine = accelerator.Allocate1D<ulong>(primesPow2LastNine.Length);
DevicePrimesPow2LastNine.View.CopyFromCPU(primesPow2LastNine);
}
public void Clear()
{
if (IsDisposed)
{
return;
}
DevicePrimesDefault.Dispose();
DevicePrimesLastOne.Dispose();
DevicePrimesLastSeven.Dispose();
@@ -384,44 +325,123 @@ public sealed class PrimeTester
DevicePrimesPow2LastSeven.Dispose();
DevicePrimesPow2LastThree.Dispose();
DevicePrimesPow2LastNine.Dispose();
IsDisposed = true;
Accelerator.Dispose();
Context.Dispose();
}
}
internal static class GpuKernelState
{
// 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.
private static readonly ConcurrentQueue<PooledContext> Pool = new();
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static KernelState GetOrCreate(Accelerator accelerator)
internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
{
var lazy = States.GetOrAdd(
accelerator,
acc => new Lazy<KernelState>(() => new KernelState(acc), LazyThreadSafetyMode.ExecutionAndPublication));
return lazy.Value;
if (!Pool.TryDequeue(out var ctx))
{
ctx = new PooledContext(minBufferCapacity);
}
else
{
ctx.EnsureCapacity(minBufferCapacity);
}
public static void Clear(Accelerator accelerator)
{
if (States.TryRemove(accelerator, out var lazy))
{
if (lazy.IsValueCreated)
{
var state = lazy.Value;
state.Clear();
}
return new PrimeTesterGpuContextLease(ctx);
}
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)
{
GpuKernelState.Clear(accelerator);
PrimeTesterGpuContextPool.Clear(accelerator);
}
internal static void DisposeGpuContexts()