From 918195939008f03656d4fb73631aacaec2965efb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marek=20Lema=C5=84czyk?= Date: Fri, 10 Oct 2025 23:12:49 +0200 Subject: [PATCH] Accelerate Montgomery odd power initialization with GPU batches --- PerfectNumbers.Core/Gpu/GpuKernelPool.cs | 52 ++++++++++ .../Gpu/MontgomeryOddPowerGpu.cs | 94 +++++++++++++++++++ PerfectNumbers.Core/ULongExtensions.cs | 30 ++++-- 3 files changed, 168 insertions(+), 8 deletions(-) create mode 100644 PerfectNumbers.Core/Gpu/MontgomeryOddPowerGpu.cs diff --git a/PerfectNumbers.Core/Gpu/GpuKernelPool.cs b/PerfectNumbers.Core/Gpu/GpuKernelPool.cs index 6b022aa..1fcee0e 100644 --- a/PerfectNumbers.Core/Gpu/GpuKernelPool.cs +++ b/PerfectNumbers.Core/Gpu/GpuKernelPool.cs @@ -37,6 +37,7 @@ public sealed class KernelContainer ResidueAutomatonArgs, ArrayView, ArrayView1D>? IncrementalOrder; public Action, ArrayView1D>? Pow2ModOrder; + public Action, ArrayView1D, ArrayView1D, int>? MontgomeryOddPowers; // Optional device buffer with small divisor cycles (<= 4M). Index = divisor, value = cycle length. public MemoryBuffer1D? SmallCycles; @@ -189,6 +190,57 @@ public ref struct GpuKernelLease(IDisposable limiter, GpuContextLease gpu, Kerne } } + public readonly Action, ArrayView1D, ArrayView1D, int> MontgomeryOddPowerKernel + { + get + { + var accel = Accelerator; + lock (_gpu!.Value.KernelInitLock) + { + return KernelContainer.InitOnce(ref _kernels.MontgomeryOddPowers, () => + { + var loaded = accel.LoadAutoGroupedStreamKernel, ArrayView1D, ArrayView1D, int>(MontgomeryOddPowerKernelScan); + var kernel = KernelUtil.GetKernel(loaded); + return kernel.CreateLauncherDelegate, ArrayView1D, ArrayView1D, int>>(); + }); + } + } + } + + private static void MontgomeryOddPowerKernelScan( + Index1D index, + ArrayView1D divisors, + ArrayView1D counts, + ArrayView1D results, + int stride) + { + int tableIndex = index.X; + int baseOffset = tableIndex * stride; + MontgomeryDivisorData divisor = divisors[tableIndex]; + int oddPowerCount = counts[tableIndex]; + if (oddPowerCount <= 0) + { + return; + } + + ulong montgomeryTwo = divisor.MontgomeryTwo; + results[baseOffset] = montgomeryTwo; + if (oddPowerCount == 1) + { + return; + } + + ulong modulus = divisor.Modulus; + ulong nPrime = divisor.NPrime; + ulong square = divisor.MontgomeryTwoSquared; + ulong current = montgomeryTwo; + for (int i = 1; i < oddPowerCount; i++) + { + current = ULongExtensions.MontgomeryMultiply(current, square, modulus, nPrime); + results[baseOffset + i] = current; + } + } + // TODO: Plumb the small-cycles device buffer into all kernels that can benefit // (some already accept it). Consider a compact type (byte/ushort) for memory footprint. diff --git a/PerfectNumbers.Core/Gpu/MontgomeryOddPowerGpu.cs b/PerfectNumbers.Core/Gpu/MontgomeryOddPowerGpu.cs new file mode 100644 index 0000000..d1a9efb --- /dev/null +++ b/PerfectNumbers.Core/Gpu/MontgomeryOddPowerGpu.cs @@ -0,0 +1,94 @@ +using System; +using System.Runtime.InteropServices; +using ILGPU.Runtime; + +namespace PerfectNumbers.Core.Gpu; + +internal static class MontgomeryOddPowerGpu +{ + public static bool TryCompute(in MontgomeryDivisorData divisor, int oddPowerCount, Span destination) + { + if (oddPowerCount <= 0) + { + return true; + } + + Span divisorSpan = stackalloc MontgomeryDivisorData[1]; + divisorSpan[0] = divisor; + Span countSpan = stackalloc int[1]; + countSpan[0] = oddPowerCount; + Span resultSpan = stackalloc ulong[PerfectNumberConstants.MaxOddPowersCount]; + + if (!TryComputeBatch(divisorSpan, countSpan, resultSpan, PerfectNumberConstants.MaxOddPowersCount)) + { + return false; + } + + resultSpan[..oddPowerCount].CopyTo(destination); + return true; + } + + public static bool TryComputeBatch( + ReadOnlySpan divisors, + ReadOnlySpan oddPowerCounts, + Span results, + int stride) + { + int divisorCount = divisors.Length; + if (divisorCount == 0) + { + return true; + } + + if (oddPowerCounts.Length < divisorCount) + { + throw new ArgumentException( + "oddPowerCounts span must contain at least as many elements as divisors span.", + nameof(oddPowerCounts)); + } + + if (results.Length < divisorCount * stride) + { + throw new ArgumentException( + "The result span is shorter than the required batched odd power storage.", + nameof(results)); + } + + if (GpuContextPool.ForceCpu) + { + return false; + } + + var lease = GpuKernelPool.GetKernel(useGpuOrder: true); + var execution = lease.EnterExecutionScope(); + var accelerator = lease.Accelerator; + var stream = lease.Stream; + + using var divisorBuffer = accelerator.Allocate1D(divisorCount); + using var countBuffer = accelerator.Allocate1D(divisorCount); + using var resultBuffer = accelerator.Allocate1D(divisorCount * stride); + + try + { + divisorBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(divisors), divisorCount); + countBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(oddPowerCounts), divisorCount); + resultBuffer.MemSetToZero(); + + var kernel = lease.MontgomeryOddPowerKernel; + kernel(stream, divisorCount, divisorBuffer.View, countBuffer.View, resultBuffer.View, stride); + stream.Synchronize(); + + resultBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(results), divisorCount * stride); + return true; + } + catch + { + return false; + } + finally + { + execution.Dispose(); + lease.Dispose(); + } + } +} diff --git a/PerfectNumbers.Core/ULongExtensions.cs b/PerfectNumbers.Core/ULongExtensions.cs index fa0d030..8c08577 100644 --- a/PerfectNumbers.Core/ULongExtensions.cs +++ b/PerfectNumbers.Core/ULongExtensions.cs @@ -535,21 +535,35 @@ public static class ULongExtensions [MethodImpl(MethodImplOptions.AggressiveInlining)] private static ulong[] InitializeMontgomeryOddPowersCpu(in MontgomeryDivisorData divisor, ulong modulus, ulong nPrime, int oddPowerCount) { - ulong[] oddPowers = oddPowerCount < PerfectNumberConstants.PooledArrayThreshold ? new ulong[oddPowerCount] : ArrayPool.Shared.Rent(oddPowerCount); - - oddPowers[0] = divisor.MontgomeryTwo; + ulong[] oddPowers = oddPowerCount < PerfectNumberConstants.PooledArrayThreshold + ? new ulong[oddPowerCount] + : ArrayPool.Shared.Rent(oddPowerCount); + + if (oddPowerCount == 0) + { + return oddPowers; + } + + Span destination = oddPowers.AsSpan(0, oddPowerCount); + bool computedOnGpu = oddPowerCount <= PerfectNumberConstants.MaxOddPowersCount + && MontgomeryOddPowerGpu.TryCompute(divisor, oddPowerCount, destination); + if (computedOnGpu) + { + return oddPowers; + } + + destination[0] = divisor.MontgomeryTwo; if (oddPowerCount == 1) { return oddPowers; } - ulong previous, - square = divisor.MontgomeryTwoSquared; - + ulong square = divisor.MontgomeryTwoSquared; + ulong current = destination[0]; for (int i = 1; i < oddPowerCount; i++) { - previous = oddPowers[i - 1]; - oddPowers[i] = previous.MontgomeryMultiply(square, modulus, nPrime); + current = current.MontgomeryMultiply(square, modulus, nPrime); + destination[i] = current; } return oddPowers;