netstack: further optimize checksumming

- Restructure calculateChecksumNoASM() to consistently handle misaligned
  inputs: special-case inputs less than 8 bytes in size, otherwise align the
  buffer address to 8 bytes before checksumming in bulk. This avoids
  significant performance degradation on misaligned inputs due to fallback to
  unrolledCalculateChecksum().

- Sum up to 64 bytes per loop iteration in calculateChecksumNoASM(), as in
  unrolledCalculateChecksum(). This significantly improves performance for
  large inputs.

- Don't plumb carry between blocks in calculateChecksumNoASM(). On x86, this is
  significantly faster because many instructions clobber the carry flag, so
  saving it requires extra overhead; on arm64, this has minimal effect on the
  generated code (the compiler inserts an extra useless copy from the zero
  register to a temporary).

- Replace the folding of uint64 accumulator to uint16 checksum with a
  branchless equivalent, which is ~2ns/checksum faster on average due to
  avoiding an unpredictable loop.

- Fix calculateChecksumNoASM() on big-endian architectures.

- Delete unrolledCalculateChecksum(), which is now unused.

Before these changes, calculateChecksumAMD64() is approximately twice as fast
as calculateChecksumNoASM(). After these changes, calculateChecksumNoASM() is
faster than calculateChecksumAMD64():

```
goos: linux
goarch: amd64
pkg: pkg/tcpip/checksum/checksum
cpu: Intel(R) Xeon(R) CPU @ 2.80GHz
                           │ checksum_baseline │         checksum_experiment         │
                           │      sec/op       │   sec/op     vs base                │
Checksum/checksum_64-48            10.50n ± 0%   10.10n ± 0%   -3.86% (p=0.000 n=10)
Checksum/checksum_128-48           13.72n ± 0%   12.16n ± 0%  -11.37% (p=0.000 n=10)
Checksum/checksum_256-48           22.80n ± 0%   16.31n ± 0%  -28.48% (p=0.000 n=10)
Checksum/checksum_512-48           39.92n ± 0%   24.96n ± 0%  -37.47% (p=0.000 n=10)
Checksum/checksum_1024-48          81.34n ± 0%   44.26n ± 2%  -45.58% (p=0.000 n=10)
Checksum/checksum_1500-48         119.70n ± 1%   65.05n ± 2%  -45.65% (p=0.000 n=10)
Checksum/checksum_2048-48         157.30n ± 0%   86.65n ± 0%  -44.91% (p=0.000 n=10)
Checksum/checksum_4096-48          309.5n ± 0%   172.4n ± 3%  -44.30% (p=0.000 n=10)
Checksum/checksum_8192-48          615.0n ± 0%   348.0n ± 0%  -43.42% (p=0.000 n=10)
Checksum/checksum_16384-48        1222.5n ± 0%   690.8n ± 0%  -43.49% (p=0.000 n=10)
Checksum/checksum_32767-48         2.447µ ± 0%   1.381µ ± 0%  -43.55% (p=0.000 n=10)
Checksum/checksum_32768-48         2.442µ ± 0%   1.381µ ± 0%  -43.48% (p=0.000 n=10)
Checksum/checksum_65535-48         4.883µ ± 0%   2.933µ ± 0%  -39.94% (p=0.000 n=10)
Checksum/checksum_65536-48         4.877µ ± 0%   2.935µ ± 0%  -39.82% (p=0.000 n=10)
geomean                            260.6n        161.9n       -37.88%
```

Thus, use calculateChecksumNoASM() on all architectures and rename it to
calculateChecksum().

---

Note: As of this writing, calculateChecksumNoASM() is still ~7% slower than an
equivalently-optimized x86 assembly implementation, for three main reasons:

1. The Go compiler emits separate MOV memory-to-register and ADD/ADC
   instructions rather than using memory operands for the latter.

2. For statements of the form `buf = buf[1:]`, updating the slice data pointer
   takes ~5 instructions instead of 1, because the compiler avoids incrementing
   the pointer if the resulting slice length is 0 (such that the pointer would
   point outside of the underlying array). This overhead can be avoided by
   using uintptr and runtime.KeepAlive(), but doing so worsens problem (1):
   replacing `*(*uint64(unsafe.Pointer(&buf[8])))` with
   `*(*uint64(unsafe.Pointer(addr+8)))` causes the compiler to split each MOV
   into LEA + MOV.

