netstack: faster checksum

This is a roll-forward of cl/540620826. The differences are:

- Added Tony's test that showed breakage on ARM64
- Fixed the 64-bit non-AMD64 checksum implementation
  - Handles odd buffers correctly
  - Doesn't call the unrolled checksum impl anymore
  - Re-ordered calculateChecksum to reduce indentation

Otherwise this is the same CL. The only changes are in checksum_test.go and
checksum_noasm_unsafe.go.

PiperOrigin-RevId: 547890206
This commit is contained in:
Kevin Krakauer
2023-07-13 12:51:44 -07:00
committed by gVisor bot
parent e7c1bc6baf
commit 91b023d95c
7 changed files with 373 additions and 16 deletions
+6 -1
View File
@@ -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"],
)
+4 -4
View File
@@ -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)
}
}
+24
View File
@@ -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
// +build amd64
package checksum
// 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)
+138
View File
@@ -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
@@ -0,0 +1,80 @@
// 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"
)
// 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 {
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 := uint(((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
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 += uint(buf[len(buf)-1])
buf = buf[:len(buf)-1]
}
// 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)
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)
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
}
+118 -8
View File
@@ -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)
@@ -120,6 +144,43 @@ func TestChecksum(t *testing.T) {
}
}
// TestIncrementalChecksum tests for breakages of Checksummer as described in
// b/289284842.
func TestIncrementalChecksum(t *testing.T) {
buf := []byte{
0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x31,
0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c,
0x3d, 0x3e, 0x3f, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46, 0x47,
0x48, 0x49, 0x4a, 0x4b, 0x4c, 0x4d, 0x4e, 0x4f, 0x50, 0x51, 0x52,
0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5a, 0x5b, 0x5c, 0x5d,
0x5e, 0x5f, 0x60, 0x61, 0x62, 0x63,
}
// Go through buf and check that checksum(buf[:end]) is equivalent to
// an incremental checksum of two chunks of buf[:end].
for end := 2; end <= len(buf); end++ {
for start := 1; start < end; start++ {
t.Run(fmt.Sprintf("end=%d start=%d", end, start), func(t *testing.T) {
var cs Checksumer
cs.Add(buf[:end])
csum := cs.Checksum()
cs = Checksumer{}
cs.Add(buf[:start])
cs.Add(buf[start:end])
csumIncremental := cs.Checksum()
if want := old(buf[:end], 0); csum != want {
t.Fatalf("checksum is wrong: got %x, expected %x", csum, want)
}
if csum != csumIncremental {
t.Errorf("checksums should be the same: %x %x", csum, csumIncremental)
}
})
}
}
}
func BenchmarkChecksum(b *testing.B) {
var bufSizes = []int{64, 128, 256, 512, 1024, 1500, 2048, 4096, 8192, 16384, 32767, 32768, 65535, 65536}
@@ -127,8 +188,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 +217,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 +249,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)
}
+3 -3
View File
@@ -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)
}