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

+14 -14
View File
@@ -95,7 +95,7 @@ public sealed class DivisorCycleCache
int[]? rentedMissing = null;
Span<int> missingBuffer = length <= StackBufferThreshold
? stackalloc int[length]
: new Span<int>(rentedMissing = ArrayPool<int>.Shared.Rent(length), 0, length);
: new Span<int>(rentedMissing = ThreadStaticPools.IntPool.Rent(length), 0, length);
int missingCount = 0;
@@ -139,7 +139,7 @@ public sealed class DivisorCycleCache
if (rentedMissing is not null)
{
ArrayPool<int>.Shared.Return(rentedMissing, clearArray: false);
ThreadStaticPools.IntPool.Return(rentedMissing, clearArray: false);
}
}
@@ -198,12 +198,12 @@ public sealed class DivisorCycleCache
ulong[]? rentedDivisors = null;
Span<ulong> workDivisors = count <= StackBufferThreshold
? stackalloc ulong[count]
: new Span<ulong>(rentedDivisors = ArrayPool<ulong>.Shared.Rent(count), 0, count);
: new Span<ulong>(rentedDivisors = ThreadStaticPools.UlongPool.Rent(count), 0, count);
ulong[]? rentedResults = null;
Span<ulong> workResults = count <= StackBufferThreshold
? stackalloc ulong[count]
: new Span<ulong>(rentedResults = ArrayPool<ulong>.Shared.Rent(count), 0, count);
: new Span<ulong>(rentedResults = ThreadStaticPools.UlongPool.Rent(count), 0, count);
for (int i = 0; i < count; i++)
{
@@ -219,8 +219,8 @@ public sealed class DivisorCycleCache
if (rentedDivisors is not null)
{
ArrayPool<ulong>.Shared.Return(rentedDivisors, clearArray: false);
ArrayPool<ulong>.Shared.Return(rentedResults!, clearArray: false);
ThreadStaticPools.UlongPool.Return(rentedDivisors, clearArray: false);
ThreadStaticPools.UlongPool.Return(rentedResults!, clearArray: false);
}
}
@@ -242,22 +242,22 @@ public sealed class DivisorCycleCache
ulong[]? rentedPow = null;
Span<ulong> powSpan = length <= StackBufferThreshold
? stackalloc ulong[length]
: new Span<ulong>(rentedPow = ArrayPool<ulong>.Shared.Rent(length), 0, length);
: new Span<ulong>(rentedPow = ThreadStaticPools.UlongPool.Rent(length), 0, length);
ulong[]? rentedOrder = null;
Span<ulong> orderSpan = length <= StackBufferThreshold
? stackalloc ulong[length]
: new Span<ulong>(rentedOrder = ArrayPool<ulong>.Shared.Rent(length), 0, length);
: new Span<ulong>(rentedOrder = ThreadStaticPools.UlongPool.Rent(length), 0, length);
ulong[]? rentedResult = null;
Span<ulong> resultSpan = length <= StackBufferThreshold
? stackalloc ulong[length]
: new Span<ulong>(rentedResult = ArrayPool<ulong>.Shared.Rent(length), 0, length);
: new Span<ulong>(rentedResult = ThreadStaticPools.UlongPool.Rent(length), 0, length);
byte[]? rentedStatus = null;
Span<byte> statusSpan = length <= StackBufferThreshold
? stackalloc byte[length]
: new Span<byte>(rentedStatus = ArrayPool<byte>.Shared.Rent(length), 0, length);
: new Span<byte>(rentedStatus = ThreadStaticPools.BytePool.Rent(length), 0, length);
ref ulong divisorRef = ref MemoryMarshal.GetReference(divisors);
divisorBuffer.View.CopyFromCPU(stream, ref divisorRef, length);
@@ -306,10 +306,10 @@ public sealed class DivisorCycleCache
resultSpan.CopyTo(destination);
if (rentedPow is not null)
{
ArrayPool<ulong>.Shared.Return(rentedPow, clearArray: false);
ArrayPool<ulong>.Shared.Return(rentedOrder!, clearArray: false);
ArrayPool<ulong>.Shared.Return(rentedResult!, clearArray: false);
ArrayPool<byte>.Shared.Return(rentedStatus!, clearArray: false);
ThreadStaticPools.UlongPool.Return(rentedPow, clearArray: false);
ThreadStaticPools.UlongPool.Return(rentedOrder!, clearArray: false);
ThreadStaticPools.UlongPool.Return(rentedResult!, clearArray: false);
ThreadStaticPools.BytePool.Return(rentedStatus!, clearArray: false);
}
gpuLease.Dispose();
@@ -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();
+1 -1
View File
@@ -797,7 +797,7 @@ public class MersenneDivisorCycles
ulong start = 3UL;
var locker = new object();
var pool = ArrayPool<ulong>.Shared;
var pool = ThreadStaticPools.UlongPool;
ulong batchSizeUL = (ulong)batchSize, d, end;
int count = (int)Math.Min(batchSizeUL, maxDivisor), i, idx;
ulong[] divisors, outCycles, validDivisors;
+9 -7
View File
@@ -101,7 +101,8 @@ public sealed class MersenneNumberTester(
return;
}
UInt128[] qsBuffer = ArrayPool<UInt128>.Shared.Rent((int)limit);
ArrayPool<UInt128> uInt128Pool = ThreadStaticPools.UInt128Pool;
UInt128[] qsBuffer = uInt128Pool.Rent((int)limit);
int idx = 0;
UInt128 k2 = 1UL;
UInt128 q2;
@@ -140,7 +141,7 @@ public sealed class MersenneNumberTester(
if (idx == 0)
{
ArrayPool<UInt128>.Shared.Return(qsBuffer);
ThreadStaticPools.UInt128Pool.Return(qsBuffer);
return;
}
@@ -156,20 +157,21 @@ public sealed class MersenneNumberTester(
int batchSize = GpuConstants.ScanBatchSize;
ulong divMul = (ulong)((((UInt128)1 << 64) - 1UL) / exponent) + 1UL;
int offset = 0;
var gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool;
while (offset < idx)
{
int count = Math.Min(batchSize, idx - offset);
var qBuffer = accelerator.Allocate1D<PerfectNumbers.Core.Gpu.GpuUInt128>(count);
var qBuffer = accelerator.Allocate1D<GpuUInt128>(count);
var orderBuffer = accelerator.Allocate1D<ulong>(count);
// Convert to device-friendly GpuUInt128 on the fly
var tmp = System.Buffers.ArrayPool<PerfectNumbers.Core.Gpu.GpuUInt128>.Shared.Rent(count);
var tmp = gpuUInt128Pool.Rent(count);
for (int i = 0; i < count; i++)
{
tmp[i] = (PerfectNumbers.Core.Gpu.GpuUInt128)qs[offset + i];
tmp[i] = (GpuUInt128)qs[offset + i];
}
// Copy only the populated elements to the device buffer to avoid overrun on pooled arrays.
qBuffer.View.CopyFromCPU(stream, ref tmp[0], count);
System.Buffers.ArrayPool<PerfectNumbers.Core.Gpu.GpuUInt128>.Shared.Return(tmp);
gpuUInt128Pool.Return(tmp);
orderKernel(stream, count, exponent, divMul, qBuffer.View, orderBuffer.View);
stream.Synchronize();
@@ -195,7 +197,7 @@ public sealed class MersenneNumberTester(
}
gpuLease.Dispose();
ArrayPool<UInt128>.Shared.Return(qs);
uInt128Pool.Return(qs);
}
public bool IsMersennePrime(ulong exponent)
@@ -3,6 +3,7 @@ using System.Runtime.CompilerServices;
using System.Numerics;
using UInt128 = System.UInt128;
using ILGPU;
using System.Buffers;
namespace PerfectNumbers.Core;
@@ -75,6 +76,7 @@ public static class MersennePrimeFactorTester
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static UInt128 OrderOf2ModPrime(UInt128 q, CancellationToken ct)
{
lock (_orderCache)
{
// TODO: Replace this lock with the lock-free order cache once the divisor-cycle snapshot exposes
@@ -86,6 +88,8 @@ public static class MersennePrimeFactorTester
}
}
ArrayPool<UInt128> uInt128Pool = ThreadStaticPools.UInt128Pool;
UInt128 phi = q - 1;
UInt128[] factors = phi <= ulong.MaxValue
? Array.ConvertAll(Factor64((ulong)phi, ct), x => (UInt128)x)
@@ -101,28 +105,32 @@ public static class MersennePrimeFactorTester
prime = factors[i];
if (ct.IsCancellationRequested)
{
uInt128Pool.Return(factors);
return 0;
}
// TODO: Swap this `%`/`/` loop for the Math.DivRem-based reducer highlighted in the
// order benchmarks so we avoid repeated 128-bit division while trimming the order.
while (order % prime == 0)
{
if (ct.IsCancellationRequested)
{
return 0;
}
// TODO: Swap this `%`/`/` loop for the Math.DivRem-based reducer highlighted in the
// order benchmarks so we avoid repeated 128-bit division while trimming the order.
while (order % prime == 0)
{
if (ct.IsCancellationRequested)
{
uInt128Pool.Return(factors);
return 0;
}
candidate = order / prime;
if (ModPowWithCancellation(UInt128Numbers.Two, candidate, q, ct) == UInt128.One)
{
order = candidate;
}
else
{
break;
}
}
candidate = order / prime;
if (ModPowWithCancellation(UInt128Numbers.Two, candidate, q, ct) == UInt128.One)
{
order = candidate;
}
else
{
break;
}
}
uInt128Pool.Return(factors);
}
// TODO: Integrate divisor-cycle data here so repeated order refinement uses cached cycle
@@ -136,6 +144,7 @@ public static class MersennePrimeFactorTester
}
}
uInt128Pool.Return(factors);
return order;
}
@@ -327,9 +336,7 @@ public static class MersennePrimeFactorTester
queue.Enqueue(m / d);
}
var result = new UInt128[dict.Count]; // TODO: Rent this array from ArrayPool<UInt128> so
// wide-factorization batches stop allocating fresh
// buffers per candidate.
var result = ThreadStaticPools.UInt128Pool.Rent(dict.Count);
int index = 0;
foreach (var key in dict.Keys)
{
+2 -2
View File
@@ -50,12 +50,12 @@ namespace PerfectNumbers.Core
FactorEntry[]? source = Factors;
if (source is null || Count == 0)
{
FactorEntry[] local = ArrayPool<FactorEntry>.Shared.Rent(1);
FactorEntry[] local = ThreadStaticPools.FactorEntryPool.Rent(1);
local[0] = new FactorEntry(prime, 1, true);
return Rent(local, 1UL, true, 1, false);
}
FactorEntry[] extended = ArrayPool<FactorEntry>.Shared.Rent(Count + 1);
FactorEntry[] extended = ThreadStaticPools.FactorEntryPool.Rent(Count + 1);
Array.Copy(source, 0, extended, 0, Count);
extended[Count] = new FactorEntry(prime, 1, true);
Span<FactorEntry> span = extended.AsSpan(0, Count + 1);
+3 -2
View File
@@ -42,7 +42,8 @@ public sealed class PrimeCache
index = ~index;
}
EInteger[] primesArray = ArrayPool<EInteger>.Shared.Rent(_primes.Count);
ArrayPool<EInteger> eIntegerPool = ThreadStaticPools.EIntegerPool;
EInteger[] primesArray = eIntegerPool.Rent(_primes.Count);
try
{
Span<EInteger> primes = new(primesArray);
@@ -64,7 +65,7 @@ public sealed class PrimeCache
}
finally
{
ArrayPool<EInteger>.Shared.Return(primesArray);
eIntegerPool.Return(primesArray);
}
}
+37
View File
@@ -3,11 +3,24 @@ using System.Collections.Generic;
using System.Runtime.CompilerServices;
using PerfectNumbers.Core.Cpu;
using PerfectNumbers.Core.Gpu;
using PeterO.Numbers;
namespace PerfectNumbers.Core
{
public readonly struct ThreadStaticPools
{
[ThreadStatic]
private static ArrayPool<EInteger>? _eintegerPool;
public static ArrayPool<EInteger> EIntegerPool
{
[MethodImpl(MethodImplOptions.AggressiveInlining)]
get
{
return _eintegerPool ??= ArrayPool<EInteger>.Create();
}
}
[ThreadStatic]
private static ArrayPool<FactorEntry>? _factorEntryPool;
@@ -32,6 +45,18 @@ namespace PerfectNumbers.Core
}
}
[ThreadStatic]
private static ArrayPool<GpuUInt128>? _gpuUInt128Pool;
public static ArrayPool<GpuUInt128> GpuUInt128Pool
{
[MethodImpl(MethodImplOptions.AggressiveInlining)]
get
{
return _gpuUInt128Pool ??= ArrayPool<GpuUInt128>.Create();
}
}
[ThreadStatic]
private static MersenneCpuDivisorScanSession? _mersenneCpuDivisorSession;
@@ -55,6 +80,18 @@ namespace PerfectNumbers.Core
_mersenneCpuDivisorSession = session;
}
[ThreadStatic]
private static ArrayPool<UInt128>? _uint128Pool;
public static ArrayPool<UInt128> UInt128Pool
{
[MethodImpl(MethodImplOptions.AggressiveInlining)]
get
{
return _uint128Pool ??= ArrayPool<UInt128>.Create();
}
}
[ThreadStatic]
private static ArrayPool<ulong>? _ulongPool;