netstack: faster checksum

tcp_benchmark shows we spend significant time checksumming -- after Syscall,
checksumming is the most time-consuming function.

AMD64 gets a 77% faster checksum via assmebly implementation.

Other 64bit machines get a new checksum function, which runs 38% faster than
current checksumming on my machine and 34% faster on an ARM test machine.

tcp_benchmark is pretty noisy here, but generally shows a 2-5% decrease in CPU
usage or a small boost to throughput:

```
                                              │ /tmp/old.log │            /tmp/new.log             │
                                              │     Mb/s     │     Mb/s      vs base               │
TCP/role=server/host-gso=false/host-gro=false    142.0 ± 28%    154.5 ± 24%       ~ (p=0.459 n=40)
TCP/role=client/host-gso=false/host-gro=false   1.794k ±  1%   1.840k ±  1%  +2.59% (p=0.000 n=40)
geomean                                          504.7          533.2        +5.65%

                                              │ /tmp/old.log │            /tmp/new.log             │
                                              │   cpu-time   │   cpu-time    vs base               │
TCP/role=server/host-gso=false/host-gro=false   244.0m ± 25%   256.3m ± 21%       ~ (p=0.529 n=40)
TCP/role=client/host-gso=false/host-gro=false    2.251 ±  1%    2.242 ±  1%       ~ (p=0.079 n=40)
geomean                                         741.1m         758.1m        +2.30%
```

So it seems checksumming is not the bottleneck in tcp_benchmark. This may be
different in other environments, especially those where we cannot rely on host
receive checksum offload.

PiperOrigin-RevId: 540620826

pkg/tcpip/checksum/BUILD

@@ -13,6 +13,11 @@ go_test(
 
 go_library(
     name = "checksum",
-    srcs = ["checksum.go"],
+    srcs = [
+        "checksum.go",
+        "checksum_amd64.go",
+        "checksum_amd64.s",
+        "checksum_noasm_unsafe.go",
+    ],
     visibility = ["//visibility:public"],
 )

pkg/tcpip/checksum/checksum.go

@@ -146,12 +146,12 @@ func unrolledCalculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bo
 }
 
 // Checksum calculates the checksum (as defined in RFC 1071) of the bytes in the
-// given byte array. This function uses an optimized unrolled version of the
-// checksum algorithm.
+// given byte array. This function uses an optimized version of the checksum
+// algorithm.
 //
 // The initial checksum must have been computed on an even number of bytes.
 func Checksum(buf []byte, initial uint16) uint16 {
-	s, _ := unrolledCalculateChecksum(buf, false, initial)
+	s, _ := calculateChecksum(buf, false, initial)
 	return s
 }
 
@@ -164,7 +164,7 @@ type Checksumer struct {
 // Add adds b to checksum.
 func (c *Checksumer) Add(b []byte) {
 	if len(b) > 0 {
-		c.sum, c.odd = unrolledCalculateChecksum(b, c.odd, c.sum)
+		c.sum, c.odd = calculateChecksum(b, c.odd, c.sum)
 	}
 }
 

pkg/tcpip/checksum/checksum_amd64.go

@@ -0,0 +1,21 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build amd64
+// +build amd64
+
+package checksum
+
+// calculateChecksum is defined in assembly.
+func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool)

pkg/tcpip/checksum/checksum_amd64.s

