Merge pull request #9500 from kevinGC:checksumalign

PiperOrigin-RevId: 573287057

pkg/tcpip/checksum/BUILD

@@ -17,6 +17,8 @@ go_library(
         "checksum.go",
         "checksum_amd64.go",
         "checksum_amd64.s",
+        "checksum_mips64.go",
+        "checksum_noasm.go",
         "checksum_noasm_unsafe.go",
     ],
     visibility = ["//visibility:public"],

pkg/tcpip/checksum/checksum_amd64.go

@@ -17,8 +17,11 @@
 
 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.
-//
-// calculateChecksum is defined in assembly.
-func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool)
+func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
+	return calculateChecksumAMD64(buf, odd, initial)
+}

pkg/tcpip/checksum/checksum_amd64.s

@@ -27,7 +27,7 @@
 // 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
+TEXT ·calculateChecksumAMD64(SB),NOSPLIT|NOFRAME,$0-35
   // Store arguments in registers.
   MOVW initial+26(FP), AX
   MOVQ buf_len+8(FP), BX
@@ -75,6 +75,30 @@ swaporder:
   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 {

pkg/tcpip/checksum/checksum_mips64.go

@@ -0,0 +1,24 @@
+// 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

@@ -0,0 +1,24 @@
+// 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

@@ -12,9 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//go:build !amd64
-// +build !amd64
-
 package checksum
 
 import (
@@ -25,10 +22,16 @@ import (
 
 // 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) {
-	// Note: we can probably remove unrolledCalculateChecksum altogether,
-	// but I don't have any 32 bit machines to benchmark on.
-	if bits.UintSize != 64 {
+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)
 	}
 
@@ -37,12 +40,12 @@ func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
 
 	// It doesn't matter what endianness we use, only that it's
 	// consistent throughout the calculation. See RFC 1071 1.2.B.
-	acc := uint(((initial & 0xff00) >> 8) | ((initial & 0x00ff) << 8))
+	acc := uint64(((initial & 0xff00) >> 8) | ((initial & 0x00ff) << 8))
 
 	// Account for initial having been calculated over an odd number of
 	// bytes.
 	if odd {
-		acc += uint(buf[0]) << 8
+		acc += uint64(buf[0]) << 8
 		buf = buf[1:]
 	}
 
@@ -50,21 +53,28 @@ func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
 	// so, the final byte is a big endian most significant byte.
 	odd = len(buf)%2 != 0
 	if odd {
-		acc += uint(buf[len(buf)-1])
+		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.
-	var carry uint
 	for len(buf) >= 8 {
-		acc, carry = bits.Add(acc, *(*uint)(unsafe.Pointer(&buf[0])), carry)
+		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.Add(acc, uint(*(*uint16)(unsafe.Pointer(&buf[0]))), carry)
+		acc, carry = bits.Add64(acc, uint64(*(*uint16)(unsafe.Pointer(&buf[0]))), carry)
 		buf = buf[2:]
 	}
 	acc += carry
@@ -78,3 +88,7 @@ func calculateChecksum(buf []byte, odd bool, initial uint16) (uint16, bool) {
 	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

@@ -19,11 +19,8 @@ package checksum
 import (
 	"bytes"
 	"fmt"
-	"math"
-	"math/bits"
 	"math/rand"
 	"testing"
-	"unsafe"
 )
 
 func TestChecksumer(t *testing.T) {
@@ -121,10 +118,8 @@ func TestChecksum(t *testing.T) {
 		1024,
 	}
 	type testCase struct {
-		buf      []byte
-		initial  uint16
-		csumOrig uint16
-		csumNew  uint16
+		buf     []byte
+		initial uint16
 	}
 	testCases := make([]testCase, 100000)
 	// Ensure same buffer generation for test consistency.
@@ -135,12 +130,37 @@ func TestChecksum(t *testing.T) {
 		rnd.Read(testCases[i].buf)
 	}
 
-	for i := range testCases {
-		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)
-		}
+	checkSumImpls := []struct {
+		fn   func([]byte, uint16) uint16
+		name string
+	}{
+		{old, "checksum_old"},
+		{unrolled, "unrolled"},
+		{bitsLib, "bitslib"},
+		{Checksum, "checksum"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("buf size %d", len(tc.buf)), func(t *testing.T) {
+			// Also test different offsets into the buffers. This
+			// tests the correctess of optimizations dealing with
+			// non-64-bit aligned numbers.
+			for offset := 0; offset < 8; offset++ {
+				t.Run(fmt.Sprintf("offset %d", offset), func(t *testing.T) {
+					if offset > len(tc.buf) {
+						t.Skip("offset is greater than buffer size")
+					}
+					buf := tc.buf[offset:]
+					for i := 0; i < len(checkSumImpls)-1; i++ {
+						first := checkSumImpls[i].fn(buf, tc.initial)
+						second := checkSumImpls[i+1].fn(buf, tc.initial)
+						if first != second {
+							t.Fatalf("for (buf = 0x%x, initial = 0x%x) checksum %q does not match %q: got: 0x%x and 0x%x", buf, tc.initial, checkSumImpls[i].name, checkSumImpls[i+1].name, first, second)
+						}
+					}
+				})
+			}
+		})
 	}
 }
 
@@ -256,43 +276,6 @@ func unrolled(buf []byte, initial uint16) uint16 {
 }
 
 func bitsLib(buf []byte, initial uint16) uint16 {
-	s, _ := bitsAdd(buf, false, initial)
+	s, _ := calculateChecksumNoASM(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)
-}