3. The faster form of reduce() described in that function is implementable in
   assembly.

Thus, we can add the optimized x86 assembly implementation back if desired.

---

`GOOS=linux GOARCH=mips64 go test -exec qemu-mips64 ./...` fails before this
CL, since the "bitslib" implementation, i.e. calculateChecksumNoASM(), is
broken on big-endian architectures. After this CL:
```
ok      gvisor.dev/gvisor/pkg/tcpip/checksum    125.970s
```

PiperOrigin-RevId: 574264329

pkg/tcpip/checksum/BUILD

@@ -15,11 +15,7 @@ go_library(
     name = "checksum",
     srcs = [
         "checksum.go",
-        "checksum_amd64.go",
-        "checksum_amd64.s",
-        "checksum_mips64.go",
-        "checksum_noasm.go",
-        "checksum_noasm_unsafe.go",
+        "checksum_unsafe.go",
     ],
     visibility = ["//visibility:public"],
 )

pkg/tcpip/checksum/checksum.go

@@ -30,121 +30,6 @@ func Put(b []byte, xsum uint16) {
 	binary.BigEndian.PutUint16(b, xsum)
 }
 
-func unrolledCalculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
-	v := uint32(initial)
-
-	if odd {
-		v += uint32(buf[0])
-		buf = buf[1:]
-	}
-
-	l := len(buf)
-	odd = l&1 != 0
-	if odd {
-		l--
-		v += uint32(buf[l]) << 8
-	}
-	for (l - 64) >= 0 {
-		i := 0
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		i += 16
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		i += 16
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		i += 16
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		buf = buf[64:]
-		l = l - 64
-	}
-	if (l - 32) >= 0 {
-		i := 0
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		i += 16
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		buf = buf[32:]
-		l = l - 32
-	}
-	if (l - 16) >= 0 {
-		i := 0
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		v += (uint32(buf[i+8]) << 8) + uint32(buf[i+9])
-		v += (uint32(buf[i+10]) << 8) + uint32(buf[i+11])
-		v += (uint32(buf[i+12]) << 8) + uint32(buf[i+13])
-		v += (uint32(buf[i+14]) << 8) + uint32(buf[i+15])
-		buf = buf[16:]
-		l = l - 16
-	}
-	if (l - 8) >= 0 {
-		i := 0
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		v += (uint32(buf[i+4]) << 8) + uint32(buf[i+5])
-		v += (uint32(buf[i+6]) << 8) + uint32(buf[i+7])
-		buf = buf[8:]
-		l = l - 8
-	}
-	if (l - 4) >= 0 {
-		i := 0
-		v += (uint32(buf[i]) << 8) + uint32(buf[i+1])
-		v += (uint32(buf[i+2]) << 8) + uint32(buf[i+3])
-		buf = buf[4:]
-		l = l - 4
-	}
-
-	// At this point since l was even before we started unrolling
-	// there can be only two bytes left to add.
-	if l != 0 {
-		v += (uint32(buf[0]) << 8) + uint32(buf[1])
-	}
-
-	return Combine(uint16(v), uint16(v>>16)), odd
-}
-
 // Checksum calculates the checksum (as defined in RFC 1071) of the bytes in the
 // given byte array. This function uses an optimized version of the checksum
 // algorithm.

pkg/tcpip/checksum/checksum_amd64.go

@@ -1,27 +0,0 @@
-// 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
-
-// calculateChecksumAMD64 is defined in assembly.
-func calculateChecksumAMD64(buf []byte, odd bool, initial uint16) (uint16, bool)
-
-// Note: odd indicates whether initial is a partial checksum over an odd number
-// of bytes.
-func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
-	return calculateChecksumAMD64(buf, odd, initial)
-}

pkg/tcpip/checksum/checksum_amd64.s