@@ -0,0 +1,138 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build amd64
+// +build amd64
+
+#include "textflag.h"
+
+// calculateChecksum computes the checksum of a slice, taking into account a
+// previously computed initial value and whether the first byte is a lower or
+// upper byte.
+//
+// It utilizes byte order independence and parallel summation as described in
+// RFC 1071 1.2.
+//
+// The best way to understand this function is to understand
+// checksum_noasm_unsafe.go first, which implements largely the same logic.
+// Using assembly speeds things up via ADC (add with carry).
+TEXT ·calculateChecksum(SB),NOSPLIT|NOFRAME,$0-35
+  // Store arguments in registers.
+  MOVW initial+26(FP), AX
+  MOVQ buf_len+8(FP), BX
+  MOVQ buf_base+0(FP), CX
+  XORQ R8, R8
+  MOVB odd+24(FP), R8
+
+  // Account for a previous odd number of bytes.
+  //
+  // if odd {
+  //   initial += buf[0]
+  //   buf = buf[1:]
+  // }
+  CMPB R8, $0
+  JE newlyodd
+  XORQ R9, R9
+  MOVB (CX), R9
+  ADDW R9, AX
+  ADCW $0, AX
+  INCQ CX
+  DECQ BX
+
+  // See whether we're checksumming an odd number of bytes. If so, the final
+  // byte is a big endian most significant byte, and so needs to be shifted.
+  //
+  // odd = buf_len%2 != 0
+  // if odd {
+  //   buf_len--
+  //   initial += buf[buf_len]<<8
+  // }
+newlyodd:
+  XORQ R8, R8
+  TESTQ $1, BX
+  JZ swaporder
+  MOVB $1, R8
+  DECQ BX
+  XORQ R10, R10
+  MOVB (CX)(BX*1), R10
+  SHLQ $8, R10
+  ADDW R10, AX
+  ADCW $0, AX
+
+swaporder:
+  // Load initial in network byte order.
+  BSWAPQ AX
+  SHRQ $48, AX
+
+  // Accumulate 8 bytes at a time.
+  //
+  // while buf_len >= 8 {
+  //   acc, carry = acc + *(uint64 *)(buf) + carry
+  //   buf_len -= 8
+  //   buf = buf[8:]
+  // }
+  // acc += carry
+  JMP addcond
+addloop:
+  ADDQ (CX), AX
+  ADCQ $0, AX
+  SUBQ $8, BX
+  ADDQ $8, CX
+addcond:
+  CMPQ BX, $8
+  JAE addloop
+
+  // TODO(krakauer): We can do 4 byte accumulation too.
+
+  // Accumulate the rest 2 bytes at a time.
+  //
+  // while buf_len > 0 {
+  //   acc, carry = acc + *(uint16 *)(buf)
+  //   buf_len -= 2
+  //   buf = buf[2:]
+  // }
+  JMP slowaddcond
+slowaddloop:
+  XORQ DX, DX
+  MOVW (CX), DX
+  ADDQ DX, AX
+  ADCQ $0, AX
+  SUBQ $2, BX
+  ADDQ $2, CX
+slowaddcond:
+  CMPQ BX, $2
+  JAE slowaddloop
+
+  // Fold into 16 bits.
+  //
+  // for acc > math.MaxUint16 {
+  //   acc = (acc & 0xffff) + acc>>16
+  // }
+  JMP foldcond
+foldloop:
+  MOVQ AX, DX
+  ANDQ $0xffff, DX
+  SHRQ $16, AX
+  ADDQ DX, AX
+  // We don't need ADC because folding will take care of it
+foldcond:
+  CMPQ AX, $0xffff
+  JA foldloop
+
+  // Return the checksum in host byte order.
+  BSWAPQ AX
+  SHRQ $48, AX
+  MOVW AX, ret+32(FP)
+  MOVB R8, ret1+34(FP)
+  RET

pkg/tcpip/checksum/checksum_noasm_unsafe.go

@@ -0,0 +1,62 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !amd64
+// +build !amd64
+
+package checksum
+
+import (
+	"math"
+	"math/bits"
+	"unsafe"
+)
+
+func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
+	if bits.UintSize == 64 {
+		// Utilize byte order independence and parallel summation as
+		// described in RFC 1071 1.2.
+
+		// Initialize the accumulator and account for odd byte input.
+		acc := uint(initial)
+		if odd {
+			acc += uint(buf[0])
+			buf = buf[1:]
+		}
+		// It doesn't matter what endianness we use, only that it's
+		// consistent throughout the calculation. See RFC ?.
+		acc = ((acc & 0xff00) >> 8) | ((acc & 0x00ff) << 8)
+
+		// Compute the checksum.
+		remaining := len(buf)
+		var carry uint
+		for remaining >= 8 {
+			acc, carry = bits.Add(acc, *(*uint)(unsafe.Pointer(&buf[0])), carry)
+			remaining -= 8
+			buf = buf[8:]
+		}
+		acc += carry
+
+		// Fold the checksum into 16 bits.
+		for acc > math.MaxUint16 {
+			acc = (acc & 0xffff) + acc>>16
+		}
+
+		// Swap back to little endian and let unrolledCalculateChecksum
+		// handle the remaining bytes.
+		acc = ((acc & 0xff00) >> 8) | ((acc & 0x00ff) << 8)
+		return unrolledCalculateChecksum(buf, false, uint16(acc))
+	}
+	return unrolledCalculateChecksum(buf, odd, initial)
+}

pkg/tcpip/checksum/checksum_test.go

