From e46d1b20fa920dc8b701a535dc7847e1225d490a Mon Sep 17 00:00:00 2001
From: Asriel Camora <asriel@camora.dev>
Date: Tue, 4 Jul 2023 08:36:41 +0200
Subject: [PATCH] Use rqrt in child score calculation

Changes sqrt -> rsqrt and div->mul. Reduces latency by 23 -> 8 cycles.
---
 Solver/Crafty/Intrinsics.cs | 15 +++++++++++++++
 Solver/Crafty/Solver.cs     |  9 ++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/Solver/Crafty/Intrinsics.cs b/Solver/Crafty/Intrinsics.cs
index 7894ae3..a1baef0 100644
--- a/Solver/Crafty/Intrinsics.cs
+++ b/Solver/Crafty/Intrinsics.cs
@@ -110,6 +110,21 @@ internal static class Intrinsics
             NthBitSetScalar(value, n);
     }
 
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static Vector<float> ReciprocalSqrt(Vector<float> data)
+    {
+        if (Avx.IsSupported && Vector<float>.Count >= Vector256<float>.Count)
+            return Avx.ReciprocalSqrt(data.AsVector256()).AsVector();
+
+        if (Sse.IsSupported && Vector<float>.Count >= Vector128<float>.Count)
+            return Sse.ReciprocalSqrt(data.AsVector128()).AsVector();
+
+        Span<float> result = stackalloc float[Vector<float>.Count];
+        for (var i = 0; i < Vector<float>.Count; ++i)
+            result[i] = MathF.ReciprocalSqrtEstimate(data[i]);
+        return new(result);
+    }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public static void CASMax(ref float location, float newValue)
diff --git a/Solver/Crafty/Solver.cs b/Solver/Crafty/Solver.cs
index df2864c..e10ac48 100644
--- a/Solver/Crafty/Solver.cs
+++ b/Solver/Crafty/Solver.cs
@@ -101,7 +101,7 @@ public class Solver
         var length = children.Length;
         var vecLength = Vector<float>.Count;
 
-        var C = Config.ExplorationConstant * MathF.Log(parentVisits);
+        var C = MathF.Sqrt(Config.ExplorationConstant * MathF.Log(parentVisits));
         var w = Config.MaxScoreWeightingConstant;
         var W = 1f - w;
         var CVector = new Vector<float>(C);
@@ -126,8 +126,11 @@ public class Solver
                 maxScores[j] = node.MaxScore;
             }
 
-            var exploitation = (W * (new Vector<float>(scoreSums) / new Vector<float>(visits))) + (w * new Vector<float>(maxScores));
-            var exploration = Vector.SquareRoot(CVector / new Vector<float>(visits));
+            var s = new Vector<float>(scoreSums);
+            var m = new Vector<float>(maxScores);
+            var v = new Vector<float>(visits);
+            var exploitation = (W * (s / v)) + (w * m);
+            var exploration = CVector * Intrinsics.ReciprocalSqrt(v);
             var evalScores = exploitation + exploration;
 
             var idx = Intrinsics.HMaxIndex(evalScores, iterCount);