Add a small cache for integer square root

math/BUILD.bazel

@@ -5,15 +5,19 @@ go_library(
     srcs = ["math_helper.go"],
     importpath = "github.com/prysmaticlabs/prysm/v3/math",
     visibility = ["//visibility:public"],
-    deps = ["@com_github_thomaso_mirodin_intmath//u64:go_default_library"],
+    deps = [
+        "//cache/lru:go_default_library",
+        "@com_github_thomaso_mirodin_intmath//u64:go_default_library",
+    ],
 )
 
 go_test(
     name = "go_default_test",
     size = "small",
-    srcs = ["math_helper_test.go"],
-    deps = [
-        ":go_default_library",
-        "//testing/require:go_default_library",
+    srcs = [
+        "math_helper_test.go",
+        "test_helper_test.go",
     ],
+    embed = [":go_default_library"],
+    deps = ["//testing/require:go_default_library"],
 )

math/math_helper.go

@@ -6,6 +6,7 @@ import (
 	stdmath "math"
 	"math/bits"
 
+	lruwrpr "github.com/prysmaticlabs/prysm/v3/cache/lru"
 	"github.com/thomaso-mirodin/intmath/u64"
 )
 
@@ -43,6 +44,12 @@ var squareRootTable = map[uint64]uint64{
 	4194304: 2048,
 }
 
+// The IntegerSquareRoot method is often called thousands of times during epoch processing. Adding a small
+// cache can significantly improve performance by about 80% for cache hits.
+const SQRT_CACHE_SIZE = 1000
+
+var sqrtCache = lruwrpr.New(SQRT_CACHE_SIZE)
+
 // IntegerSquareRoot defines a function that returns the
 // largest possible integer root of a number using go's standard library.
 func IntegerSquareRoot(n uint64) uint64 {
@@ -50,13 +57,21 @@ func IntegerSquareRoot(n uint64) uint64 {
 		return v
 	}
 
+	if v, ok := sqrtCache.Get(n); ok {
+		return v.(uint64)
+	}
+
 	// Golang floating point precision may be lost above 52 bits, so we use a
 	// non floating point method. u64.Sqrt is about x2.5 slower than math.Sqrt.
 	if n >= 1<<52 {
-		return u64.Sqrt(n)
+		r := u64.Sqrt(n)
+		sqrtCache.Add(n, r)
+		return r
 	}
 
-	return uint64(stdmath.Sqrt(float64(n)))
+	r := uint64(stdmath.Sqrt(float64(n)))
+	sqrtCache.Add(n, r)
+	return r
 }
 
 // CeilDiv8 divides the input number by 8

math/math_helper_test.go

@@ -78,6 +78,10 @@ func TestIntegerSquareRoot(t *testing.T) {
 			number: 4503599761588224,
 			root:   67108864,
 		},
+		{
+			number: 13803562000000000,
+			root:   117488561,
+		},
 	}
 
 	for _, testVals := range tt {
@@ -154,6 +158,7 @@ func TestMath_Mod(t *testing.T) {
 }
 
 func BenchmarkIntegerSquareRootBelow52Bits(b *testing.B) {
+	math.DisableCaches()
 	val := uint64(1 << 33)
 	for i := 0; i < b.N; i++ {
 		require.Equal(b, uint64(92681), math.IntegerSquareRoot(val))
@@ -161,6 +166,7 @@ func BenchmarkIntegerSquareRootBelow52Bits(b *testing.B) {
 }
 
 func BenchmarkIntegerSquareRootAbove52Bits(b *testing.B) {
+	math.DisableCaches()
 	val := uint64(1 << 62)
 	for i := 0; i < b.N; i++ {
 		require.Equal(b, uint64(1<<31), math.IntegerSquareRoot(val))
@@ -168,12 +174,29 @@ func BenchmarkIntegerSquareRootAbove52Bits(b *testing.B) {
 }
 
 func BenchmarkIntegerSquareRoot_WithDatatable(b *testing.B) {
+	math.DisableCaches()
 	val := uint64(1024)
 	for i := 0; i < b.N; i++ {
 		require.Equal(b, uint64(32), math.IntegerSquareRoot(val))
 	}
 }
 
+func BenchmarkIntegerSquareRoot_Cache(b *testing.B) {
+	val := uint64(1 << 62)
+	b.Run("on", func(b *testing.B) {
+		math.EnableCaches()
+		for i := 0; i < b.N; i++ {
+			require.Equal(b, uint64(1<<31), math.IntegerSquareRoot(val))
+		}
+	})
+	b.Run("off", func(b *testing.B) {
+		math.DisableCaches()
+		for i := 0; i < b.N; i++ {
+			require.Equal(b, uint64(1<<31), math.IntegerSquareRoot(val))
+		}
+	})
+}
+
 func TestCeilDiv8(t *testing.T) {
 	tests := []struct {
 		number int

math/test_helper_test.go

@@ -0,0 +1,9 @@
+package math
+
+func DisableCaches() {
+	sqrtCache.Resize(0)
+}
+
+func EnableCaches() {
+	sqrtCache.Resize(SQRT_CACHE_SIZE)
+}