Add GPU lease warm-up and shared limiter

This commit is contained in:
Marek Lemańczyk committed 2025-11-02 13:47:57 +01:00
1 parent 0e7bbae132
commit 00d4e9bde2
4 files changed
+139 -48

No files matched your search

+3 -2
View File
@@ -58,14 +58,15 @@ internal static class Program
bool startPrimeProvided = _cliArguments.StartPrimeProvided;
int threadCount = Math.Max(1, _cliArguments.ThreadCount);
int gpuPrimeBatch = Math.Max(1, _cliArguments.GpuPrimeBatch);
int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads);
UnboundedTaskScheduler.ConfigureThreadCount(threadCount);
PrimeTester.GpuBatchSize = gpuPrimeBatch;
GpuPrimeLeaseLimiter.Initialize(gpuPrimeThreads);
GpuContextPool.WarmUpPool(threadCount);
PrimeTester.WarmUpGpuKernels(threadCount);
PrimeTester.WarmUpGpuKernels(gpuPrimeThreads);
Console.WriteLine("Starting up threads...");
_ = UnboundedTaskScheduler.Instance;
int blockSize = Math.Max(1, _cliArguments.BlockSize);
int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads);
GpuKernelType kernelType = _cliArguments.KernelType;
bool useBitTransform = _cliArguments.UseBitTransform;
bool useOrder = _cliArguments.UseOrder;
+32 -12
View File
@@ -50,13 +50,23 @@ public static class GpuContextPool
// Allows callers to choose CPU/GPU per use-case, decoupled from ForceCpu.
public static GpuContextLease Rent()
{
if (GpuPool.TryDequeue(out var gpu))
{
return new GpuContextLease(gpu);
}
GpuPrimeLeaseLimiter.Enter();
// Create a new accelerator when the pool does not have one available.
return new GpuContextLease(new PooledContext());
try
{
if (GpuPool.TryDequeue(out var gpu))
{
return new GpuContextLease(gpu);
}
// Create a new accelerator when the pool does not have one available.
return new GpuContextLease(new PooledContext());
}
catch
{
GpuPrimeLeaseLimiter.Exit();
throw;
}
}
public static void WarmUpPool(int threadCount)
@@ -96,23 +106,33 @@ public static class GpuContextPool
{
ctx.Accelerator.Synchronize();
GpuPool.Enqueue(ctx);
GpuPrimeLeaseLimiter.Exit();
}
public readonly struct GpuContextLease
public struct GpuContextLease
{
private readonly PooledContext _ctx;
private PooledContext? _ctx;
private bool _disposed;
internal GpuContextLease(PooledContext ctx)
{
_ctx = ctx;
_disposed = false;
}
public readonly Context Context => _ctx.Context;
public readonly Accelerator Accelerator => _ctx.Accelerator;
public Context Context => _ctx!.Context;
public Accelerator Accelerator => _ctx!.Accelerator;
public readonly void Dispose()
public void Dispose()
{
Return(_ctx);
if (_disposed)
{
return;
}
_disposed = true;
Return(_ctx!);
_ctx = null;
}
}
}
@@ -0,0 +1,37 @@
using System.Threading;
namespace PerfectNumbers.Core.Gpu;
public static class GpuPrimeLeaseLimiter
{
private static SemaphoreSlim _semaphore = new(int.MaxValue, int.MaxValue);
private static int _currentLimit = int.MaxValue;
public static void Initialize(int limit)
{
if (limit <= 0)
{
limit = 1;
}
if (limit == _currentLimit)
{
return;
}
var previous = _semaphore;
_semaphore = new SemaphoreSlim(limit, limit);
_currentLimit = limit;
previous.Dispose();
}
public static void Enter()
{
_semaphore.Wait();
}
public static void Exit()
{
_semaphore.Release();
}
}
+67 -34
View File
@@ -159,18 +159,7 @@ public sealed class PrimeTester
return;
}
int toWarm = target - WarmedGpuLeaseCount;
var leases = new PrimeTesterGpuContextPool.PrimeTesterGpuContextLease[toWarm];
for (int i = 0; i < toWarm; i++)
{
leases[i] = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
}
for (int i = 0; i < toWarm; i++)
{
leases[i].Dispose();
}
PrimeTesterGpuContextPool.WarmUp(target, GpuBatchSize);
WarmedGpuLeaseCount = target;
}
}
@@ -229,6 +218,8 @@ public sealed class PrimeTester
internal static class PrimeTesterGpuContextPool
{
private static readonly ConcurrentQueue<PrimeTesterGpuContextLease> Pool = new();
private static readonly object WarmUpLock = new();
private static int WarmedLeaseCount;
private static readonly object SharedTablesLock = new();
private static readonly Dictionary<string, SharedHeuristicGpuTables> SharedTables = new();
@@ -267,25 +258,60 @@ public sealed class PrimeTester
}
}
internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
internal static void WarmUp(int target, int minBufferCapacity)
{
if (!Pool.TryDequeue(out var lease))
if (target <= 0)
{
lease = new PrimeTesterGpuContextLease(minBufferCapacity);
}
else
{
lease.EnsureCapacity(minBufferCapacity);
return;
}
lease.ResetReturnFlag();
return lease;
lock (WarmUpLock)
{
if (target <= WarmedLeaseCount)
{
return;
}
int toCreate = target - WarmedLeaseCount;
for (int i = 0; i < toCreate; i++)
{
Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity));
}
WarmedLeaseCount = target;
}
}
internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
{
GpuPrimeLeaseLimiter.Enter();
try
{
if (!Pool.TryDequeue(out var lease))
{
lease = new PrimeTesterGpuContextLease(minBufferCapacity);
}
else
{
lease.EnsureCapacity(minBufferCapacity);
}
lease.ResetReturnFlag();
return lease;
}
catch
{
GpuPrimeLeaseLimiter.Exit();
throw;
}
}
internal static void Return(PrimeTesterGpuContextLease lease)
{
lease.Accelerator.Synchronize();
Pool.Enqueue(lease);
GpuPrimeLeaseLimiter.Exit();
}
internal static void DisposeAll()
@@ -304,6 +330,8 @@ public sealed class PrimeTester
SharedTables.Clear();
}
WarmedLeaseCount = 0;
}
internal static void Clear(Accelerator accelerator)
@@ -313,23 +341,28 @@ public sealed class PrimeTester
return;
}
var retained = new List<PrimeTesterGpuContextLease>();
while (Pool.TryDequeue(out var lease))
lock (WarmUpLock)
{
if (lease.Accelerator == accelerator)
var retained = new List<PrimeTesterGpuContextLease>();
while (Pool.TryDequeue(out var lease))
{
lease.DisposeResources();
if (lease.Accelerator == accelerator)
{
lease.DisposeResources();
}
else
{
retained.Add(lease);
}
}
else
{
retained.Add(lease);
}
}
int retainedCount = retained.Count;
for (int i = 0; i < retainedCount; i++)
{
Pool.Enqueue(retained[i]);
int retainedCount = retained.Count;
for (int i = 0; i < retainedCount; i++)
{
Pool.Enqueue(retained[i]);
}
WarmedLeaseCount = Pool.Count;
}
}