From eaee599f2503429243b80b224c990ea58c0eccde Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Lema=C5=84czyk?= <marek@lemanczyk.eu>
Date: Tue, 28 Oct 2025 20:58:02 +0100
Subject: [PATCH] Use GPU-safe Montgomery multiply in kernels

---
 PerfectNumbers.Core/ULongExtensions.Gpu.cs | 31 +++++++++++++++++-----
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/PerfectNumbers.Core/ULongExtensions.Gpu.cs b/PerfectNumbers.Core/ULongExtensions.Gpu.cs
index 085d41c..bc062c1 100644
--- a/PerfectNumbers.Core/ULongExtensions.Gpu.cs
+++ b/PerfectNumbers.Core/ULongExtensions.Gpu.cs
@@ -46,6 +46,22 @@ public static partial class ULongExtensions
 		return state.MulMod(right, modulus);
 	}
 
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	internal static ulong MontgomeryMultiplyGpu(this ulong a, ulong b, ulong modulus, ulong nPrime)
+	{
+		ulong tLow = unchecked(a * b);
+		ulong m = unchecked(tLow * nPrime);
+		ulong mTimesModulusLow = unchecked(m * modulus);
+
+		ulong result = unchecked(a.MulHighGpu(b) + m.MulHighGpu(modulus) + (unchecked(tLow + mTimesModulusLow) < tLow ? 1UL : 0UL));
+		if (result >= modulus)
+		{
+			result -= modulus;
+		}
+
+		return result;
+	}
+
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	public static ulong Pow2ModWindowedGpu(this ulong exponent, ulong modulus)
 	{
@@ -204,14 +220,15 @@ public static partial class ULongExtensions
 		ulong modulus = divisor.Modulus;
 		ulong nPrime = divisor.NPrime;
 		ulong montgomeryResult = Pow2MontgomeryModWindowedGpuMontgomeryResult(divisor, exponent);
-		return MontgomeryMultiply(montgomeryResult, 1UL, modulus, nPrime);
+		return MontgomeryMultiplyGpu(montgomeryResult, 1UL, modulus, nPrime);
 	}
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	private static ulong Pow2MontgomeryModWindowedGpuMontgomeryResult(in MontgomeryDivisorData divisor, ulong exponent)
 	{
 		ulong modulus = divisor.Modulus;
-		// Reuse the optimized MontgomeryMultiply extension highlighted in MontgomeryMultiplyBenchmarks so the GPU path matches the fastest CPU reduction.
+		// Reuse the GPU-compatible MontgomeryMultiply helper highlighted in MontgomeryMultiplyBenchmarks so the accelerator path
+		// mirrors the fastest CPU reduction without relying on UInt128 intrinsics.
 		// We shouldn't hit it in production code. If we do, we're doing something wrong.
 		// if (exponent == 0UL)
 		// {
@@ -233,7 +250,7 @@ public static partial class ULongExtensions
 		while (index >= 0)
 		{
 			ulong currentBit = (exponent >> index) & 1UL;
-			ulong squared = MontgomeryMultiply(result, result, modulus, nPrime);
+			ulong squared = MontgomeryMultiplyGpu(result, result, modulus, nPrime);
 			bool processWindow = currentBit != 0UL;
 
 			result = processWindow ? result : squared;
@@ -253,11 +270,11 @@ public static partial class ULongExtensions
 			int windowLength = processWindow ? index - windowStart + 1 : 0;
 			for (int square = 0; square < windowLength; square++)
 			{
-				result = MontgomeryMultiply(result, result, modulus, nPrime);
+				result = MontgomeryMultiplyGpu(result, result, modulus, nPrime);
 			}
 
 			result = processWindow
-					? MontgomeryMultiply(
+					? MontgomeryMultiplyGpu(
 							result,
 							ComputeMontgomeryOddPowerGpu(
 									(exponent >> windowStart) & ((1UL << windowLength) - 1UL),
@@ -330,7 +347,7 @@ public static partial class ULongExtensions
 		while (remaining != 0UL)
 		{
 			ulong bit = remaining & 1UL;
-			ulong multiplied = MontgomeryMultiply(power, baseValue, modulus, nPrime);
+			ulong multiplied = MontgomeryMultiplyGpu(power, baseValue, modulus, nPrime);
 			ulong mask = (ulong)-(long)bit;
 			power = (power & ~mask) | (multiplied & mask);
 
@@ -340,7 +357,7 @@ public static partial class ULongExtensions
 				break;
 			}
 
-			baseValue = MontgomeryMultiply(baseValue, baseValue, modulus, nPrime);
+			baseValue = MontgomeryMultiplyGpu(baseValue, baseValue, modulus, nPrime);
 		}
 
 		return power;