ConcurrentDictionary removed. Kernel made dependent on the shared GPU context.

This commit is contained in:
mlemanczyk committed 2025-11-07 10:16:33 +01:00
1 parent 9537addf38
commit 119690086a
11 files changed
+93 -65

No files matched your search

@@ -349,7 +349,7 @@ public class GpuKernelCompilationTests
private static void CompileGpuKernelPoolKernels(Accelerator accelerator)
{
_ = accelerator;
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
try
{
_ = lease.OrderKernel;
@@ -363,7 +363,7 @@ public class GpuKernelCompilationTests
lease.Dispose();
}
var orderLease = GpuKernelPool.GetKernel();
var orderLease = GpuKernelPool.Rent();
try
{
_ = orderLease.OrderKernel;
+50 -14
View File
@@ -16,7 +16,7 @@ public sealed class DivisorCycleCache
private const byte ByteOne = 1;
private const int StackBufferThreshold = 128;
private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>> _gpuKernelCache = new(AcceleratorReferenceComparer.Instance);
private readonly Action<AcceleratorStream, Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> _gpuKernel = LoadKernel(SharedGpuContext.Accelerator);
private volatile ulong[] _snapshot;
private volatile bool _useGpuGeneration = true;
private readonly int _divisorCyclesBatchSize;
@@ -224,20 +224,55 @@ public sealed class DivisorCycleCache
}
}
private const int DefaultCapacity = 128;
[ThreadStatic]
private static MemoryBuffer1D<ulong, Stride1D.Dense>? divisorBuffer;
[ThreadStatic]
private static MemoryBuffer1D<ulong, Stride1D.Dense>? powBuffer;
[ThreadStatic]
private static MemoryBuffer1D<ulong, Stride1D.Dense>? orderBuffer;
[ThreadStatic]
private static MemoryBuffer1D<ulong, Stride1D.Dense>? resultBuffer;
[ThreadStatic]
private static MemoryBuffer1D<byte, Stride1D.Dense>? statusBuffer;
private static void EnsureCapacity(int requiredCapacity)
{
if (requiredCapacity <= (divisorBuffer?.Length ?? 0))
{
return;
}
if (divisorBuffer != null)
{
divisorBuffer.Dispose();
powBuffer!.Dispose();
orderBuffer!.Dispose();
resultBuffer!.Dispose();
statusBuffer!.Dispose();
}
divisorBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
powBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
orderBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
resultBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
statusBuffer = SharedGpuContext.Accelerator.Allocate1D<byte>(requiredCapacity);
}
private void ComputeCyclesGpuCore(ReadOnlySpan<ulong> divisors, Span<ulong> destination)
{
int length = divisors.Length;
var gpuLease = GpuKernelPool.GetKernel();
var gpuLease = GpuKernelPool.Rent();
Accelerator accelerator = gpuLease.Accelerator;
var stream = gpuLease.Stream;
var kernel = _gpuKernelCache.GetOrAdd(accelerator, LoadKernel);
using MemoryBuffer1D<ulong, Stride1D.Dense> divisorBuffer = accelerator.Allocate1D<ulong>(length);
using MemoryBuffer1D<ulong, Stride1D.Dense> powBuffer = accelerator.Allocate1D<ulong>(length);
using MemoryBuffer1D<ulong, Stride1D.Dense> orderBuffer = accelerator.Allocate1D<ulong>(length);
using MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = accelerator.Allocate1D<ulong>(length);
using MemoryBuffer1D<byte, Stride1D.Dense> statusBuffer = accelerator.Allocate1D<byte>(length);
var kernel = _gpuKernel;
EnsureCapacity(length);
ulong[]? rentedPow = null;
Span<ulong> powSpan = length <= StackBufferThreshold
@@ -260,7 +295,7 @@ public sealed class DivisorCycleCache
: new Span<byte>(rentedStatus = ThreadStaticPools.BytePool.Rent(length), 0, length);
ref ulong divisorRef = ref MemoryMarshal.GetReference(divisors);
divisorBuffer.View.CopyFromCPU(stream, ref divisorRef, length);
divisorBuffer!.View.CopyFromCPU(stream, ref divisorRef, length);
for (int i = 0; i < length; i++)
{
@@ -277,10 +312,10 @@ public sealed class DivisorCycleCache
statusSpan[i] = ByteZero;
}
powBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(powSpan), length);
orderBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(orderSpan), length);
resultBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(resultSpan), length);
statusBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(statusSpan), length);
powBuffer!.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(powSpan), length);
orderBuffer!.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(orderSpan), length);
resultBuffer!.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(resultSpan), length);
statusBuffer!.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(statusSpan), length);
int pending;
do
@@ -303,6 +338,7 @@ public sealed class DivisorCycleCache
resultBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(resultSpan), length);
stream.Synchronize();
resultSpan.CopyTo(destination);
if (rentedPow is not null)
{
+7 -16
View File
@@ -83,7 +83,7 @@ public readonly struct ResiduePrimeViews(
public readonly ArrayView1D<ulong, Stride1D.Dense> LastSevenPow2 = lastSevenPow2;
}
public readonly struct SmallPrimeFactorTables(
public readonly struct SmallPrimeFactorViews(
MemoryBuffer1D<uint, Stride1D.Dense> primes,
MemoryBuffer1D<ulong, Stride1D.Dense> squares,
int count)
@@ -184,18 +184,13 @@ public class GpuKernelPool
return new ResiduePrimeViews(deviceLastOne.View, deviceLastSeven.View, deviceLastOnePow2.View, deviceLastSevenPow2.View);
}
public static SmallPrimeFactorTables EnsureSmallPrimeFactorTables(KernelContainer kernels, Accelerator accelerator, AcceleratorStream stream)
public static SmallPrimeFactorViews EnsureSmallPrimeFactorTables(KernelContainer kernels, Accelerator accelerator, AcceleratorStream stream)
{
if (kernels.SmallPrimeFactorsPrimes is { } primeBuffer && kernels.SmallPrimeFactorsSquares is { } squareBuffer)
{
return new SmallPrimeFactorTables(primeBuffer, squareBuffer, (int)primeBuffer.Length);
return new SmallPrimeFactorViews(primeBuffer, squareBuffer, (int)primeBuffer.Length);
}
// if (kernels.SmallPrimeFactorsPrimes is { } existingPrimes && kernels.SmallPrimeFactorsSquares is { } existingSquares)
// {
// return new SmallPrimeFactorTables(existingPrimes, existingSquares, (int)existingPrimes.Length);
// }
var hostPrimes = PrimesGenerator.SmallPrimes;
var hostSquares = PrimesGenerator.SmallPrimesPow2;
@@ -207,17 +202,13 @@ public class GpuKernelPool
kernels.SmallPrimeFactorsPrimes = devicePrimes;
kernels.SmallPrimeFactorsSquares = deviceSquares;
return new SmallPrimeFactorTables(devicePrimes, deviceSquares, hostPrimes.Length);
return new SmallPrimeFactorViews(devicePrimes, deviceSquares, hostPrimes.Length);
}
public static GpuKernelLease GetKernel()
public static GpuKernelLease Rent()
{
GpuPrimeWorkLimiter.Acquire();
var gpu = Rent();
var gpu = GpuContextPool.Rent();
var stream = gpu.Accelerator.CreateStream();
var kernels = GetKernels(gpu.Accelerator, stream);
return GpuKernelLease.Rent(gpu, stream, kernels);
@@ -229,7 +220,7 @@ public class GpuKernelPool
/// <param name="action">Action to run with (Accelerator, Stream).</param>
public static void Run(Action<Accelerator, AcceleratorStream> action)
{
var lease = GetKernel();
var lease = Rent();
var accelerator = lease.Accelerator;
var stream = lease.Stream;
action(accelerator, stream);
@@ -95,7 +95,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
ulong processedCount;
ulong lastProcessed;
var gpuLease = GpuKernelPool.GetKernel();
var gpuLease = GpuKernelPool.Rent();
var accelerator = gpuLease.Accelerator;
var stream = gpuLease.Stream;
@@ -501,7 +501,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
internal DivisorScanSession(MersenneNumberDivisorByDivisorGpuTester owner)
{
_owner = owner;
_lease = GpuKernelPool.GetKernel();
_lease = GpuKernelPool.Rent();
_accelerator = _lease.Accelerator;
_stream = _lease.Stream;
_capacity = Math.Max(1, owner._gpuBatchSize);
@@ -16,7 +16,7 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
{
throw new NotImplementedException($"GPU incremental scanning requires the device cycle heuristics implementation (kernel: {_kernelType}, GPU order: {_useGpuOrder}).");
var gpuLease = GpuKernelPool.GetKernel();
var gpuLease = GpuKernelPool.Rent();
var accelerator = gpuLease.Accelerator;
var stream = gpuLease.Stream;
int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy
@@ -9,7 +9,7 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime)
{
var gpuLease = GpuKernelPool.GetKernel();
var gpuLease = GpuKernelPool.Rent();
var accelerator = gpuLease.Accelerator;
var stream = gpuLease.Stream;
int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy and avoids TDR
@@ -198,7 +198,7 @@ internal static partial class PrimeOrderGpuHeuristics
try
{
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
Accelerator accelerator = lease.Accelerator;
AcceleratorStream stream = lease.Stream;
@@ -365,7 +365,7 @@ internal static partial class PrimeOrderGpuHeuristics
{
order = 0UL;
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
Accelerator accelerator = lease.Accelerator;
AcceleratorStream stream = lease.Stream;
@@ -526,7 +526,7 @@ internal static partial class PrimeOrderGpuHeuristics
private static bool TryComputeOnGpu(ReadOnlySpan<ulong> exponents, ulong prime, in MontgomeryDivisorData divisorData, Span<ulong> results)
{
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
Accelerator accelerator = lease.Accelerator;
AcceleratorStream stream = lease.Stream;
var kernel = GetPow2ModKernel(accelerator);
@@ -581,7 +581,7 @@ internal static partial class PrimeOrderGpuHeuristics
GpuUInt128[]? rentedExponents = null;
GpuUInt128[]? rentedResults = null;
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
ArrayPool<GpuUInt128> gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool;
try
+1 -1
View File
@@ -804,7 +804,7 @@ public class MersenneDivisorCycles
divisors = pool.Rent(count);
outCycles = pool.Rent(count);
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
var accelerator = lease.Accelerator;
var stream = lease.Stream;
var loaded = accelerator.LoadAutoGroupedStreamKernel<
+1 -1
View File
@@ -147,7 +147,7 @@ public sealed class MersenneNumberTester(
UInt128[] qs = qsBuffer;
var gpuLease = GpuKernelPool.GetKernel();
var gpuLease = GpuKernelPool.Rent();
var accelerator = gpuLease.Accelerator;
var stream = gpuLease.Stream;
var orderKernel = gpuLease.OrderKernel!;
@@ -19,14 +19,14 @@ internal static partial class PrimeOrderCalculator
Span<int> exponentBuffer = exponentBufferArray.AsSpan(0, GpuSmallPrimeFactorSlots);
remaining = value;
var lease = GpuKernelPool.GetKernel();
var lease = GpuKernelPool.Rent();
Accelerator accelerator = lease.Accelerator;
var stream = lease.Stream;
KernelContainer kernels = lease.Kernels;
ScratchBuffer scratch = GpuScratchBufferPool.Rent(accelerator, GpuSmallPrimeFactorSlots, 0);
SmallPrimeFactorTables tables = GpuKernelPool.EnsureSmallPrimeFactorTables(kernels, accelerator, stream);
SmallPrimeFactorViews tables = GpuKernelPool.EnsureSmallPrimeFactorTables(kernels, accelerator, stream);
if (kernels.SmallPrimeFactorsPrimes is null) throw new NullReferenceException($"{nameof(kernels.SmallPrimeFactorsPrimes)} is null");
if (kernels.SmallPrimeFactorsSquares is null) throw new NullReferenceException($"{nameof(kernels.SmallPrimeFactorsSquares)} is null");
@@ -100,7 +100,7 @@ internal static partial class PrimeOrderCalculator
factorSpan[i] = factors[i].Value;
}
GpuKernelLease lease = GpuKernelPool.GetKernel();
GpuKernelLease lease = GpuKernelPool.Rent();
Accelerator accelerator = lease.Accelerator;
var stream = lease.Stream;
+21 -20
View File
@@ -440,7 +440,7 @@ public static partial class ULongExtensions
private static class Pow2MontgomeryGpuExecutor
{
private static readonly ConcurrentDictionary<Accelerator, Pow2MontgomeryKernelGroup> KernelCache = new();
private static readonly Pow2MontgomeryKernelGroup Kernel = Pow2MontgomeryKernelGroup.Create(SharedGpuContext.Accelerator);
public static ulong ExecuteKeep(ulong exponent, in MontgomeryDivisorData divisor)
{
@@ -452,24 +452,25 @@ public static partial class ULongExtensions
return Execute(exponent, divisor, useKeepKernel: false);
}
[ThreadStatic]
private static MemoryBuffer1D<ulong, Stride1D.Dense>? _exponentBuffer;
[ThreadStatic]
private static MemoryBuffer1D<ulong, Stride1D.Dense>? _resultBuffer;
private static ulong Execute(ulong exponent, in MontgomeryDivisorData divisor, bool useKeepKernel)
{
ulong result = 0UL;
GpuKernelLease lease = GpuKernelPool.GetKernel();
GpuKernelLease lease = GpuKernelPool.Rent();
Accelerator accelerator = lease.Accelerator;
AcceleratorStream stream = lease.Stream;
// Keep this commented out. It should never happen in production code.
// if (accelerator.AcceleratorType == AcceleratorType.CPU)
// {
// return 0UL;
// }
var kernelGroup = Kernel;
var kernelGroup = KernelCache.GetOrAdd(accelerator, static accel => Pow2MontgomeryKernelGroup.Create(accel));
EnsureCapacity();
var exponentBuffer = accelerator.Allocate1D<ulong>(1);
var resultBuffer = accelerator.Allocate1D<ulong>(1);
MemoryBuffer1D<ulong, Stride1D.Dense> exponentBuffer = _exponentBuffer!;
MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = _resultBuffer!;
exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1);
// We don't need to worry about any left-overs here.
@@ -483,19 +484,19 @@ public static partial class ULongExtensions
resultBuffer.Dispose();
exponentBuffer.Dispose();
lease.Dispose();
// Keep this commented. We don't want to catch any exceptions. All should crash the scanner.
// catch (AcceleratorException)
// {
// return 0UL;
// }
// catch (NotSupportedException)
// {
// return 0UL;
// }
// Intentionally avoid exception handling here; any accelerator failure should crash the scanner.
return result;
}
private static void EnsureCapacity()
{
if (_exponentBuffer == null)
{
Accelerator accelerator = SharedGpuContext.Accelerator;
_exponentBuffer = accelerator.Allocate1D<ulong>(1);
_resultBuffer = accelerator.Allocate1D<ulong>(1);
}
}
private sealed class Pow2MontgomeryKernelGroup
{
internal Pow2MontgomeryKernelGroup(