uint256: optimize Sqrt

[AaronChen0]

1. Make x.IsUint64() a special case.
2. If x.IsUint64() is false, then

z2.Div(x, z1)

is equal to

z2.Clear()
udivrem(z2[:], x[:], z1, nil)

3. Lt is faster than Cmp.
4. The first div can be replaced with a right shift. Rsh is 10x cheaper than Div.

Benchmark

goos: linux
goarch: amd64
pkg: github.com/holiman/uint256
cpu: AMD Ryzen 7 7735H with Radeon Graphics         
                       │     old     │                 new                 │
                       │   sec/op    │   sec/op     vs base                │
Sqrt/single/uint256-16   241.0n ± 1%   172.4n ± 1%  -28.46% (p=0.000 n=10)
Sqrt/single/big-16       994.1n ± 1%   988.9n ± 1%        ~ (p=0.089 n=10)
geomean                  489.5n        412.9n       -15.64%

[codecov]

Codecov Report

Attention: Patch coverage is 80.00000% with 5 lines in your changes are missing coverage. Please review.

Project coverage is 99.69%. Comparing base (70cbe2b) to head (6f34595).

Additional details and impacted files

@@             Coverage Diff             @@
##            master     #174      +/-   ##
===========================================
- Coverage   100.00%   99.69%   -0.31%     
===========================================
  Files            5        5              
  Lines         1632     1639       +7     
===========================================
+ Hits          1632     1634       +2     
- Misses           0        5       +5     

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (70cbe2b) to head (11d7be8).
Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff            @@
##            master      #174   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            5         5           
  Lines         1632      1640    +8     
=========================================
+ Hits          1632      1640    +8     

[AaronChen0]

Follow your suggestion.

[holiman]

We can simplify even more, by moving the division to the end of the loop, and lose the flag.
For some reason, my linter wants to do further changes, but the CI-linter doesn't seem to complain.

diff --git a/uint256.go b/uint256.go
index 493f44a..7b93ee9 100644
--- a/uint256.go
+++ b/uint256.go
@@ -369,9 +369,9 @@ func umul(x, y *Int, res *[8]uint64) {
 func (z *Int) Mul(x, y *Int) *Int {
 	var (
 		carry0, carry1, carry2 uint64
-		res1, res2 uint64
-		x0, x1, x2, x3 = x[0], x[1], x[2], x[3]
-		y0, y1, y2, y3 = y[0], y[1], y[2], y[3]
+		res1, res2             uint64
+		x0, x1, x2, x3         = x[0], x[1], x[2], x[3]
+		y0, y1, y2, y3         = y[0], y[1], y[2], y[3]
 	)
 
 	carry0, z[0] = bits.Mul64(x0, y0)
@@ -1283,14 +1283,14 @@ func (z *Int) Sqrt(x *Int) *Int {
 	if x.IsUint64() {
 		var (
 			x0 uint64 = x.Uint64()
-			z1 uint64 = 1 << ((bits.Len64(x0)+1)/2)
+			z1 uint64 = 1 << ((bits.Len64(x0) + 1) / 2)
 			z2 uint64
 		)
 		if x0 < 2 {
 			return z.SetUint64(x0)
 		}
 		for {
-			z2 = (z1 + x0 / z1) >> 1
+			z2 = (z1 + x0/z1) >> 1
 			if z2 >= z1 {
 				return z.SetUint64(z1)
 			}
@@ -1306,27 +1306,24 @@ func (z *Int) Sqrt(x *Int) *Int {
 
 	// We can do the first division outside the loop
 	z2.Rsh(x, uint(x.BitLen()+1)/2) // The first div is equal to a right shift
-	first := true
 
 	for {
-		// z2.Div(x, z1) -- x > MaxUint64, x > z1 > 0
-		if !first { // first division was done outside the loop
-			z2.Clear()
-			udivrem(z2[:], x[:], z1, nil)
-		}
-		first = false
 		z2.Add(z2, z1)
-		
+
 		// z2 = z2.Rsh(z2, 1) -- the code below does a 1-bit rsh faster
-		z2[0] = (z2[0] >> 1) | z2[1] << 63
-		z2[1] = (z2[1] >> 1) | z2[2] << 63
-		z2[2] = (z2[2] >> 1) | z2[3] << 63
+		z2[0] = (z2[0] >> 1) | z2[1]<<63
+		z2[1] = (z2[1] >> 1) | z2[2]<<63
+		z2[2] = (z2[2] >> 1) | z2[3]<<63
 		z2[3] >>= 1
 
 		if !z2.Lt(z1) {
 			return z.Set(z1)
 		}
 		z1.Set(z2)
+		// Next iteration of the loop
+		// z2.Div(x, z1) -- x > MaxUint64, x > z1 > 0
+		z2.Clear()
+		udivrem(z2[:], x[:], z1, nil)
 	}
 }
 

[AaronChen0]

Good observation!

[holiman]

goos: linux
goarch: amd64
pkg: github.com/holiman/uint256
cpu: 12th Gen Intel(R) Core(TM) i7-1270P
                      │ bench.sqrt.1 │             bench.sqrt.4             │
                      │    sec/op    │    sec/op     vs base                │
Sqrt/single/uint256-8   650.5n ± 42%   391.6n ± 11%  -39.80% (p=0.002 n=10)

👍