@@ -1,162 +0,0 @@
-// 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 ·calculateChecksumAMD64(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
-
-  // Handle any bytes that aren't 64-bit aligned. If the buffer starts at an
-  // odd address, we just live with the alignment because doing otherwise
-  // messes up the endianness expected by the below.
-  //
-  // while buf_len > 0 && buf_base%8 != 0 {
-  //   acc, carry = acc + *(uint16 *)(buf)
-  //   buf_len -= 2
-  //   buf = buf[2:]
-  // }
-  JMP unalignedaddcond
-unalignedaddloop:
-  XORQ DX, DX
-  MOVW (CX), DX
-  ADDQ DX, AX
-  ADCQ $0, AX
-  SUBQ $2, BX
-  ADDQ $2, CX
-unalignedaddcond:
-  CMPQ BX, $2
-  JLE addcond
-  TESTQ $7, CX
-  JZ addcond
-  JMP unalignedaddloop
-
-  // 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_mips64.go

@@ -1,24 +0,0 @@
-// 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 mips64
-// +build mips64
-
-package checksum
-
-// Note: odd indicates whether initial is a partial checksum over an odd number
-// of bytes.
-func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
-	return unrolledCalculateChecksum(buf, odd, initial)
-}

pkg/tcpip/checksum/checksum_noasm.go

@@ -1,24 +0,0 @@
-// 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 && !mips64
-// +build !amd64,!mips64
-
-package checksum
-
-// Note: odd indicates whether initial is a partial checksum over an odd number
-// of bytes.
-func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
-	return calculateChecksumNoASM(buf, odd, initial)
-}

pkg/tcpip/checksum/checksum_noasm_unsafe.go

@@ -1,94 +0,0 @@
-// 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.
-
-package checksum
-
-import (
-	"math"
-	"math/bits"
-	"unsafe"
-)
-
-// Note: odd indicates whether initial is a partial checksum over an odd number
-// of bytes.
-func calculateChecksumNoASM(buf []byte, odd bool, initial uint16) (uint16, bool) {
-	// Fall back to slower checksum if we're not on a 64 bit machine or if
-	// this optimization will result in misaligned accesses. Calculating
-	// the checksum starting at an odd address messes up the endianness
-	// expected by the below.
-	var oddOffset uintptr
-	if odd {
-		oddOffset = 1
-	}
-	if bits.UintSize != 64 || (sliceAddr(buf)+oddOffset)%2 != 0 {
-		return unrolledCalculateChecksum(buf, odd, initial)
-	}
-
-	// Utilize byte order independence and parallel summation as
-	// described in RFC 1071 1.2.
-
-	// It doesn't matter what endianness we use, only that it's
-	// consistent throughout the calculation. See RFC 1071 1.2.B.
-	acc := uint64(((initial & 0xff00) >> 8) | ((initial & 0x00ff) << 8))
-
-	// Account for initial having been calculated over an odd number of
-	// bytes.
-	if odd {
-		acc += uint64(buf[0]) << 8
-		buf = buf[1:]
-	}
-
-	// See whether we're checksumming an odd number of bytes. If
-	// so, the final byte is a big endian most significant byte.
-	odd = len(buf)%2 != 0
-	if odd {
-		acc += uint64(buf[len(buf)-1])
-		buf = buf[:len(buf)-1]
-	}
-
-	// Deal with unaligned bytes. We're guaranteed at this point that buf
-	// points to an even address.
-	var carry uint64
-	for sliceAddr(buf)%8 != 0 && len(buf) >= 2 {
-		acc, carry = bits.Add64(acc, uint64(*(*uint16)(unsafe.Pointer(&buf[0]))), carry)
-		buf = buf[2:]
-	}
-
-	// Compute the checksum 8 bytes at a time.
-	for len(buf) >= 8 {
-		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[0])), carry)
-		buf = buf[8:]
-	}
-
-	// Compute the remainder 2 bytes at a time. We are guaranteed that
-	// len(buf) is even due to the above handling of odd-length buffers.
-	for len(buf) > 0 {
-		acc, carry = bits.Add64(acc, uint64(*(*uint16)(unsafe.Pointer(&buf[0]))), carry)
-		buf = buf[2:]
-	}
-	acc += carry
-
-	// Fold the checksum into 16 bits.
-	for acc > math.MaxUint16 {
-		acc = (acc & 0xffff) + acc>>16
-	}
-
-	// Swap the byte order before returning.
-	acc = ((acc & 0xff00) >> 8) | ((acc & 0x00ff) << 8)
-	return uint16(acc), odd
-}
-
-func sliceAddr(buf []byte) uintptr {
-	return uintptr(unsafe.Pointer(unsafe.SliceData(buf)))
-}

