From d0b1639a5dcc73429ac4a9a046a3ea0deaed03a1 Mon Sep 17 00:00:00 2001 From: mlemanczyk Date: Sun, 9 Nov 2025 21:52:16 +0100 Subject: [PATCH] GPU stabilization improvements to prevent calculation crashes --- .../Pow2ModBenchmarks.cs | 2 +- .../Pow2MontgomeryModBenchmarks.cs | 2 +- ...yModWindowedKeepMontgomeryGpuBenchmarks.cs | 4 +- .../PrimeTesterBenchmarks.cs | 10 +- EvenPerfectBitScanner/Program.cs | 6 +- .../Gpu/GpuKernelCompilationTests.cs | 38 +- .../Gpu/NttGpuMathTests.cs | 6 +- .../Gpu/NttMontgomeryTests.cs | 5 +- PerfectNumbers.Core/DivisorCycleCache.cs | 24 +- PerfectNumbers.Core/Gpu/AcceleratorPool.cs | 33 + .../Accelerators/Pow2MontgomeryAccelerator.cs | 74 ++ .../Gpu/ExponentRemainderStepperGpu.cs | 17 +- PerfectNumbers.Core/Gpu/GpuKernelPool.cs | 153 +-- .../Gpu/GpuScratchBufferPool.cs | 41 +- .../Gpu/GpuStaticTableInitializer.cs | 20 +- .../Gpu/HeuristicGpuDivisorTables.cs | 253 ++-- .../Gpu/InputOutputScratchBuffers.cs | 126 ++ .../Gpu/Kernels/NttMontgomeryKernels.cs | 2 +- .../Gpu/Kernels/Pow2MontgomeryKernels.cs | 3 +- .../Gpu/Kernels/PrimeOrderKernels.cs | 24 +- ...MersenneNumberDivisorByDivisorGpuTester.cs | 87 +- .../Gpu/MersenneNumberDivisorGpuTester.cs | 57 +- .../Gpu/MersenneNumberIncrementalGpuTester.cs | 8 +- .../Gpu/MersenneNumberLucasLehmerGpuTester.cs | 201 +-- .../Gpu/MersenneNumberOrderGpuTester.cs | 14 +- .../Gpu/MersenneNumberResidueGpuTester.cs | 301 +++-- PerfectNumbers.Core/Gpu/NttGpuMath.cs | 515 ++++---- .../Gpu/Pow2MontgomeryGpuCalculator.cs | 83 ++ .../Gpu/PrimeOrderGpuHeuristics.cs | 1102 +++++++++-------- PerfectNumbers.Core/Gpu/SharedGpuContext.cs | 27 +- .../HeuristicCombinedPrimeTester.cs | 186 ++- PerfectNumbers.Core/HeuristicPrimeTester.cs | 2 +- PerfectNumbers.Core/MersenneDivisorCycles.cs | 47 +- .../MersenneNumberDivisorByDivisorTester.cs | 2 +- PerfectNumbers.Core/MersenneNumberTester.cs | 25 +- PerfectNumbers.Core/PerfectNumberConstants.cs | 5 +- .../PrimeOrderCalculator.Cpu.cs | 50 +- .../PrimeOrderCalculator.Gpu.cs | 28 +- PerfectNumbers.Core/PrimeTester.cs | 367 +++--- PerfectNumbers.Core/ULongExtensions.Gpu.cs | 132 +- 40 files changed, 2318 insertions(+), 1764 deletions(-) create mode 100644 PerfectNumbers.Core/Gpu/AcceleratorPool.cs create mode 100644 PerfectNumbers.Core/Gpu/Accelerators/Pow2MontgomeryAccelerator.cs create mode 100644 PerfectNumbers.Core/Gpu/InputOutputScratchBuffers.cs create mode 100644 PerfectNumbers.Core/Gpu/Pow2MontgomeryGpuCalculator.cs diff --git a/EvenPerfectBitScanner.Benchmarks/Pow2ModBenchmarks.cs b/EvenPerfectBitScanner.Benchmarks/Pow2ModBenchmarks.cs index 5bd59d7..e515ecb 100644 --- a/EvenPerfectBitScanner.Benchmarks/Pow2ModBenchmarks.cs +++ b/EvenPerfectBitScanner.Benchmarks/Pow2ModBenchmarks.cs @@ -101,7 +101,7 @@ public class Pow2ModBenchmarks for (int i = 0; i < SampleCount; i++) { - checksum ^= exponents[i].Pow2MontgomeryModWindowedGpu(divisors[i], keepMontgomery: false); + checksum ^= exponents[i].Pow2MontgomeryModWindowedConvertGpu(divisors[i]); } return checksum; diff --git a/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModBenchmarks.cs b/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModBenchmarks.cs index 37946e8..b9323ef 100644 --- a/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModBenchmarks.cs +++ b/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModBenchmarks.cs @@ -77,7 +77,7 @@ public class Pow2MontgomeryModBenchmarks for (int i = 0; i < SampleCount; i++) { - checksum ^= exponents[i].Pow2MontgomeryModWindowedGpu(divisors[i], keepMontgomery: false); + checksum ^= exponents[i].Pow2MontgomeryModWindowedConvertGpu(divisors[i]); } return checksum; diff --git a/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModWindowedKeepMontgomeryGpuBenchmarks.cs b/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModWindowedKeepMontgomeryGpuBenchmarks.cs index 781c33b..38b1d5b 100644 --- a/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModWindowedKeepMontgomeryGpuBenchmarks.cs +++ b/EvenPerfectBitScanner.Benchmarks/Pow2MontgomeryModWindowedKeepMontgomeryGpuBenchmarks.cs @@ -34,7 +34,7 @@ public class Pow2MontgomeryModWindowedKeepMontgomeryGpuBenchmarks for (int i = 0; i < cases.Length; i++) { ref readonly Pow2MontgomeryModWindowedBenchmarkCase current = ref cases[i]; - checksum ^= current.Exponent.Pow2MontgomeryModWindowedGpu(current.Divisor, keepMontgomery: false); + checksum ^= current.Exponent.Pow2MontgomeryModWindowedConvertGpu(current.Divisor); } return checksum; @@ -52,7 +52,7 @@ public class Pow2MontgomeryModWindowedKeepMontgomeryGpuBenchmarks for (int i = 0; i < cases.Length; i++) { ref readonly Pow2MontgomeryModWindowedBenchmarkCase current = ref cases[i]; - checksum ^= current.Exponent.Pow2MontgomeryModWindowedGpu(current.Divisor, keepMontgomery: true); + checksum ^= current.Exponent.Pow2MontgomeryModWindowedKeepGpu(current.Divisor); } return checksum; diff --git a/EvenPerfectBitScanner.Benchmarks/PrimeTesterBenchmarks.cs b/EvenPerfectBitScanner.Benchmarks/PrimeTesterBenchmarks.cs index 0cf9227..8cbc69c 100644 --- a/EvenPerfectBitScanner.Benchmarks/PrimeTesterBenchmarks.cs +++ b/EvenPerfectBitScanner.Benchmarks/PrimeTesterBenchmarks.cs @@ -30,7 +30,7 @@ public class PrimeTesterBenchmarks _ = Prime.Numbers.IsPrime(_candidates[0]); } - [Benchmark] + // [Benchmark] public int HeuristicCpu() { HeuristicPrimeTester tester = _heuristicTester; @@ -48,7 +48,7 @@ public class PrimeTesterBenchmarks return primeCount; } - [Benchmark] + // [Benchmark] public int HeuristicGpu() { HeuristicPrimeTester tester = _heuristicTester; @@ -66,7 +66,7 @@ public class PrimeTesterBenchmarks return primeCount; } - [Benchmark] + // [Benchmark] public int HeuristicCombinedGpu() { ulong[] values = _candidates; @@ -83,7 +83,7 @@ public class PrimeTesterBenchmarks return primeCount; } - [Benchmark] + // [Benchmark] public int HeuristicCombinedCpu() { ulong[] values = _candidates; @@ -134,7 +134,7 @@ public class PrimeTesterBenchmarks return primeCount; } - [Benchmark] + // [Benchmark] public int OpenNumericPrimes() { ulong[] values = _candidates; diff --git a/EvenPerfectBitScanner/Program.cs b/EvenPerfectBitScanner/Program.cs index 26a5b9d..0666824 100644 --- a/EvenPerfectBitScanner/Program.cs +++ b/EvenPerfectBitScanner/Program.cs @@ -58,12 +58,14 @@ internal static class Program bool startPrimeProvided = _cliArguments.StartPrimeProvided; int threadCount = Math.Max(1, _cliArguments.ThreadCount); int gpuPrimeBatch = Math.Max(1, _cliArguments.GpuPrimeBatch); + // int gpuPrimeThreads = SharedGpuContext.Device.MaxNumThreads; int gpuPrimeThreads = Math.Max(1, _cliArguments.GpuPrimeThreads); UnboundedTaskScheduler.ConfigureThreadCount(threadCount); PrimeTester.GpuBatchSize = gpuPrimeBatch; GpuPrimeWorkLimiter.SetLimit(gpuPrimeThreads); - PrimeTester.WarmUpGpuKernels(gpuPrimeThreads); - GpuScratchBufferPool.WarmUp(gpuPrimeThreads << 4, 64, 1024); + PrimeTester.WarmUpGpuKernels(1024); + // PrimeTester.WarmUpGpuKernels(gpuPrimeThreads >> 4); + GpuScratchBufferPool.WarmUp(gpuPrimeThreads >> 4, 64, 128); Console.WriteLine("Starting up threads..."); _ = UnboundedTaskScheduler.Instance; int blockSize = Math.Max(1, _cliArguments.BlockSize); diff --git a/PerfectNumbers.Core.Tests/Gpu/GpuKernelCompilationTests.cs b/PerfectNumbers.Core.Tests/Gpu/GpuKernelCompilationTests.cs index 839a0fb..919ff7e 100644 --- a/PerfectNumbers.Core.Tests/Gpu/GpuKernelCompilationTests.cs +++ b/PerfectNumbers.Core.Tests/Gpu/GpuKernelCompilationTests.cs @@ -348,32 +348,40 @@ public class GpuKernelCompilationTests private static void CompileGpuKernelPoolKernels(Accelerator accelerator) { - _ = accelerator; - GpuPrimeWorkLimiter.Acquire(); + + GpuPrimeWorkLimiter.Acquire(); + AcceleratorStream stream = accelerator.CreateStream(); try - { - _ = GpuKernelPool.Kernels.Order; - _ = GpuKernelPool.Kernels.Incremental; - _ = GpuKernelPool.Kernels.Pow2Mod; - _ = GpuKernelPool.Kernels.IncrementalOrder; - _ = GpuKernelPool.Kernels.Pow2ModOrder; + { + var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); + _ = kernels.Order; + _ = kernels.Incremental; + _ = kernels.Pow2Mod; + _ = kernels.IncrementalOrder; + _ = kernels.Pow2ModOrder; } finally - { + { + stream.Synchronize(); + stream.Dispose(); GpuPrimeWorkLimiter.Release(); } - GpuPrimeWorkLimiter.Acquire(); + stream = accelerator.CreateStream(); + GpuPrimeWorkLimiter.Acquire(); try { - _ = GpuKernelPool.Kernels.Order; - _ = GpuKernelPool.Kernels.Incremental; - _ = GpuKernelPool.Kernels.Pow2Mod; - _ = GpuKernelPool.Kernels.IncrementalOrder; - _ = GpuKernelPool.Kernels.Pow2ModOrder; + var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); + _ = kernels.Order; + _ = kernels.Incremental; + _ = kernels.Pow2Mod; + _ = kernels.IncrementalOrder; + _ = kernels.Pow2ModOrder; } finally { + stream.Synchronize(); + stream.Dispose(); GpuPrimeWorkLimiter.Release(); } } diff --git a/PerfectNumbers.Core.Tests/Gpu/NttGpuMathTests.cs b/PerfectNumbers.Core.Tests/Gpu/NttGpuMathTests.cs index 3e83968..234e7db 100644 --- a/PerfectNumbers.Core.Tests/Gpu/NttGpuMathTests.cs +++ b/PerfectNumbers.Core.Tests/Gpu/NttGpuMathTests.cs @@ -150,7 +150,7 @@ public class NttGpuMathTests var modulus = new GpuUInt128(0UL, 17UL); var primitiveRoot = new GpuUInt128(0UL, 3UL); - var accelerator = SharedGpuContext.Accelerator; + var accelerator = SharedGpuContext.CreateAccelerator(); var stream = accelerator.CreateStream(); var method = typeof(NttGpuMath).GetMethod("GetSquareCache", BindingFlags.NonPublic | BindingFlags.Static); @@ -170,6 +170,7 @@ public class NttGpuMathTests twiddlesInvMont.CopyToCPU(stream, inverse); stream.Synchronize(); stream.Dispose(); + accelerator.Dispose(); for (int i = 0; i < count; i++) { @@ -340,7 +341,7 @@ public class NttGpuMathTests values[1] = new GpuUInt128(0UL, 2UL); var modulus = new GpuUInt128(0UL, 17UL); var primitiveRoot = new GpuUInt128(0UL, 3UL); - Accelerator accelerator = SharedGpuContext.Accelerator; + Accelerator accelerator = SharedGpuContext.CreateAccelerator(); var stream = accelerator.CreateStream(); var buffer = accelerator.Allocate1D(values.Length); buffer.View.CopyFromCPU(stream, ref values[0], values.Length); @@ -348,6 +349,7 @@ public class NttGpuMathTests stream.Synchronize(); stream.Dispose(); buffer.Dispose(); + accelerator.Dispose(); // Verify that the square cache now holds entries for the accelerator. var field = typeof(NttGpuMath).GetField("SquareCache", BindingFlags.NonPublic | BindingFlags.Static)!; diff --git a/PerfectNumbers.Core.Tests/Gpu/NttMontgomeryTests.cs b/PerfectNumbers.Core.Tests/Gpu/NttMontgomeryTests.cs index 2348296..3e21079 100644 --- a/PerfectNumbers.Core.Tests/Gpu/NttMontgomeryTests.cs +++ b/PerfectNumbers.Core.Tests/Gpu/NttMontgomeryTests.cs @@ -33,20 +33,21 @@ public class NttMontgomeryTests NttGpuMath.Square(cpu, modulus, primitiveRoot); // Device execution - var accelerator = SharedGpuContext.Accelerator; + var accelerator = SharedGpuContext.CreateAccelerator(); var stream = accelerator.CreateStream(); using var buffer = accelerator.Allocate1D(gpu.Length); buffer.View.CopyFromCPU(stream, ref gpu[0], gpu.Length); NttGpuMath.SquareDevice(accelerator, stream, buffer.View, modulus, primitiveRoot); buffer.View.CopyToCPU(stream, ref gpu[0], gpu.Length); stream.Synchronize(); + stream.Dispose(); + accelerator.Dispose(); for (int i = 0; i < cpu.Length; i++) { gpu[i].Should().Be(cpu[i]); } - stream.Dispose(); } [Fact] diff --git a/PerfectNumbers.Core/DivisorCycleCache.cs b/PerfectNumbers.Core/DivisorCycleCache.cs index bba7044..0077181 100644 --- a/PerfectNumbers.Core/DivisorCycleCache.cs +++ b/PerfectNumbers.Core/DivisorCycleCache.cs @@ -15,8 +15,8 @@ public sealed class DivisorCycleCache private const byte ByteZero = 0; private const byte ByteOne = 1; private const int StackBufferThreshold = 128; - - private readonly Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> _gpuKernel = LoadKernel(SharedGpuContext.Accelerator); + private readonly Accelerator _accelerator; + private readonly Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> _gpuKernel; private volatile ulong[] _snapshot; private volatile bool _useGpuGeneration = true; private readonly int _divisorCyclesBatchSize; @@ -40,6 +40,8 @@ public sealed class DivisorCycleCache private DivisorCycleCache(int divisorCyclesBatchSize) { _divisorCyclesBatchSize = Math.Max(1, divisorCyclesBatchSize); + _accelerator = SharedGpuContext.Device.CreateAccelerator(SharedGpuContext.Context); + _gpuKernel = LoadKernel(_accelerator); _snapshot = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot(); } @@ -241,7 +243,7 @@ public sealed class DivisorCycleCache [ThreadStatic] private static MemoryBuffer1D? statusBuffer; - private static void EnsureCapacity(int requiredCapacity) + private void EnsureCapacity(int requiredCapacity) { if (requiredCapacity <= (divisorBuffer?.Length ?? 0)) { @@ -257,11 +259,15 @@ public sealed class DivisorCycleCache statusBuffer!.Dispose(); } - divisorBuffer = SharedGpuContext.Accelerator.Allocate1D(requiredCapacity); - powBuffer = SharedGpuContext.Accelerator.Allocate1D(requiredCapacity); - orderBuffer = SharedGpuContext.Accelerator.Allocate1D(requiredCapacity); - resultBuffer = SharedGpuContext.Accelerator.Allocate1D(requiredCapacity); - statusBuffer = SharedGpuContext.Accelerator.Allocate1D(requiredCapacity); + Accelerator accelerator = _accelerator; + // lock(accelerator) + { + divisorBuffer = accelerator.Allocate1D(requiredCapacity); + powBuffer = accelerator.Allocate1D(requiredCapacity); + orderBuffer = accelerator.Allocate1D(requiredCapacity); + resultBuffer = accelerator.Allocate1D(requiredCapacity); + statusBuffer = accelerator.Allocate1D(requiredCapacity); + } } private void ComputeCyclesGpuCore(ReadOnlySpan divisors, Span destination) @@ -269,7 +275,7 @@ public sealed class DivisorCycleCache int length = divisors.Length; GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; + Accelerator accelerator = _accelerator; var stream = accelerator.CreateStream(); var kernel = _gpuKernel; EnsureCapacity(length); diff --git a/PerfectNumbers.Core/Gpu/AcceleratorPool.cs b/PerfectNumbers.Core/Gpu/AcceleratorPool.cs new file mode 100644 index 0000000..e6b1b5d --- /dev/null +++ b/PerfectNumbers.Core/Gpu/AcceleratorPool.cs @@ -0,0 +1,33 @@ +using ILGPU; +using ILGPU.Runtime; + +namespace PerfectNumbers.Core.Gpu +{ + public struct AcceleratorPool(int capacity) + { + internal readonly Accelerator[] Accelerators = [.. Enumerable.Range(1, capacity).Select(_ => { + return SharedGpuContext.CreateAccelerator(); + })]; + + private int _index; + + public Accelerator Rent() + { + Atomic.CompareExchange(ref _index, capacity, 0); + var index = Atomic.Add(ref _index, 1); + if (index >= capacity) + { + index -= capacity; + } + + return Accelerators[index]; + } + + public void Return(Accelerator accelerator) + { + // Intentionally left empty - there's nothing to do here. We're not really renting anything. + } + + public static readonly AcceleratorPool Shared = new(PerfectNumberConstants.RollingAccelerators); + } +} \ No newline at end of file diff --git a/PerfectNumbers.Core/Gpu/Accelerators/Pow2MontgomeryAccelerator.cs b/PerfectNumbers.Core/Gpu/Accelerators/Pow2MontgomeryAccelerator.cs new file mode 100644 index 0000000..eb3f6bf --- /dev/null +++ b/PerfectNumbers.Core/Gpu/Accelerators/Pow2MontgomeryAccelerator.cs @@ -0,0 +1,74 @@ +using System.Runtime.CompilerServices; +using ILGPU; +using ILGPU.Runtime; +using PerfectNumbers.Core.Gpu.Kernels; + +namespace PerfectNumbers.Core.Gpu.Accelerators; + +public sealed class Pow2MontgomeryAccelerator +{ + #region Pool + + [ThreadStatic] + private static Queue? _pool; + + private static AcceleratorPool _accelerators = new(PerfectNumberConstants.RollingAccelerators); + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Pow2MontgomeryAccelerator Rent() + { + var pool = _pool ??= []; + if (!pool.TryDequeue(out var gpu)) + { + Accelerator accelerator = _accelerators.Rent(); + gpu = new(accelerator); + } + + return gpu; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static void Return(Pow2MontgomeryAccelerator gpu) + { + gpu.Stream!.Dispose(); + gpu.Stream = null; + _pool!.Enqueue(gpu); + } + + #endregion + + public readonly Accelerator Accelerator; + public readonly MemoryBuffer1D Input; + public readonly MemoryBuffer1D Output; + public AcceleratorStream? Stream; + + public readonly Action, MontgomeryDivisorData, ArrayView1D> ConvertToStandardKernel; + public readonly Action, MontgomeryDivisorData, ArrayView1D> KeepMontgomeryKernel; + + public Pow2MontgomeryAccelerator(Accelerator accelerator) + { + Accelerator = accelerator; + Input = accelerator.Allocate1D(1); + Output = accelerator.Allocate1D(1); + Stream = accelerator.CreateStream(); + + var keepKernel = accelerator.LoadAutoGroupedStreamKernel, MontgomeryDivisorData, ArrayView1D>(Pow2MontgomeryKernels.Pow2MontgomeryKernelKeepMontgomery); + var keepLauncher = KernelUtil.GetKernel(keepKernel); + + KeepMontgomeryKernel = keepLauncher.CreateLauncherDelegate, MontgomeryDivisorData, ArrayView1D>>(); + + var convertKernel = accelerator.LoadAutoGroupedStreamKernel, MontgomeryDivisorData, ArrayView1D>(Pow2MontgomeryKernels.Pow2MontgomeryKernelConvertToStandard); + var convertLauncher = KernelUtil.GetKernel(convertKernel); + + ConvertToStandardKernel = convertLauncher.CreateLauncherDelegate, MontgomeryDivisorData, ArrayView1D>>(); + } + + public void Dispose() + { + Input.Dispose(); + Output.Dispose(); + Stream?.Dispose(); + // Accelerator is shared with other threads + // Accelerator.Dispose(); + } +} diff --git a/PerfectNumbers.Core/Gpu/ExponentRemainderStepperGpu.cs b/PerfectNumbers.Core/Gpu/ExponentRemainderStepperGpu.cs index d322406..67261da 100644 --- a/PerfectNumbers.Core/Gpu/ExponentRemainderStepperGpu.cs +++ b/PerfectNumbers.Core/Gpu/ExponentRemainderStepperGpu.cs @@ -24,17 +24,14 @@ internal struct ExponentRemainderStepperGpu _currentResidue = GpuUInt128.One; } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool MatchesDivisor(in GpuDivisorPartialData divisor) - { - return _modulus == divisor.Modulus; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public readonly bool MatchesDivisor(in GpuDivisorPartialData divisor) => _modulus == divisor.Modulus; - /// - /// Clears the cached residue so the stepper can be reused for another divisor. - /// Callers must reinitialize the stepper before consuming any residues after a reset. - /// - [MethodImpl(MethodImplOptions.AggressiveInlining)] + /// + /// Clears the cached residue so the stepper can be reused for another divisor. + /// Callers must reinitialize the stepper before consuming any residues after a reset. + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] public void Reset() { PreviousExponent = 0UL; diff --git a/PerfectNumbers.Core/Gpu/GpuKernelPool.cs b/PerfectNumbers.Core/Gpu/GpuKernelPool.cs index b70ae34..a7cc055 100644 --- a/PerfectNumbers.Core/Gpu/GpuKernelPool.cs +++ b/PerfectNumbers.Core/Gpu/GpuKernelPool.cs @@ -5,7 +5,8 @@ namespace PerfectNumbers.Core.Gpu; public class GpuKernelPool { - public static readonly KernelContainer Kernels = CreateKernels(); + [ThreadStatic] + private static Dictionary? _pool; public static T InitOnce(ref T? target, Func valueFactory) where T : class { @@ -18,14 +19,19 @@ public class GpuKernelPool return Interlocked.CompareExchange(ref target, newValue, null) ?? newValue; } - private static KernelContainer CreateKernels() + internal static KernelContainer GetOrAddKernels(Accelerator accelerator, AcceleratorStream stream) { + var pool = _pool ??= []; + if (pool.TryGetValue(accelerator, out var kernelContainer)) + { + return kernelContainer; + } + var kernels = new KernelContainer(); - var stream = SharedGpuContext.Accelerator.CreateStream(); InitOnce(ref kernels.Order, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView>(OrderKernels.OrderKernelScan); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView>(OrderKernels.OrderKernelScan); var kernel = KernelUtil.GetKernel(loaded); @@ -34,7 +40,7 @@ public class GpuKernelPool InitOnce(ref kernels.Pow2Mod, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(Pow2ModKernels.Pow2ModKernelScan); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(Pow2ModKernels.Pow2ModKernelScan); var kernel = KernelUtil.GetKernel(loaded); @@ -43,7 +49,7 @@ public class GpuKernelPool InitOnce(ref kernels.Incremental, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView1D>(IncrementalKernels.IncrementalKernelScan); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D>(IncrementalKernels.IncrementalKernelScan); var kernel = KernelUtil.GetKernel(loaded); @@ -52,7 +58,7 @@ public class GpuKernelPool InitOnce(ref kernels.IncrementalOrder, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView1D>(IncrementalKernels.IncrementalOrderKernelScan); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D>(IncrementalKernels.IncrementalOrderKernelScan); var kernel = KernelUtil.GetKernel(loaded); @@ -62,7 +68,7 @@ public class GpuKernelPool InitOnce(ref kernels.Pow2ModOrder, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView1D>(Pow2ModKernels.Pow2ModOrderKernelScan); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D>(Pow2ModKernels.Pow2ModOrderKernelScan); var kernel = KernelUtil.GetKernel(loaded); @@ -71,7 +77,7 @@ public class GpuKernelPool InitOnce(ref kernels.SpecialMax, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, int, MontgomeryDivisorData, ArrayView1D, ArrayView1D>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel); + var loaded = accelerator.LoadAutoGroupedStreamKernel, int, MontgomeryDivisorData, ArrayView1D, ArrayView1D>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel); var kernel = KernelUtil.GetKernel(loaded); @@ -80,126 +86,132 @@ public class GpuKernelPool InitOnce(ref kernels.SmallPrimeFactor, () => { - var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, int, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(SmallPrimeFactorKernels.SmallPrimeFactorKernelScan); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, int, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(SmallPrimeFactorKernels.SmallPrimeFactorKernelScan); var kernel = KernelUtil.GetKernel(loaded); return kernel.CreateLauncherDelegate, ArrayView1D, int, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>>(); }); - PreloadStaticTables(kernels, stream); + PreloadStaticTables(accelerator, kernels, stream); - stream.Synchronize(); - stream.Dispose(); + pool[accelerator] = kernels; return kernels; } - private static KernelContainer CreateKernels(AcceleratorStream stream) + internal static void PreloadStaticTables(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream) { - KernelContainer kernels = new(); - PreloadStaticTables(kernels, stream); - return Kernels; + EnsureSmallCyclesOnDevice(accelerator, kernels, stream); + EnsureSmallPrimesOnDevice(accelerator, kernels, stream); + EnsureSmallPrimeFactorTables(accelerator, kernels, stream); } - internal static void PreloadStaticTables(KernelContainer kernels, AcceleratorStream stream) - { - _ = EnsureSmallCyclesOnDevice(kernels, stream); - _ = EnsureSmallPrimesOnDevice(kernels, stream); - _ = EnsureSmallPrimeFactorTables(kernels, stream); - } - - public static ArrayView1D GetSmallCyclesOnDevice() => Kernels.SmallCycles!; + public static ArrayView1D GetSmallCyclesOnDevice(KernelContainer kernels) => kernels.SmallCycles!; // Ensures the small cycles table is uploaded to the device for the given accelerator. // Returns the ArrayView to pass into kernels (when kernels are extended to accept it). - private static ArrayView1D EnsureSmallCyclesOnDevice(KernelContainer kernels, AcceleratorStream stream) + private static void EnsureSmallCyclesOnDevice(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream) { - if (kernels.SmallCycles is { } buffer) + if (kernels.SmallCycles is { }) { - return buffer.View; + return; + } + + var host = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot(); + MemoryBuffer1D? device; + + // lock (accelerator) + { + device = accelerator.Allocate1D(host.Length); } - var host = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot(); // TODO: Preload this device buffer during startup and keep it immutable so we can delete the lock above in favor of the preloaded snapshot. - var device = SharedGpuContext.Accelerator.Allocate1D(host.Length); device.View.CopyFromCPU(stream, host); kernels.SmallCycles = device; - return device.View; } - public static ResiduePrimeViews GetSmallPrimesOnDevice() => new( - Kernels.SmallPrimesLastOne!, - Kernels.SmallPrimesLastSeven!, - Kernels.SmallPrimesPow2LastOne!, - Kernels.SmallPrimesPow2LastSeven! + public static ResiduePrimeViews GetSmallPrimesOnDevice(KernelContainer kernels) => new( + kernels.SmallPrimesLastOne!, + kernels.SmallPrimesLastSeven!, + kernels.SmallPrimesPow2LastOne!, + kernels.SmallPrimesPow2LastSeven! ); - private static ResiduePrimeViews EnsureSmallPrimesOnDevice(KernelContainer kernels, AcceleratorStream stream) + private static void EnsureSmallPrimesOnDevice(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream) { - if (kernels.SmallPrimesLastOne is { } lastOne && - kernels.SmallPrimesLastSeven is { } lastSeven && - kernels.SmallPrimesPow2LastOne is { } lastOnePow2 && - kernels.SmallPrimesPow2LastSeven is { } lastSevenPow2) + if (kernels.SmallPrimesLastOne is { } && + kernels.SmallPrimesLastSeven is { } && + kernels.SmallPrimesPow2LastOne is { } && + kernels.SmallPrimesPow2LastSeven is { }) { - return new ResiduePrimeViews(lastOne.View, lastSeven.View, lastOnePow2.View, lastSevenPow2.View); + return; } - if (kernels.SmallPrimesLastOne is { } existingLastOne && - kernels.SmallPrimesLastSeven is { } existingLastSeven && - kernels.SmallPrimesPow2LastOne is { } existingLastOnePow2 && - kernels.SmallPrimesPow2LastSeven is { } existingLastSevenPow2) + if (kernels.SmallPrimesLastOne is { } && + kernels.SmallPrimesLastSeven is { } && + kernels.SmallPrimesPow2LastOne is { } && + kernels.SmallPrimesPow2LastSeven is { }) { - return new ResiduePrimeViews(existingLastOne.View, existingLastSeven.View, existingLastOnePow2.View, existingLastSevenPow2.View); + return; } - Accelerator accelerator = SharedGpuContext.Accelerator; var hostLastOne = PrimesGenerator.SmallPrimesLastOne; var hostLastSeven = PrimesGenerator.SmallPrimesLastSeven; var hostLastOnePow2 = PrimesGenerator.SmallPrimesPow2LastOne; var hostLastSevenPow2 = PrimesGenerator.SmallPrimesPow2LastSeven; - var deviceLastOne = accelerator.Allocate1D(hostLastOne.Length); - deviceLastOne.View.CopyFromCPU(stream, hostLastOne); - var deviceLastSeven = accelerator.Allocate1D(hostLastSeven.Length); - deviceLastSeven.View.CopyFromCPU(stream, hostLastSeven); - var deviceLastOnePow2 = accelerator.Allocate1D(hostLastOnePow2.Length); - deviceLastOnePow2.View.CopyFromCPU(stream, hostLastOnePow2); - var deviceLastSevenPow2 = accelerator.Allocate1D(hostLastSevenPow2.Length); - deviceLastSevenPow2.View.CopyFromCPU(stream, hostLastSevenPow2); + MemoryBuffer1D? deviceLastOne; + MemoryBuffer1D? deviceLastSeven; + MemoryBuffer1D? deviceLastOnePow2; + MemoryBuffer1D? deviceLastSevenPow2; + // lock (accelerator) + // { + deviceLastOne = accelerator.Allocate1D(hostLastOne.Length); + deviceLastSeven = accelerator.Allocate1D(hostLastSeven.Length); + deviceLastOnePow2 = accelerator.Allocate1D(hostLastOnePow2.Length); + deviceLastSevenPow2 = accelerator.Allocate1D(hostLastSevenPow2.Length); + // } + + deviceLastOne.View.CopyFromCPU(stream, hostLastOne); + deviceLastSeven.View.CopyFromCPU(stream, hostLastSeven); + deviceLastOnePow2.View.CopyFromCPU(stream, hostLastOnePow2); + deviceLastSevenPow2.View.CopyFromCPU(stream, hostLastSevenPow2); kernels.SmallPrimesLastOne = deviceLastOne; kernels.SmallPrimesLastSeven = deviceLastSeven; kernels.SmallPrimesPow2LastOne = deviceLastOnePow2; kernels.SmallPrimesPow2LastSeven = deviceLastSevenPow2; - - return new ResiduePrimeViews(deviceLastOne.View, deviceLastSeven.View, deviceLastOnePow2.View, deviceLastSevenPow2.View); } - public static SmallPrimeFactorViews GetSmallPrimeFactorTables() => new( - Kernels.SmallPrimeFactorsPrimes!, - Kernels.SmallPrimeFactorsSquares! + public static SmallPrimeFactorViews GetSmallPrimeFactorTables(KernelContainer kernels) => new( + kernels.SmallPrimeFactorsPrimes!, + kernels.SmallPrimeFactorsSquares! ); - private static SmallPrimeFactorViews EnsureSmallPrimeFactorTables(KernelContainer kernels, AcceleratorStream stream) + private static void EnsureSmallPrimeFactorTables(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream) { - if (kernels.SmallPrimeFactorsPrimes is { } primeBuffer && kernels.SmallPrimeFactorsSquares is { } squareBuffer) + if (kernels.SmallPrimeFactorsPrimes is { } && kernels.SmallPrimeFactorsSquares is { }) { - return new SmallPrimeFactorViews(primeBuffer, squareBuffer); + return; } - Accelerator accelerator = SharedGpuContext.Accelerator; var hostPrimes = PrimesGenerator.SmallPrimes; var hostSquares = PrimesGenerator.SmallPrimesPow2; int length = hostPrimes.Length; - var devicePrimes = accelerator.Allocate1D(length); + MemoryBuffer1D? devicePrimes; + MemoryBuffer1D? deviceSquares; + + // lock (accelerator) + { + devicePrimes = accelerator.Allocate1D(length); + deviceSquares = accelerator.Allocate1D(length); + } + devicePrimes.View.CopyFromCPU(stream, hostPrimes); - var deviceSquares = accelerator.Allocate1D(length); deviceSquares.View.CopyFromCPU(stream, hostSquares); kernels.SmallPrimeFactorsPrimes = devicePrimes; kernels.SmallPrimeFactorsSquares = deviceSquares; - - return new SmallPrimeFactorViews(devicePrimes, deviceSquares); } /// @@ -209,11 +221,12 @@ public class GpuKernelPool public static void Run(Action action) { GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + var accelerator = SharedGpuContext.CreateAccelerator(); var stream = accelerator.CreateStream(); action(accelerator, stream); stream.Synchronize(); stream.Dispose(); + accelerator.Dispose(); GpuPrimeWorkLimiter.Release(); } } diff --git a/PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs b/PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs index 3b97d80..dad48c0 100644 --- a/PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs +++ b/PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs @@ -22,7 +22,7 @@ public sealed class ScratchBuffer _accelerator = accelerator; // lock (accelerator) - // { + { SmallPrimeFactorPrimeSlots = accelerator.Allocate1D(smallPrimeFactorSlotCount); SmallPrimeFactorExponentSlots = accelerator.Allocate1D(smallPrimeFactorSlotCount); SmallPrimeFactorCountSlot = accelerator.Allocate1D(1); @@ -31,7 +31,7 @@ public sealed class ScratchBuffer SpecialMaxFactors = accelerator.Allocate1D(specialMaxFactorCapacity); SpecialMaxCandidates = accelerator.Allocate1D(specialMaxFactorCapacity); SpecialMaxResult = accelerator.Allocate1D(1); - // } + } } public void ResizeSmallPrimeFactorSlots(int newSize) @@ -41,8 +41,11 @@ public sealed class ScratchBuffer SmallPrimeFactorPrimeSlots.Dispose(); SmallPrimeFactorExponentSlots.Dispose(); - SmallPrimeFactorPrimeSlots = _accelerator.Allocate1D(newSize); - SmallPrimeFactorExponentSlots = _accelerator.Allocate1D(newSize); + // lock(_accelerator) + { + SmallPrimeFactorPrimeSlots = _accelerator.Allocate1D(newSize); + SmallPrimeFactorExponentSlots = _accelerator.Allocate1D(newSize); + } } } @@ -53,8 +56,11 @@ public sealed class ScratchBuffer SpecialMaxFactors.Dispose(); SpecialMaxCandidates.Dispose(); - SpecialMaxFactors = _accelerator.Allocate1D(newSize); - SpecialMaxCandidates = _accelerator.Allocate1D(newSize); + // lock(_accelerator) + { + SpecialMaxFactors = _accelerator.Allocate1D(newSize); + SpecialMaxCandidates = _accelerator.Allocate1D(newSize); + } } } } @@ -63,6 +69,7 @@ public static class GpuScratchBufferPool { [ThreadStatic] private static Dictionary>? _pools; + // private static readonly ConcurrentDictionary> _pools = new(); private const int DefaultSmallPrimeFactorSlotCount = 64; // From PrimeOrderCalculator.Gpu.cs private const int DefaultSpecialMaxFactorCapacity = 1024; // A reasonable default, will be resized if needed @@ -89,13 +96,13 @@ public static class GpuScratchBufferPool { if (buffer.SmallPrimeFactorPrimeSlots.Length < smallPrimeFactorSlotCount || buffer.SpecialMaxFactors.Length < specialMaxFactorCapacity) { - throw new InvalidOperationException(); + // throw new InvalidOperationException(); // lock (accelerator) - // { + { // Console.WriteLine($"Resizing GPU scratch buffer from pool ({buffer.SmallPrimeFactorPrimeSlots.Length} / {smallPrimeFactorSlotCount}), ({buffer.SpecialMaxFactors.Length}/{specialMaxFactorCapacity})"); - // buffer.ResizeSmallPrimeFactorSlots(smallPrimeFactorSlotCount); - // buffer.ResizeSpecialMaxFactors(specialMaxFactorCapacity); - // } + buffer.ResizeSmallPrimeFactorSlots(smallPrimeFactorSlotCount); + buffer.ResizeSpecialMaxFactors(specialMaxFactorCapacity); + } } return buffer; } @@ -113,12 +120,12 @@ public static class GpuScratchBufferPool public static void WarmUp(int count, int smallPrimeFactorSlots, int specialMaxFactorCapacity) { - Accelerator accelerator = SharedGpuContext.Accelerator; + // Accelerator accelerator = SharedGpuContext.Accelerator; - for (int i = 0; i < count; i++) - { - var buffer = Rent(accelerator, smallPrimeFactorSlots, specialMaxFactorCapacity); - Return(buffer); - } + // for (int i = 0; i < count; i++) + // { + // var buffer = Rent(accelerator, smallPrimeFactorSlots, specialMaxFactorCapacity); + // Return(buffer); + // } } } \ No newline at end of file diff --git a/PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs b/PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs index 4922a11..0f44b69 100644 --- a/PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs +++ b/PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs @@ -1,25 +1,15 @@ -using System.Collections.Concurrent; -using System.Collections.Generic; using System.Runtime.CompilerServices; -using ILGPU; using ILGPU.Runtime; -using PerfectNumbers.Core; namespace PerfectNumbers.Core.Gpu; internal static class GpuStaticTableInitializer { - private static readonly ConcurrentDictionary Initialized = new(AcceleratorReferenceComparer.Instance); - - internal static void EnsureStaticTables(KernelContainer kernels, Accelerator accelerator, AcceleratorStream stream) - { - Initialized.GetOrAdd(accelerator, acc => - { - PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(kernels, acc, stream); - GpuKernelPool.PreloadStaticTables(kernels, stream); - PrimeOrderGpuHeuristics.PreloadStaticTables( acc, stream); - return 0; - }); + internal static void EnsureStaticTables(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream) + { + // PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(accelerator); + GpuKernelPool.PreloadStaticTables(accelerator, kernels, stream); + // PrimeOrderGpuHeuristics.PreloadStaticTables(accelerator, stream); } private sealed class AcceleratorReferenceComparer : IEqualityComparer diff --git a/PerfectNumbers.Core/Gpu/HeuristicGpuDivisorTables.cs b/PerfectNumbers.Core/Gpu/HeuristicGpuDivisorTables.cs index d295417..778dc7b 100644 --- a/PerfectNumbers.Core/Gpu/HeuristicGpuDivisorTables.cs +++ b/PerfectNumbers.Core/Gpu/HeuristicGpuDivisorTables.cs @@ -6,141 +6,134 @@ namespace PerfectNumbers.Core.Gpu; internal readonly struct HeuristicGpuDivisorTables { - public static readonly HeuristicGpuDivisorTables Shared = default; + public readonly ArrayView1D CombinedDivisorsEnding1; + public readonly ArrayView1D CombinedDivisorsEnding3; + public readonly ArrayView1D CombinedDivisorsEnding7; + public readonly ArrayView1D CombinedDivisorsEnding9; + public readonly ArrayView1D CombinedDivisorSquaresEnding1; + public readonly ArrayView1D CombinedDivisorSquaresEnding3; + public readonly ArrayView1D CombinedDivisorSquaresEnding7; + public readonly ArrayView1D CombinedDivisorSquaresEnding9; + public readonly ArrayView1D GroupADivisors; + public readonly ArrayView1D GroupADivisorSquares; + public readonly ArrayView1D GroupBDivisorsEnding1; + public readonly ArrayView1D GroupBDivisorSquaresEnding1; + public readonly ArrayView1D GroupBDivisorsEnding7; + public readonly ArrayView1D GroupBDivisorSquaresEnding7; + public readonly ArrayView1D GroupBDivisorsEnding9; + public readonly ArrayView1D GroupBDivisorSquaresEnding9; - public readonly ArrayView1D CombinedDivisorsEnding1; - public readonly ArrayView1D CombinedDivisorsEnding3; - public readonly ArrayView1D CombinedDivisorsEnding7; - public readonly ArrayView1D CombinedDivisorsEnding9; - public readonly ArrayView1D CombinedDivisorSquaresEnding1; - public readonly ArrayView1D CombinedDivisorSquaresEnding3; - public readonly ArrayView1D CombinedDivisorSquaresEnding7; - public readonly ArrayView1D CombinedDivisorSquaresEnding9; - public readonly ArrayView1D GroupADivisors; - public readonly ArrayView1D GroupADivisorSquares; - public readonly ArrayView1D GroupBDivisorsEnding1; - public readonly ArrayView1D GroupBDivisorSquaresEnding1; - public readonly ArrayView1D GroupBDivisorsEnding7; - public readonly ArrayView1D GroupBDivisorSquaresEnding7; - public readonly ArrayView1D GroupBDivisorsEnding9; - public readonly ArrayView1D GroupBDivisorSquaresEnding9; + public HeuristicGpuDivisorTables( + ArrayView1D combinedDivisorsEnding1, + ArrayView1D combinedDivisorsEnding3, + ArrayView1D combinedDivisorsEnding7, + ArrayView1D combinedDivisorsEnding9, + ArrayView1D combinedDivisorSquaresEnding1, + ArrayView1D combinedDivisorSquaresEnding3, + ArrayView1D combinedDivisorSquaresEnding7, + ArrayView1D combinedDivisorSquaresEnding9, + ArrayView1D groupADivisors, + ArrayView1D groupADivisorSquares, + ArrayView1D groupBDivisorsEnding1, + ArrayView1D groupBDivisorSquaresEnding1, + ArrayView1D groupBDivisorsEnding7, + ArrayView1D groupBDivisorSquaresEnding7, + ArrayView1D groupBDivisorsEnding9, + ArrayView1D groupBDivisorSquaresEnding9) + { + CombinedDivisorsEnding1 = combinedDivisorsEnding1; + CombinedDivisorsEnding3 = combinedDivisorsEnding3; + CombinedDivisorsEnding7 = combinedDivisorsEnding7; + CombinedDivisorsEnding9 = combinedDivisorsEnding9; + CombinedDivisorSquaresEnding1 = combinedDivisorSquaresEnding1; + CombinedDivisorSquaresEnding3 = combinedDivisorSquaresEnding3; + CombinedDivisorSquaresEnding7 = combinedDivisorSquaresEnding7; + CombinedDivisorSquaresEnding9 = combinedDivisorSquaresEnding9; + GroupADivisors = groupADivisors; + GroupADivisorSquares = groupADivisorSquares; + GroupBDivisorsEnding1 = groupBDivisorsEnding1; + GroupBDivisorSquaresEnding1 = groupBDivisorSquaresEnding1; + GroupBDivisorsEnding7 = groupBDivisorsEnding7; + GroupBDivisorSquaresEnding7 = groupBDivisorSquaresEnding7; + GroupBDivisorsEnding9 = groupBDivisorsEnding9; + GroupBDivisorSquaresEnding9 = groupBDivisorSquaresEnding9; + } - public HeuristicGpuDivisorTables( - ArrayView1D combinedDivisorsEnding1, - ArrayView1D combinedDivisorsEnding3, - ArrayView1D combinedDivisorsEnding7, - ArrayView1D combinedDivisorsEnding9, - ArrayView1D combinedDivisorSquaresEnding1, - ArrayView1D combinedDivisorSquaresEnding3, - ArrayView1D combinedDivisorSquaresEnding7, - ArrayView1D combinedDivisorSquaresEnding9, - ArrayView1D groupADivisors, - ArrayView1D groupADivisorSquares, - ArrayView1D groupBDivisorsEnding1, - ArrayView1D groupBDivisorSquaresEnding1, - ArrayView1D groupBDivisorsEnding7, - ArrayView1D groupBDivisorSquaresEnding7, - ArrayView1D groupBDivisorsEnding9, - ArrayView1D groupBDivisorSquaresEnding9) - { - CombinedDivisorsEnding1 = combinedDivisorsEnding1; - CombinedDivisorsEnding3 = combinedDivisorsEnding3; - CombinedDivisorsEnding7 = combinedDivisorsEnding7; - CombinedDivisorsEnding9 = combinedDivisorsEnding9; - CombinedDivisorSquaresEnding1 = combinedDivisorSquaresEnding1; - CombinedDivisorSquaresEnding3 = combinedDivisorSquaresEnding3; - CombinedDivisorSquaresEnding7 = combinedDivisorSquaresEnding7; - CombinedDivisorSquaresEnding9 = combinedDivisorSquaresEnding9; - GroupADivisors = groupADivisors; - GroupADivisorSquares = groupADivisorSquares; - GroupBDivisorsEnding1 = groupBDivisorsEnding1; - GroupBDivisorSquaresEnding1 = groupBDivisorSquaresEnding1; - GroupBDivisorsEnding7 = groupBDivisorsEnding7; - GroupBDivisorSquaresEnding7 = groupBDivisorSquaresEnding7; - GroupBDivisorsEnding9 = groupBDivisorsEnding9; - GroupBDivisorSquaresEnding9 = groupBDivisorSquaresEnding9; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ArrayView1D SelectCombined(byte ending) + { + return ending switch + { + 1 => CombinedDivisorsEnding1, + 3 => CombinedDivisorsEnding3, + 7 => CombinedDivisorsEnding7, + 9 => CombinedDivisorsEnding9, + _ => ArrayView1D.Empty, + }; + } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ArrayView1D SelectCombined(byte ending) - { - return ending switch - { - 1 => CombinedDivisorsEnding1, - 3 => CombinedDivisorsEnding3, - 7 => CombinedDivisorsEnding7, - 9 => CombinedDivisorsEnding9, - _ => ArrayView1D.Empty, - }; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ArrayView1D SelectCombinedSquares(byte ending) + { + return ending switch + { + 1 => CombinedDivisorSquaresEnding1, + 3 => CombinedDivisorSquaresEnding3, + 7 => CombinedDivisorSquaresEnding7, + 9 => CombinedDivisorSquaresEnding9, + _ => ArrayView1D.Empty, + }; + } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ArrayView1D SelectCombinedSquares(byte ending) - { - return ending switch - { - 1 => CombinedDivisorSquaresEnding1, - 3 => CombinedDivisorSquaresEnding3, - 7 => CombinedDivisorSquaresEnding7, - 9 => CombinedDivisorSquaresEnding9, - _ => ArrayView1D.Empty, - }; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ArrayView1D SelectGroupB(byte ending) + { + return ending switch + { + 1 => GroupBDivisorsEnding1, + 7 => GroupBDivisorsEnding7, + 9 => GroupBDivisorsEnding9, + _ => ArrayView1D.Empty, + }; + } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ArrayView1D SelectGroupB(byte ending) - { - return ending switch - { - 1 => GroupBDivisorsEnding1, - 7 => GroupBDivisorsEnding7, - 9 => GroupBDivisorsEnding9, - _ => ArrayView1D.Empty, - }; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ArrayView1D SelectGroupBSquares(byte ending) + { + return ending switch + { + 1 => GroupBDivisorSquaresEnding1, + 7 => GroupBDivisorSquaresEnding7, + 9 => GroupBDivisorSquaresEnding9, + _ => ArrayView1D.Empty, + }; + } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ArrayView1D SelectGroupBSquares(byte ending) - { - return ending switch - { - 1 => GroupBDivisorSquaresEnding1, - 7 => GroupBDivisorSquaresEnding7, - 9 => GroupBDivisorSquaresEnding9, - _ => ArrayView1D.Empty, - }; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ArrayView1D SelectDivisors(HeuristicGpuDivisorTableKind kind, byte ending) + { + return kind switch + { + HeuristicGpuDivisorTableKind.GroupA => GroupADivisors, + HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorsEnding1, + HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorsEnding7, + HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorsEnding9, + HeuristicGpuDivisorTableKind.Combined => SelectCombined(ending), + _ => ArrayView1D.Empty, + }; + } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ArrayView1D SelectDivisors(HeuristicGpuDivisorTableKind kind, byte ending) - { - return kind switch - { - HeuristicGpuDivisorTableKind.GroupA => GroupADivisors, - HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorsEnding1, - HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorsEnding7, - HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorsEnding9, - HeuristicGpuDivisorTableKind.Combined => SelectCombined(ending), - _ => ArrayView1D.Empty, - }; - } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ArrayView1D SelectDivisorSquares(HeuristicGpuDivisorTableKind kind, byte ending) - { - return kind switch - { - HeuristicGpuDivisorTableKind.GroupA => GroupADivisorSquares, - HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorSquaresEnding1, - HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorSquaresEnding7, - HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorSquaresEnding9, - HeuristicGpuDivisorTableKind.Combined => SelectCombinedSquares(ending), - _ => ArrayView1D.Empty, - }; - } - - internal static void InitializeShared(in HeuristicGpuDivisorTables tables) - { - Unsafe.AsRef(in Shared) = tables; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ArrayView1D SelectDivisorSquares(HeuristicGpuDivisorTableKind kind, byte ending) + { + return kind switch + { + HeuristicGpuDivisorTableKind.GroupA => GroupADivisorSquares, + HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorSquaresEnding1, + HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorSquaresEnding7, + HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorSquaresEnding9, + HeuristicGpuDivisorTableKind.Combined => SelectCombinedSquares(ending), + _ => ArrayView1D.Empty, + }; + } } diff --git a/PerfectNumbers.Core/Gpu/InputOutputScratchBuffers.cs b/PerfectNumbers.Core/Gpu/InputOutputScratchBuffers.cs new file mode 100644 index 0000000..39f409a --- /dev/null +++ b/PerfectNumbers.Core/Gpu/InputOutputScratchBuffers.cs @@ -0,0 +1,126 @@ +using System.Collections.Concurrent; +using ILGPU; +using ILGPU.Runtime; + +namespace PerfectNumbers.Core.Gpu; + +public sealed class GpuInputOutputScratchBuffer + where TInput : unmanaged + where TOutput: unmanaged +{ + private readonly Accelerator _accelerator; + public MemoryBuffer1D Input { get; private set; } + public MemoryBuffer1D Output { get; private set; } + + public GpuInputOutputScratchBuffer(Accelerator accelerator,int inputCapacity, int outputCapacity) + { + _accelerator = accelerator; + // lock (accelerator) + { + Input = accelerator.Allocate1D(inputCapacity); + Output = accelerator.Allocate1D(outputCapacity); + } + } + + public void EnsureCapacity(int inputCapacity, int outputCapacity) + { + if (Input.Length >= inputCapacity && Output.Length >= outputCapacity) + { + return; + } + + Accelerator accelerator = _accelerator; + if (Input.Length >= inputCapacity && Output.Length >= outputCapacity) + { + return; + } + + // lock (accelerator) + { + if (Input.Length < inputCapacity) + { + Input.Dispose(); + Input = accelerator.Allocate1D(inputCapacity); + } + + if (Output.Length < outputCapacity) + { + Output.Dispose(); + Output = accelerator.Allocate1D(outputCapacity); + } + } + } + + public void Dispose() + { + Input.Dispose(); + Output.Dispose(); + } +} + +public static class GpuInputOutputScratchBufferPool + where TInput : unmanaged + where TOutput : unmanaged +{ + private const int WarmUpCapacity = 20480; + + [ThreadStatic] + private static Dictionary>>? _pool; + + private static Queue> GetPools(Accelerator accelerator) + { + var pool = _pool ??= []; + + if (!pool.TryGetValue(accelerator, out var bufferQueue)) + { + bufferQueue = []; + pool[accelerator] = bufferQueue; + } + + return bufferQueue; + } + + // private static readonly ConcurrentDictionary> _pools = new(); + + private const int DefaultSmallPrimeFactorSlotCount = 64; // From PrimeOrderCalculator.Gpu.cs + private const int DefaultSpecialMaxFactorCapacity = 1024; // A reasonable default, will be resized if needed + + public static GpuInputOutputScratchBuffer Rent(Accelerator accelerator, int inputCapacity, int outputCapacity) + { + var bufferQueue = GetPools(accelerator); + if (bufferQueue.TryDequeue(out var buffer)) + { + if (buffer.Input.Length < inputCapacity || buffer.Output.Length < outputCapacity) + { + // lock (accelerator) + // { + // Console.WriteLine($"Resizing GPU scratch buffer from pool ({buffer.SmallPrimeFactorPrimeSlots.Length} / {smallPrimeFactorSlotCount}), ({buffer.SpecialMaxFactors.Length}/{specialMaxFactorCapacity})"); + buffer.EnsureCapacity(inputCapacity, outputCapacity); + // } + } + + return buffer; + } + + return new GpuInputOutputScratchBuffer(accelerator, inputCapacity, outputCapacity); + } + + public static void Return(Accelerator accelerator, GpuInputOutputScratchBuffer buffer) => GetPools(accelerator).Enqueue(buffer); + + public static void WarmUp(int count, int inputCapacity, int outputCapacity) + { + // List> rented = new(count); + // for (int i = 0; i < count; i++) + // { + // rented.Add(Rent(inputCapacity, outputCapacity)); + // } + + // for (int i = 0; i < count; i++) + // { + // Return(rented[i]); + // } + + // rented.Clear(); + // rented.Capacity = 1; + } +} \ No newline at end of file diff --git a/PerfectNumbers.Core/Gpu/Kernels/NttMontgomeryKernels.cs b/PerfectNumbers.Core/Gpu/Kernels/NttMontgomeryKernels.cs index ea914b3..0ce580b 100644 --- a/PerfectNumbers.Core/Gpu/Kernels/NttMontgomeryKernels.cs +++ b/PerfectNumbers.Core/Gpu/Kernels/NttMontgomeryKernels.cs @@ -30,7 +30,7 @@ internal static class NttMontgomeryKernels data[i2] = new GpuUInt128(0UL, diff); } - public static void ToMont64Kernel(Index1D index, AcceleratorStream stream, ArrayView data, ulong modulus, ulong nPrime, ulong r2) + public static void ToMont64Kernel(Index1D index, ArrayView data, ulong modulus, ulong nPrime, ulong r2) { ulong a = data[index].Low; ulong v = MontMul64(a, r2, modulus, nPrime); diff --git a/PerfectNumbers.Core/Gpu/Kernels/Pow2MontgomeryKernels.cs b/PerfectNumbers.Core/Gpu/Kernels/Pow2MontgomeryKernels.cs index 5b64914..306b2b8 100644 --- a/PerfectNumbers.Core/Gpu/Kernels/Pow2MontgomeryKernels.cs +++ b/PerfectNumbers.Core/Gpu/Kernels/Pow2MontgomeryKernels.cs @@ -1,8 +1,7 @@ using ILGPU; using ILGPU.Runtime; -using PerfectNumbers.Core; -namespace PerfectNumbers.Core.Gpu; +namespace PerfectNumbers.Core.Gpu.Kernels; internal static class Pow2MontgomeryKernels { diff --git a/PerfectNumbers.Core/Gpu/Kernels/PrimeOrderKernels.cs b/PerfectNumbers.Core/Gpu/Kernels/PrimeOrderKernels.cs index 044e510..5c5ae14 100644 --- a/PerfectNumbers.Core/Gpu/Kernels/PrimeOrderKernels.cs +++ b/PerfectNumbers.Core/Gpu/Kernels/PrimeOrderKernels.cs @@ -1,10 +1,7 @@ -using System; using System.Numerics; using System.Runtime.CompilerServices; using ILGPU; -using ILGPU.Algorithms.Random; using ILGPU.Runtime; -using PerfectNumbers.Core; namespace PerfectNumbers.Core.Gpu; @@ -1171,13 +1168,20 @@ internal static partial class PrimeOrderGpuHeuristics } private static Action, GpuUInt128, ArrayView1D> GetPow2ModWideKernel(Accelerator accelerator) - { - return Pow2ModKernelWideCache.GetOrAdd(accelerator, static accel => - { - var loaded = accel.LoadAutoGroupedStreamKernel, GpuUInt128, ArrayView1D>(Pow2ModKernelWide); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, GpuUInt128, ArrayView1D>>(); - }); + { + var pool = _pow2ModWideKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel, GpuUInt128, ArrayView1D>(Pow2ModKernelWide); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, GpuUInt128, ArrayView1D>>(); + pool[accelerator] = cached; + return cached; } /// This kernel always sets the result of the corresponding element. Callers don't need to clear the output buffers. diff --git a/PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs b/PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs index 176369f..6137b29 100644 --- a/PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs +++ b/PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs @@ -22,22 +22,30 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv // so the synchronization fields from the previous implementation remain commented out here. // private readonly object _sync = new(); private ulong _divisorLimit; - // private bool _isConfigured; + // private bool _isConfigured; - private readonly ConcurrentDictionary, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>> _kernelCache = new(AcceleratorReferenceComparer.Instance); - private readonly ConcurrentBag _resourcePool = new(); - private readonly ConcurrentBag _sessionPool = new(); + [ThreadStatic] + private static Dictionary, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>>? _kernel; + + private readonly ConcurrentBag _resourcePool = []; + private readonly ConcurrentBag _sessionPool = []; - private Action, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView> GetKernel(Accelerator accelerator) => - _kernelCache.GetOrAdd(accelerator, acc => + private static Action, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView> GetKernel(Accelerator accelerator) + { + var pool = _kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>(DivisorByDivisorKernels.CheckKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>>(); - }); + return cached; + } + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>(DivisorByDivisorKernels.CheckKernel); + var kernel = KernelUtil.GetKernel(loaded); + cached = kernel.CreateLauncherDelegate, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>>(); + pool[accelerator] = cached; + return cached; + } - public int GpuBatchSize + public int GpuBatchSize { get => _gpuBatchSize; set => _gpuBatchSize = value < 1 ? 1 : value; @@ -96,7 +104,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv ulong lastProcessed; GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); // Monitor.Enter(gpuLease.ExecutionLock); @@ -502,7 +510,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv { _owner = owner; GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; + Accelerator accelerator = AcceleratorPool.Shared.Rent(); _accelerator = accelerator; _stream = accelerator.CreateStream(); _capacity = Math.Max(1, owner._gpuBatchSize); @@ -514,18 +522,23 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv private void EnsureExecutionResourcesLocked(int requiredCapacity) { + Accelerator accelerator = _accelerator; + if (_kernel == null) { - _kernel = _owner.GetKernel(_accelerator); + _kernel = GetKernel(_accelerator); } if (_divisorBuffer == null) - { - _divisorBuffer = _accelerator.Allocate1D(1); - _offsetBuffer = _accelerator.Allocate1D(1); - _countBuffer = _accelerator.Allocate1D(1); - _cycleBuffer = _accelerator.Allocate1D(1); - _firstHitBuffer = _accelerator.Allocate1D(1); + { + // lock(accelerator) + { + _divisorBuffer = accelerator.Allocate1D(1); + _offsetBuffer = accelerator.Allocate1D(1); + _countBuffer = accelerator.Allocate1D(1); + _cycleBuffer = accelerator.Allocate1D(1); + _firstHitBuffer = accelerator.Allocate1D(1); + } } bool allocateBuffers = _exponentsBuffer == null; @@ -547,10 +560,14 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv if (allocateBuffers) { int desiredCapacity = requiredCapacity > _capacity ? requiredCapacity : _capacity; - _capacity = desiredCapacity; - _exponentsBuffer = _accelerator.Allocate1D(_capacity); - _hitBuffer = _accelerator.Allocate1D(_capacity); - _hostBuffer = ThreadStaticPools.UlongPool.Rent(_capacity); + _capacity = desiredCapacity; + // lock (accelerator) + { + _exponentsBuffer = accelerator.Allocate1D(desiredCapacity); + _hitBuffer = accelerator.Allocate1D(desiredCapacity); + } + + _hostBuffer = ThreadStaticPools.UlongPool.Rent(desiredCapacity); } } @@ -769,15 +786,19 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv } private void AllocateDeviceBuffers(Accelerator accelerator, int capacity) - { - _divisorDataBuffer = accelerator.Allocate1D(capacity); - _offsetBuffer = accelerator.Allocate1D(capacity); - _countBuffer = accelerator.Allocate1D(capacity); - _cycleBuffer = accelerator.Allocate1D(capacity); - _exponentBuffer = accelerator.Allocate1D(capacity); - _hitsBuffer = accelerator.Allocate1D(capacity); - _hitIndexBuffer = accelerator.Allocate1D(1); - _deviceCapacity = capacity; + { + // lock (accelerator) + { + _divisorDataBuffer = accelerator.Allocate1D(capacity); + _offsetBuffer = accelerator.Allocate1D(capacity); + _countBuffer = accelerator.Allocate1D(capacity); + _cycleBuffer = accelerator.Allocate1D(capacity); + _exponentBuffer = accelerator.Allocate1D(capacity); + _hitsBuffer = accelerator.Allocate1D(capacity); + _hitIndexBuffer = accelerator.Allocate1D(1); + } + + _deviceCapacity = capacity; } private void RentHostArrays(int capacity) diff --git a/PerfectNumbers.Core/Gpu/MersenneNumberDivisorGpuTester.cs b/PerfectNumbers.Core/Gpu/MersenneNumberDivisorGpuTester.cs index fb233dc..630d655 100644 --- a/PerfectNumbers.Core/Gpu/MersenneNumberDivisorGpuTester.cs +++ b/PerfectNumbers.Core/Gpu/MersenneNumberDivisorGpuTester.cs @@ -10,16 +10,29 @@ namespace PerfectNumbers.Core.Gpu; /// public sealed class MersenneNumberDivisorGpuTester { - private readonly ConcurrentDictionary>> _kernelCache = new(); - private readonly ConcurrentDictionary> _resultBuffers = new(); + [ThreadStatic] + private static Dictionary>>? _resultBuffers; - private Action> GetKernel(Accelerator accelerator) => - _kernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel>(DivisorKernels.Kernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate>>(); - }); + private static Queue> GetQueue(Accelerator accelerator) + { + var pool = _resultBuffers ??= []; + if (!pool.TryGetValue(accelerator, out var bufferQueue)) + { + bufferQueue = []; + pool[accelerator] = bufferQueue; + } + + return bufferQueue; + } + + private static Action> GetDivisorKernel(Accelerator accelerator) + { + var loaded = accelerator.LoadAutoGroupedStreamKernel>(DivisorKernels.Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + return kernel.CreateLauncherDelegate>>(); + } public static void BuildDivisorCandidates() { @@ -43,18 +56,30 @@ public sealed class MersenneNumberDivisorGpuTester public bool IsDivisible(ulong exponent, in ReadOnlyGpuUInt128 divisor) { - var accelerator = SharedGpuContext.Accelerator; + // var accelerator = SharedGpuContext.CreateAccelerator(); + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var kernel = GetKernel(accelerator); - var resultBuffer = _resultBuffers.GetOrAdd(accelerator, acc => acc.Allocate1D(1)); + Queue> resultBuffers = GetQueue(accelerator); + if (!resultBuffers.TryDequeue(out var resultBuffer)) + { + // lock(accelerator) + { + resultBuffer = accelerator.Allocate1D(1); + } + } + // There is no point in clearing this buffer. We always override item [0] and never use it beyond item [0] // resultBuffer.MemSetToZero(stream); - kernel(stream, 1, exponent, divisor, resultBuffer.View); - Span result = stackalloc byte[1]; - resultBuffer.View.CopyToCPU(stream, ref result[0], 1); + var divisorKernel = GetDivisorKernel(accelerator); + divisorKernel(stream, 1, exponent, divisor, resultBuffer.View); + byte result = 0; + resultBuffer.View.CopyToCPU(stream, ref result, 1); stream.Synchronize(); - bool divisible = result[0] != 0; + bool divisible = result != 0; + resultBuffers.Enqueue(resultBuffer); stream.Dispose(); + // AcceleratorPool.Shared.Return(accelerator); + // accelerator.Dispose(); return divisible; } diff --git a/PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs b/PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs index 09ca4ca..fbed3c1 100644 --- a/PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs +++ b/PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs @@ -17,14 +17,14 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u throw new NotImplementedException($"GPU incremental scanning requires the device cycle heuristics implementation (kernel: {_kernelType}, GPU order: {_useGpuOrder})."); GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy UInt128 kStart = 1UL; ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL; byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters - KernelContainer kernels = GpuKernelPool.Kernels; + KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); var pow2Kernel = kernels.Pow2Mod!; var incKernel = kernels.Incremental!; ulong step10 = (exponent.Mod10() << 1).Mod10(); @@ -32,11 +32,11 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u ulong step3 = ((exponent % 3UL) << 1) % 3UL; ulong step5 = ((exponent % 5UL) << 1) % 5UL; GpuUInt128 twoPGpu = (GpuUInt128)twoP; - var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(); + var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(kernels); ResiduePrimeViews primeViews = default; if (_kernelType == GpuKernelType.Pow2Mod) { - primeViews = GpuKernelPool.GetSmallPrimesOnDevice(); + primeViews = GpuKernelPool.GetSmallPrimesOnDevice(kernels); } var orderBuffer = accelerator.Allocate1D(batchSize); diff --git a/PerfectNumbers.Core/Gpu/MersenneNumberLucasLehmerGpuTester.cs b/PerfectNumbers.Core/Gpu/MersenneNumberLucasLehmerGpuTester.cs index 57b6add..59c7e42 100644 --- a/PerfectNumbers.Core/Gpu/MersenneNumberLucasLehmerGpuTester.cs +++ b/PerfectNumbers.Core/Gpu/MersenneNumberLucasLehmerGpuTester.cs @@ -1,7 +1,6 @@ using System.Buffers; using System.Collections.Concurrent; using ILGPU; -using ILGPU.Algorithms; using ILGPU.Runtime; using UInt128 = System.UInt128; @@ -9,60 +8,125 @@ namespace PerfectNumbers.Core.Gpu; public class MersenneNumberLucasLehmerGpuTester { - private readonly ConcurrentDictionary>> KernelCache = new(); - private readonly ConcurrentDictionary>> AddSmallKernelCache = new(); - private readonly ConcurrentDictionary>> SubSmallKernelCache = new(); - private readonly ConcurrentDictionary>> ReduceKernelCache = new(); - private readonly ConcurrentDictionary, ArrayView>> IsZeroKernelCache = new(); - private readonly ConcurrentDictionary, ArrayView, ArrayView>> BatchKernelCache = new(); + [ThreadStatic] + private static Dictionary>>? _lucasLehmerKernel; - private Action> GetKernel(Accelerator accelerator) => - KernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.Kernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate>>(); - }); + [ThreadStatic] + private static Dictionary>>? _addSmallKernel; - private Action> GetAddSmallKernel(Accelerator accelerator) => - AddSmallKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.AddSmallKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate>>(); - }); + [ThreadStatic] + private static Dictionary>>? _subSmallKernel; + + [ThreadStatic] + private static Dictionary>>? _reduceKernel; - private Action> GetSubSmallKernel(Accelerator accelerator) => - SubSmallKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.SubtractSmallKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate>>(); - }); + [ThreadStatic] + private static Dictionary, ArrayView>>? _isZeroKernel; - private Action> GetReduceKernel(Accelerator accelerator) => - ReduceKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.ReduceModMersenneKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate>>(); - }); + [ThreadStatic] + private static Dictionary, ArrayView, ArrayView>>? _lucasLehmerBatchKernel; - private Action, ArrayView> GetIsZeroKernel(Accelerator accelerator) => - IsZeroKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView>(LucasLehmerKernels.IsZeroKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView>>(); - }); + private static Action> GetLucasLehmerKernel(Accelerator accelerator) + { + var pool = _lucasLehmerKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } - private Action, ArrayView, ArrayView> GetBatchKernel(Accelerator accelerator) => - BatchKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView, ArrayView>(LucasLehmerKernels.KernelBatch); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView, ArrayView>>(); - }); + var loaded = accelerator.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action> GetAddSmallKernel(Accelerator accelerator) + { + var pool = _addSmallKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.AddSmallKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action> GetSubSmallKernel(Accelerator accelerator) + { + var pool = _subSmallKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.SubtractSmallKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action> GetReduceKernel(Accelerator accelerator) + { + var pool = _reduceKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel>(LucasLehmerKernels.ReduceModMersenneKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, ArrayView> GetIsZeroKernel(Accelerator accelerator) + { + var pool = _isZeroKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView>(LucasLehmerKernels.IsZeroKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ArrayView>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, ArrayView, ArrayView> GetLucasLehmerBatchKernel(Accelerator accelerator) + { + var pool = _lucasLehmerBatchKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView, ArrayView>(LucasLehmerKernels.KernelBatch); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ArrayView, ArrayView>>(); + pool[accelerator] = cached; + return cached; + } // Configurable LL slice size to keep kernels short. Default 32. public int SliceSize = 32; @@ -97,12 +161,12 @@ public class MersenneNumberLucasLehmerGpuTester } else { - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var kernel = GetKernel(accelerator); var modulus = new GpuUInt128(((UInt128)1 << (int)exponent) - 1UL); // TODO: Cache these Mersenne moduli per exponent so LL GPU runs skip rebuilding them every launch. - var buffer = accelerator.Allocate1D(1); - kernel(stream, 1, exponent, modulus, buffer.View); + var buffer = accelerator.Allocate1D(1); + var lucasLehmerKernel = GetLucasLehmerKernel(accelerator); + lucasLehmerKernel(stream, 1, exponent, modulus, buffer.View); stream.Synchronize(); // TODO: Replace this with .CopyToCPU with pre-allocated buffer using stream. Synchronize on the stream after copy. result = buffer.GetAsArray1D()[0].IsZero; @@ -130,9 +194,8 @@ public class MersenneNumberLucasLehmerGpuTester } } - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var kernel = GetBatchKernel(accelerator); var expBuffer = accelerator.Allocate1D(count); ArrayPool ulongPool = ThreadStaticPools.UlongPool; @@ -152,8 +215,9 @@ public class MersenneNumberLucasLehmerGpuTester var modBuffer = accelerator.Allocate1D(count); modBuffer.View.CopyFromCPU(stream, ref modulusArray[0], count); - var stateBuffer = accelerator.Allocate1D(count); - kernel(stream, count, expBuffer.View, modBuffer.View, stateBuffer.View); + var stateBuffer = accelerator.Allocate1D(count); + var lucasLehmerBatchKernel = GetLucasLehmerBatchKernel(accelerator); + lucasLehmerBatchKernel(stream, count, expBuffer.View, modBuffer.View, stateBuffer.View); stateBuffer.View.CopyToCPU(stream, ref residues[0], count); stream.Synchronize(); @@ -185,13 +249,12 @@ public class MersenneNumberLucasLehmerGpuTester private bool IsMersennePrimeNtt(ulong exponent, GpuUInt128 nttMod, GpuUInt128 primitiveRoot, bool runOnGpu) { - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var addKernel = GetAddSmallKernel(accelerator); var subKernel = GetSubSmallKernel(accelerator); - var reduceKernel = GetReduceKernel(accelerator); - var zeroKernel = GetIsZeroKernel(accelerator); + var reduceKernel = GetReduceKernel(accelerator); + var isZeroKernel = GetIsZeroKernel(accelerator); int limbCount = (int)((exponent + 127UL) / 128UL); int target = limbCount * 2; @@ -201,8 +264,9 @@ public class MersenneNumberLucasLehmerGpuTester transformLength <<= 1; } var stateBuffer = accelerator.Allocate1D(transformLength); - stateBuffer.MemSetToZero(stream); - addKernel(stream, 1, 4UL, stateBuffer.View); + stateBuffer.MemSetToZero(stream); + var addSmallKernel = GetAddSmallKernel(accelerator); + addSmallKernel(stream, 1, 4UL, stateBuffer.View); // TODO(LL-SLICE): Slice Lucas–Lehmer iterations into short batches to // avoid long-running kernels and TDR. Example: process 8–64 iterations @@ -227,11 +291,12 @@ public class MersenneNumberLucasLehmerGpuTester } reduceKernel(stream, 1, exponent, stateBuffer.View); - var resultBuffer = accelerator.Allocate1D(1); - zeroKernel(stream, 1, stateBuffer.View, resultBuffer.View); - byte[] result = new byte[1]; + var resultBuffer = accelerator.Allocate1D(1); + isZeroKernel(stream, 1, stateBuffer.View, resultBuffer.View); + Span result = stackalloc byte[1]; resultBuffer.View.CopyToCPU(stream, result); - stream.Synchronize(); + stream.Synchronize(); + bool isPrime = result[0] != 0; resultBuffer.Dispose(); stateBuffer.Dispose(); @@ -504,6 +569,4 @@ public class MersenneNumberLucasLehmerGpuTester return result; } - - } diff --git a/PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs b/PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs index 372546d..bdfe669 100644 --- a/PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs +++ b/PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs @@ -10,28 +10,34 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime) { GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + // var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy and avoids TDR UInt128 kStart = UInt128.One; ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL; byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters - KernelContainer kernels = GpuKernelPool.Kernels; + KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); var kernel = _kernelType switch { GpuKernelType.Pow2Mod => kernels.Pow2ModOrder!, _ => kernels.IncrementalOrder!, }; - var foundBuffer = accelerator.Allocate1D(1); // TODO: Replace this allocation with the pooled device buffer from GpuOrderKernelBenchmarks so repeated scans reuse the pinned staging memory instead of allocating per run. + MemoryBuffer1D? foundBuffer; + // lock (accelerator) + { + foundBuffer = accelerator.Allocate1D(1); // TODO: Replace this allocation with the pooled device buffer from GpuOrderKernelBenchmarks so repeated scans reuse the pinned staging memory instead of allocating per run. + + } exponent.Mod10_8_5_3Steps(out ulong step10, out ulong step8, out ulong step5, out ulong step3); UInt128 remaining; int currentSize; int found = 0; - var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(); + var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(kernels); while (kStart <= maxK && Volatile.Read(ref isPrime)) { remaining = maxK - kStart + 1UL; diff --git a/PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs b/PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs index 1c2f5ea..a7232ba 100644 --- a/PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs +++ b/PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs @@ -1,10 +1,5 @@ -using PerfectNumbers.Core.Gpu; -using System; -using System.Buffers; using System.Collections.Concurrent; -using System.Collections.Generic; using System.Runtime.InteropServices; -using System.Threading; using System.Runtime.CompilerServices; using ILGPU; using ILGPU.Runtime; @@ -13,174 +8,208 @@ namespace PerfectNumbers.Core.Gpu; public class MersenneNumberResidueGpuTester(bool useGpuOrder) { - private readonly bool _useGpuOrder = useGpuOrder; - private readonly ConcurrentDictionary, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>> _pow2ModKernelCache = new(AcceleratorReferenceComparer.Instance); - private readonly ConcurrentDictionary> _resourcePools = new(AcceleratorReferenceComparer.Instance); + private readonly bool _useGpuOrder = useGpuOrder; - // TODO: Integrate MersenneDivisorCycles.Shared to consult cycle lengths for small q (<= 4M) and fast-reject - // candidates without launching heavy kernels. Consider a device-visible constant buffer per-accelerator - // via GpuKernelPool to avoid host round-trips. - // TODO: When cycles are missing for larger q values, compute only the single required cycle on the device - // requested by the caller, skip queuing additional block generation, and keep the snapshot cache untouched. + [ThreadStatic] + private static Dictionary, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>>? _pow2ModKernel; - // GPU residue variant: check 2^p % q == 1 for q = 2*p*k + 1. - public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime) - { - GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + [ThreadStatic] + private static Queue? _resourcePools; - // Monitor.Enter(gpuLease.ExecutionLock); + // TODO: Integrate MersenneDivisorCycles.Shared to consult cycle lengths for small q (<= 4M) and fast-reject + // candidates without launching heavy kernels. Consider a device-visible constant buffer per-accelerator + // via GpuKernelPool to avoid host round-trips. + // TODO: When cycles are missing for larger q values, compute only the single required cycle on the device + // requested by the caller, skip queuing additional block generation, and keep the snapshot cache untouched. - var stream = accelerator.CreateStream(); - var kernel = GetPow2ModKernel(accelerator); + // GPU residue variant: check 2^p % q == 1 for q = 2*p*k + 1. + public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime) + { + GpuPrimeWorkLimiter.Acquire(); + var accelerator = AcceleratorPool.Shared.Rent(); + + // Monitor.Enter(gpuLease.ExecutionLock); + + var stream = accelerator.CreateStream(); + var kernel = GetPow2ModKernel(accelerator); var resources = RentResources(accelerator, GpuConstants.ScanBatchSize); - var gpuKernels = GpuKernelPool.Kernels; + var gpuKernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); - var orderBuffer = resources.OrderBuffer; - ulong[] orderArray = resources.OrderArray; + var orderBuffer = resources.OrderBuffer; + ulong[] orderArray = resources.OrderArray; int batchSize = resources.Capacity; - // Ensure device has small cycles and primes tables for in-kernel lookup + // Ensure device has small cycles and primes tables for in-kernel lookup // TODO: Pass-in the stream down the path - var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(); - ResiduePrimeViews primeViews = GpuKernelPool.GetSmallPrimesOnDevice(); + var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(gpuKernels); + ResiduePrimeViews primeViews = GpuKernelPool.GetSmallPrimesOnDevice(gpuKernels); - UInt128 kStart = 1UL; - UInt128 limit = maxK + UInt128.One; - byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; - twoP.Mod10_8_5_3(out ulong step10, out ulong step8, out ulong step5, out ulong step3); - step10 = step10.Mod10(); + UInt128 kStart = 1UL; + UInt128 limit = maxK + UInt128.One; + byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; + twoP.Mod10_8_5_3(out ulong step10, out ulong step8, out ulong step5, out ulong step3); + step10 = step10.Mod10(); - GpuUInt128 twoPGpu = (GpuUInt128)twoP; - UInt128 batchSize128 = (UInt128)batchSize; - UInt128 fullBatchStep = twoP * batchSize128; - UInt128 q = twoP * kStart + UInt128.One; + GpuUInt128 twoPGpu = (GpuUInt128)twoP; + UInt128 batchSize128 = (UInt128)batchSize; + UInt128 fullBatchStep = twoP * batchSize128; + UInt128 q = twoP * kStart + UInt128.One; - while (kStart < limit && Volatile.Read(ref isPrime)) - { - UInt128 remaining = limit - kStart; - int currentSize = remaining > batchSize128 ? batchSize : (int)remaining; - Span orders = orderArray.AsSpan(0, currentSize); - ref ulong ordersRef = ref MemoryMarshal.GetReference(orders); + while (kStart < limit && Volatile.Read(ref isPrime)) + { + UInt128 remaining = limit - kStart; + int currentSize = remaining > batchSize128 ? batchSize : (int)remaining; + Span orders = orderArray.AsSpan(0, currentSize); + ref ulong ordersRef = ref MemoryMarshal.GetReference(orders); - q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3); - var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5); + q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3); + var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5); - kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL, - kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2); + kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL, + kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2); - - orderBuffer.View.CopyToCPU(stream, ref ordersRef, currentSize); + orderBuffer.View.CopyToCPU(stream, ref ordersRef, currentSize); stream.Synchronize(); - if (!Volatile.Read(ref isPrime)) - { - break; - } + if (!Volatile.Read(ref isPrime)) + { + break; + } - bool compositeFound = false; - for (int i = 0; i < currentSize; i++) - { - if (orders[i] != 0UL) - { - Volatile.Write(ref isPrime, false); - compositeFound = true; - break; - } - } + bool compositeFound = false; + for (int i = 0; i < currentSize; i++) + { + if (orders[i] != 0UL) + { + Volatile.Write(ref isPrime, false); + compositeFound = true; + break; + } + } - if (compositeFound) - { - break; - } + if (compositeFound) + { + break; + } - UInt128 processed = (UInt128)currentSize; - kStart += processed; - if (kStart < limit) - { - if (processed == batchSize128) - { - q += fullBatchStep; - } - else - { - q += twoP * processed; // TODO: Swap this multiplication for the shared cycle stepping helper once residue cycles are mandatory so we can reuse the cached remainder ladder instead of performing UInt128 multiply operations. - } - } - } + UInt128 processed = (UInt128)currentSize; + kStart += processed; + if (kStart < limit) + { + if (processed == batchSize128) + { + q += fullBatchStep; + } + else + { + q += twoP * processed; // TODO: Swap this multiplication for the shared cycle stepping helper once residue cycles are mandatory so we can reuse the cached remainder ladder instead of performing UInt128 multiply operations. + } + } + } - ReturnResources(accelerator, resources); - stream.Dispose(); + ReturnResources(resources); + stream.Dispose(); // Monitor.Exit(gpuLease.ExecutionLock); gpuKernels.Dispose(); GpuPrimeWorkLimiter.Release(); - } + } - private Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> GetPow2ModKernel(Accelerator accelerator) => - _pow2ModKernelCache.GetOrAdd( - accelerator, - static acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(Pow2ModKernels.Pow2ModKernelScan); - return (Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>)Delegate.CreateDelegate( - typeof(Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>), - loaded.Target, - loaded.Method); - }); + private static Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> GetPow2ModKernel(Accelerator accelerator) + { - private ResidueResources RentResources(Accelerator accelerator, int capacity) - { - var bag = _resourcePools.GetOrAdd(accelerator, static _ => []); - while (bag.TryTake(out var resources)) - { - if (resources.Capacity >= capacity) + var pool = _pow2ModKernel ?? []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + cached = CreatePow2ModKernel(accelerator); + pool[accelerator] = cached; + return cached; + } + + private static Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> CreatePow2ModKernel(Accelerator accelerator) + { + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(Pow2ModKernels.Pow2ModKernelScan); + + var kernel = (Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>)Delegate.CreateDelegate( + typeof(Action, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>), + loaded.Target, + loaded.Method); + return kernel; + } + + + private static ResidueResources RentResources(Accelerator accelerator, int capacity) + { + var resourcePools = _resourcePools ??= new(); + if (resourcePools.TryDequeue(out var bag)) + { + bag.EnsureCapacity(accelerator, capacity); + return bag; + } + + return new ResidueResources(accelerator, capacity); + } + + private static void ReturnResources(ResidueResources resources) + { + var resourcePools = _resourcePools ??= new(); + resourcePools.Enqueue(resources); + } + + private sealed class ResidueResources + { + internal ResidueResources(Accelerator accelerator, int capacity) + { + Capacity = Math.Max(1, capacity); + // lock (accelerator) { - // The caller always overwrites the required elements. We don't need to worry about clearing the output buffer. - // resources.OrderBuffer.MemSetToZero(); - return resources; - } + OrderBuffer = accelerator.Allocate1D(Capacity); + } - resources.Dispose(); - } + OrderArray = ThreadStaticPools.UlongPool.Rent(Capacity); + } - var created = new ResidueResources(accelerator, capacity); - // The caller always overwrites the required elements. We don't need to worry about clearing the output buffer. - // created.OrderBuffer.MemSetToZero(); - return created; - } + internal void EnsureCapacity(Accelerator accelerator, int capacity) + { + if (OrderBuffer.Length >= capacity) + { + return; + } - private void ReturnResources(Accelerator accelerator, ResidueResources resources) => - _resourcePools.GetOrAdd(accelerator, static _ => new ConcurrentBag()).Add(resources); + OrderBuffer.Dispose(); + ThreadStaticPools.UlongPool.Return(OrderArray); - private sealed class ResidueResources - { - internal ResidueResources(Accelerator accelerator, int capacity) - { - Capacity = Math.Max(1, capacity); - OrderBuffer = accelerator.Allocate1D(Capacity); - OrderArray = ThreadStaticPools.UlongPool.Rent(Capacity); - } + // lock (accelerator) + { + OrderBuffer = accelerator.Allocate1D(capacity); + } - internal MemoryBuffer1D OrderBuffer { get; } + OrderArray = ThreadStaticPools.UlongPool.Rent(capacity); + Capacity = capacity; + } - internal ulong[] OrderArray { get; } + internal MemoryBuffer1D OrderBuffer { get; private set; } - internal int Capacity { get; } + internal ulong[] OrderArray { get; private set; } - public void Dispose() - { - OrderBuffer.Dispose(); - ThreadStaticPools.UlongPool.Return(OrderArray, clearArray: false); - } - } + internal int Capacity { get; private set; } - private sealed class AcceleratorReferenceComparer : IEqualityComparer - { - public static AcceleratorReferenceComparer Instance { get; } = new(); + public void Dispose() + { + OrderBuffer.Dispose(); + ThreadStaticPools.UlongPool.Return(OrderArray, clearArray: false); + } + } - public bool Equals(Accelerator? x, Accelerator? y) => ReferenceEquals(x, y); + private sealed class AcceleratorReferenceComparer : IEqualityComparer + { + public static AcceleratorReferenceComparer Instance { get; } = new(); - public int GetHashCode(Accelerator obj) => RuntimeHelpers.GetHashCode(obj); - } + public bool Equals(Accelerator? x, Accelerator? y) => ReferenceEquals(x, y); + + public int GetHashCode(Accelerator obj) => RuntimeHelpers.GetHashCode(obj); + } } \ No newline at end of file diff --git a/PerfectNumbers.Core/Gpu/NttGpuMath.cs b/PerfectNumbers.Core/Gpu/NttGpuMath.cs index 0f4b167..1ef9301 100644 --- a/PerfectNumbers.Core/Gpu/NttGpuMath.cs +++ b/PerfectNumbers.Core/Gpu/NttGpuMath.cs @@ -29,41 +29,62 @@ public static class NttGpuMath public static NttBackend GpuTransformBackend { get; set; } = NttBackend.Reference; // Controls modular reduction strategy inside staged NTT paths. public static ModReductionMode ReductionMode { get; set; } = ModReductionMode.Auto; - - private static readonly ConcurrentDictionary, ArrayView, GpuUInt128>> MulKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - private static readonly ConcurrentDictionary, int, int, int, ArrayView, GpuUInt128>> StageKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - private static readonly ConcurrentDictionary, GpuUInt128, GpuUInt128>> ScaleKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + [ThreadStatic] + private static Dictionary, ArrayView, GpuUInt128>>? _mulKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - private static readonly ConcurrentDictionary, int, int, int, ArrayView, ulong, ulong>> StageMontKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, int, int, int, ArrayView, ulong, ulong, ulong, ulong>> StageBarrett128KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, ulong, ulong, ulong>> ToMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, ulong, ulong>> FromMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, GpuUInt128, ulong, ulong, ulong, ulong>> ScaleBarrett128KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, int, int, int, ArrayView, GpuUInt128>> SquareMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, ulong, ulong, ulong>> ScaleMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, ArrayView, int, GpuUInt128, GpuUInt128>> ForwardKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128>> InverseKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static readonly ConcurrentDictionary, int>> BitReverseKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - - private static Action, int, int, int, ArrayView, GpuUInt128> GetStageKernel(Accelerator accelerator) => - StageKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, int, int, int, ArrayView, GpuUInt128>(NttButterflyKernels.StageKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, int, int, int, ArrayView, GpuUInt128>>(); - }); + [ThreadStatic] + private static Dictionary, int, int, int, ArrayView, GpuUInt128>>? _stageKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. - private sealed class SquareCacheEntry + [ThreadStatic] + private static Dictionary, GpuUInt128, GpuUInt128>>? _scaleKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, int, int, int, ArrayView, ulong, ulong>>? _stageMontKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, int, int, int, ArrayView, ulong, ulong, ulong, ulong>>? _stageBarrett128Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, ulong, ulong, ulong>>? _toMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, ulong, ulong>>? _fromMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, GpuUInt128, ulong, ulong, ulong, ulong>>? _scaleBarrett128Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static readonly Dictionary, int, int, int, ArrayView, GpuUInt128>>? _squareMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, ulong, ulong, ulong>>? _scaleMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, ArrayView, int, GpuUInt128, GpuUInt128>>? _forwardKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128>>? _inverseKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + [ThreadStatic] + private static Dictionary, int>>? _bitReverseKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup. + + private static Action, int, int, int, ArrayView, GpuUInt128> GetStageKernel(Accelerator accelerator) + { + var pool = _stageKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel, int, int, int, ArrayView, GpuUInt128>(NttButterflyKernels.StageKernel); + var kernel = KernelUtil.GetKernel(loaded); + cached = kernel.CreateLauncherDelegate, int, int, int, ArrayView, GpuUInt128>>(); + pool[accelerator] = cached; + return cached; + } + + private sealed class SquareCacheEntry { public MemoryBuffer1D Buffer { get; } public GpuUInt128 Root { get; } @@ -266,7 +287,7 @@ public static class NttGpuMath private readonly record struct SquareCacheKey(int Length, GpuUInt128 Modulus, GpuUInt128 PrimitiveRoot); - private static readonly ConcurrentDictionary> SquareCache = new(); + private static readonly ConcurrentDictionary SquareCache = []; // TODO(NTT-OPT): Add a TwiddleCache for stage-wise Cooley–Tukey NTT. // Design: @@ -276,9 +297,8 @@ public static class NttGpuMath private static SquareCacheEntry GetSquareCache(Accelerator accelerator, AcceleratorStream stream, int length, GpuUInt128 modulus, GpuUInt128 primitiveRoot) { - var perAccel = SquareCache.GetOrAdd(accelerator, _ => new ConcurrentDictionary()); var key = new SquareCacheKey(length, modulus, primitiveRoot); - return perAccel.GetOrAdd(key, _ => new SquareCacheEntry(accelerator, stream, length, modulus, primitiveRoot)); + return SquareCache.GetOrAdd(key, _ => new SquareCacheEntry(accelerator, stream, length, modulus, primitiveRoot)); } private static readonly ulong[] SmallPrimes = GenerateSmallPrimes(1000); @@ -346,147 +366,197 @@ public static class NttGpuMath { if (accelerator is null) { - foreach (var perAccel in SquareCache.Values) + foreach (var entry in SquareCache.Values) { - foreach (var entry in perAccel.Values) - { - entry.Dispose(); - } + entry.Dispose(); } SquareCache.Clear(); - MulKernelCache.Clear(); - StageKernelCache.Clear(); - ScaleKernelCache.Clear(); - StageMontKernelCache.Clear(); - StageBarrett128KernelCache.Clear(); - ToMont64KernelCache.Clear(); - FromMont64KernelCache.Clear(); - SquareMont64KernelCache.Clear(); - ScaleMont64KernelCache.Clear(); - ForwardKernelCache.Clear(); - InverseKernelCache.Clear(); - BitReverseKernelCache.Clear(); // TODO(NTT-OPT): Clear TwiddleCache once introduced. return; } - if (SquareCache.TryRemove(accelerator, out var cache)) - { - foreach (var entry in cache.Values) - { - entry.Dispose(); - } - } + var squareCache = SquareCache; + foreach (var entry in squareCache.Values) + { + entry.Dispose(); + } - MulKernelCache.TryRemove(accelerator, out _); - StageKernelCache.TryRemove(accelerator, out _); - ScaleKernelCache.TryRemove(accelerator, out _); - StageMontKernelCache.TryRemove(accelerator, out _); - StageBarrett128KernelCache.TryRemove(accelerator, out _); - ToMont64KernelCache.TryRemove(accelerator, out _); - FromMont64KernelCache.TryRemove(accelerator, out _); - SquareMont64KernelCache.TryRemove(accelerator, out _); - ScaleMont64KernelCache.TryRemove(accelerator, out _); - ForwardKernelCache.TryRemove(accelerator, out _); - InverseKernelCache.TryRemove(accelerator, out _); - BitReverseKernelCache.TryRemove(accelerator, out _); // TODO(NTT-OPT): Remove per-accelerator twiddle buffers once added. } - private static Action, int> GetBitReverseKernel(Accelerator accelerator) => - BitReverseKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, int>(NttTransformKernels.BitReverseKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, int>>(); - }); - - private static Action, ArrayView, int, GpuUInt128, GpuUInt128> GetForwardKernel(Accelerator accelerator) => - ForwardKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView, int, GpuUInt128, GpuUInt128>(NttTransformKernels.ForwardKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView, int, GpuUInt128, GpuUInt128>>(); - }); - - private static Action, ulong, ulong> GetFromMont64Kernel(Accelerator accelerator) => - FromMont64KernelCache.GetOrAdd(accelerator, acc => + private static Action, int> GetBitReverseKernel(Accelerator accelerator) + { + var pool = _bitReverseKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, ulong, ulong>(NttMontgomeryKernels.FromMont64Kernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ulong, ulong>>(); - }); + return cached; + } - private static Action, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128> GetInverseKernel(Accelerator accelerator) => - InverseKernelCache.GetOrAdd(accelerator, acc => - { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128>(NttTransformKernels.InverseKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128>>(); - }); + var loaded = accelerator.LoadAutoGroupedStreamKernel, int>(NttTransformKernels.BitReverseKernel); + + var kernel = KernelUtil.GetKernel(loaded); - private static Action, ArrayView, GpuUInt128> GetMulKernel(Accelerator accelerator) => - MulKernelCache.GetOrAdd(accelerator, valueFactory: acc => + cached = kernel.CreateLauncherDelegate, int>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, ArrayView, int, GpuUInt128, GpuUInt128> GetForwardKernel(Accelerator accelerator) + { + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView, int, GpuUInt128, GpuUInt128>(NttTransformKernels.ForwardKernel); + var kernel = KernelUtil.GetKernel(loaded); + return kernel.CreateLauncherDelegate, ArrayView, int, GpuUInt128, GpuUInt128>>(); + } + + private static Action, ulong, ulong> GetFromMont64Kernel(Accelerator accelerator) + { + var pool = _fromMont64Kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, ArrayView, GpuUInt128>(NttPointwiseKernels.MulKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView, GpuUInt128>>(); - }); + return cached; + } - private static Action, GpuUInt128, GpuUInt128> GetScaleKernel(Accelerator accelerator) => - ScaleKernelCache.GetOrAdd(accelerator, acc => + var loaded = accelerator.LoadAutoGroupedStreamKernel, ulong, ulong>(NttMontgomeryKernels.FromMont64Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ulong, ulong>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128> GetInverseKernel(Accelerator accelerator) + { + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128>(NttTransformKernels.InverseKernel); + var kernel = KernelUtil.GetKernel(loaded); + return kernel.CreateLauncherDelegate, ArrayView, int, GpuUInt128, GpuUInt128, GpuUInt128>>(); + } + + private static Action, ArrayView, GpuUInt128> GetMulKernel(Accelerator accelerator) + { + var pool = _mulKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, GpuUInt128, GpuUInt128>(NttPointwiseKernels.ScaleKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, GpuUInt128, GpuUInt128>>(); - }); + return cached; + } - private static Action, GpuUInt128, ulong, ulong, ulong, ulong> GetScaleBarrett128Kernel(Accelerator accelerator) => - ScaleBarrett128KernelCache.GetOrAdd(accelerator, acc => + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView, GpuUInt128>(NttPointwiseKernels.MulKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ArrayView, GpuUInt128>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, GpuUInt128, GpuUInt128> GetScaleKernel(Accelerator accelerator) + { + var pool = _scaleKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, GpuUInt128, ulong, ulong, ulong, ulong>(NttBarrettKernels.ScaleBarrett128Kernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, GpuUInt128, ulong, ulong, ulong, ulong>>(); - }); + return cached; + } - private static Action, ulong, ulong, ulong> GetScaleMont64Kernel(Accelerator accelerator) => - ScaleMont64KernelCache.GetOrAdd(accelerator, acc => + var loaded = accelerator.LoadAutoGroupedStreamKernel, GpuUInt128, GpuUInt128>(NttPointwiseKernels.ScaleKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, GpuUInt128, GpuUInt128>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, GpuUInt128, ulong, ulong, ulong, ulong> GetScaleBarrett128Kernel(Accelerator accelerator) + { + var pool = _scaleBarrett128Kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, ulong, ulong, ulong>(NttMontgomeryKernels.ScaleMont64Kernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ulong, ulong, ulong>>(); - }); + return cached; + } - private static Action, int, int, int, ArrayView, ulong, ulong, ulong, ulong> GetStageBarrett128Kernel(Accelerator accelerator) => - StageBarrett128KernelCache.GetOrAdd(accelerator, acc => + var loaded = accelerator.LoadAutoGroupedStreamKernel, GpuUInt128, ulong, ulong, ulong, ulong>(NttBarrettKernels.ScaleBarrett128Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, GpuUInt128, ulong, ulong, ulong, ulong>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, ulong, ulong, ulong> GetScaleMont64Kernel(Accelerator accelerator) + { + var pool = _scaleMont64Kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, int, int, int, ArrayView, ulong, ulong, ulong, ulong>(NttBarrettKernels.StageBarrett128Kernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, int, int, int, ArrayView, ulong, ulong, ulong, ulong>>(); - }); + return cached; + } - private static Action, int, int, int, ArrayView, ulong, ulong> GetStageMontKernel(Accelerator accelerator) => - StageMontKernelCache.GetOrAdd(accelerator, acc => + var loaded = accelerator.LoadAutoGroupedStreamKernel, ulong, ulong, ulong>(NttMontgomeryKernels.ScaleMont64Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ulong, ulong, ulong>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, int, int, int, ArrayView, ulong, ulong, ulong, ulong> GetStageBarrett128Kernel(Accelerator accelerator) + { + var pool = _stageBarrett128Kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, int, int, int, ArrayView, ulong, ulong>(NttMontgomeryKernels.StageMontKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, int, int, int, ArrayView, ulong, ulong>>(); - }); + return cached; + } - private static Action, ulong, ulong, ulong> GetToMont64Kernel(Accelerator accelerator) => - ToMont64KernelCache.GetOrAdd(accelerator, acc => + var loaded = accelerator.LoadAutoGroupedStreamKernel, int, int, int, ArrayView, ulong, ulong, ulong, ulong>(NttBarrettKernels.StageBarrett128Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, int, int, int, ArrayView, ulong, ulong, ulong, ulong>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, int, int, int, ArrayView, ulong, ulong> GetStageMontKernel(Accelerator accelerator) + { + var pool = _stageMontKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) { - var loaded = acc.LoadAutoGroupedStreamKernel, int>(NttTransformKernels.BitReverseKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ulong, ulong, ulong>>(); - }); + return cached; + } - // TODO(NTT-OPT): Mirror the forward stage-wise design for the inverse - // transform. Precompute inverse twiddles and a normalization factor - // (n^-1 mod m). Launch one short kernel per stage to avoid TDR. + var loaded = accelerator.LoadAutoGroupedStreamKernel, int, int, int, ArrayView, ulong, ulong>(NttMontgomeryKernels.StageMontKernel); - public static void BitReverse(Span values) + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, int, int, int, ArrayView, ulong, ulong>>(); + pool[accelerator] = cached; + return cached; + } + + private static Action, ulong, ulong, ulong> GetToMont64Kernel(Accelerator accelerator) + { + var pool = _toMont64Kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel, ulong, ulong, ulong>(NttMontgomeryKernels.ToMont64Kernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ulong, ulong, ulong>>(); + pool[accelerator] = cached; + return cached; + } + + // TODO(NTT-OPT): Mirror the forward stage-wise design for the inverse + // transform. Precompute inverse twiddles and a normalization factor + // (n^-1 mod m). Launch one short kernel per stage to avoid TDR. + + public static void BitReverse(Span values) { int n = values.Length; int bits = (int)Math.Log2(n); @@ -504,15 +574,15 @@ public static class NttGpuMath { int n = values.Length; int bits = (int)Math.Log2(n); - var accelerator = SharedGpuContext.Accelerator; - var kernel = GetBitReverseKernel(accelerator); + var accelerator = AcceleratorPool.Shared.Rent(); var pool = ThreadStaticPools.GpuUInt128Pool; var array = pool.Rent(n); values.CopyTo(array); var buffer = accelerator.Allocate1D(n); var stream = accelerator.CreateStream(); - buffer.View.CopyFromCPU(stream, ref array[0], n); - kernel(stream, n, buffer.View, bits); + buffer.View.CopyFromCPU(stream, ref array[0], n); + var bitReverseKernel = GetBitReverseKernel(accelerator); + bitReverseKernel(stream, n, buffer.View, bits); buffer.View.CopyToCPU(stream, ref array[0], n); stream.Synchronize(); array.AsSpan(0, n).CopyTo(values); @@ -566,9 +636,8 @@ public static class NttGpuMath private static void ForwardGpuReference(Span values, GpuUInt128 modulus, GpuUInt128 primitiveRoot) { int n = (int)values.Length; - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var kernel = GetForwardKernel(accelerator); var pool = ThreadStaticPools.GpuUInt128Pool; var inputArray = pool.Rent(n); var outputArray = pool.Rent(n); @@ -579,8 +648,9 @@ public static class NttGpuMath UInt128 modValue = modulus; UInt128 expBase = (modValue - 1UL) / (ulong)n; var root = new GpuUInt128(primitiveRoot); - root.ModPow((ulong)expBase, modulus); - kernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, root); + root.ModPow((ulong)expBase, modulus); + var forwardKernel = GetForwardKernel(accelerator); + forwardKernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, root); outputBuffer.View.CopyToCPU(stream, ref outputArray[0], n); stream.Synchronize(); outputArray.AsSpan(0, n).CopyTo(values); @@ -662,9 +732,8 @@ public static class NttGpuMath private static void InverseGpuReference(Span values, GpuUInt128 modulus, GpuUInt128 primitiveRoot) { int n = values.Length; - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var kernel = GetInverseKernel(accelerator); var pool = ThreadStaticPools.GpuUInt128Pool; var inputArray = pool.Rent(n); var outputArray = pool.Rent(n); @@ -687,15 +756,17 @@ public static class NttGpuMath } var nInv = new GpuUInt128((UInt128)n); - if (modulus.High == 0UL) - { - nInv.ModInv(modulus.Low); - } - else - { - nInv.ModInv(modulus); - } - kernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, rootInv, nInv); + if (modulus.High == 0UL) + { + nInv.ModInv(modulus.Low); + } + else + { + nInv.ModInv(modulus); + } + + var inverseKernel = GetInverseKernel(accelerator); + inverseKernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, rootInv, nInv); outputBuffer.View.CopyToCPU(stream, ref outputArray[0], n); stream.Synchronize(); outputArray.AsSpan(0, n).CopyTo(values); @@ -709,7 +780,7 @@ public static class NttGpuMath private static void ForwardGpuStaged(Span values, GpuUInt128 modulus, GpuUInt128 primitiveRoot) { int n = values.Length; - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); var buffer = accelerator.Allocate1D(n); buffer.View.CopyFromCPU(stream, ref values[0], n); @@ -724,7 +795,7 @@ public static class NttGpuMath private static void InverseGpuStaged(Span values, GpuUInt128 modulus, GpuUInt128 primitiveRoot) { int n = values.Length; - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); var buffer = accelerator.Allocate1D(n); buffer.View.CopyFromCPU(stream, ref values[0], n); @@ -737,24 +808,24 @@ public static class NttGpuMath } private static void ForwardDevice(Accelerator accelerator, AcceleratorStream stream, ArrayView data, int length, GpuUInt128 modulus, SquareCacheEntry cache) - { - var bitKernel = GetBitReverseKernel(accelerator); - bitKernel(stream, length, data, cache.Bits); + { + var bitReverseKernel = GetBitReverseKernel(accelerator); + bitReverseKernel(stream, length, data, cache.Bits); if (cache.UseMontgomery64) - { - var toMont = GetToMont64Kernel(accelerator); - toMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64); + { + var toMont64Kernel = GetToMont64Kernel(accelerator); + toMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64); var stageKernel = GetStageMontKernel(accelerator); - int butterflies = length >> 1; - for (int s = 0; s < cache.Bits; s++) - { - int len = 1 << (s + 1); - int half = len >> 1; - stageKernel( stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesMont!.View, cache.ModulusLow, cache.MontNPrime64); - } - var fromMont = GetFromMont64Kernel(accelerator); - fromMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64); + int butterflies = length >> 1; + for (int s = 0; s < cache.Bits; s++) + { + int len = 1 << (s + 1); + int half = len >> 1; + stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesMont!.View, cache.ModulusLow, cache.MontNPrime64); + } + var fromMont64Kernel = GetFromMont64Kernel(accelerator); + fromMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64); } else if (cache.UseBarrett128) { @@ -783,51 +854,55 @@ public static class NttGpuMath } private static void InverseDevice(Accelerator accelerator, AcceleratorStream stream, ArrayView data, int length, GpuUInt128 modulus, SquareCacheEntry cache) - { - var bitKernel = GetBitReverseKernel(accelerator); - bitKernel(stream, length, data, cache.Bits); + { + var bitReverseKernel = GetBitReverseKernel(accelerator); + bitReverseKernel(stream, length, data, cache.Bits); if (cache.UseMontgomery64) - { - var toMont = GetToMont64Kernel(accelerator); - toMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64); + { + var toMont64Kernel = GetToMont64Kernel(accelerator); + toMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64); var stageKernel = GetStageMontKernel(accelerator); int butterflies = length >> 1; - for (int s = 0; s < cache.Bits; s++) - { - int len = 1 << (s + 1); - int half = len >> 1; - stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInvMont!.View, cache.ModulusLow, cache.MontNPrime64); - } - var scaleKernel = GetScaleMont64Kernel(accelerator); - scaleKernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontNInvR64); - var fromMont = GetFromMont64Kernel(accelerator); - fromMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64); + for (int s = 0; s < cache.Bits; s++) + { + int len = 1 << (s + 1); + int half = len >> 1; + stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInvMont!.View, cache.ModulusLow, cache.MontNPrime64); + } + + var scaleMont64Kernel = GetScaleMont64Kernel(accelerator); + scaleMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontNInvR64); + + var fromMont64Kernel = GetFromMont64Kernel(accelerator); + fromMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64); } else if (cache.UseBarrett128) { var stageKernel = GetStageBarrett128Kernel(accelerator); int butterflies = length >> 1; - for (int s = 0; s < cache.Bits; s++) - { - int len = 1 << (s + 1); - int half = len >> 1; - stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow); - } - var scaleKernel = GetScaleBarrett128Kernel(accelerator); - scaleKernel(stream, length, data, cache.NInv, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow); + for (int s = 0; s < cache.Bits; s++) + { + int len = 1 << (s + 1); + int half = len >> 1; + stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow); + } + + var scaleBarrett128Kernel = GetScaleBarrett128Kernel(accelerator); + scaleBarrett128Kernel(stream, length, data, cache.NInv, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow); } else { var stageKernel = GetStageKernel(accelerator); int butterflies = length >> 1; - for (int s = 0; s < cache.Bits; s++) - { - int len = 1 << (s + 1); - int half = len >> 1; - stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, modulus); - } - var scaleKernel = GetScaleKernel(accelerator); + for (int s = 0; s < cache.Bits; s++) + { + int len = 1 << (s + 1); + int half = len >> 1; + stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, modulus); + } + + var scaleKernel = GetScaleKernel(accelerator); scaleKernel(stream, length, data, cache.NInv, modulus); } @@ -917,17 +992,16 @@ public static class NttGpuMath public static void SquareDevice(Accelerator accelerator, AcceleratorStream stream, ArrayView values, GpuUInt128 modulus, GpuUInt128 primitiveRoot) { var cache = GetSquareCache(accelerator, stream, (int)values.Length, modulus, primitiveRoot); - ForwardDevice(accelerator, stream, values, (int)values.Length, modulus, cache); - var mul = GetMulKernel(accelerator); - mul(stream, (int)values.Length, values, values, modulus); + ForwardDevice(accelerator, stream, values, (int)values.Length, modulus, cache); + var mulKernel = GetMulKernel(accelerator); + mulKernel(stream, (int)values.Length, values, values, modulus); InverseDevice(accelerator, stream, values, (int)values.Length, modulus, cache); } public static void PointwiseMultiply(Span left, Span right, GpuUInt128 modulus) { - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var kernel = GetMulKernel(accelerator); var pool = ThreadStaticPools.GpuUInt128Pool; var leftArray = pool.Rent(left.Length); var rightArray = pool.Rent(right.Length); @@ -936,8 +1010,9 @@ public static class NttGpuMath var leftBuffer = accelerator.Allocate1D(left.Length); var rightBuffer = accelerator.Allocate1D(right.Length); leftBuffer.View.CopyFromCPU(stream, ref leftArray[0], left.Length); - rightBuffer.View.CopyFromCPU(stream, ref rightArray[0], right.Length); - kernel(stream, left.Length, leftBuffer.View, rightBuffer.View, modulus); + rightBuffer.View.CopyFromCPU(stream, ref rightArray[0], right.Length); + var mulKernel = GetMulKernel(accelerator); + mulKernel(stream, left.Length, leftBuffer.View, rightBuffer.View, modulus); leftBuffer.View.CopyToCPU(stream, ref leftArray[0], left.Length); rightBuffer.View.CopyToCPU(stream, ref rightArray[0], right.Length); stream.Synchronize(); diff --git a/PerfectNumbers.Core/Gpu/Pow2MontgomeryGpuCalculator.cs b/PerfectNumbers.Core/Gpu/Pow2MontgomeryGpuCalculator.cs new file mode 100644 index 0000000..3838a71 --- /dev/null +++ b/PerfectNumbers.Core/Gpu/Pow2MontgomeryGpuCalculator.cs @@ -0,0 +1,83 @@ +using System.Runtime.CompilerServices; +using ILGPU; +using ILGPU.Runtime; +using PerfectNumbers.Core.Gpu.Accelerators; + +namespace PerfectNumbers.Core.Gpu; + +public static class Pow2MontgomeryGpuCalculator +{ + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ulong CalculateKeep(ulong exponent, in MontgomeryDivisorData divisor) + { + GpuPrimeWorkLimiter.Acquire(); + + var gpu = Pow2MontgomeryAccelerator.Rent(); + Accelerator? accelerator = gpu.Accelerator; + AcceleratorStream stream = gpu.Stream!; + MemoryBuffer1D exponentBuffer = gpu.Input; + MemoryBuffer1D resultBuffer = gpu.Output; + + exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1); + // We don't need to worry about any left-overs here. + // resultBuffer.MemSetToZero(stream); + + var kernel = gpu.KeepMontgomeryKernel; + kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View); + + ulong result = 0UL; + resultBuffer.View.CopyToCPU(stream, ref result, 1); + stream.Synchronize(); + + Pow2MontgomeryAccelerator.Return(gpu); + GpuPrimeWorkLimiter.Release(); + return result; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ulong CalculateConvert(ulong exponent, in MontgomeryDivisorData divisor) + { + GpuPrimeWorkLimiter.Acquire(); + + var gpu = Pow2MontgomeryAccelerator.Rent(); + AcceleratorStream stream = gpu.Stream!; + MemoryBuffer1D exponentBuffer = gpu.Input; + MemoryBuffer1D resultBuffer = gpu.Output; + + exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1); + // We don't need to worry about any left-overs here. + // resultBuffer.MemSetToZero(stream); + + var kernel = gpu.ConvertToStandardKernel; + kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View); + + ulong result = 0UL; + resultBuffer.View.CopyToCPU(stream, ref result, 1); + stream.Synchronize(); + + Pow2MontgomeryAccelerator.Return(gpu); + GpuPrimeWorkLimiter.Release(); + return result; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ulong CalculateConvert(Pow2MontgomeryAccelerator gpu, ulong exponent, in MontgomeryDivisorData divisor) + { + AcceleratorStream stream = gpu.Stream!; + MemoryBuffer1D exponentBuffer = gpu.Input; + MemoryBuffer1D resultBuffer = gpu.Output; + + exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1); + // We don't need to worry about any left-overs here. + // resultBuffer.MemSetToZero(stream); + + var kernel = gpu.ConvertToStandardKernel; + kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View); + + ulong result = 0UL; + resultBuffer.View.CopyToCPU(stream, ref result, 1); + stream.Synchronize(); + + return result; + } +} diff --git a/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs b/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs index 6052a19..18e3837 100644 --- a/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs +++ b/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs @@ -8,428 +8,481 @@ namespace PerfectNumbers.Core.Gpu; internal enum GpuPow2ModStatus { - Success, - Overflow, - Unavailable, + Success, + Overflow, + Unavailable, } internal static partial class PrimeOrderGpuHeuristics { - internal static void PreloadStaticTables(Accelerator accelerator, AcceleratorStream stream) - { - _ = GetSmallPrimeDeviceCache(accelerator, stream); - } + internal static void PreloadStaticTables(AcceleratorStream stream) + { + // _ = GetSmallPrimeDeviceCache(stream); + } - private static readonly ConcurrentDictionary OverflowedPrimes = new(); - private static readonly ConcurrentDictionary OverflowedPrimesWide = new(); - private static readonly ConcurrentDictionary, MontgomeryDivisorData, ArrayView1D>> Pow2ModKernelCache = new(); - private static readonly ConcurrentDictionary, GpuUInt128, ArrayView1D>> Pow2ModKernelWideCache = new(); - private static readonly ConcurrentDictionary, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>> PartialFactorKernelCache = new(); - private static readonly ConcurrentDictionary OrderKernelCache = new(); - private static readonly ConcurrentDictionary SmallPrimeDeviceCaches = new(); + private static readonly ConcurrentDictionary OverflowedPrimes = new(); + private static readonly ConcurrentDictionary OverflowedPrimesWide = new(); - public readonly struct OrderKernelConfig(ulong previousOrder, byte hasPreviousOrder, uint smallFactorLimit, int maxPowChecks, int mode) - { - public readonly ulong PreviousOrder = previousOrder; - public readonly byte HasPreviousOrder = hasPreviousOrder; - public readonly uint SmallFactorLimit = smallFactorLimit; - public readonly int MaxPowChecks = maxPowChecks; - public readonly int Mode = mode; - } + [ThreadStatic] + private static Dictionary, MontgomeryDivisorData, ArrayView1D>>? _pow2ModKernel; - public readonly struct OrderKernelBuffers( - ArrayView1D phiFactors, - ArrayView1D phiExponents, - ArrayView1D workFactors, - ArrayView1D workExponents, - ArrayView1D candidates, - ArrayView1D stackIndex, - ArrayView1D stackExponent, - ArrayView1D stackProduct, - ArrayView1D result, - ArrayView1D status) - { - public readonly ArrayView1D PhiFactors = phiFactors; - public readonly ArrayView1D PhiExponents = phiExponents; - public readonly ArrayView1D WorkFactors = workFactors; - public readonly ArrayView1D WorkExponents = workExponents; - public readonly ArrayView1D Candidates = candidates; - public readonly ArrayView1D StackIndex = stackIndex; - public readonly ArrayView1D StackExponent = stackExponent; - public readonly ArrayView1D StackProduct = stackProduct; - public readonly ArrayView1D Result = result; - public readonly ArrayView1D Status = status; - } + [ThreadStatic] + private static Dictionary, GpuUInt128, ArrayView1D>>? _pow2ModWideKernel; - private delegate void OrderKernelLauncher( - AcceleratorStream stream, - Index1D extent, - ulong prime, - OrderKernelConfig config, - MontgomeryDivisorData divisor, - ArrayView1D primes, - ArrayView1D squares, - int primeCount, - OrderKernelBuffers buffers); + [ThreadStatic] + private static Dictionary, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>>? _partialFactorKernel; - private sealed class SmallPrimeDeviceCache - { - public MemoryBuffer1D? Primes; - public MemoryBuffer1D? Squares; - public int Count; - } + [ThreadStatic] + private static Dictionary? _orderKernel; - private const int WideStackThreshold = 8; - private static PrimeOrderGpuCapability s_capability = PrimeOrderGpuCapability.Default; + [ThreadStatic] + private static Dictionary? _smallPrimeDeviceCache; - private const int Pow2WindowSizeBits = 8; - private const int Pow2WindowOddPowerCount = 1 << (Pow2WindowSizeBits - 1); - private const ulong Pow2WindowFallbackThreshold = 32UL; - private const int HeuristicCandidateLimit = 512; - private const int HeuristicStackCapacity = 256; + public readonly struct OrderKernelConfig(ulong previousOrder, byte hasPreviousOrder, uint smallFactorLimit, int maxPowChecks, int mode) + { + public readonly ulong PreviousOrder = previousOrder; + public readonly byte HasPreviousOrder = hasPreviousOrder; + public readonly uint SmallFactorLimit = smallFactorLimit; + public readonly int MaxPowChecks = maxPowChecks; + public readonly int Mode = mode; + } - private const int GpuSmallPrimeFactorSlots = 64; + public readonly struct OrderKernelBuffers( + ArrayView1D phiFactors, + ArrayView1D phiExponents, + ArrayView1D workFactors, + ArrayView1D workExponents, + ArrayView1D candidates, + ArrayView1D stackIndex, + ArrayView1D stackExponent, + ArrayView1D stackProduct, + ArrayView1D result, + ArrayView1D status) + { + public readonly ArrayView1D PhiFactors = phiFactors; + public readonly ArrayView1D PhiExponents = phiExponents; + public readonly ArrayView1D WorkFactors = workFactors; + public readonly ArrayView1D WorkExponents = workExponents; + public readonly ArrayView1D Candidates = candidates; + public readonly ArrayView1D StackIndex = stackIndex; + public readonly ArrayView1D StackExponent = stackExponent; + public readonly ArrayView1D StackProduct = stackProduct; + public readonly ArrayView1D Result = result; + public readonly ArrayView1D Status = status; + } - internal static ConcurrentDictionary OverflowRegistry => OverflowedPrimes; - internal static ConcurrentDictionary OverflowRegistryWide => OverflowedPrimesWide; + private delegate void OrderKernelLauncher( + AcceleratorStream stream, + Index1D extent, + ulong prime, + OrderKernelConfig config, + MontgomeryDivisorData divisor, + ArrayView1D primes, + ArrayView1D squares, + int primeCount, + OrderKernelBuffers buffers); - internal static void OverrideCapabilitiesForTesting(PrimeOrderGpuCapability capability) - { - s_capability = capability; - } + private sealed class SmallPrimeDeviceCache + { + public MemoryBuffer1D? Primes; + public MemoryBuffer1D? Squares; + public int Count; + } - internal static void ResetCapabilitiesForTesting() - { - s_capability = PrimeOrderGpuCapability.Default; - } + private const int WideStackThreshold = 8; + private static PrimeOrderGpuCapability s_capability = PrimeOrderGpuCapability.Default; - private static SmallPrimeDeviceCache GetSmallPrimeDeviceCache(Accelerator accelerator, AcceleratorStream stream) - { - return SmallPrimeDeviceCaches.GetOrAdd(accelerator, acc => - { - uint[] primes = PrimesGenerator.SmallPrimes; - ulong[] squares = PrimesGenerator.SmallPrimesPow2; - var primeBuffer = acc.Allocate1D(primes.Length); - primeBuffer.View.CopyFromCPU(stream, primes); - var squareBuffer = acc.Allocate1D(squares.Length); - squareBuffer.View.CopyFromCPU(stream, squares); - return new SmallPrimeDeviceCache - { - Primes = primeBuffer, - Squares = squareBuffer, - Count = primes.Length, - }; - }); - } + private const int Pow2WindowSizeBits = 8; + private const int Pow2WindowOddPowerCount = 1 << (Pow2WindowSizeBits - 1); + private const ulong Pow2WindowFallbackThreshold = 32UL; + private const int HeuristicCandidateLimit = 512; + private const int HeuristicStackCapacity = 256; - private static Action, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> GetPartialFactorKernel(Accelerator accelerator) - { - return PartialFactorKernelCache.GetOrAdd(accelerator, static accel => - { - var loaded = accel.LoadAutoGroupedStreamKernel, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(PartialFactorKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>>(); - }); - } + private const int GpuSmallPrimeFactorSlots = 64; - public static bool TryPartialFactor( - ulong value, - uint limit, - Span primeTargets, - Span exponentTargets, - out int factorCount, - out ulong remaining, - out bool fullyFactored) - { - factorCount = 0; - remaining = value; - fullyFactored = false; + internal static ConcurrentDictionary OverflowRegistry => OverflowedPrimes; + internal static ConcurrentDictionary OverflowRegistryWide => OverflowedPrimesWide; - if (primeTargets.Length == 0 || exponentTargets.Length == 0) - { - return false; - } + internal static void OverrideCapabilitiesForTesting(PrimeOrderGpuCapability capability) + { + s_capability = capability; + } - primeTargets.Clear(); - exponentTargets.Clear(); + internal static void ResetCapabilitiesForTesting() + { + s_capability = PrimeOrderGpuCapability.Default; + } - if (!TryLaunchPartialFactorKernel( - value, - limit, - primeTargets, - exponentTargets, - out int extracted, - out ulong leftover, - out bool kernelFullyFactored)) - { - return false; - } + private static SmallPrimeDeviceCache GetSmallPrimeDeviceCache(Accelerator accelerator, AcceleratorStream stream) + { + var pool = _smallPrimeDeviceCache ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } - int capacity = Math.Min(primeTargets.Length, exponentTargets.Length); - if (extracted > capacity) - { - factorCount = 0; - remaining = value; - fullyFactored = false; - return false; - } + uint[] primes = PrimesGenerator.SmallPrimes; + ulong[] squares = PrimesGenerator.SmallPrimesPow2; + MemoryBuffer1D? primeBuffer; + MemoryBuffer1D? squareBuffer; + // lock (accelerator) + { + primeBuffer = accelerator.Allocate1D(primes.Length); + squareBuffer = accelerator.Allocate1D(squares.Length); + } + + primeBuffer.View.CopyFromCPU(stream, primes); + squareBuffer.View.CopyFromCPU(stream, squares); + cached = new SmallPrimeDeviceCache + { + Primes = primeBuffer, + Squares = squareBuffer, + Count = primes.Length, + }; - factorCount = extracted; - remaining = leftover; - fullyFactored = kernelFullyFactored && leftover == 1UL; - return true; - } + pool[accelerator] = cached; + return cached; + } - private static bool TryLaunchPartialFactorKernel( - ulong value, - uint limit, - Span primeTargets, - Span exponentTargets, - out int factorCount, - out ulong remaining, - out bool fullyFactored) - { - factorCount = 0; - remaining = value; - fullyFactored = false; + private static Action, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D> GetPartialFactorKernel(Accelerator accelerator) + { + var pool = _partialFactorKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } - try - { - GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; - AcceleratorStream stream = accelerator.CreateStream(); + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>(PartialFactorKernel); - var kernel = GetPartialFactorKernel(accelerator); - SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator, stream); + var kernel = KernelUtil.GetKernel(loaded); - // TODO: We should create / reallocate these buffer only once or if the new required length is bigger than capacity. - var factorBuffer = accelerator.Allocate1D(primeTargets.Length); - var exponentBuffer = accelerator.Allocate1D(exponentTargets.Length); - var countBuffer = accelerator.Allocate1D(1); - var remainingBuffer = accelerator.Allocate1D(1); - var fullyFactoredBuffer = accelerator.Allocate1D(1); + cached = kernel.CreateLauncherDelegate, ArrayView1D, int, int, ulong, uint, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D, ArrayView1D>>(); - // TODO: There is no need to clear these buffers because the kernel will always assign values within the required bounds. + pool[accelerator] = cached; + return cached; + } + + public static bool TryPartialFactor( + ulong value, + uint limit, + Span primeTargets, + Span exponentTargets, + out int factorCount, + out ulong remaining, + out bool fullyFactored) + { + factorCount = 0; + remaining = value; + fullyFactored = false; + + if (primeTargets.Length == 0 || exponentTargets.Length == 0) + { + return false; + } + + primeTargets.Clear(); + exponentTargets.Clear(); + + if (!TryLaunchPartialFactorKernel( + value, + limit, + primeTargets, + exponentTargets, + out int extracted, + out ulong leftover, + out bool kernelFullyFactored)) + { + return false; + } + + int capacity = Math.Min(primeTargets.Length, exponentTargets.Length); + if (extracted > capacity) + { + factorCount = 0; + remaining = value; + fullyFactored = false; + return false; + } + + factorCount = extracted; + remaining = leftover; + fullyFactored = kernelFullyFactored && leftover == 1UL; + return true; + } + + private static bool TryLaunchPartialFactorKernel( + ulong value, + uint limit, + Span primeTargets, + Span exponentTargets, + out int factorCount, + out ulong remaining, + out bool fullyFactored) + { + factorCount = 0; + remaining = value; + fullyFactored = false; + + try + { + GpuPrimeWorkLimiter.Acquire(); + Accelerator accelerator = AcceleratorPool.Shared.Rent(); + AcceleratorStream stream = accelerator.CreateStream(); + + SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator, stream); + + // TODO: We should create / reallocate these buffer only once or if the new required length is bigger than capacity. + MemoryBuffer1D? factorBuffer; + MemoryBuffer1D? exponentBuffer; + MemoryBuffer1D? countBuffer; + MemoryBuffer1D? remainingBuffer; + MemoryBuffer1D? fullyFactoredBuffer; + // lock(accelerator) + { + factorBuffer = accelerator.Allocate1D(primeTargets.Length); + exponentBuffer = accelerator.Allocate1D(exponentTargets.Length); + countBuffer = accelerator.Allocate1D(1); + remainingBuffer = accelerator.Allocate1D(1); + fullyFactoredBuffer = accelerator.Allocate1D(1); + } + + // TODO: There is no need to clear these buffers because the kernel will always assign values within the required bounds. // factorBuffer.MemSetToZero(stream); - // exponentBuffer.MemSetToZero(stream); - // countBuffer.MemSetToZero(stream); - // remainingBuffer.MemSetToZero(stream); - // fullyFactoredBuffer.MemSetToZero(stream); + // exponentBuffer.MemSetToZero(stream); + // countBuffer.MemSetToZero(stream); + // remainingBuffer.MemSetToZero(stream); + // fullyFactoredBuffer.MemSetToZero(stream); - kernel( - stream, - 1, - cache.Primes!.View, - cache.Squares!.View, - cache.Count, - primeTargets.Length, - value, - limit, - factorBuffer.View, - exponentBuffer.View, - countBuffer.View, - remainingBuffer.View, - fullyFactoredBuffer.View); + var partialFactorKernel = GetPartialFactorKernel(accelerator); + + partialFactorKernel( + stream, + 1, + cache.Primes!.View, + cache.Squares!.View, + cache.Count, + primeTargets.Length, + value, + limit, + factorBuffer.View, + exponentBuffer.View, + countBuffer.View, + remainingBuffer.View, + fullyFactoredBuffer.View); - countBuffer.View.CopyToCPU(stream, ref factorCount, 1); - factorBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(primeTargets), primeTargets.Length); - exponentBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(exponentTargets), exponentTargets.Length); - remainingBuffer.View.CopyToCPU(stream, ref remaining, 1); + countBuffer.View.CopyToCPU(stream, ref factorCount, 1); + factorBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(primeTargets), primeTargets.Length); + exponentBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(exponentTargets), exponentTargets.Length); + remainingBuffer.View.CopyToCPU(stream, ref remaining, 1); - byte fullyFactoredFlag = 0; - fullyFactoredBuffer.View.CopyToCPU(stream, ref fullyFactoredFlag, 1); + byte fullyFactoredFlag = 0; + fullyFactoredBuffer.View.CopyToCPU(stream, ref fullyFactoredFlag, 1); stream.Synchronize(); stream.Dispose(); fullyFactored = fullyFactoredFlag != 0; - factorCount = Math.Min(factorCount, primeTargets.Length); + factorCount = Math.Min(factorCount, primeTargets.Length); - // TODO: These buffers shouldn't be disposed but rather reused accross call with the assigned accelerator. - factorBuffer.Dispose(); - exponentBuffer.Dispose(); - countBuffer.Dispose(); - remainingBuffer.Dispose(); + // TODO: These buffers shouldn't be disposed but rather reused accross call with the assigned accelerator. + factorBuffer.Dispose(); + exponentBuffer.Dispose(); + countBuffer.Dispose(); + remainingBuffer.Dispose(); fullyFactoredBuffer.Dispose(); - GpuPrimeWorkLimiter.Release(); - return true; - } - catch (CLException ex) - { - Console.WriteLine($"GPU ERROR ({ex.Error}): {ex.Message}"); - factorCount = 0; - remaining = value; - fullyFactored = false; - return false; - } - catch (Exception ex) when (ex is AcceleratorException or InternalCompilerException or NotSupportedException or InvalidOperationException or AggregateException) - { - Console.WriteLine($"GPU ERROR: {ex.Message}"); - factorCount = 0; - remaining = value; - fullyFactored = false; - return false; - } - } + GpuPrimeWorkLimiter.Release(); + return true; + } + catch (CLException ex) + { + Console.WriteLine($"GPU ERROR ({ex.Error}): {ex.Message}"); + factorCount = 0; + remaining = value; + fullyFactored = false; + return false; + } + catch (Exception ex) when (ex is AcceleratorException or InternalCompilerException or NotSupportedException or InvalidOperationException or AggregateException) + { + Console.WriteLine($"GPU ERROR: {ex.Message}"); + factorCount = 0; + remaining = value; + fullyFactored = false; + return false; + } + } - public static GpuPow2ModStatus TryPow2Mod(ulong exponent, ulong prime, out ulong remainder, in MontgomeryDivisorData divisorData) - { - Span exponents = stackalloc ulong[1]; - Span remainders = stackalloc ulong[1]; - exponents[0] = exponent; + public static GpuPow2ModStatus TryPow2Mod(ulong exponent, ulong prime, out ulong remainder, in MontgomeryDivisorData divisorData) + { + Span exponents = stackalloc ulong[1]; + Span remainders = stackalloc ulong[1]; + exponents[0] = exponent; - GpuPow2ModStatus status = TryPow2ModBatch(exponents, prime, remainders, divisorData); - remainder = remainders[0]; - return status; - } + GpuPow2ModStatus status = TryPow2ModBatch(exponents, prime, remainders, divisorData); + remainder = remainders[0]; + return status; + } - public static GpuPow2ModStatus TryPow2ModBatch(ReadOnlySpan exponents, ulong prime, Span remainders, in MontgomeryDivisorData divisorData) - { - // This will never occur in production code - // if (exponents.Length == 0) - // { - // return GpuPow2ModStatus.Success; - // } + public static GpuPow2ModStatus TryPow2ModBatch(ReadOnlySpan exponents, ulong prime, Span remainders, in MontgomeryDivisorData divisorData) + { + // This will never occur in production code + // if (exponents.Length == 0) + // { + // return GpuPow2ModStatus.Success; + // } - // This will never occur in production code - // if (remainders.Length < exponents.Length) - // { - // throw new ArgumentException("Remainder span is shorter than the exponent span.", nameof(remainders)); - // } + // This will never occur in production code + // if (remainders.Length < exponents.Length) + // { + // throw new ArgumentException("Remainder span is shorter than the exponent span.", nameof(remainders)); + // } - // This will never occur in production code - // if (remainders.Length > exponents.Length) - // throw new ArgumentException("Remainder span is longer than the exponent span.", nameof(remainders)); + // This will never occur in production code + // if (remainders.Length > exponents.Length) + // throw new ArgumentException("Remainder span is longer than the exponent span.", nameof(remainders)); - // We need to clear more, but we save on allocations - // Span target = remainders[..exponents.Length]; - remainders.Clear(); + // We need to clear more, but we save on allocations + // Span target = remainders[..exponents.Length]; + remainders.Clear(); - // This will never occur in production code - // if (prime <= 1UL) - // { - // return GpuPow2ModStatus.Unavailable; - // } + // This will never occur in production code + // if (prime <= 1UL) + // { + // return GpuPow2ModStatus.Unavailable; + // } - ConcurrentDictionary overflowRegistry = OverflowedPrimes; + ConcurrentDictionary overflowRegistry = OverflowedPrimes; - if (overflowRegistry.ContainsKey(prime)) - { - return GpuPow2ModStatus.Overflow; - } + if (overflowRegistry.ContainsKey(prime)) + { + return GpuPow2ModStatus.Overflow; + } - PrimeOrderGpuCapability capability = s_capability; + PrimeOrderGpuCapability capability = s_capability; - if (prime.GetBitLength() > capability.ModulusBits) - { - overflowRegistry[prime] = 0; - return GpuPow2ModStatus.Overflow; - } + if (prime.GetBitLength() > capability.ModulusBits) + { + overflowRegistry[prime] = 0; + return GpuPow2ModStatus.Overflow; + } - for (int i = 0; i < exponents.Length; i++) - { - if (exponents[i].GetBitLength() > capability.ExponentBits) - { - return GpuPow2ModStatus.Overflow; - } - } + for (int i = 0; i < exponents.Length; i++) + { + if (exponents[i].GetBitLength() > capability.ExponentBits) + { + return GpuPow2ModStatus.Overflow; + } + } - bool computed = TryComputeOnGpu(exponents, prime, divisorData, remainders); - return computed ? GpuPow2ModStatus.Success : GpuPow2ModStatus.Unavailable; - } + bool computed = TryComputeOnGpu(exponents, prime, divisorData, remainders); + return computed ? GpuPow2ModStatus.Success : GpuPow2ModStatus.Unavailable; + } - public static GpuPow2ModStatus TryPow2Mod(in UInt128 exponent, in UInt128 prime, out UInt128 remainder) - { - Span exponents = stackalloc UInt128[1]; - Span remainders = stackalloc UInt128[1]; - exponents[0] = exponent; + public static GpuPow2ModStatus TryPow2Mod(in UInt128 exponent, in UInt128 prime, out UInt128 remainder) + { + Span exponents = stackalloc UInt128[1]; + Span remainders = stackalloc UInt128[1]; + exponents[0] = exponent; - GpuPow2ModStatus status = TryPow2ModBatch(exponents, prime, remainders); - remainder = remainders[0]; - return status; - } + GpuPow2ModStatus status = TryPow2ModBatch(exponents, prime, remainders); + remainder = remainders[0]; + return status; + } - public static GpuPow2ModStatus TryPow2ModBatch(ReadOnlySpan exponents, UInt128 prime, Span remainders) - { - return TryPow2ModBatchInternal(exponents, prime, remainders); - } + public static GpuPow2ModStatus TryPow2ModBatch(ReadOnlySpan exponents, UInt128 prime, Span remainders) + { + return TryPow2ModBatchInternal(exponents, prime, remainders); + } - internal static bool TryCalculateOrder( - ulong prime, - ulong? previousOrder, - PrimeOrderCalculator.PrimeOrderSearchConfig config, - in MontgomeryDivisorData divisorData, - out ulong order) - { - order = 0UL; + internal static bool TryCalculateOrder( + ulong prime, + ulong? previousOrder, + PrimeOrderCalculator.PrimeOrderSearchConfig config, + in MontgomeryDivisorData divisorData, + out ulong order) + { + order = 0UL; - GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; - AcceleratorStream stream = accelerator.CreateStream(); + GpuPrimeWorkLimiter.Acquire(); + Accelerator accelerator = AcceleratorPool.Shared.Rent(); + AcceleratorStream stream = accelerator.CreateStream(); - var kernel = GetOrderKernel(accelerator); - SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator, stream); + SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator, stream); // TODO: These buffers should be created reallocated once and assigned to an accelerator. Callers should use their own // thread static cache to prevent other threads from using taking accelerators with pre-allocated buffers if they don't // use them. - var phiFactorBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); - var phiExponentBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); - var workFactorBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); - var workExponentBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); - var candidateBuffer = accelerator.Allocate1D(HeuristicCandidateLimit); - var stackIndexBuffer = accelerator.Allocate1D(HeuristicStackCapacity); - var stackExponentBuffer = accelerator.Allocate1D(HeuristicStackCapacity); - var stackProductBuffer = accelerator.Allocate1D(HeuristicStackCapacity); - var resultBuffer = accelerator.Allocate1D(1); - var statusBuffer = accelerator.Allocate1D(1); + MemoryBuffer1D? phiFactorBuffer; + MemoryBuffer1D? phiExponentBuffer; + MemoryBuffer1D? workFactorBuffer; + MemoryBuffer1D? workExponentBuffer; + MemoryBuffer1D? candidateBuffer; + MemoryBuffer1D? stackIndexBuffer; + MemoryBuffer1D? stackExponentBuffer; + MemoryBuffer1D? stackProductBuffer; + MemoryBuffer1D? resultBuffer; + MemoryBuffer1D? statusBuffer; + + // lock(accelerator) + { + phiFactorBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); + phiExponentBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); + workFactorBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); + workExponentBuffer = accelerator.Allocate1D(GpuSmallPrimeFactorSlots); + candidateBuffer = accelerator.Allocate1D(HeuristicCandidateLimit); + stackIndexBuffer = accelerator.Allocate1D(HeuristicStackCapacity); + stackExponentBuffer = accelerator.Allocate1D(HeuristicStackCapacity); + stackProductBuffer = accelerator.Allocate1D(HeuristicStackCapacity); + resultBuffer = accelerator.Allocate1D(1); + statusBuffer = accelerator.Allocate1D(1); + } // TODO: Remove the cleaning after the order kernel is modified to always set the result. phiFactorBuffer.MemSetToZero(stream); - phiExponentBuffer.MemSetToZero(stream); - workFactorBuffer.MemSetToZero(stream); - workExponentBuffer.MemSetToZero(stream); - candidateBuffer.MemSetToZero(stream); - stackIndexBuffer.MemSetToZero(stream); - stackExponentBuffer.MemSetToZero(stream); - stackProductBuffer.MemSetToZero(stream); - resultBuffer.MemSetToZero(stream); - statusBuffer.MemSetToZero(stream); + phiExponentBuffer.MemSetToZero(stream); + workFactorBuffer.MemSetToZero(stream); + workExponentBuffer.MemSetToZero(stream); + candidateBuffer.MemSetToZero(stream); + stackIndexBuffer.MemSetToZero(stream); + stackExponentBuffer.MemSetToZero(stream); + stackProductBuffer.MemSetToZero(stream); + resultBuffer.MemSetToZero(stream); + statusBuffer.MemSetToZero(stream); - uint limit = config.SmallFactorLimit == 0 ? uint.MaxValue : config.SmallFactorLimit; - byte hasPrevious = previousOrder.HasValue ? (byte)1 : (byte)0; - ulong previousValue = previousOrder ?? 0UL; + uint limit = config.SmallFactorLimit == 0 ? uint.MaxValue : config.SmallFactorLimit; + byte hasPrevious = previousOrder.HasValue ? (byte)1 : (byte)0; + ulong previousValue = previousOrder ?? 0UL; - var kernelConfig = new OrderKernelConfig(previousValue, hasPrevious, limit, config.MaxPowChecks, (int)config.Mode); - var buffers = new OrderKernelBuffers( - phiFactorBuffer.View, - phiExponentBuffer.View, - workFactorBuffer.View, - workExponentBuffer.View, - candidateBuffer.View, - stackIndexBuffer.View, - stackExponentBuffer.View, - stackProductBuffer.View, - resultBuffer.View, - statusBuffer.View); + var kernelConfig = new OrderKernelConfig(previousValue, hasPrevious, limit, config.MaxPowChecks, (int)config.Mode); + var buffers = new OrderKernelBuffers( + phiFactorBuffer.View, + phiExponentBuffer.View, + workFactorBuffer.View, + workExponentBuffer.View, + candidateBuffer.View, + stackIndexBuffer.View, + stackExponentBuffer.View, + stackProductBuffer.View, + resultBuffer.View, + statusBuffer.View); - kernel( - stream, - 1, - prime, - kernelConfig, - divisorData, - cache.Primes!.View, - cache.Squares!.View, - cache.Count, - buffers); + var orderKernel = GetOrderKernel(accelerator); + orderKernel( + stream, + 1, + prime, + kernelConfig, + divisorData, + cache.Primes!.View, + cache.Squares!.View, + cache.Count, + buffers); - byte status = 0; - statusBuffer.View.CopyToCPU(stream, ref status, 1); - resultBuffer.View.CopyToCPU(stream, ref order, 1); + byte status = 0; + statusBuffer.View.CopyToCPU(stream, ref status, 1); + resultBuffer.View.CopyToCPU(stream, ref order, 1); stream.Synchronize(); - + stream.Dispose(); phiFactorBuffer.Dispose(); phiExponentBuffer.Dispose(); @@ -443,105 +496,115 @@ internal static partial class PrimeOrderGpuHeuristics statusBuffer.Dispose(); GpuPrimeWorkLimiter.Release(); - PrimeOrderKernelStatus kernelStatus = (PrimeOrderKernelStatus)status; - if (kernelStatus == PrimeOrderKernelStatus.Fallback) - { - return false; - } + PrimeOrderKernelStatus kernelStatus = (PrimeOrderKernelStatus)status; + if (kernelStatus == PrimeOrderKernelStatus.Fallback) + { + return false; + } - if (kernelStatus == PrimeOrderKernelStatus.PollardOverflow) - { - throw new InvalidOperationException("GPU Pollard Rho stack overflow; increase HeuristicStackCapacity."); - } + if (kernelStatus == PrimeOrderKernelStatus.PollardOverflow) + { + throw new InvalidOperationException("GPU Pollard Rho stack overflow; increase HeuristicStackCapacity."); + } - if (kernelStatus == PrimeOrderKernelStatus.FactoringFailure) - { - order = 0UL; - } + if (kernelStatus == PrimeOrderKernelStatus.FactoringFailure) + { + order = 0UL; + } - return order != 0UL; - } + return order != 0UL; + } - private static OrderKernelLauncher GetOrderKernel(Accelerator accelerator) - { - return OrderKernelCache.GetOrAdd(accelerator, static accel => - { - var loaded = accel.LoadAutoGroupedStreamKernel, ArrayView1D, int, OrderKernelBuffers>(CalculateOrderKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate(); - }); - } + private static OrderKernelLauncher GetOrderKernel(Accelerator accelerator) + { + var pool = _orderKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } - private static GpuPow2ModStatus TryPow2ModBatchInternal(ReadOnlySpan exponents, UInt128 prime, Span remainders) - { - if (exponents.Length == 0) - { - return GpuPow2ModStatus.Success; - } + var loaded = accelerator.LoadAutoGroupedStreamKernel, ArrayView1D, int, OrderKernelBuffers>(CalculateOrderKernel); + var kernel = KernelUtil.GetKernel(loaded); + cached = kernel.CreateLauncherDelegate(); + pool[accelerator] = cached; + return cached; + } - if (remainders.Length < exponents.Length) - { - throw new ArgumentException("Remainder span is shorter than the exponent span.", nameof(remainders)); - } + private static GpuPow2ModStatus TryPow2ModBatchInternal(ReadOnlySpan exponents, UInt128 prime, Span remainders) + { + if (exponents.Length == 0) + { + return GpuPow2ModStatus.Success; + } - Span target = remainders.Slice(0, exponents.Length); - target.Clear(); + if (remainders.Length < exponents.Length) + { + throw new ArgumentException("Remainder span is shorter than the exponent span.", nameof(remainders)); + } - ConcurrentDictionary overflowRegistryWide = OverflowedPrimesWide; + Span target = remainders[..exponents.Length]; + target.Clear(); - if (overflowRegistryWide.ContainsKey(prime)) - { - return GpuPow2ModStatus.Overflow; - } + ConcurrentDictionary overflowRegistryWide = OverflowedPrimesWide; - PrimeOrderGpuCapability capability = s_capability; + if (overflowRegistryWide.ContainsKey(prime)) + { + return GpuPow2ModStatus.Overflow; + } - if (prime.GetBitLength() > capability.ModulusBits) - { - overflowRegistryWide[prime] = 0; - return GpuPow2ModStatus.Overflow; - } + PrimeOrderGpuCapability capability = s_capability; - for (int i = 0; i < exponents.Length; i++) - { - if (exponents[i].GetBitLength() > capability.ExponentBits) - { - return GpuPow2ModStatus.Overflow; - } - } + if (prime.GetBitLength() > capability.ModulusBits) + { + overflowRegistryWide[prime] = 0; + return GpuPow2ModStatus.Overflow; + } - bool computed = TryComputeOnGpuWide(exponents, prime, target); - return computed ? GpuPow2ModStatus.Success : GpuPow2ModStatus.Unavailable; - } + for (int i = 0; i < exponents.Length; i++) + { + if (exponents[i].GetBitLength() > capability.ExponentBits) + { + return GpuPow2ModStatus.Overflow; + } + } - private static bool TryComputeOnGpu(ReadOnlySpan exponents, ulong prime, in MontgomeryDivisorData divisorData, Span results) - { - GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; - AcceleratorStream? stream = accelerator.CreateStream(); - var kernel = GetPow2ModKernel(accelerator); + bool computed = TryComputeOnGpuWide(exponents, prime, target); + return computed ? GpuPow2ModStatus.Success : GpuPow2ModStatus.Unavailable; + } + + private static bool TryComputeOnGpu(ReadOnlySpan exponents, ulong prime, in MontgomeryDivisorData divisorData, Span results) + { + GpuPrimeWorkLimiter.Acquire(); + Accelerator accelerator = AcceleratorPool.Shared.Rent(); + AcceleratorStream? stream = accelerator.CreateStream(); // TODO: These buffers should be allocated once per accelerator and only reallocated when the new length exceeds capacity. // Modify the callers to use their own pool of buffers per accelerator, so that other threads don't use the accelerators // with out preallocated buffers. Share the pool with Pow2ModWide kernel. - var exponentBuffer = accelerator.Allocate1D(exponents.Length); - var remainderBuffer = accelerator.Allocate1D(exponents.Length); + MemoryBuffer1D? exponentBuffer; + MemoryBuffer1D? remainderBuffer; + // lock(accelerator) + { + exponentBuffer = accelerator.Allocate1D(exponents.Length); + remainderBuffer = accelerator.Allocate1D(exponents.Length); + } - exponentBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(exponents), exponents.Length); + exponentBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(exponents), exponents.Length); // Pow2Mod kernel always assigns the required output elements so we don't need to worry about clearing these. // remainderBuffer.MemSetToZero(stream); + var pow2ModKernel = GetPow2ModKernel(accelerator); try { - kernel(stream, exponents.Length, exponentBuffer.View, divisorData, remainderBuffer.View); - remainderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(results), exponents.Length); + pow2ModKernel(stream, exponents.Length, exponentBuffer.View, divisorData, remainderBuffer.View); + remainderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(results), exponents.Length); stream.Synchronize(); stream.Dispose(); stream = null; - exponentBuffer.Dispose(); - remainderBuffer.Dispose(); - GpuPrimeWorkLimiter.Release(); + exponentBuffer.Dispose(); + remainderBuffer.Dispose(); + GpuPrimeWorkLimiter.Release(); } catch (Exception) { @@ -553,127 +616,136 @@ internal static partial class PrimeOrderGpuHeuristics throw; } - return true; - } + return true; + } - private static void ComputePow2ModCpu(ReadOnlySpan exponents, ulong prime, in MontgomeryDivisorData divisorData, Span results) - { - int length = exponents.Length; - for (int i = 0; i < length; i++) - { - results[i] = Pow2ModCpu(exponents[i], prime, divisorData); - } - } + private static void ComputePow2ModCpu(ReadOnlySpan exponents, ulong prime, in MontgomeryDivisorData divisorData, Span results) + { + int length = exponents.Length; + for (int i = 0; i < length; i++) + { + results[i] = Pow2ModCpu(exponents[i], prime, divisorData); + } + } - private static bool TryComputeOnGpuWide(ReadOnlySpan exponents, UInt128 prime, Span results) - { - int length = exponents.Length; - if (length == 0) - { - return true; - } + private static bool TryComputeOnGpuWide(ReadOnlySpan exponents, UInt128 prime, Span results) + { + int length = exponents.Length; + if (length == 0) + { + return true; + } - GpuUInt128[]? rentedExponents = null; - GpuUInt128[]? rentedResults = null; - GpuPrimeWorkLimiter.Acquire(); + GpuUInt128[]? rentedExponents = null; + GpuUInt128[]? rentedResults = null; + GpuPrimeWorkLimiter.Acquire(); ArrayPool gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool; - try - { - Accelerator accelerator = SharedGpuContext.Accelerator; - AcceleratorStream stream = accelerator.CreateStream(); - var kernel = GetPow2ModWideKernel(accelerator); + try + { + Accelerator accelerator = AcceleratorPool.Shared.Rent(); + AcceleratorStream stream = accelerator.CreateStream(); // TODO: These buffers should be allocated once per accelerator and only reallocated when the new length exceeds capacity. // Modify the callers to use their own pool of buffers per accelerator, so that other threads don't use the accelerators // with out preallocated buffers. Share the pool with Pow2Mod kernel. - var exponentBuffer = accelerator.Allocate1D(length); - var remainderBuffer = accelerator.Allocate1D(length); + MemoryBuffer1D? exponentBuffer; + MemoryBuffer1D? remainderBuffer; + // lock(accelerator) + { + exponentBuffer = accelerator.Allocate1D(length); + remainderBuffer = accelerator.Allocate1D(length); + } Span exponentSpan = length <= WideStackThreshold - ? stackalloc GpuUInt128[length] - : new Span(rentedExponents = gpuUInt128Pool.Rent(length), 0, length); + ? stackalloc GpuUInt128[length] + : new Span(rentedExponents = gpuUInt128Pool.Rent(length), 0, length); - for (int i = 0; i < length; i++) - { - exponentSpan[i] = (GpuUInt128)exponents[i]; - } + for (int i = 0; i < length; i++) + { + exponentSpan[i] = (GpuUInt128)exponents[i]; + } - exponentBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(exponentSpan), length); + exponentBuffer.View.CopyFromCPU(stream, ref MemoryMarshal.GetReference(exponentSpan), length); // Pow2ModWide kernel always assigns the required output elements so we don't need to worry about clearing these. - // remainderBuffer.MemSetToZero(stream); + // remainderBuffer.MemSetToZero(stream); - GpuUInt128 modulus = (GpuUInt128)prime; - kernel(stream, length, exponentBuffer.View, modulus, remainderBuffer.View); + GpuUInt128 modulus = (GpuUInt128)prime; + var pow2ModWideKernel = GetPow2ModWideKernel(accelerator); + pow2ModWideKernel(stream, length, exponentBuffer.View, modulus, remainderBuffer.View); - Span resultSpan = length <= WideStackThreshold - ? stackalloc GpuUInt128[length] - : new Span(rentedResults = gpuUInt128Pool.Rent(length), 0, length); + Span resultSpan = length <= WideStackThreshold + ? stackalloc GpuUInt128[length] + : new Span(rentedResults = gpuUInt128Pool.Rent(length), 0, length); - remainderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(resultSpan), length); - stream.Synchronize(); + remainderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(resultSpan), length); + stream.Synchronize(); stream.Dispose(); - for (int i = 0; i < length; i++) - { - results[i] = (UInt128)resultSpan[i]; - } + for (int i = 0; i < length; i++) + { + results[i] = (UInt128)resultSpan[i]; + } - exponentBuffer.Dispose(); - remainderBuffer.Dispose(); - return true; - } - catch (Exception) - { - Console.WriteLine($"Exception when computing {prime} with exponents: {string.Join(", ", exponents.ToArray())}"); - throw; - } - finally - { - if (rentedExponents is not null) - { - gpuUInt128Pool.Return(rentedExponents, clearArray: false); - gpuUInt128Pool.Return(rentedResults!, clearArray: false); - } + exponentBuffer.Dispose(); + remainderBuffer.Dispose(); + return true; + } + catch (Exception) + { + Console.WriteLine($"Exception when computing {prime} with exponents: {string.Join(", ", exponents.ToArray())}"); + throw; + } + finally + { + if (rentedExponents is not null) + { + gpuUInt128Pool.Return(rentedExponents, clearArray: false); + gpuUInt128Pool.Return(rentedResults!, clearArray: false); + } GpuPrimeWorkLimiter.Release(); - } - } + } + } - private static void ComputePow2ModCpuWide(ReadOnlySpan exponents, UInt128 prime, Span results) - { - int length = exponents.Length; - for (int i = 0; i < length; i++) - { - results[i] = exponents[i].Pow2MontgomeryModWindowed(prime); - } - } + private static void ComputePow2ModCpuWide(ReadOnlySpan exponents, UInt128 prime, Span results) + { + int length = exponents.Length; + for (int i = 0; i < length; i++) + { + results[i] = exponents[i].Pow2MontgomeryModWindowed(prime); + } + } - private static ulong Pow2ModCpu(ulong exponent, ulong modulus, in MontgomeryDivisorData divisorData) - { - if (modulus <= 1UL) - { - return 0UL; - } + private static ulong Pow2ModCpu(ulong exponent, ulong modulus, in MontgomeryDivisorData divisorData) + { + if (modulus <= 1UL) + { + return 0UL; + } - return ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisorData, exponent); - } + return ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisorData, exponent); + } - private static Action, MontgomeryDivisorData, ArrayView1D> GetPow2ModKernel(Accelerator accelerator) - { - return Pow2ModKernelCache.GetOrAdd(accelerator, static accel => - { - var loaded = accel.LoadAutoGroupedStreamKernel, MontgomeryDivisorData, ArrayView1D>(Pow2ModKernel); - var kernel = KernelUtil.GetKernel(loaded); - return kernel.CreateLauncherDelegate, MontgomeryDivisorData, ArrayView1D>>(); - }); - } + private static Action, MontgomeryDivisorData, ArrayView1D> GetPow2ModKernel(Accelerator accelerator) + { + var pool = _pow2ModKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + var loaded = accelerator.LoadAutoGroupedStreamKernel, MontgomeryDivisorData, ArrayView1D>(Pow2ModKernel); + var kernel = KernelUtil.GetKernel(loaded); + cached = kernel.CreateLauncherDelegate, MontgomeryDivisorData, ArrayView1D>>(); + pool[accelerator] = cached; + return cached; + } - internal readonly record struct PrimeOrderGpuCapability(int ModulusBits, int ExponentBits) - { - public static PrimeOrderGpuCapability Default => new(128, 128); - } + internal readonly record struct PrimeOrderGpuCapability(int ModulusBits, int ExponentBits) + { + public static PrimeOrderGpuCapability Default => new(128, 128); + } } diff --git a/PerfectNumbers.Core/Gpu/SharedGpuContext.cs b/PerfectNumbers.Core/Gpu/SharedGpuContext.cs index bb62d32..07008a4 100644 --- a/PerfectNumbers.Core/Gpu/SharedGpuContext.cs +++ b/PerfectNumbers.Core/Gpu/SharedGpuContext.cs @@ -1,5 +1,4 @@ -using System; -using System.Threading; +using System.Runtime.CompilerServices; using ILGPU; using ILGPU.Runtime; @@ -7,14 +6,24 @@ namespace PerfectNumbers.Core.Gpu; public static class SharedGpuContext { - // TODO: Remove Lazy in favor for better performance with nullable fields. - private static readonly Lazy s_context = new(() => ILGPU.Context.CreateDefault(), LazyThreadSafetyMode.ExecutionAndPublication); - private static readonly Lazy s_device = new(() => s_context.Value.GetPreferredDevice(false), LazyThreadSafetyMode.ExecutionAndPublication); - private static readonly Lazy s_accelerator = new(() => s_device.Value.CreateAccelerator(s_context.Value), LazyThreadSafetyMode.ExecutionAndPublication); + public static readonly Context Context = Context.CreateDefault(); + public static readonly Device Device = Context.GetPreferredDevice(false); - internal static Context Context => s_context.Value; + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Accelerator CreateAccelerator() => Device.CreateAccelerator(Context); +} - internal static Device Device => s_device.Value; +public readonly struct GpuContext +{ + public readonly Context Context; + public readonly Device Device; - internal static Accelerator Accelerator => s_accelerator.Value; + public GpuContext() + { + var context = Context.CreateDefault(); + Context = context; + Device = context.GetPreferredDevice(false); + } + + public Accelerator CreateAccelerator() => Device.CreateAccelerator(Context); } diff --git a/PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs b/PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs index 4ea3259..20f9b16 100644 --- a/PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs +++ b/PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs @@ -2,7 +2,6 @@ using Open.Numeric.Primes; using ILGPU.Runtime; using System.Runtime.CompilerServices; using PerfectNumbers.Core.Gpu; -using System.Runtime.InteropServices; namespace PerfectNumbers.Core; @@ -41,7 +40,6 @@ public sealed class HeuristicCombinedPrimeTester CombinedDivisorsEnding9Squares = CombinedDivisorsEnding9OneAOneBSquares; } - public static void EnsureInitialized() { // Intentionally left blank. Accessing this method forces the static constructor to run. @@ -229,7 +227,7 @@ public sealed class HeuristicCombinedPrimeTester var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(1); var accelerator = gpu.Accelerator; var stream = gpu.Stream; - var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.HeuristicTrialDivisionKernel; + var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetHeuristicTrialDivisionKernel(accelerator); var flagView1D = gpu.HeuristicFlag.View; var flagView = flagView1D.AsContiguous(); @@ -267,109 +265,97 @@ public sealed class HeuristicCombinedPrimeTester return n; } - private static ReadOnlySpan GetCombinedDivisors(byte nMod10) + private static ReadOnlySpan GetCombinedDivisors(byte nMod10) => nMod10 switch { - return nMod10 switch + 1 => CombinedDivisorsEnding1, + 3 => CombinedDivisorsEnding3, + 7 => CombinedDivisorsEnding7, + 9 => CombinedDivisorsEnding9, + _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), + }; + + private static ReadOnlySpan GetCombinedDivisorSquares(byte nMod10) => nMod10 switch + { + 1 => CombinedDivisorsEnding1Squares, + 3 => CombinedDivisorsEnding3Squares, + 7 => CombinedDivisorsEnding7Squares, + 9 => CombinedDivisorsEnding9Squares, + _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), + }; + + internal static ReadOnlySpan GetCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern) => pattern switch + { + CombinedDivisorPattern.TwoAOneB => nMod10 switch { - 1 => CombinedDivisorsEnding1, - 3 => CombinedDivisorsEnding3, - 7 => CombinedDivisorsEnding7, - 9 => CombinedDivisorsEnding9, + 1 => CombinedDivisorsEnding1TwoAOneB, + 3 => CombinedDivisorsEnding3TwoAOneB, + 7 => CombinedDivisorsEnding7TwoAOneB, + 9 => CombinedDivisorsEnding9TwoAOneB, _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), - }; - } - - private static ReadOnlySpan GetCombinedDivisorSquares(byte nMod10) - { - return nMod10 switch + }, + CombinedDivisorPattern.OneAOneB => nMod10 switch { - 1 => CombinedDivisorsEnding1Squares, - 3 => CombinedDivisorsEnding3Squares, - 7 => CombinedDivisorsEnding7Squares, - 9 => CombinedDivisorsEnding9Squares, + 1 => CombinedDivisorsEnding1OneAOneB, + 3 => CombinedDivisorsEnding3OneAOneB, + 7 => CombinedDivisorsEnding7OneAOneB, + 9 => CombinedDivisorsEnding9OneAOneB, + _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), + }, + CombinedDivisorPattern.ThreeAOneB => nMod10 switch + { + 1 => CombinedDivisorsEnding1ThreeAOneB, + 3 => CombinedDivisorsEnding3ThreeAOneB, + 7 => CombinedDivisorsEnding7ThreeAOneB, + 9 => CombinedDivisorsEnding9ThreeAOneB, + _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), + }, + CombinedDivisorPattern.ThreeATwoB => nMod10 switch + { + 1 => CombinedDivisorsEnding1ThreeATwoB, + 3 => CombinedDivisorsEnding3ThreeATwoB, + 7 => CombinedDivisorsEnding7ThreeATwoB, + 9 => CombinedDivisorsEnding9ThreeATwoB, + _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), + }, + _ => throw new InvalidOperationException($"Unsupported combined divisor pattern: {pattern}."), + }; + + internal static ReadOnlySpan GetCombinedDivisorSquares(byte nMod10, CombinedDivisorPattern pattern) => pattern switch + { + CombinedDivisorPattern.TwoAOneB => nMod10 switch + { + 1 => CombinedDivisorsEnding1TwoAOneBSquares, + 3 => CombinedDivisorsEnding3TwoAOneBSquares, + 7 => CombinedDivisorsEnding7TwoAOneBSquares, + 9 => CombinedDivisorsEnding9TwoAOneBSquares, _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), - }; - } - - internal static ReadOnlySpan GetCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern) - { - return pattern switch + }, + CombinedDivisorPattern.OneAOneB => nMod10 switch { - CombinedDivisorPattern.TwoAOneB => nMod10 switch - { - 1 => CombinedDivisorsEnding1TwoAOneB, - 3 => CombinedDivisorsEnding3TwoAOneB, - 7 => CombinedDivisorsEnding7TwoAOneB, - 9 => CombinedDivisorsEnding9TwoAOneB, - _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), - }, - CombinedDivisorPattern.OneAOneB => nMod10 switch - { - 1 => CombinedDivisorsEnding1OneAOneB, - 3 => CombinedDivisorsEnding3OneAOneB, - 7 => CombinedDivisorsEnding7OneAOneB, - 9 => CombinedDivisorsEnding9OneAOneB, - _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), - }, - CombinedDivisorPattern.ThreeAOneB => nMod10 switch - { - 1 => CombinedDivisorsEnding1ThreeAOneB, - 3 => CombinedDivisorsEnding3ThreeAOneB, - 7 => CombinedDivisorsEnding7ThreeAOneB, - 9 => CombinedDivisorsEnding9ThreeAOneB, - _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), - }, - CombinedDivisorPattern.ThreeATwoB => nMod10 switch - { - 1 => CombinedDivisorsEnding1ThreeATwoB, - 3 => CombinedDivisorsEnding3ThreeATwoB, - 7 => CombinedDivisorsEnding7ThreeATwoB, - 9 => CombinedDivisorsEnding9ThreeATwoB, - _ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."), - }, - _ => throw new InvalidOperationException($"Unsupported combined divisor pattern: {pattern}."), - }; - } - - internal static ReadOnlySpan GetCombinedDivisorSquares(byte nMod10, CombinedDivisorPattern pattern) - { - return pattern switch + 1 => CombinedDivisorsEnding1OneAOneBSquares, + 3 => CombinedDivisorsEnding3OneAOneBSquares, + 7 => CombinedDivisorsEnding7OneAOneBSquares, + 9 => CombinedDivisorsEnding9OneAOneBSquares, + _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), + }, + CombinedDivisorPattern.ThreeAOneB => nMod10 switch { - CombinedDivisorPattern.TwoAOneB => nMod10 switch - { - 1 => CombinedDivisorsEnding1TwoAOneBSquares, - 3 => CombinedDivisorsEnding3TwoAOneBSquares, - 7 => CombinedDivisorsEnding7TwoAOneBSquares, - 9 => CombinedDivisorsEnding9TwoAOneBSquares, - _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), - }, - CombinedDivisorPattern.OneAOneB => nMod10 switch - { - 1 => CombinedDivisorsEnding1OneAOneBSquares, - 3 => CombinedDivisorsEnding3OneAOneBSquares, - 7 => CombinedDivisorsEnding7OneAOneBSquares, - 9 => CombinedDivisorsEnding9OneAOneBSquares, - _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), - }, - CombinedDivisorPattern.ThreeAOneB => nMod10 switch - { - 1 => CombinedDivisorsEnding1ThreeAOneBSquares, - 3 => CombinedDivisorsEnding3ThreeAOneBSquares, - 7 => CombinedDivisorsEnding7ThreeAOneBSquares, - 9 => CombinedDivisorsEnding9ThreeAOneBSquares, - _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), - }, - CombinedDivisorPattern.ThreeATwoB => nMod10 switch - { - 1 => CombinedDivisorsEnding1ThreeATwoBSquares, - 3 => CombinedDivisorsEnding3ThreeATwoBSquares, - 7 => CombinedDivisorsEnding7ThreeATwoBSquares, - 9 => CombinedDivisorsEnding9ThreeATwoBSquares, - _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), - }, - _ => throw new InvalidOperationException($"Unsupported combined divisor pattern for squares: {pattern}."), - }; - } + 1 => CombinedDivisorsEnding1ThreeAOneBSquares, + 3 => CombinedDivisorsEnding3ThreeAOneBSquares, + 7 => CombinedDivisorsEnding7ThreeAOneBSquares, + 9 => CombinedDivisorsEnding9ThreeAOneBSquares, + _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), + }, + CombinedDivisorPattern.ThreeATwoB => nMod10 switch + { + 1 => CombinedDivisorsEnding1ThreeATwoBSquares, + 3 => CombinedDivisorsEnding3ThreeATwoBSquares, + 7 => CombinedDivisorsEnding7ThreeATwoBSquares, + 9 => CombinedDivisorsEnding9ThreeATwoBSquares, + _ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."), + }, + _ => throw new InvalidOperationException($"Unsupported combined divisor pattern for squares: {pattern}."), + }; private static ulong[] BuildCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern, out ulong[] squares) { @@ -443,8 +429,6 @@ public sealed class HeuristicCombinedPrimeTester return result; } - - private static int GetGroupBStartIndex(ReadOnlySpan divisors) { int index = 0; diff --git a/PerfectNumbers.Core/HeuristicPrimeTester.cs b/PerfectNumbers.Core/HeuristicPrimeTester.cs index 7b30cc1..5b17183 100644 --- a/PerfectNumbers.Core/HeuristicPrimeTester.cs +++ b/PerfectNumbers.Core/HeuristicPrimeTester.cs @@ -244,7 +244,7 @@ public sealed class HeuristicPrimeTester var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(1); var accelerator = gpu.Accelerator; var stream = gpu.Stream; - var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.HeuristicTrialDivisionKernel; + var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetHeuristicTrialDivisionKernel(accelerator); var flagView1D = gpu.HeuristicFlag.View; var flagView = flagView1D.AsContiguous(); diff --git a/PerfectNumbers.Core/MersenneDivisorCycles.cs b/PerfectNumbers.Core/MersenneDivisorCycles.cs index 2ce1f14..3467868 100644 --- a/PerfectNumbers.Core/MersenneDivisorCycles.cs +++ b/PerfectNumbers.Core/MersenneDivisorCycles.cs @@ -48,9 +48,7 @@ public class MersenneDivisorCycles public ulong[] ExportSmallCyclesSnapshot() { EnsureSmallBuffer(); - ulong[] snapshot = new ulong[PerfectNumberConstants.MaxQForDivisorCycles + 1]; - Array.Copy(_smallCycles!, snapshot, snapshot.Length); - return snapshot; + return _smallCycles!; } public static bool CycleEqualsExponent(ulong divisor, in MontgomeryDivisorData divisorData, ulong exponent) @@ -790,6 +788,9 @@ public class MersenneDivisorCycles Task.WaitAll(tasks); } + [ThreadStatic] + private static Dictionary, ArrayView1D>>? _divisorCycleKernel; + public static void GenerateGpu(string path, ulong maxDivisor, int batchSize = 1_000_000, long skipCount = 0L, long nextPosition = 0L) { // Prepare output file with header @@ -804,15 +805,8 @@ public class MersenneDivisorCycles divisors = pool.Rent(count); outCycles = pool.Rent(count); GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var loaded = accelerator.LoadAutoGroupedStreamKernel< - Index1D, - ArrayView1D, - ArrayView1D>(DivisorCycleKernels.GpuDivisorCycleKernel); - var kernel = KernelUtil.GetKernel(loaded); - var launcher = kernel.CreateLauncherDelegate, ArrayView1D>>(); - using Stream outputStream = new FileStream(path, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.Read, BufferSize10M, useAsync: true); if (nextPosition > 0L) @@ -841,6 +835,7 @@ public class MersenneDivisorCycles skipCount--; } + var divisorCycleKernel = GetDivisorCycleKernel(accelerator); while (start <= maxDivisor) { end = Math.Min(start + batchSizeUL - 1UL, maxDivisor); @@ -859,10 +854,15 @@ public class MersenneDivisorCycles // Use GpuKernelPool to get a kernel and context // GpuKernelPool.Run((accelerator, stream) => - var bufferDiv = accelerator.Allocate1D(validDivisors); - var bufferCycle = accelerator.Allocate1D(idx); + MemoryBuffer1D? bufferDiv; + MemoryBuffer1D? bufferCycle; + // lock(accelerator) + { + bufferDiv = accelerator.Allocate1D(validDivisors); + bufferCycle = accelerator.Allocate1D(idx); + } - launcher( + divisorCycleKernel( stream, idx, bufferDiv.View, @@ -900,6 +900,25 @@ public class MersenneDivisorCycles GpuPrimeWorkLimiter.Release(); } + private static Action, ArrayView1D> GetDivisorCycleKernel(Accelerator accelerator) + { + var pool = _divisorCycleKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + var loaded = accelerator.LoadAutoGroupedStreamKernel< + Index1D, + ArrayView1D, + ArrayView1D>(DivisorCycleKernels.GpuDivisorCycleKernel); + + var kernel = KernelUtil.GetKernel(loaded); + + cached = kernel.CreateLauncherDelegate, ArrayView1D>>(); + pool[accelerator] = cached; + return cached; + } public static ulong CalculateCycleLengthGpu(ulong divisor) => DivisorCycleKernels.CalculateCycleLengthGpu(divisor); diff --git a/PerfectNumbers.Core/MersenneNumberDivisorByDivisorTester.cs b/PerfectNumbers.Core/MersenneNumberDivisorByDivisorTester.cs index c01c62e..bcd9bda 100644 --- a/PerfectNumbers.Core/MersenneNumberDivisorByDivisorTester.cs +++ b/PerfectNumbers.Core/MersenneNumberDivisorByDivisorTester.cs @@ -293,7 +293,7 @@ public static class MersenneNumberDivisorByDivisorTester tasks[taskIndex++] = Task.Factory.StartNew( () => { - startGate.Wait(); + // startGate.Wait(); Console.WriteLine($"Task started for range {rangeStart}"); for (int index = rangeStart; index < rangeEnd; index++) diff --git a/PerfectNumbers.Core/MersenneNumberTester.cs b/PerfectNumbers.Core/MersenneNumberTester.cs index 17ade5e..c7b746a 100644 --- a/PerfectNumbers.Core/MersenneNumberTester.cs +++ b/PerfectNumbers.Core/MersenneNumberTester.cs @@ -33,7 +33,7 @@ public sealed class MersenneNumberTester( private readonly bool _useGpuLucas = useGpuLucas; private readonly bool _useGpuScan = useGpuScan; // device for pow2mod/incremental scanning private readonly bool _useGpuOrder = useGpuOrder; // device for order computations - private static readonly ConcurrentDictionary OrderCache = new(); + private static readonly ConcurrentDictionary _orderCache = []; // TODO: Swap this ConcurrentDictionary for the pooled dictionary variant highlighted in // Pow2MontgomeryModBenchmarks once order warmups reuse deterministic divisor-cycle snapshots; the pooled approach removed // the locking overhead when scanning p >= 138M in those benchmarks. @@ -89,13 +89,13 @@ public sealed class MersenneNumberTester( } } - if (!shouldCheck || !cacheEnabled || (cacheEnabled && OrderCache.ContainsKey(q))) + if (!shouldCheck || !cacheEnabled || (cacheEnabled && _orderCache.ContainsKey(q))) { continue; } ulong ord = q.CalculateOrder(); - OrderCache[q] = ord; + _orderCache[q] = ord; } return; @@ -131,7 +131,7 @@ public sealed class MersenneNumberTester( } } - if (!shouldCheck2 || (cacheEnabled && OrderCache.ContainsKey(q2))) + if (!shouldCheck2 || (cacheEnabled && _orderCache.ContainsKey(q2))) { continue; } @@ -148,9 +148,9 @@ public sealed class MersenneNumberTester( UInt128[] qs = qsBuffer; GpuPrimeWorkLimiter.Acquire(); - var accelerator = SharedGpuContext.Accelerator; + var accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - var orderKernel = GpuKernelPool.Kernels.Order!; + var orderKernel = GpuKernelPool.GetOrAddKernels(accelerator, stream).Order!; // Guard long-running kernels by chunking the warm-up across batches. // Reuse the same ScanBatchSize knob used for scanning. @@ -158,11 +158,18 @@ public sealed class MersenneNumberTester( ulong divMul = (ulong)((((UInt128)1 << 64) - 1UL) / exponent) + 1UL; int offset = 0; var gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool; + MemoryBuffer1D? qBuffer; + MemoryBuffer1D? orderBuffer; + while (offset < idx) { int count = Math.Min(batchSize, idx - offset); - var qBuffer = accelerator.Allocate1D(count); - var orderBuffer = accelerator.Allocate1D(count); + // lock (accelerator) + { + qBuffer = accelerator.Allocate1D(count); + orderBuffer = accelerator.Allocate1D(count); + } + // Convert to device-friendly GpuUInt128 on the fly var tmp = gpuUInt128Pool.Rent(count); for (int i = 0; i < count; i++) @@ -187,7 +194,7 @@ public sealed class MersenneNumberTester( if (useOrderCache) { - OrderCache[qs[offset + i]] = order; + _orderCache[qs[offset + i]] = order; } } diff --git a/PerfectNumbers.Core/PerfectNumberConstants.cs b/PerfectNumbers.Core/PerfectNumberConstants.cs index 3e37bd3..a8bf7db 100644 --- a/PerfectNumbers.Core/PerfectNumberConstants.cs +++ b/PerfectNumbers.Core/PerfectNumberConstants.cs @@ -1,3 +1,5 @@ +using PerfectNumbers.Core.Gpu; + namespace PerfectNumbers.Core; public static class PerfectNumberConstants @@ -8,7 +10,8 @@ public static class PerfectNumberConstants public const int MaxQForDivisorCycles = 4_000_000; public static readonly uint PrimesLimit = 1_000_000; //(ulong)Array.MaxLength;// 1_000_000; public const int MaxOddPowersCount = 128; - public const int PooledArrayThreshold = 64; + public const int PooledArrayThreshold = 64; + public static readonly int RollingAccelerators = 128; //SharedGpuContext.Device.MaxNumThreadsPerGroup; // TODO: Load these limits from the benchmark-driven configuration so CPU and GPU scans stay aligned with the optimal // divisor-cycle datasets we generate offline. // TODO: Promote these magic numbers into a runtime profile derived from EvenPerfectBitScanner.Benchmarks so we can retune diff --git a/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs b/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs index 07ce6ef..23c298a 100644 --- a/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs +++ b/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs @@ -3,8 +3,8 @@ using System.Diagnostics; using System.Numerics; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; -using ILGPU; using PerfectNumbers.Core.Gpu; +using PerfectNumbers.Core.Gpu.Accelerators; namespace PerfectNumbers.Core; @@ -52,8 +52,10 @@ internal static partial class PrimeOrderCalculator ulong phi = prime - 1UL; + Pow2MontgomeryAccelerator gpu = Pow2MontgomeryAccelerator.Rent(); if (IsGpuHeuristicDevice && PrimeOrderGpuHeuristics.TryCalculateOrder(prime, previousOrder, config, divisorData, out ulong gpuOrder)) { + Pow2MontgomeryAccelerator.Return(gpu); return gpuOrder; } @@ -62,17 +64,19 @@ internal static partial class PrimeOrderCalculator ulong result; if (phiFactors.Factors is null) { - result = CalculateByFactorizationCpu(prime, divisorData); + result = CalculateByFactorizationCpu(gpu, prime, divisorData); phiFactors.Dispose(); return result; } - result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors); + result = RunHeuristicPipelineCpu(gpu, prime, previousOrder, config, divisorData, phi, phiFactors); phiFactors.Dispose(); + Pow2MontgomeryAccelerator.Return(gpu); return result; } private static ulong RunHeuristicPipelineCpu( + Pow2MontgomeryAccelerator gpu, ulong prime, ulong? previousOrder, in PrimeOrderSearchConfig config, @@ -85,25 +89,30 @@ internal static partial class PrimeOrderCalculator return phi; } - ulong candidateOrder = InitializeStartingOrderCpu(prime, phi, divisorData); + GpuPrimeWorkLimiter.Acquire(); + ulong candidateOrder = InitializeStartingOrderCpu(gpu, prime, phi, divisorData); candidateOrder = ExponentLoweringCpu(candidateOrder, prime, phiFactors, divisorData); - if (TryConfirmOrderCpu(prime, candidateOrder, divisorData, config, out PartialFactorResult? orderFactors)) + if (TryConfirmOrderCpu(gpu, prime, candidateOrder, divisorData, config, out PartialFactorResult? orderFactors)) { + GpuPrimeWorkLimiter.Release(); return candidateOrder; } if (config.Mode == PrimeOrderMode.Strict) { orderFactors?.Dispose(); - return CalculateByFactorizationCpu(prime, divisorData); + GpuPrimeWorkLimiter.Release(); + return CalculateByFactorizationCpu(gpu, prime, divisorData); } if (TryHeuristicFinishCpu(prime, candidateOrder, previousOrder, divisorData, config, phiFactors, orderFactors, out ulong order)) { + GpuPrimeWorkLimiter.Release(); return order; } + GpuPrimeWorkLimiter.Release(); return candidateOrder; } @@ -185,13 +194,13 @@ internal static partial class PrimeOrderCalculator return true; } - private static ulong InitializeStartingOrderCpu(ulong prime, ulong phi, in MontgomeryDivisorData divisorData) + private static ulong InitializeStartingOrderCpu(Pow2MontgomeryAccelerator gpu, ulong prime, ulong phi, in MontgomeryDivisorData divisorData) { ulong order = phi; if ((prime & 7UL) == 1UL || (prime & 7UL) == 7UL) { ulong half = phi >> 1; - if (Pow2EqualsOneCpu(half, prime, divisorData)) + if (Pow2EqualsOneCpu(gpu, half, prime, divisorData)) { order = half; } @@ -319,6 +328,7 @@ internal static partial class PrimeOrderCalculator } private static bool TryConfirmOrderCpu( + Pow2MontgomeryAccelerator gpu, ulong prime, ulong order, in MontgomeryDivisorData divisorData, @@ -332,7 +342,7 @@ internal static partial class PrimeOrderCalculator return false; } - if (!Pow2EqualsOneCpu(order, prime, divisorData)) + if (!Pow2EqualsOneCpu(gpu, order, prime, divisorData)) { return false; } @@ -1021,11 +1031,15 @@ internal static partial class PrimeOrderCalculator return false; } - private static bool Pow2EqualsOneCpu(ulong exponent, ulong prime, in MontgomeryDivisorData divisorData) + private static ulong _pow2GpuHits; + private static ulong _pow2CpuHits; + private static bool Pow2EqualsOneCpu(Pow2MontgomeryAccelerator gpu, ulong exponent, ulong prime, in MontgomeryDivisorData divisorData) { - if (IsGpuPow2Allowed) + if (exponent <= (PerfectNumberConstants.MaxQForDivisorCycles) || prime <= (PerfectNumberConstants.MaxQForDivisorCycles)) { - ulong remainder = exponent.Pow2MontgomeryModWindowedGpu(divisorData, false); + // Atomic.Add(ref _pow2GpuHits, 1); + // Console.WriteLine($"pow2 GPU Hits {_pow2GpuHits}"); + ulong remainder = exponent.Pow2MontgomeryModWindowedConvertGpu(gpu, divisorData); // GpuPow2ModStatus status = PrimeOrderGpuHeuristics.TryPow2Mod(exponent, prime, out ulong remainder, divisorData); // if (status == GpuPow2ModStatus.Success) // { @@ -1033,6 +1047,8 @@ internal static partial class PrimeOrderCalculator // } } + // Atomic.Add(ref _pow2CpuHits, 1); + // Console.WriteLine($"pow2 CPU Hits {_pow2CpuHits}"); return exponent.Pow2MontgomeryModWindowedCpu(divisorData, keepMontgomery: false) == 1UL; } @@ -1638,7 +1654,7 @@ internal static partial class PrimeOrderCalculator exponents = newExponentSpan; } - private static ulong CalculateByFactorizationCpu(ulong prime, in MontgomeryDivisorData divisorData) + private static ulong CalculateByFactorizationCpu(Pow2MontgomeryAccelerator gpu, ulong prime, in MontgomeryDivisorData divisorData) { ulong phi = prime - 1UL; Dictionary counts = new(capacity: 8); @@ -1653,6 +1669,8 @@ internal static partial class PrimeOrderCalculator ulong order = phi; int entryCount = entries.Count; + + GpuPrimeWorkLimiter.Acquire(); for (int i = 0; i < entryCount; i++) { ulong primeFactor = entries[i].Key; @@ -1665,7 +1683,10 @@ internal static partial class PrimeOrderCalculator } ulong candidate = order / primeFactor; - if (Pow2EqualsOneCpu(candidate, prime, divisorData)) + + // if (candidate.Pow2MontgomeryModWindowedCpu(divisorData, keepMontgomery: false) == 1UL) + if (candidate.Pow2MontgomeryModWindowedConvertGpu(gpu, divisorData) == 1UL) + // if (Pow2EqualsOneCpu(candidate, prime, divisorData)) { order = candidate; continue; @@ -1675,6 +1696,7 @@ internal static partial class PrimeOrderCalculator } } + GpuPrimeWorkLimiter.Release(); return order; } diff --git a/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs b/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs index 19a9e90..24b066a 100644 --- a/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs +++ b/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs @@ -1,6 +1,7 @@ using System.Runtime.InteropServices; using ILGPU; using ILGPU.Runtime; +using ILGPU.Runtime.Cuda; using PerfectNumbers.Core.Gpu; namespace PerfectNumbers.Core; @@ -20,18 +21,18 @@ internal static partial class PrimeOrderCalculator remaining = value; GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; + Accelerator accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); - KernelContainer kernels = GpuKernelPool.Kernels; + KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); ScratchBuffer scratch = GpuScratchBufferPool.Rent(accelerator, GpuSmallPrimeFactorSlots, 0); - SmallPrimeFactorViews tables = GpuKernelPool.GetSmallPrimeFactorTables(); + SmallPrimeFactorViews tables = GpuKernelPool.GetSmallPrimeFactorTables(kernels); 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"); - var kernel = GpuKernelPool.Kernels.SmallPrimeFactor!; + var kernel = kernels.SmallPrimeFactor!; ArrayView1D smallPrimeFactorCountSlotView = scratch.SmallPrimeFactorCountSlot.View; ArrayView1D smallPrimeFactorRemainingSlotView = scratch.SmallPrimeFactorRemainingSlot.View; @@ -53,10 +54,11 @@ internal static partial class PrimeOrderCalculator smallPrimeFactorRemainingSlotView); + // Span factorCountTemp = stackalloc int[1]; factorCount = 0; smallPrimeFactorCountSlotView.CopyToCPU(stream, ref factorCount, 1); stream.Synchronize(); - + // factorCount = factorCountTemp[0]; factorCount = Math.Min(factorCount, GpuSmallPrimeFactorSlots); if (factorCount > 0) @@ -66,7 +68,10 @@ internal static partial class PrimeOrderCalculator smallPrimeFactorExponentSlotsView.CopyToCPU(stream, ref MemoryMarshal.GetReference(exponentBuffer), factorCount); } + // Span remainingTemp = stackalloc ulong[1]; + smallPrimeFactorRemainingSlotView.CopyToCPU(stream, ref remaining, 1); + // remaining = remainingTemp[0]; stream.Synchronize(); stream.Dispose(); @@ -102,16 +107,16 @@ internal static partial class PrimeOrderCalculator } GpuPrimeWorkLimiter.Acquire(); - Accelerator accelerator = SharedGpuContext.Accelerator; - + Accelerator accelerator = AcceleratorPool.Shared.Rent(); var stream = accelerator.CreateStream(); + var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream); ScratchBuffer scratch = GpuScratchBufferPool.Rent(accelerator, 0, factorCount); ArrayView1D specialMaxFactorsView = scratch.SpecialMaxFactors.View; specialMaxFactorsView.SubView(0, factorCount).CopyFromCPU(stream, ref MemoryMarshal.GetReference(factorSpan), factorCount); - var kernel = GpuKernelPool.Kernels.SpecialMax!; + var kernel = kernels.SpecialMax!; ArrayView1D specialMaxResultView = scratch.SpecialMaxResult.View; @@ -125,9 +130,10 @@ internal static partial class PrimeOrderCalculator scratch.SpecialMaxCandidates.View, specialMaxResultView); - Span result = stackalloc ulong[1]; + // Span result = stackalloc ulong[1]; + ulong result = 0UL; // specialMaxResultView.CopyToCPU(stream, ref result, 1); - specialMaxResultView.CopyToCPU(stream, in result); + specialMaxResultView.CopyToCPU(stream, ref result, 1); stream.Synchronize(); stream.Dispose(); @@ -135,7 +141,7 @@ internal static partial class PrimeOrderCalculator // return result != 0; GpuScratchBufferPool.Return(scratch); GpuPrimeWorkLimiter.Release(); - return result[0] != 0; + return result != 0; } } diff --git a/PerfectNumbers.Core/PrimeTester.cs b/PerfectNumbers.Core/PrimeTester.cs index 3b49dcb..9799a4f 100644 --- a/PerfectNumbers.Core/PrimeTester.cs +++ b/PerfectNumbers.Core/PrimeTester.cs @@ -99,9 +99,10 @@ public sealed class PrimeTester var gpu = PrimeTesterGpuContextPool.Rent(1); var accelerator = gpu.Accelerator; var stream = gpu.Stream; - var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.Kernel; + var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetKernel(accelerator); ulong value = n; + // Span flag = stackalloc byte[1]; byte flag = 0; var input = gpu.Input; @@ -166,7 +167,7 @@ public sealed class PrimeTester var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize); var accelerator = gpu.Accelerator; var stream = gpu.Stream; - var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.Kernel; + var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetKernel(accelerator); int totalLength = values.Length; int batchSize = GpuBatchSize; @@ -215,83 +216,52 @@ public sealed class PrimeTester internal static class PrimeTesterGpuContextPool { - private static readonly ConcurrentQueue Pool = new(); + [ThreadStatic] + private static Queue? Pool; private static readonly object WarmUpLock = new(); private static int WarmedLeaseCount; - private static readonly object SharedTablesLock = new(); - private static readonly Dictionary SharedTables = new(AcceleratorReferenceComparer.Instance); - private static SharedHeuristicGpuTables RentSharedTables(Accelerator accelerator, AcceleratorStream stream) - { - lock (SharedTablesLock) - { - if (!SharedTables.TryGetValue(accelerator, out var tables)) - { - tables = new SharedHeuristicGpuTables(accelerator, stream); - SharedTables.Add(accelerator, tables); - } - - tables.AddReference(); - return tables; - } - } - - private static void ReleaseSharedTables(SharedHeuristicGpuTables tables) - { - if (tables is null) - { - return; - } - - lock (SharedTablesLock) - { - tables.Release(); - } - } - - internal static void EnsureStaticTables(KernelContainer kernels, Accelerator accelerator, AcceleratorStream stream) - { - var tables = RentSharedTables(accelerator, stream); - var heuristicTables = tables.CreateHeuristicDivisorTables(); - HeuristicGpuDivisorTables.InitializeShared(in heuristicTables); - ReleaseSharedTables(tables); - } + // internal static void EnsureStaticTables(Accelerator accelerator) + // { + // var tables = CreateSharedTables(accelerator); + // } internal static void WarmUp(int target, int minBufferCapacity) { - if (target <= 0) - { - return; - } + // if (target <= 0) + // { + // return; + // } - lock (WarmUpLock) - { - if (target <= WarmedLeaseCount) - { - return; - } + // lock (WarmUpLock) + // { + // if (target <= WarmedLeaseCount) + // { + // return; + // } - Accelerator accelerator = SharedGpuContext.Accelerator; - AcceleratorStream stream = accelerator.CreateStream(); - KernelContainer kernels = GpuKernelPool.Kernels; + // Accelerator accelerator = SharedGpuContext.Accelerator; + // AcceleratorStream stream = accelerator.CreateStream(); + // KernelContainer kernels = GpuKernelPool.Kernels; - GpuStaticTableInitializer.EnsureStaticTables(kernels, accelerator, stream); + // GpuStaticTableInitializer.EnsureStaticTables(kernels, stream); - int toCreate = target - WarmedLeaseCount; - for (int i = 0; i < toCreate; i++) - { - Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity, stream)); - } + // int toCreate = target - WarmedLeaseCount; + // for (int i = 0; i < toCreate; i++) + // { + // Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity, stream)); + // } - stream.Synchronize(); - stream.Dispose(); - WarmedLeaseCount = target; - } + // stream.Synchronize(); + // stream.Dispose(); + // WarmedLeaseCount = target; + // } } internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1) { - if (!Pool.TryDequeue(out var lease)) + var pool = Pool ??= new(); + if (!pool.TryDequeue(out var lease)) { lease = new PrimeTesterGpuContextLease(minBufferCapacity); } @@ -306,74 +276,37 @@ public sealed class PrimeTester internal static void Return(PrimeTesterGpuContextLease lease) { - Pool.Enqueue(lease); + var pool = Pool ??= new(); + pool.Enqueue(lease); } internal static void DisposeAll() { - while (Pool.TryDequeue(out var lease)) + if (Pool != null) { - lease.DisposeResources(); - } - - lock (SharedTablesLock) - { - foreach (var entry in SharedTables.Values) + while (Pool.TryDequeue(out var lease)) { - entry.Dispose(); + lease.DisposeResources(); } - - SharedTables.Clear(); - var emptyTables = default(HeuristicGpuDivisorTables); - HeuristicGpuDivisorTables.InitializeShared(in emptyTables); } WarmedLeaseCount = 0; } - internal static void Clear(Accelerator accelerator) + internal static void Clear() { - if (accelerator is null) + if (Pool != null) { - return; - } - - lock (WarmUpLock) - { - var retained = new List(); - while (Pool.TryDequeue(out var lease)) + lock (WarmUpLock) { - if (lease.Accelerator == accelerator) + var retained = new List(); + while (Pool.TryDequeue(out var lease)) { lease.DisposeResources(); } - else - { - retained.Add(lease); - } - } - int retainedCount = retained.Count; - for (int i = 0; i < retainedCount; i++) - { - Pool.Enqueue(retained[i]); - } - - WarmedLeaseCount = Pool.Count; - } - - lock (SharedTablesLock) - { - if (SharedTables.Remove(accelerator, out var tables)) - { - tables.Dispose(); - } - - if (SharedTables.Count == 0) - { - var emptyTables = default(HeuristicGpuDivisorTables); - HeuristicGpuDivisorTables.InitializeShared(in emptyTables); - } + WarmedLeaseCount = Pool.Count; + } } } @@ -390,9 +323,38 @@ public sealed class PrimeTester get => _stream ?? Accelerator.CreateStream(); } - public static readonly Action, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView> Kernel = KernelUtil.GetKernel(SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>(PrimeTesterKernels.SmallPrimeSieveKernel)).CreateLauncherDelegate, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>>(); + [ThreadStatic] + public static Dictionary, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>>? Kernel; + + internal static Action, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView> GetKernel(Accelerator accelerator) + { + var pool = Kernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + cached = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>(PrimeTesterKernels.SmallPrimeSieveKernel)).CreateLauncherDelegate, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>>(); + pool[accelerator] = cached; + return cached; + } + + [ThreadStatic] + public static Dictionary, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>>? HeuristicTrialDivisionKernel; + + internal static Action, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables> GetHeuristicTrialDivisionKernel(Accelerator accelerator) + { + var pool = HeuristicTrialDivisionKernel ??= []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + cached = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel)).CreateLauncherDelegate, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>>(); + pool[accelerator] = cached; + return cached; + } - public static readonly Action, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables> HeuristicTrialDivisionKernel = KernelUtil.GetKernel(SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel)).CreateLauncherDelegate, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>>(); public MemoryBuffer1D Input = null!; public MemoryBuffer1D Output = null!; @@ -431,37 +393,27 @@ public sealed class PrimeTester internal PrimeTesterGpuContextLease(int minBufferCapacity) { AcceleratorContext = SharedGpuContext.Context; - Accelerator = SharedGpuContext.Accelerator; - AcceleratorStream stream = SharedGpuContext.Accelerator.CreateStream(); - var kernels = GpuKernelPool.Kernels; - GpuStaticTableInitializer.EnsureStaticTables(kernels, Accelerator, stream); + Accelerator accelerator = AcceleratorPool.Shared.Rent(); + // lock (_accelerators) + // { + // } - _sharedTables = RentSharedTables(Accelerator, stream); + Accelerator = accelerator; + AcceleratorStream stream = accelerator.CreateStream(); + // var kernels = _kernels[index]; //GpuKernelPool.Kernels; + // GpuStaticTableInitializer.EnsureStaticTables(accelerator, kernels, stream); - _heuristicDivisorTables = _sharedTables.CreateHeuristicDivisorTables(); - HeuristicGpuDivisorTables.InitializeShared(in _heuristicDivisorTables); + var sharedTables = SharedHeuristicGpuTables.EnsureStaticTables(accelerator, stream); + _sharedTables = sharedTables; + + var heuristicDivisorTables = sharedTables.CreateHeuristicDivisorTables(); + _heuristicDivisorTables = heuristicDivisorTables; EnsureCapacity(minBufferCapacity); stream.Synchronize(); stream.Dispose(); } - internal PrimeTesterGpuContextLease(int minBufferCapacity, AcceleratorStream stream) - { - AcceleratorContext = SharedGpuContext.Context; - Accelerator = SharedGpuContext.Accelerator; - // GpuStaticTableInitializer.EnsureStaticTables(Accelerator); - // Kernel = Accelerator.LoadAutoGroupedStreamKernel, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView, ArrayView>(PrimeTesterKernels.SmallPrimeSieveKernel); - // HeuristicTrialDivisionKernel = Accelerator.LoadAutoGroupedStreamKernel, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel); - - _sharedTables = RentSharedTables(Accelerator, stream); - - _heuristicDivisorTables = _sharedTables.CreateHeuristicDivisorTables(); - HeuristicGpuDivisorTables.InitializeShared(in _heuristicDivisorTables); - - EnsureCapacity(minBufferCapacity); - } - public void EnsureCapacity(int minCapacity) { if (minCapacity <= BufferCapacity) @@ -473,14 +425,20 @@ public sealed class PrimeTester Output?.Dispose(); HeuristicFlag?.Dispose(); - Input = Accelerator.Allocate1D(minCapacity); - Output = Accelerator.Allocate1D(minCapacity); - HeuristicFlag = Accelerator.Allocate1D(minCapacity); + // lock (Accelerator) + { + Input = Accelerator.Allocate1D(minCapacity); + Output = Accelerator.Allocate1D(minCapacity); + HeuristicFlag = Accelerator.Allocate1D(minCapacity); + } + BufferCapacity = minCapacity; } public void Dispose() { + _stream?.Dispose(); + _stream = null; Return(this); } @@ -491,7 +449,7 @@ public sealed class PrimeTester HeuristicFlag?.Dispose(); _stream?.Dispose(); _stream = null; - ReleaseSharedTables(_sharedTables); + // ReleaseSharedTables(_sharedTables); // These resources are shared between GPU leases // Accelerator.Dispose(); // AcceleratorContext.Dispose(); @@ -515,49 +473,49 @@ public sealed class PrimeTester private sealed class SharedHeuristicGpuTables { + private static readonly ConcurrentDictionary _sharedTables = new(20_480, PerfectNumberConstants.RollingAccelerators); + private int _referenceCount; + internal static SharedHeuristicGpuTables EnsureStaticTables(Accelerator accelerator, AcceleratorStream stream) + => _sharedTables.GetOrAdd(accelerator, _ => new SharedHeuristicGpuTables(accelerator, stream)); + internal SharedHeuristicGpuTables(Accelerator accelerator, AcceleratorStream stream) { - var primesDefault = DivisorGenerator.SmallPrimes; - DevicePrimesDefault = accelerator.Allocate1D(primesDefault.Length); - DevicePrimesDefault.View.CopyFromCPU(stream, primesDefault); - var primesLastOne = DivisorGenerator.SmallPrimesLastOne; - DevicePrimesLastOne = accelerator.Allocate1D(primesLastOne.Length); - DevicePrimesLastOne.View.CopyFromCPU(stream, primesLastOne); - var primesLastSeven = DivisorGenerator.SmallPrimesLastSeven; - DevicePrimesLastSeven = accelerator.Allocate1D(primesLastSeven.Length); - DevicePrimesLastSeven.View.CopyFromCPU(stream, primesLastSeven); - var primesLastThree = DivisorGenerator.SmallPrimesLastThree; - DevicePrimesLastThree = accelerator.Allocate1D(primesLastThree.Length); - DevicePrimesLastThree.View.CopyFromCPU(stream, primesLastThree); - var primesLastNine = DivisorGenerator.SmallPrimesLastNine; - DevicePrimesLastNine = accelerator.Allocate1D(primesLastNine.Length); - DevicePrimesLastNine.View.CopyFromCPU(stream, primesLastNine); - var primesPow2Default = DivisorGenerator.SmallPrimesPow2; - DevicePrimesPow2Default = accelerator.Allocate1D(primesPow2Default.Length); - DevicePrimesPow2Default.View.CopyFromCPU(stream, primesPow2Default); - var primesPow2LastOne = DivisorGenerator.SmallPrimesPow2LastOne; - DevicePrimesPow2LastOne = accelerator.Allocate1D(primesPow2LastOne.Length); - DevicePrimesPow2LastOne.View.CopyFromCPU(stream, primesPow2LastOne); - var primesPow2LastSeven = DivisorGenerator.SmallPrimesPow2LastSeven; - DevicePrimesPow2LastSeven = accelerator.Allocate1D(primesPow2LastSeven.Length); - DevicePrimesPow2LastSeven.View.CopyFromCPU(stream, primesPow2LastSeven); - var primesPow2LastThree = DivisorGenerator.SmallPrimesPow2LastThree; - DevicePrimesPow2LastThree = accelerator.Allocate1D(primesPow2LastThree.Length); - DevicePrimesPow2LastThree.View.CopyFromCPU(stream, primesPow2LastThree); - var primesPow2LastNine = DivisorGenerator.SmallPrimesPow2LastNine; - DevicePrimesPow2LastNine = accelerator.Allocate1D(primesPow2LastNine.Length); + + // lock (accelerator) + { + DevicePrimesDefault = accelerator.Allocate1D(primesDefault.Length); + DevicePrimesLastOne = accelerator.Allocate1D(primesLastOne.Length); + DevicePrimesLastSeven = accelerator.Allocate1D(primesLastSeven.Length); + DevicePrimesLastThree = accelerator.Allocate1D(primesLastThree.Length); + DevicePrimesLastNine = accelerator.Allocate1D(primesLastNine.Length); + DevicePrimesPow2Default = accelerator.Allocate1D(primesPow2Default.Length); + DevicePrimesPow2LastOne = accelerator.Allocate1D(primesPow2LastOne.Length); + DevicePrimesPow2LastSeven = accelerator.Allocate1D(primesPow2LastSeven.Length); + DevicePrimesPow2LastThree = accelerator.Allocate1D(primesPow2LastThree.Length); + DevicePrimesPow2LastNine = accelerator.Allocate1D(primesPow2LastNine.Length); + } + + DevicePrimesDefault.View.CopyFromCPU(stream, primesDefault); + DevicePrimesLastOne.View.CopyFromCPU(stream, primesLastOne); + DevicePrimesLastSeven.View.CopyFromCPU(stream, primesLastSeven); + DevicePrimesLastThree.View.CopyFromCPU(stream, primesLastThree); + DevicePrimesLastNine.View.CopyFromCPU(stream, primesLastNine); + DevicePrimesPow2Default.View.CopyFromCPU(stream, primesPow2Default); + DevicePrimesPow2LastOne.View.CopyFromCPU(stream, primesPow2LastOne); + DevicePrimesPow2LastSeven.View.CopyFromCPU(stream, primesPow2LastSeven); + DevicePrimesPow2LastThree.View.CopyFromCPU(stream, primesPow2LastThree); DevicePrimesPow2LastNine.View.CopyFromCPU(stream, primesPow2LastNine); HeuristicPrimeSieves.EnsureInitialized(); @@ -647,22 +605,20 @@ public sealed class PrimeTester HeuristicGroupBDivisorSquaresEnding9.View); } - internal void AddReference() - { - _referenceCount++; - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal void AddReference() => Atomic.Add(ref _referenceCount, 1); - internal void Release() - { - if (_referenceCount > 0) - { - _referenceCount--; - } - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal void Release() => Atomic.Add(ref _referenceCount, -1); private static MemoryBuffer1D CopySpanToDevice(Accelerator accelerator, AcceleratorStream stream, ReadOnlySpan span) { - var buffer = accelerator.Allocate1D(span.Length); + MemoryBuffer1D? buffer; + // lock (accelerator) + { + buffer = accelerator.Allocate1D(span.Length); + } + if (!span.IsEmpty) { ref ulong sourceRef = ref MemoryMarshal.GetReference(span); @@ -674,7 +630,12 @@ public sealed class PrimeTester private static MemoryBuffer1D CopyUintSpanToDevice(Accelerator accelerator, AcceleratorStream stream, ReadOnlySpan span) { - var buffer = accelerator.Allocate1D(span.Length); + MemoryBuffer1D? buffer; + // lock (accelerator) + { + buffer = accelerator.Allocate1D(span.Length); + } + if (!span.IsEmpty) { var converted = new ulong[span.Length]; @@ -722,9 +683,9 @@ public sealed class PrimeTester } // Expose cache clearing for accelerator disposal coordination - public static void ClearGpuCaches(Accelerator accelerator) + public static void ClearGpuCaches() { - PrimeTesterGpuContextPool.Clear(accelerator); + PrimeTesterGpuContextPool.Clear(); } internal static void DisposeGpuContexts() @@ -746,6 +707,9 @@ public sealed class PrimeTester return BinaryGcd(n, m) != 1UL; } + [ThreadStatic] + private static readonly Dictionary, ArrayView>>? _kernel; + internal static void SharesFactorWithMaxExponentBatch(ReadOnlySpan values, Span results) { // TODO: Route this batch helper through the shared GPU kernel pool from @@ -755,18 +719,24 @@ public sealed class PrimeTester var accelerator = gpu.Accelerator; var stream = gpu.Stream; - var kernel = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel, ArrayView>(PrimeTesterKernels.SharesFactorKernel)).CreateLauncherDelegate, ArrayView>>(); - int length = values.Length; - var inputBuffer = accelerator.Allocate1D(length); - var resultBuffer = accelerator.Allocate1D(length); - ArrayPool pool = ThreadStaticPools.UlongPool; ulong[] temp = pool.Rent(length); values.CopyTo(temp); + + MemoryBuffer1D? inputBuffer; + MemoryBuffer1D? resultBuffer; + + // lock (accelerator) + { + inputBuffer = accelerator.Allocate1D(length); + resultBuffer = accelerator.Allocate1D(length); + } + inputBuffer.View.CopyFromCPU(stream, ref temp[0], length); + var kernel = GetSharesFactorKernel(accelerator); kernel(stream, length, inputBuffer.View, resultBuffer.View); - resultBuffer.View.CopyToCPU(stream, ref results[0], length); + resultBuffer.View.CopyToCPU(stream, in results); stream.Synchronize(); pool.Return(temp, clearArray: false); resultBuffer.Dispose(); @@ -774,6 +744,19 @@ public sealed class PrimeTester gpu.Dispose(); } + private static Action, ArrayView> GetSharesFactorKernel(Accelerator accelerator) + { + var pool = _kernel ?? []; + if (pool.TryGetValue(accelerator, out var cached)) + { + return cached; + } + + cached = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel, ArrayView>(PrimeTesterKernels.SharesFactorKernel)).CreateLauncherDelegate, ArrayView>>(); + pool[accelerator] = cached; + return cached; + } + [MethodImpl(MethodImplOptions.AggressiveInlining)] private static ulong BinaryGcd(ulong u, ulong v) { diff --git a/PerfectNumbers.Core/ULongExtensions.Gpu.cs b/PerfectNumbers.Core/ULongExtensions.Gpu.cs index 91e1564..01d3f76 100644 --- a/PerfectNumbers.Core/ULongExtensions.Gpu.cs +++ b/PerfectNumbers.Core/ULongExtensions.Gpu.cs @@ -1,10 +1,10 @@ using System; using System.Collections.Concurrent; using System.Runtime.CompilerServices; -using ILGPU; using ILGPU.Algorithms; using ILGPU.Runtime; using PerfectNumbers.Core.Gpu; +using PerfectNumbers.Core.Gpu.Accelerators; namespace PerfectNumbers.Core; @@ -315,25 +315,16 @@ public static partial class ULongExtensions } [MethodImpl(MethodImplOptions.AggressiveInlining)] - public static ulong Pow2MontgomeryModWindowedGpu(this ulong exponent, in MontgomeryDivisorData divisor, bool keepMontgomery) - { - // We shouldn't hit it in production code. If we do, we're doing something wrong. - // ulong modulus = divisor.Modulus; - // if (exponent == 0UL) - // { - // return divisor.MontgomeryOne; - // } + public static ulong Pow2MontgomeryModWindowedKeepGpu(this ulong exponent, in MontgomeryDivisorData divisor) + => Pow2MontgomeryGpuCalculator.CalculateKeep(exponent, divisor); - // We barely ever hit it in production code. It's e.g. 772 calls out of billions - // if (exponent <= Pow2WindowFallbackThreshold) - // { - // return Pow2MontgomeryModSingleBit(exponent, divisor); - // } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ulong Pow2MontgomeryModWindowedConvertGpu(this ulong exponent, in MontgomeryDivisorData divisor) + => Pow2MontgomeryGpuCalculator.CalculateConvert(exponent, divisor); - return keepMontgomery - ? Pow2MontgomeryGpuExecutor.ExecuteKeep(exponent, divisor) - : Pow2MontgomeryGpuExecutor.ExecuteConvert(exponent, divisor); - } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ulong Pow2MontgomeryModWindowedConvertGpu(this ulong exponent, Pow2MontgomeryAccelerator gpu, in MontgomeryDivisorData divisor) + => Pow2MontgomeryGpuCalculator.CalculateConvert(gpu, exponent, divisor); [MethodImpl(MethodImplOptions.AggressiveInlining)] private static void InitializeStandardOddPowers(ulong modulus, Span oddPowers) @@ -387,16 +378,7 @@ public static partial class ULongExtensions } [MethodImpl(MethodImplOptions.AggressiveInlining)] - private static int GetPortableBitLengthGpu(ulong value) - { - // Keep this commented out. It will never happen in production code. - // if (value == 0UL) - // { - // return 0; - // } - - return 64 - XMath.LeadingZeroCount(value); - } + private static int GetPortableBitLengthGpu(ulong value) => 64 - XMath.LeadingZeroCount(value); [MethodImpl(MethodImplOptions.AggressiveInlining)] private static int GetWindowSizeGpu(int n) @@ -429,102 +411,12 @@ public static partial class ULongExtensions public static ulong Pow2MontgomeryModWithCycleGpu(this ulong exponent, ulong cycleLength, in MontgomeryDivisorData divisor) { ulong rotationCount = exponent % cycleLength; - return Pow2MontgomeryGpuExecutor.ExecuteConvert(rotationCount, divisor); + return Pow2MontgomeryGpuCalculator.CalculateConvert(rotationCount, divisor); } [MethodImpl(MethodImplOptions.AggressiveInlining)] public static ulong Pow2MontgomeryModFromCycleRemainderGpu(this ulong reducedExponent, in MontgomeryDivisorData divisor) { - return Pow2MontgomeryGpuExecutor.ExecuteConvert(reducedExponent, divisor); - } - - private static class Pow2MontgomeryGpuExecutor - { - private static readonly Pow2MontgomeryKernelGroup Kernel = Pow2MontgomeryKernelGroup.Create(SharedGpuContext.Accelerator); - - public static ulong ExecuteKeep(ulong exponent, in MontgomeryDivisorData divisor) - { - return Execute(exponent, divisor, useKeepKernel: true); - } - - public static ulong ExecuteConvert(ulong exponent, in MontgomeryDivisorData divisor) - { - return Execute(exponent, divisor, useKeepKernel: false); - } - - [ThreadStatic] - private static MemoryBuffer1D? _exponentBuffer; - [ThreadStatic] - private static MemoryBuffer1D? _resultBuffer; - - private static ulong Execute(ulong exponent, in MontgomeryDivisorData divisor, bool useKeepKernel) - { - ulong result = 0UL; - - GpuPrimeWorkLimiter.Acquire(); - - Accelerator accelerator = SharedGpuContext.Accelerator; - AcceleratorStream stream = accelerator.CreateStream(); - var kernelGroup = Kernel; - - EnsureCapacity(); - - MemoryBuffer1D exponentBuffer = _exponentBuffer!; - MemoryBuffer1D resultBuffer = _resultBuffer!; - - exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1); - // We don't need to worry about any left-overs here. - // resultBuffer.MemSetToZero(stream); - - var kernel = useKeepKernel ? kernelGroup.KeepMontgomery : kernelGroup.ConvertToStandard; - kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View); - resultBuffer.View.CopyToCPU(stream, ref result, 1); - stream.Synchronize(); - - stream.Dispose(); - resultBuffer.Dispose(); - exponentBuffer.Dispose(); - GpuPrimeWorkLimiter.Release(); - return result; - } - - private static void EnsureCapacity() - { - if (_exponentBuffer == null) - { - Accelerator accelerator = SharedGpuContext.Accelerator; - _exponentBuffer = accelerator.Allocate1D(1); - _resultBuffer = accelerator.Allocate1D(1); - } - } - - private sealed class Pow2MontgomeryKernelGroup - { - internal Pow2MontgomeryKernelGroup( - Action, MontgomeryDivisorData, ArrayView1D> keepMontgomery, - Action, MontgomeryDivisorData, ArrayView1D> convertToStandard) - { - KeepMontgomery = keepMontgomery; - ConvertToStandard = convertToStandard; - } - - internal Action, MontgomeryDivisorData, ArrayView1D> KeepMontgomery { get; } - - internal Action, MontgomeryDivisorData, ArrayView1D> ConvertToStandard { get; } - - internal static Pow2MontgomeryKernelGroup Create(Accelerator accelerator) - { - var keepKernel = accelerator.LoadAutoGroupedStreamKernel, MontgomeryDivisorData, ArrayView1D>(Pow2MontgomeryKernels.Pow2MontgomeryKernelKeepMontgomery); - var keepLauncher = KernelUtil.GetKernel(keepKernel); - var keepDelegate = keepLauncher.CreateLauncherDelegate, MontgomeryDivisorData, ArrayView1D>>(); - - var convertKernel = accelerator.LoadAutoGroupedStreamKernel, MontgomeryDivisorData, ArrayView1D>(Pow2MontgomeryKernels.Pow2MontgomeryKernelConvertToStandard); - var convertLauncher = KernelUtil.GetKernel(convertKernel); - var convertDelegate = convertLauncher.CreateLauncherDelegate, MontgomeryDivisorData, ArrayView1D>>(); - - return new Pow2MontgomeryKernelGroup(keepDelegate, convertDelegate); - } - } - + return Pow2MontgomeryGpuCalculator.CalculateConvert(reducedExponent, divisor); } }