53 lines
1.8 KiB
C#
53 lines
1.8 KiB
C#
using System.Collections.Concurrent;
|
|
using ILGPU.Runtime;
|
|
using Open.Collections;
|
|
|
|
namespace PerfectNumbers.Core.Gpu.Accelerators;
|
|
|
|
public static class AcceleratorStreamPool
|
|
{
|
|
private static readonly Dictionary<Accelerator, ConcurrentQueue<AcceleratorStream>> _streams = new(PerfectNumberConstants .RollingAccelerators);
|
|
// private static readonly ConcurrentDictionary<Accelerator, AcceleratorStream> _streams = new(AcceleratorPool.Shared.Accelerators.ToDictionary(static x => x, CreateStream));
|
|
|
|
private static SemaphoreSlim CreateLock() => new(PerfectNumberConstants.ThreadsByStream);
|
|
private static AcceleratorStream CreateStream(Accelerator accelerator) => accelerator.CreateStream();
|
|
|
|
private static readonly Dictionary<Accelerator, SemaphoreSlim> _locks = new(AcceleratorPool.Shared.Accelerators.ToDictionary(x => x, static x => CreateLock()));
|
|
|
|
public static AcceleratorStream Rent(Accelerator accelerator)
|
|
{
|
|
var streamLock = _locks[accelerator];
|
|
// var streamLock = _locks.GetOrAdd(accelerator, static _ => CreateLock());
|
|
|
|
streamLock.Wait();
|
|
var queue = _streams[accelerator];
|
|
|
|
return queue.TryDequeue(out var stream)
|
|
? stream
|
|
: accelerator.CreateStream();
|
|
}
|
|
|
|
public static void Return(AcceleratorStream stream)
|
|
{
|
|
Accelerator accelerator = stream.Accelerator;
|
|
var streamLock = _locks[accelerator];
|
|
var queue = _streams[accelerator];
|
|
queue.Enqueue(stream);
|
|
streamLock.Release();
|
|
}
|
|
|
|
public static void WarmUp(Accelerator accelerator)
|
|
{
|
|
_streams.Add(accelerator, new());
|
|
_ = _locks.TryAdd(accelerator, CreateLock());
|
|
|
|
// ConcurrentQueue<AcceleratorStream> streamQueue = new();
|
|
// var stream = accelerator.CreateStream();
|
|
// // var stream = accelerator.CreateStream();
|
|
// streamQueue.Enqueue(stream);
|
|
// _streams[accelerator] = streamQueue;
|
|
// _locks[accelerator] = CreateLock();
|
|
// stream.Synchronize();
|
|
}
|
|
}
|