ArrayPool<T> replaced with ThreadStaticPools

This commit is contained in:
mlemanczyk committed 2025-11-07 00:46:16 +01:00
1 parent c004e140db
commit 9537addf38
12 files changed
+123 -75

No files matched your search

@@ -39,9 +39,10 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
}
var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);
// Avoid giant stack allocations that can trigger StackOverflow when batchSize is large.
// Rent a reusable array from the shared pool instead.
ulong[] orderArray = ArrayPool<ulong>.Shared.Rent(batchSize);
// Avoid giant stack allocations that can trigger StackOverflow when batchSize is large.
// Rent a reusable array from the shared pool instead.
ArrayPool<ulong> ulongPool = ThreadStaticPools.UlongPool;
ulong[] orderArray = ulongPool.Rent(batchSize);
UInt128 remaining;
int currentSize;
int i;
@@ -120,7 +121,7 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
kStart += (UInt128)currentSize;
}
ArrayPool<ulong>.Shared.Return(orderArray);
ulongPool.Return(orderArray);
orderBuffer.Dispose();
gpuLease.Dispose();
}
@@ -138,11 +138,13 @@ public class MersenneNumberLucasLehmerGpuTester
var kernel = GetBatchKernel(accelerator);
var expBuffer = accelerator.Allocate1D<ulong>(count);
ulong[] expArray = ArrayPool<ulong>.Shared.Rent(count);
ArrayPool<ulong> ulongPool = ThreadStaticPools.UlongPool;
ulong[] expArray = ulongPool.Rent(count);
exponents.CopyTo(expArray);
expBuffer.View.CopyFromCPU(stream, ref expArray[0], count);
GpuUInt128[] modulusArray = ArrayPool<GpuUInt128>.Shared.Rent(count);
ArrayPool<GpuUInt128> gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool;
GpuUInt128[] modulusArray = gpuUInt128Pool.Rent(count);
for (int i = 0; i < count; i++)
{
// TODO: Replace this per-exponent shift with a small shared table (one entry per supported exponent < 128)
@@ -161,8 +163,8 @@ public class MersenneNumberLucasLehmerGpuTester
stateBuffer.Dispose();
modBuffer.Dispose();
expBuffer.Dispose();
ArrayPool<GpuUInt128>.Shared.Return(modulusArray);
ArrayPool<ulong>.Shared.Return(expArray);
gpuUInt128Pool.Return(modulusArray);
ulongPool.Return(expArray);
stream.Dispose();
gpu.Dispose();
}
@@ -175,13 +177,14 @@ public class MersenneNumberLucasLehmerGpuTester
throw new ArgumentException("Result span too small", nameof(results));
}
GpuUInt128[] buffer = ArrayPool<GpuUInt128>.Shared.Rent(count);
ArrayPool<GpuUInt128> gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool;
GpuUInt128[] buffer = gpuUInt128Pool.Rent(count);
ComputeResidues(exponents, buffer.AsSpan(0, count));
for (int i = 0; i < count; i++)
{
results[i] = buffer[i].IsZero;
}
ArrayPool<GpuUInt128>.Shared.Return(buffer);
gpuUInt128Pool.Return(buffer);
}
private bool IsMersennePrimeNtt(ulong exponent, GpuUInt128 nttMod, GpuUInt128 primitiveRoot, bool runOnGpu)
@@ -172,7 +172,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
{
Capacity = Math.Max(1, capacity);
OrderBuffer = accelerator.Allocate1D<ulong>(Capacity);
OrderArray = ArrayPool<ulong>.Shared.Rent(Capacity);
OrderArray = ThreadStaticPools.UlongPool.Rent(Capacity);
}
internal MemoryBuffer1D<ulong, Stride1D.Dense> OrderBuffer { get; }
@@ -184,7 +184,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
public void Dispose()
{
OrderBuffer.Dispose();
ArrayPool<ulong>.Shared.Return(OrderArray, clearArray: false);
ThreadStaticPools.UlongPool.Return(OrderArray, clearArray: false);
}
}
+7 -7
View File
@@ -115,7 +115,7 @@ public static class NttGpuMath
// Precompute stage twiddle factors (forward and inverse) and upload to GPU.
Bits = (int)Math.Log2(length);
int twiddleCount = length - 1;
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var forward = pool.Rent(twiddleCount);
var inverse = pool.Rent(twiddleCount);
StageOffsets = new int[Bits];
@@ -507,7 +507,7 @@ public static class NttGpuMath
var gpu = GpuContextPool.Rent();
var accelerator = gpu.Accelerator;
var kernel = GetBitReverseKernel(accelerator);
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var array = pool.Rent(n);
values.CopyTo(array);
var buffer = accelerator.Allocate1D<GpuUInt128>(n);
@@ -572,7 +572,7 @@ public static class NttGpuMath
var accelerator = gpu.Accelerator;
var stream = accelerator.CreateStream();
var kernel = GetForwardKernel(accelerator);
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var inputArray = pool.Rent(n);
var outputArray = pool.Rent(n);
values.CopyTo(inputArray);
@@ -671,7 +671,7 @@ public static class NttGpuMath
var accelerator = gpu.Accelerator;
var stream = accelerator.CreateStream();
var kernel = GetInverseKernel(accelerator);
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var inputArray = pool.Rent(n);
var outputArray = pool.Rent(n);
values.CopyTo(inputArray);
@@ -857,7 +857,7 @@ public static class NttGpuMath
n <<= 1;
}
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var aBuffer = pool.Rent(n);
var bBuffer = pool.Rent(n);
var a = aBuffer.AsSpan(0, n);
@@ -894,7 +894,7 @@ public static class NttGpuMath
n <<= 1;
}
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var aBuffer = pool.Rent(n);
var bBuffer = pool.Rent(n);
var a = aBuffer.AsSpan(0, n);
@@ -943,7 +943,7 @@ public static class NttGpuMath
var accelerator = gpu.Accelerator;
var stream = accelerator.CreateStream();
var kernel = GetMulKernel(accelerator);
var pool = ArrayPool<GpuUInt128>.Shared;
var pool = ThreadStaticPools.GpuUInt128Pool;
var leftArray = pool.Rent(left.Length);
var rightArray = pool.Rent(right.Length);
left.CopyTo(leftArray);
@@ -582,6 +582,7 @@ internal static partial class PrimeOrderGpuHeuristics
GpuUInt128[]? rentedExponents = null;
GpuUInt128[]? rentedResults = null;
var lease = GpuKernelPool.GetKernel();
ArrayPool<GpuUInt128> gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool;
try
{
@@ -594,9 +595,9 @@ internal static partial class PrimeOrderGpuHeuristics
var exponentBuffer = accelerator.Allocate1D<GpuUInt128>(length);
var remainderBuffer = accelerator.Allocate1D<GpuUInt128>(length);
Span<GpuUInt128> exponentSpan = length <= WideStackThreshold
Span<GpuUInt128> exponentSpan = length <= WideStackThreshold
? stackalloc GpuUInt128[length]
: new Span<GpuUInt128>(rentedExponents = ArrayPool<GpuUInt128>.Shared.Rent(length), 0, length);
: new Span<GpuUInt128>(rentedExponents = gpuUInt128Pool.Rent(length), 0, length);
for (int i = 0; i < length; i++)
{
@@ -613,7 +614,7 @@ internal static partial class PrimeOrderGpuHeuristics
Span<GpuUInt128> resultSpan = length <= WideStackThreshold
? stackalloc GpuUInt128[length]
: new Span<GpuUInt128>(rentedResults = ArrayPool<GpuUInt128>.Shared.Rent(length), 0, length);
: new Span<GpuUInt128>(rentedResults = gpuUInt128Pool.Rent(length), 0, length);
remainderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(resultSpan), length);
stream.Synchronize();
@@ -636,12 +637,8 @@ internal static partial class PrimeOrderGpuHeuristics
{
if (rentedExponents is not null)
{
ArrayPool<GpuUInt128>.Shared.Return(rentedExponents, clearArray: false);
}
if (rentedResults is not null)
{
ArrayPool<GpuUInt128>.Shared.Return(rentedResults, clearArray: false);
gpuUInt128Pool.Return(rentedExponents, clearArray: false);
gpuUInt128Pool.Return(rentedResults!, clearArray: false);
}
lease.Dispose();