ArrayPool<T> replaced with ThreadStaticPools
This commit is contained in:
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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in new issue
Block a user