perf: Speedup Dec.Sqrt() (cosmos#16141)

osmosis-labs · May 15, 2023 · fb8ff07 · fb8ff07
1 parent 793cbe5
commit fb8ff07
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 13 deletions.
diff --git a/math/CHANGELOG.md b/math/CHANGELOG.md
@@ -43,6 +43,7 @@ Ref: https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.j
 ### Improvements
 
 * [#15768](https://github.com/cosmos/cosmos-sdk/pull/15768) Removed the second call to the `init` method for the global variable `grand`.
+* [#16141](https://github.com/cosmos/cosmos-sdk/pull/16141) Speedup `LegacyDec.ApproxRoot` and `LegacyDec.ApproxSqrt`.
 
 ## [math/v1.0.0](https://github.com/cosmos/cosmos-sdk/releases/tag/math/v1.0.0) - 2023-03-23
 

diff --git a/math/dec.go b/math/dec.go
@@ -220,6 +220,7 @@ func (d LegacyDec) LTE(d2 LegacyDec) bool      { return (d.i).Cmp(d2.i) <= 0 }
 func (d LegacyDec) Neg() LegacyDec             { return LegacyDec{new(big.Int).Neg(d.i)} } // reverse the decimal sign
 func (d LegacyDec) NegMut() LegacyDec          { d.i.Neg(d.i); return d }                  // reverse the decimal sign, mutable
 func (d LegacyDec) Abs() LegacyDec             { return LegacyDec{new(big.Int).Abs(d.i)} } // absolute value
+func (d LegacyDec) AbsMut() LegacyDec          { d.i.Abs(d.i); return d }                  // absolute value, mutable
 func (d LegacyDec) Set(d2 LegacyDec) LegacyDec { d.i.Set(d2.i); return d }                 // set to existing dec value
 func (d LegacyDec) Clone() LegacyDec           { return LegacyDec{new(big.Int).Set(d.i)} } // clone new dec
 
@@ -442,24 +443,38 @@ func (d LegacyDec) ApproxRoot(root uint64) (guess LegacyDec, err error) {
 		return absRoot.NegMut(), err
 	}
 
-	if root == 1 || d.IsZero() || d.Equal(LegacyOneDec()) {
+	// One decimal, that we invalidate later. Helps us save a heap allocation.
+	scratchOneDec := LegacyOneDec()
+	if root == 1 || d.IsZero() || d.Equal(scratchOneDec) {
 		return d, nil
 	}
 
 	if root == 0 {
-		return LegacyOneDec(), nil
+		return scratchOneDec, nil
 	}
 
-	guess, delta := LegacyOneDec(), LegacyOneDec()
+	guess, delta := scratchOneDec, LegacyOneDec()
+	smallestDec := LegacySmallestDec()
 
-	for iter := 0; delta.Abs().GT(LegacySmallestDec()) && iter < maxApproxRootIterations; iter++ {
-		prev := guess.Power(root - 1)
+	for iter := 0; delta.AbsMut().GT(smallestDec) && iter < maxApproxRootIterations; iter++ {
+		// Set prev = guess^{root - 1}, with an optimization for sqrt
+		// where root=2 => prev = guess. (And thus no extra heap allocations)
+		prev := guess
+		if root != 2 {
+			prev = guess.Power(root - 1)
+		}
 		if prev.IsZero() {
-			prev = LegacySmallestDec()
+			prev = smallestDec
 		}
 		delta.Set(d).QuoMut(prev)
 		delta.SubMut(guess)
-		delta.QuoInt64Mut(int64(root))
+		// delta = delta / root.
+		// We optimize for sqrt, where root=2 => delta = delta >> 1
+		if root == 2 {
+			delta.i.Rsh(delta.i, 1)
+		} else {
+			delta.QuoInt64Mut(int64(root))
+		}
 
 		guess.AddMut(delta)
 	}

diff --git a/math/dec_test.go b/math/dec_test.go
@@ -453,12 +453,16 @@ func (s *decimalTestSuite) TestApproxSqrt() {
 		input    math.LegacyDec
 		expected math.LegacyDec
 	}{
-		{math.LegacyOneDec(), math.LegacyOneDec()},                                                     // 1.0 => 1.0
-		{math.LegacyNewDecWithPrec(25, 2), math.LegacyNewDecWithPrec(5, 1)},                            // 0.25 => 0.5
-		{math.LegacyNewDecWithPrec(4, 2), math.LegacyNewDecWithPrec(2, 1)},                             // 0.09 => 0.3
-		{math.LegacyNewDecFromInt(math.NewInt(9)), math.LegacyNewDecFromInt(math.NewInt(3))},           // 9 => 3
-		{math.LegacyNewDecFromInt(math.NewInt(-9)), math.LegacyNewDecFromInt(math.NewInt(-3))},         // -9 => -3
-		{math.LegacyNewDecFromInt(math.NewInt(2)), math.LegacyNewDecWithPrec(1414213562373095049, 18)}, // 2 => 1.414213562373095049
+		{math.LegacyOneDec(), math.LegacyOneDec()},                                 // 1.0 => 1.0
+		{math.LegacyNewDecWithPrec(25, 2), math.LegacyNewDecWithPrec(5, 1)},        // 0.25 => 0.5
+		{math.LegacyNewDecWithPrec(4, 2), math.LegacyNewDecWithPrec(2, 1)},         // 0.09 => 0.3
+		{math.LegacyNewDec(9), math.LegacyNewDecFromInt(math.NewInt(3))},           // 9 => 3
+		{math.LegacyNewDec(-9), math.LegacyNewDecFromInt(math.NewInt(-3))},         // -9 => -3
+		{math.LegacyNewDec(2), math.LegacyNewDecWithPrec(1414213562373095049, 18)}, // 2 => 1.414213562373095049
+		{ // 2^127 - 1 => 13043817825332782212.3495718062525083688 which rounds to 13043817825332782212.3495718062525083689
+			math.LegacyNewDec(2).Power(127).Sub(math.LegacyOneDec()),
+			math.LegacyMustNewDecFromStr("13043817825332782212.349571806252508369"),
+		},
 	}
 
 	for i, tc := range testCases {
@@ -651,6 +655,20 @@ func BenchmarkLegacyQuoTruncateMut(b *testing.B) {
 	sink = (interface{})(nil)
 }
 
+func BenchmarkLegacySqrtOnMersennePrime(b *testing.B) {
+	b1 := math.LegacyNewDec(2).Power(127).Sub(math.LegacyOneDec())
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		sink, _ = b1.ApproxSqrt()
+	}
+
+	if sink == nil {
+		b.Fatal("Benchmark did not run")
+	}
+	sink = (interface{})(nil)
+}
+
 func BenchmarkLegacyQuoRoundupMut(b *testing.B) {
 	b1 := math.LegacyNewDec(17e2 + 8371)
 	b2 := math.LegacyNewDec(4371)