pkg/tcpip/checksum/checksum_test.go

@@ -135,8 +135,6 @@ func TestChecksum(t *testing.T) {
 		name string
 	}{
 		{old, "checksum_old"},
-		{unrolled, "unrolled"},
-		{bitsLib, "bitslib"},
 		{Checksum, "checksum"},
 	}
 
@@ -209,8 +207,6 @@ func BenchmarkChecksum(b *testing.B) {
 		name string
 	}{
 		{old, "checksum_old"},
-		{unrolled, "unrolled"},
-		{bitsLib, "bitslib"},
 		{Checksum, "checksum"},
 	}
 
@@ -269,13 +265,3 @@ func oldCalculateChecksum(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, _ := calculateChecksumNoASM(buf, false, initial)
-	return s
-}

pkg/tcpip/checksum/checksum_unsafe.go

@@ -0,0 +1,182 @@
+// 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.
+
+package checksum
+
+import (
+	"encoding/binary"
+	"math/bits"
+	"unsafe"
+)
+
+// Note: odd indicates whether initial is a partial checksum over an odd number
+// of bytes.
+func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
+	// Use a larger-than-uint16 accumulator to benefit from parallel summation
+	// as described in RFC 1071 1.2.C.
+	acc := uint64(initial)
+
+	// Handle an odd number of previously-summed bytes, and get the return
+	// value for odd.
+	if odd {
+		acc += uint64(buf[0])
+		buf = buf[1:]
+	}
+	odd = len(buf)&1 != 0
+
+	// Aligning &buf[0] below is much simpler if len(buf) >= 8; special-case
+	// smaller bufs.
+	if len(buf) < 8 {
+		if len(buf) >= 4 {
+			acc += (uint64(buf[0]) << 8) + uint64(buf[1])
+			acc += (uint64(buf[2]) << 8) + uint64(buf[3])
+			buf = buf[4:]
+		}
+		if len(buf) >= 2 {
+			acc += (uint64(buf[0]) << 8) + uint64(buf[1])
+			buf = buf[2:]
+		}
+		if len(buf) >= 1 {
+			acc += uint64(buf[0]) << 8
+			// buf = buf[1:] is skipped because it's unused and nogo will
+			// complain.
+		}
+		return reduce(acc), odd
+	}
+
+	// On little-endian architectures, multi-byte loads from buf will load
+	// bytes in the wrong order. Rather than byte-swap after each load (slow),
+	// we byte-swap the accumulator before summing any bytes and byte-swap it
+	// back before returning, which still produces the correct result as
+	// described in RFC 1071 1.2.B "Byte Order Independence".
+	//
+	// acc is at most a uint16 + a uint8, so its upper 32 bits must be 0s. We
+	// preserve this property by byte-swapping only the lower 32 bits of acc,
+	// so that additions to acc performed during alignment can't overflow.
+	acc = uint64(bswapIfLittleEndian32(uint32(acc)))
+
+	// Align &buf[0] to an 8-byte boundary.
+	bswapped := false
+	if sliceAddr(buf)&1 != 0 {
+		// Compute the rest of the partial checksum with bytes swapped, and
+		// swap back before returning; see the last paragraph of
+		// RFC 1071 1.2.B.
+		acc = uint64(bits.ReverseBytes32(uint32(acc)))
+		bswapped = true
+		// No `<< 8` here due to the byte swap we just did.
+		acc += uint64(bswapIfLittleEndian16(uint16(buf[0])))
+		buf = buf[1:]
+	}
+	if sliceAddr(buf)&2 != 0 {
+		acc += uint64(*(*uint16)(unsafe.Pointer(&buf[0])))
+		buf = buf[2:]
+	}
+	if sliceAddr(buf)&4 != 0 {
+		acc += uint64(*(*uint32)(unsafe.Pointer(&buf[0])))
+		buf = buf[4:]
+	}
+
+	// Sum 64 bytes at a time. Beyond this point, additions to acc may
+	// overflow, so we have to handle carrying.
+	for len(buf) >= 64 {
+		var carry uint64
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[0])), 0)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[8])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[16])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[24])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[32])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[40])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[48])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[56])), carry)
+		acc, _ = bits.Add64(acc, 0, carry)
+		buf = buf[64:]
+	}
+
+	// Sum the remaining 0-63 bytes.
+	if len(buf) >= 32 {
+		var carry uint64
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[0])), 0)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[8])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[16])), carry)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[24])), carry)
+		acc, _ = bits.Add64(acc, 0, carry)
+		buf = buf[32:]
+	}
+	if len(buf) >= 16 {
+		var carry uint64
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[0])), 0)
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[8])), carry)
+		acc, _ = bits.Add64(acc, 0, carry)
+		buf = buf[16:]
+	}
+	if len(buf) >= 8 {
+		var carry uint64
+		acc, carry = bits.Add64(acc, *(*uint64)(unsafe.Pointer(&buf[0])), 0)
+		acc, _ = bits.Add64(acc, 0, carry)
+		buf = buf[8:]
+	}
+	if len(buf) >= 4 {
+		var carry uint64
+		acc, carry = bits.Add64(acc, uint64(*(*uint32)(unsafe.Pointer(&buf[0]))), 0)
+		acc, _ = bits.Add64(acc, 0, carry)
+		buf = buf[4:]
+	}
+	if len(buf) >= 2 {
+		var carry uint64
+		acc, carry = bits.Add64(acc, uint64(*(*uint16)(unsafe.Pointer(&buf[0]))), 0)
+		acc, _ = bits.Add64(acc, 0, carry)
+		buf = buf[2:]
+	}
+	if len(buf) >= 1 {
+		// bswapIfBigEndian16(buf[0]) == bswapIfLittleEndian16(buf[0]<<8).
+		var carry uint64
+		acc, carry = bits.Add64(acc, uint64(bswapIfBigEndian16(uint16(buf[0]))), 0)
+		acc, _ = bits.Add64(acc, 0, carry)
+		// buf = buf[1:] is skipped because it's unused and nogo will complain.
+	}
+
+	// Reduce the checksum to 16 bits and undo byte swaps before returning.
+	acc16 := bswapIfLittleEndian16(reduce(acc))
+	if bswapped {
+		acc16 = bits.ReverseBytes16(acc16)
+	}
+	return acc16, odd
+}
+
+func reduce(acc uint64) uint16 {
+	// Ideally we would do:
+	//   return uint16(acc>>48) +' uint16(acc>>32) +' uint16(acc>>16) +' uint16(acc)
+	// for more instruction-level parallelism; however, there is no
+	// bits.Add16().
+	acc = (acc >> 32) + (acc & 0xffff_ffff)  // at most 0x1_ffff_fffe
+	acc32 := uint32(acc>>32 + acc)           // at most 0xffff_ffff
+	acc32 = (acc32 >> 16) + (acc32 & 0xffff) // at most 0x1_fffe
+	return uint16(acc32>>16 + acc32)         // at most 0xffff
+}
+
+func bswapIfLittleEndian32(val uint32) uint32 {
+	return binary.BigEndian.Uint32((*[4]byte)(unsafe.Pointer(&val))[:])
+}
+
+func bswapIfLittleEndian16(val uint16) uint16 {
+	return binary.BigEndian.Uint16((*[2]byte)(unsafe.Pointer(&val))[:])
+}
+
+func bswapIfBigEndian16(val uint16) uint16 {
+	return binary.LittleEndian.Uint16((*[2]byte)(unsafe.Pointer(&val))[:])
+}
+
+func sliceAddr(buf []byte) uintptr {
+	return uintptr(unsafe.Pointer(unsafe.SliceData(buf)))
+}