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
+111 -20

No files matched your search

+3 -2
View File
@@ -58,14 +58,15 @@ internal static class Program
bool startPrimeProvided = _cliArguments.StartPrimeProvided; bool startPrimeProvided = _cliArguments.StartPrimeProvided;
int threadCount = Math.Max(1, _cliArguments.ThreadCount); int threadCount = Math.Max(1, _cliArguments.ThreadCount);
int gpuPrimeBatch = Math.Max(1, _cliArguments.GpuPrimeBatch); int gpuPrimeBatch = Math.Max(1, _cliArguments.GpuPrimeBatch);
int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads);
UnboundedTaskScheduler.ConfigureThreadCount(threadCount); UnboundedTaskScheduler.ConfigureThreadCount(threadCount);
PrimeTester.GpuBatchSize = gpuPrimeBatch; PrimeTester.GpuBatchSize = gpuPrimeBatch;
GpuPrimeLeaseLimiter.Initialize(gpuPrimeThreads);
GpuContextPool.WarmUpPool(threadCount); GpuContextPool.WarmUpPool(threadCount);
PrimeTester.WarmUpGpuKernels(threadCount); PrimeTester.WarmUpGpuKernels(gpuPrimeThreads);
Console.WriteLine("Starting up threads..."); Console.WriteLine("Starting up threads...");
_ = UnboundedTaskScheduler.Instance; _ = UnboundedTaskScheduler.Instance;
int blockSize = Math.Max(1, _cliArguments.BlockSize); int blockSize = Math.Max(1, _cliArguments.BlockSize);
int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads);
GpuKernelType kernelType = _cliArguments.KernelType; GpuKernelType kernelType = _cliArguments.KernelType;
bool useBitTransform = _cliArguments.UseBitTransform; bool useBitTransform = _cliArguments.UseBitTransform;
bool useOrder = _cliArguments.UseOrder; bool useOrder = _cliArguments.UseOrder;
+26 -6
View File
@@ -49,6 +49,10 @@ public static class GpuContextPool
// Allows callers to choose CPU/GPU per use-case, decoupled from ForceCpu. // Allows callers to choose CPU/GPU per use-case, decoupled from ForceCpu.
public static GpuContextLease Rent() public static GpuContextLease Rent()
{
GpuPrimeLeaseLimiter.Enter();
try
{ {
if (GpuPool.TryDequeue(out var gpu)) if (GpuPool.TryDequeue(out var gpu))
{ {
@@ -58,6 +62,12 @@ public static class GpuContextPool
// Create a new accelerator when the pool does not have one available. // Create a new accelerator when the pool does not have one available.
return new GpuContextLease(new PooledContext()); return new GpuContextLease(new PooledContext());
} }
catch
{
GpuPrimeLeaseLimiter.Exit();
throw;
}
}
public static void WarmUpPool(int threadCount) public static void WarmUpPool(int threadCount)
{ {
@@ -96,23 +106,33 @@ public static class GpuContextPool
{ {
ctx.Accelerator.Synchronize(); ctx.Accelerator.Synchronize();
GpuPool.Enqueue(ctx); 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) internal GpuContextLease(PooledContext ctx)
{ {
_ctx = ctx; _ctx = ctx;
_disposed = false;
} }
public readonly Context Context => _ctx.Context; public Context Context => _ctx!.Context;
public readonly Accelerator Accelerator => _ctx.Accelerator; 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();
}
}
+45 -12
View File
@@ -159,18 +159,7 @@ public sealed class PrimeTester
return; return;
} }
int toWarm = target - WarmedGpuLeaseCount; PrimeTesterGpuContextPool.WarmUp(target, GpuBatchSize);
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();
}
WarmedGpuLeaseCount = target; WarmedGpuLeaseCount = target;
} }
} }
@@ -229,6 +218,8 @@ public sealed class PrimeTester
internal static class PrimeTesterGpuContextPool internal static class PrimeTesterGpuContextPool
{ {
private static readonly ConcurrentQueue<PrimeTesterGpuContextLease> Pool = new(); 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 object SharedTablesLock = new();
private static readonly Dictionary<string, SharedHeuristicGpuTables> SharedTables = new(); private static readonly Dictionary<string, SharedHeuristicGpuTables> SharedTables = new();
@@ -267,7 +258,35 @@ public sealed class PrimeTester
} }
} }
internal static void WarmUp(int target, int minBufferCapacity)
{
if (target <= 0)
{
return;
}
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) internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
{
GpuPrimeLeaseLimiter.Enter();
try
{ {
if (!Pool.TryDequeue(out var lease)) if (!Pool.TryDequeue(out var lease))
{ {
@@ -281,11 +300,18 @@ public sealed class PrimeTester
lease.ResetReturnFlag(); lease.ResetReturnFlag();
return lease; return lease;
} }
catch
{
GpuPrimeLeaseLimiter.Exit();
throw;
}
}
internal static void Return(PrimeTesterGpuContextLease lease) internal static void Return(PrimeTesterGpuContextLease lease)
{ {
lease.Accelerator.Synchronize(); lease.Accelerator.Synchronize();
Pool.Enqueue(lease); Pool.Enqueue(lease);
GpuPrimeLeaseLimiter.Exit();
} }
internal static void DisposeAll() internal static void DisposeAll()
@@ -304,6 +330,8 @@ public sealed class PrimeTester
SharedTables.Clear(); SharedTables.Clear();
} }
WarmedLeaseCount = 0;
} }
internal static void Clear(Accelerator accelerator) internal static void Clear(Accelerator accelerator)
@@ -313,6 +341,8 @@ public sealed class PrimeTester
return; return;
} }
lock (WarmUpLock)
{
var retained = new List<PrimeTesterGpuContextLease>(); var retained = new List<PrimeTesterGpuContextLease>();
while (Pool.TryDequeue(out var lease)) while (Pool.TryDequeue(out var lease))
{ {
@@ -331,6 +361,9 @@ public sealed class PrimeTester
{ {
Pool.Enqueue(retained[i]); Pool.Enqueue(retained[i]);
} }
WarmedLeaseCount = Pool.Count;
}
} }
internal sealed class PrimeTesterGpuContextLease : IDisposable internal sealed class PrimeTesterGpuContextLease : IDisposable