@@ -19,8 +19,11 @@ package checksum
 import (
 	"bytes"
 	"fmt"
+	"math"
+	"math/bits"
 	"math/rand"
 	"testing"
+	"unsafe"
 )
 
 func TestChecksumer(t *testing.T) {
@@ -95,7 +98,28 @@ func TestChecksumer(t *testing.T) {
 }
 
 func TestChecksum(t *testing.T) {
-	var bufSizes = []int{0, 1, 2, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128, 255, 256, 257, 1023, 1024}
+	var bufSizes = []int{
+		0,
+		1,
+		2,
+		3,
+		4,
+		7,
+		8,
+		15,
+		16,
+		31,
+		32,
+		63,
+		64,
+		127,
+		128,
+		255,
+		256,
+		257,
+		1023,
+		1024,
+	}
 	type testCase struct {
 		buf      []byte
 		initial  uint16
@@ -112,7 +136,7 @@ func TestChecksum(t *testing.T) {
 	}
 
 	for i := range testCases {
-		testCases[i].csumOrig = Old(testCases[i].buf, testCases[i].initial)
+		testCases[i].csumOrig = old(testCases[i].buf, testCases[i].initial)
 		testCases[i].csumNew = Checksum(testCases[i].buf, testCases[i].initial)
 		if got, want := testCases[i].csumNew, testCases[i].csumOrig; got != want {
 			t.Fatalf("new checksum for (buf = %x, initial = %d) does not match old got: %d, want: %d", testCases[i].buf, testCases[i].initial, got, want)
@@ -127,8 +151,10 @@ func BenchmarkChecksum(b *testing.B) {
 		fn   func([]byte, uint16) uint16
 		name string
 	}{
-		{Old, fmt.Sprintf("checksum_old")},
-		{Checksum, fmt.Sprintf("checksum")},
+		{old, "checksum_old"},
+		{unrolled, "unrolled"},
+		{bitsLib, "bitslib"},
+		{Checksum, "checksum"},
 	}
 
 	for _, csumImpl := range checkSumImpls {
@@ -154,18 +180,18 @@ func BenchmarkChecksum(b *testing.B) {
 	}
 }
 
-// Old calculates the checksum (as defined in RFC 1071) of the bytes in
+// old calculates the checksum (as defined in RFC 1071) of the bytes in
 // the given byte array. This function uses a non-optimized implementation. Its
 // only retained for reference and to use as a benchmark/test. Most code should
 // use the header.Checksum function.
 //
 // The initial checksum must have been computed on an even number of bytes.
-func Old(buf []byte, initial uint16) uint16 {
-	s, _ := calculateChecksum(buf, false, uint32(initial))
+func old(buf []byte, initial uint16) uint16 {
+	s, _ := oldCalculateChecksum(buf, false, uint32(initial))
 	return s
 }
 
-func calculateChecksum(buf []byte, odd bool, initial uint32) (uint16, bool) {
+func oldCalculateChecksum(buf []byte, odd bool, initial uint32) (uint16, bool) {
 	v := initial
 
 	if odd {
@@ -186,3 +212,50 @@ func calculateChecksum(buf []byte, odd bool, initial uint32) (uint16, bool) {
 
 	return Combine(uint16(v), uint16(v>>16)), odd
 }
+
+func unrolled(buf []byte, initial uint16) uint16 {
+	s, _ := unrolledCalculateChecksum(buf, false, initial)
+	return s
+}
+
+func bitsLib(buf []byte, initial uint16) uint16 {
+	s, _ := bitsAdd(buf, false, initial)
+	return s
+}
+
+// bitsAdd is copied from checksum_noasm_unsafe.go so that it can be
+// benchmarked.
+func bitsAdd(buf []byte, odd bool, initial uint16) (uint16, bool) {
+	if bits.UintSize == 64 {
+		// Initialize the accumulator and account for odd byte input.
+		acc := uint(initial)
+		if odd {
+			acc += uint(buf[0])
+			buf = buf[1:]
+		}
+		// It doesn't matter what endianness we use, only that it's
+		// consistent throughout the calculation. See RFC ?.
+		acc = ((acc & 0xff00) >> 8) | ((acc & 0x00ff) << 8)
+
+		// Compute the checksum.
+		remaining := len(buf)
+		var carry uint
+		for remaining >= 8 {
+			acc, carry = bits.Add(acc, *(*uint)(unsafe.Pointer(&buf[0])), carry)
+			remaining -= 8
+			buf = buf[8:]
+		}
+		acc += carry
+
+		// Fold the checksum into 16 bits.
+		for acc > math.MaxUint16 {
+			acc = (acc & 0xffff) + acc>>16
+		}
+
+		// Swap back to little endian and let unrolledCalculateChecksum
+		// handle the remaining bytes.
+		acc = ((acc & 0xff00) >> 8) | ((acc & 0x00ff) << 8)
+		return unrolledCalculateChecksum(buf, false, uint16(acc))
+	}
+	return unrolledCalculateChecksum(buf, odd, initial)
+}

pkg/tcpip/header/checksum.go

@@ -32,9 +32,9 @@ func PseudoHeaderChecksum(protocol tcpip.TransportProtocolNumber, srcAddr tcpip.
 	xsum = checksum.Checksum(dstAddr.AsSlice(), xsum)
 
 	// Add the length portion of the checksum to the pseudo-checksum.
-	tmp := make([]byte, 2)
-	binary.BigEndian.PutUint16(tmp, totalLen)
-	xsum = checksum.Checksum(tmp, xsum)
+	var tmp [2]byte
+	binary.BigEndian.PutUint16(tmp[:], totalLen)
+	xsum = checksum.Checksum(tmp[:], xsum)
 
 	return checksum.Checksum([]byte{0, uint8(protocol)}, xsum)
 }