From 6c0f22c66ea171c5107796d8849c02e7f6fa899d Mon Sep 17 00:00:00 2001 From: Kevin Krakauer Date: Thu, 15 Jun 2023 10:25:01 -0700 Subject: [PATCH] netstack: faster checksum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 7 +- pkg/tcpip/checksum/checksum.go | 8 +- pkg/tcpip/checksum/checksum_amd64.go | 21 +++ pkg/tcpip/checksum/checksum_amd64.s | 138 ++++++++++++++++++++ pkg/tcpip/checksum/checksum_noasm_unsafe.go | 62 +++++++++ pkg/tcpip/checksum/checksum_test.go | 89 +++++++++++-- pkg/tcpip/header/checksum.go | 6 +- 7 files changed, 315 insertions(+), 16 deletions(-) create mode 100644 pkg/tcpip/checksum/checksum_amd64.go create mode 100644 pkg/tcpip/checksum/checksum_amd64.s create mode 100644 pkg/tcpip/checksum/checksum_noasm_unsafe.go diff --git a/pkg/tcpip/checksum/BUILD b/pkg/tcpip/checksum/BUILD index 26dd9a5cc..5dffbfd0a 100644 --- a/pkg/tcpip/checksum/BUILD +++ b/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"], ) diff --git a/pkg/tcpip/checksum/checksum.go b/pkg/tcpip/checksum/checksum.go index 7f01877fc..fa6e8bc3a 100644 --- a/pkg/tcpip/checksum/checksum.go +++ b/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) } } diff --git a/pkg/tcpip/checksum/checksum_amd64.go b/pkg/tcpip/checksum/checksum_amd64.go new file mode 100644 index 000000000..3be3afefa --- /dev/null +++ b/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) diff --git a/pkg/tcpip/checksum/checksum_amd64.s b/pkg/tcpip/checksum/checksum_amd64.s new file mode 100644 index 000000000..c2a97a647 --- /dev/null +++ b/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 diff --git a/pkg/tcpip/checksum/checksum_noasm_unsafe.go b/pkg/tcpip/checksum/checksum_noasm_unsafe.go new file mode 100644 index 000000000..b67455cf1 --- /dev/null +++ b/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) +} diff --git a/pkg/tcpip/checksum/checksum_test.go b/pkg/tcpip/checksum/checksum_test.go index d462bb11c..777790fe0 100644 --- a/pkg/tcpip/checksum/checksum_test.go +++ b/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) +} diff --git a/pkg/tcpip/header/checksum.go b/pkg/tcpip/header/checksum.go index cde74e730..1cdda6b90 100644 --- a/pkg/tcpip/header/checksum.go +++ b/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) }