crypto/aes: speed up AES by reducing allocations

marten-seemann · neild · commit 2c150f401bae · 2024-03-25T14:58:19.000Z
By embedding the arrays into aesCipher directly, we can save a few allocations when creating a new AES cipher. This also avoids a lot of pointer chasing when encrypting and decrypting, leading to 3-4% faster performance. Fixes #65507. name old time/op new time/op delta Encrypt/AES-128-16 4.70ns ± 3% 4.46ns ± 2% -5.08% (p=0.000 n=10+10) Encrypt/AES-192-16 4.90ns ± 3% 4.71ns ± 2% -3.98% (p=0.000 n=10+10) Encrypt/AES-256-16 5.18ns ± 3% 4.91ns ± 1% -5.15% (p=0.000 n=10+10) Decrypt/AES-128-16 4.51ns ± 2% 4.37ns ± 1% -3.01% (p=0.000 n=10+10) Decrypt/AES-192-16 4.77ns ± 2% 4.63ns ± 2% -3.05% (p=0.000 n=10+10) Decrypt/AES-256-16 5.10ns ± 2% 4.93ns ± 3% -3.27% (p=0.000 n=10+10) Expand/AES-128-16 52.5ns ± 4% 55.9ns ± 3% +6.58% (p=0.000 n=10+10) Expand/AES-192-16 45.6ns ± 5% 48.1ns ± 2% +5.56% (p=0.000 n=10+10) Expand/AES-256-16 69.2ns ± 5% 71.7ns ± 3% +3.61% (p=0.006 n=9+10) CreateCipher/AES-128-16 130ns ± 3% 84ns ± 3% -35.13% (p=0.000 n=10+10) CreateCipher/AES-192-16 128ns ± 2% 78ns ± 3% -38.82% (p=0.000 n=10+10) CreateCipher/AES-256-16 156ns ± 3% 105ns ± 4% -32.33% (p=0.000 n=10+10) name old speed new speed delta Encrypt/AES-128-16 3.40GB/s ± 2% 3.59GB/s ± 2% +5.35% (p=0.000 n=10+10) Encrypt/AES-192-16 3.27GB/s ± 3% 3.40GB/s ± 2% +4.13% (p=0.000 n=10+10) Encrypt/AES-256-16 3.09GB/s ± 2% 3.26GB/s ± 1% +5.42% (p=0.000 n=10+10) Decrypt/AES-128-16 3.55GB/s ± 2% 3.66GB/s ± 1% +3.09% (p=0.000 n=10+10) Decrypt/AES-192-16 3.35GB/s ± 2% 3.46GB/s ± 2% +3.14% (p=0.000 n=10+10) Decrypt/AES-256-16 3.14GB/s ± 2% 3.24GB/s ± 3% +3.39% (p=0.000 n=10+10) name old alloc/op new alloc/op delta Encrypt/AES-128-16 0.00B 0.00B ~ (all equal) Encrypt/AES-192-16 0.00B 0.00B ~ (all equal) Encrypt/AES-256-16 0.00B 0.00B ~ (all equal) Decrypt/AES-128-16 0.00B 0.00B ~ (all equal) Decrypt/AES-192-16 0.00B 0.00B ~ (all equal) Decrypt/AES-256-16 0.00B 0.00B ~ (all equal) Expand/AES-128-16 0.00B 0.00B ~ (all equal) Expand/AES-192-16 0.00B 0.00B ~ (all equal) Expand/AES-256-16 0.00B 0.00B ~ (all equal) CreateCipher/AES-128-16 448B ± 0% 512B ± 0% +14.29% (p=0.000 n=10+10) CreateCipher/AES-192-16 512B ± 0% 512B ± 0% ~ (all equal) CreateCipher/AES-256-16 576B ± 0% 512B ± 0% -11.11% (p=0.000 n=10+10) name old allocs/op new allocs/op delta Encrypt/AES-128-16 0.00 0.00 ~ (all equal) Encrypt/AES-192-16 0.00 0.00 ~ (all equal) Encrypt/AES-256-16 0.00 0.00 ~ (all equal) Decrypt/AES-128-16 0.00 0.00 ~ (all equal) Decrypt/AES-192-16 0.00 0.00 ~ (all equal) Decrypt/AES-256-16 0.00 0.00 ~ (all equal) Expand/AES-128-16 0.00 0.00 ~ (all equal) Expand/AES-192-16 0.00 0.00 ~ (all equal) Expand/AES-256-16 0.00 0.00 ~ (all equal) CreateCipher/AES-128-16 4.00 ± 0% 1.00 ± 0% -75.00% (p=0.000 n=10+10) CreateCipher/AES-192-16 4.00 ± 0% 1.00 ± 0% -75.00% (p=0.000 n=10+10) CreateCipher/AES-256-16 4.00 ± 0% 1.00 ± 0% -75.00% (p=0.000 n=10+10) Change-Id: I0ea0b21cf84b11b6a5fc7c6ace144390eb55438b Reviewed-on: https://go-review.googlesource.com/c/go/+/561080 Reviewed-by: Damien Neil <dneil@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by: Jorropo <jorropo.pgm@gmail.com> Reviewed-by: Roland Shoemaker <roland@golang.org>
diff --git a/src/crypto/aes/aes_gcm.go b/src/crypto/aes/aes_gcm.go
@@ -45,7 +45,7 @@ var _ gcmAble = (*aesCipherGCM)(nil)
 // NewGCM returns the AES cipher wrapped in Galois Counter Mode. This is only
 // called by [crypto/cipher.NewGCM] via the gcmAble interface.
 func (c *aesCipherGCM) NewGCM(nonceSize, tagSize int) (cipher.AEAD, error) {
-	g := &gcmAsm{ks: c.enc, nonceSize: nonceSize, tagSize: tagSize}
+	g := &gcmAsm{ks: c.enc[:c.l], nonceSize: nonceSize, tagSize: tagSize}
 	gcmAesInit(&g.productTable, g.ks)
 	return g, nil
 }
diff --git a/src/crypto/aes/aes_test.go b/src/crypto/aes/aes_test.go
@@ -345,7 +345,12 @@ func mustPanic(t *testing.T, msg string, f func()) {
 }
 
 func BenchmarkEncrypt(b *testing.B) {
-	tt := encryptTests[0]
+	b.Run("AES-128", func(b *testing.B) { benchmarkEncrypt(b, encryptTests[1]) })
+	b.Run("AES-192", func(b *testing.B) { benchmarkEncrypt(b, encryptTests[2]) })
+	b.Run("AES-256", func(b *testing.B) { benchmarkEncrypt(b, encryptTests[3]) })
+}
+
+func benchmarkEncrypt(b *testing.B, tt CryptTest) {
 	c, err := NewCipher(tt.key)
 	if err != nil {
 		b.Fatal("NewCipher:", err)
@@ -359,7 +364,12 @@ func BenchmarkEncrypt(b *testing.B) {
 }
 
 func BenchmarkDecrypt(b *testing.B) {
-	tt := encryptTests[0]
+	b.Run("AES-128", func(b *testing.B) { benchmarkDecrypt(b, encryptTests[1]) })
+	b.Run("AES-192", func(b *testing.B) { benchmarkDecrypt(b, encryptTests[2]) })
+	b.Run("AES-256", func(b *testing.B) { benchmarkDecrypt(b, encryptTests[3]) })
+}
+
+func benchmarkDecrypt(b *testing.B, tt CryptTest) {
 	c, err := NewCipher(tt.key)
 	if err != nil {
 		b.Fatal("NewCipher:", err)
@@ -373,11 +383,30 @@ func BenchmarkDecrypt(b *testing.B) {
 }
 
 func BenchmarkExpand(b *testing.B) {
-	tt := encryptTests[0]
-	n := len(tt.key) + 28
-	c := &aesCipher{make([]uint32, n), make([]uint32, n)}
+	b.Run("AES-128", func(b *testing.B) { benchmarkExpand(b, encryptTests[1]) })
+	b.Run("AES-192", func(b *testing.B) { benchmarkExpand(b, encryptTests[2]) })
+	b.Run("AES-256", func(b *testing.B) { benchmarkExpand(b, encryptTests[3]) })
+}
+
+func benchmarkExpand(b *testing.B, tt CryptTest) {
+	c := &aesCipher{l: uint8(len(tt.key) + 28)}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		expandKey(tt.key, c.enc, c.dec)
+		expandKey(tt.key, c.enc[:c.l], c.dec[:c.l])
+	}
+}
+
+func BenchmarkCreateCipher(b *testing.B) {
+	b.Run("AES-128", func(b *testing.B) { benchmarkCreateCipher(b, encryptTests[1]) })
+	b.Run("AES-192", func(b *testing.B) { benchmarkCreateCipher(b, encryptTests[2]) })
+	b.Run("AES-256", func(b *testing.B) { benchmarkCreateCipher(b, encryptTests[3]) })
+}
+
+func benchmarkCreateCipher(b *testing.B, tt CryptTest) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		if _, err := NewCipher(tt.key); err != nil {
+			b.Fatal(err)
+		}
 	}
 }
diff --git a/src/crypto/aes/cipher.go b/src/crypto/aes/cipher.go
@@ -16,8 +16,9 @@ const BlockSize = 16
 
 // A cipher is an instance of AES encryption using a particular key.
 type aesCipher struct {
-	enc []uint32
-	dec []uint32
+	l   uint8 // only this length of the enc and dec array is actually used
+	enc [28 + 32]uint32
+	dec [28 + 32]uint32
 }
 
 type KeySizeError int
@@ -47,9 +48,8 @@ func NewCipher(key []byte) (cipher.Block, error) {
 // newCipherGeneric creates and returns a new cipher.Block
 // implemented in pure Go.
 func newCipherGeneric(key []byte) (cipher.Block, error) {
-	n := len(key) + 28
-	c := aesCipher{make([]uint32, n), make([]uint32, n)}
-	expandKeyGo(key, c.enc, c.dec)
+	c := aesCipher{l: uint8(len(key) + 28)}
+	expandKeyGo(key, c.enc[:c.l], c.dec[:c.l])
 	return &c, nil
 }
 
@@ -65,7 +65,7 @@ func (c *aesCipher) Encrypt(dst, src []byte) {
 	if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {
 		panic("crypto/aes: invalid buffer overlap")
 	}
-	encryptBlockGo(c.enc, dst, src)
+	encryptBlockGo(c.enc[:c.l], dst, src)
 }
 
 func (c *aesCipher) Decrypt(dst, src []byte) {
@@ -78,5 +78,5 @@ func (c *aesCipher) Decrypt(dst, src []byte) {
 	if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {
 		panic("crypto/aes: invalid buffer overlap")
 	}
-	decryptBlockGo(c.dec, dst, src)
+	decryptBlockGo(c.dec[:c.l], dst, src)
 }
diff --git a/src/crypto/aes/cipher_asm.go b/src/crypto/aes/cipher_asm.go
@@ -44,8 +44,9 @@ func newCipher(key []byte) (cipher.Block, error) {
 	if !supportsAES {
 		return newCipherGeneric(key)
 	}
-	n := len(key) + 28
-	c := aesCipherAsm{aesCipher{make([]uint32, n), make([]uint32, n)}}
+	// Note that under certain circumstances, we only return the inner aesCipherAsm.
+	// This avoids an unnecessary allocation of the aesCipher struct.
+	c := aesCipherGCM{aesCipherAsm{aesCipher{l: uint8(len(key) + 28)}}}
 	var rounds int
 	switch len(key) {
 	case 128 / 8:
@@ -60,9 +61,9 @@ func newCipher(key []byte) (cipher.Block, error) {
 
 	expandKeyAsm(rounds, &key[0], &c.enc[0], &c.dec[0])
 	if supportsAES && supportsGFMUL {
-		return &aesCipherGCM{c}, nil
+		return &c, nil
 	}
-	return &c, nil
+	return &c.aesCipherAsm, nil
 }
 
 func (c *aesCipherAsm) BlockSize() int { return BlockSize }
@@ -78,7 +79,7 @@ func (c *aesCipherAsm) Encrypt(dst, src []byte) {
 	if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {
 		panic("crypto/aes: invalid buffer overlap")
 	}
-	encryptBlockAsm(len(c.enc)/4-1, &c.enc[0], &dst[0], &src[0])
+	encryptBlockAsm(int(c.l)/4-1, &c.enc[0], &dst[0], &src[0])
 }
 
 func (c *aesCipherAsm) Decrypt(dst, src []byte) {
@@ -92,7 +93,7 @@ func (c *aesCipherAsm) Decrypt(dst, src []byte) {
 	if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {
 		panic("crypto/aes: invalid buffer overlap")
 	}
-	decryptBlockAsm(len(c.dec)/4-1, &c.dec[0], &dst[0], &src[0])
+	decryptBlockAsm(int(c.l)/4-1, &c.dec[0], &dst[0], &src[0])
 }
 
 // expandKey is used by BenchmarkExpand to ensure that the asm implementation
diff --git a/src/crypto/aes/gcm_ppc64x.go b/src/crypto/aes/gcm_ppc64x.go
@@ -57,7 +57,7 @@ func counterCryptASM(nr int, out, in []byte, counter *[gcmBlockSize]byte, key *u
 // called by [crypto/cipher.NewGCM] via the gcmAble interface.
 func (c *aesCipherAsm) NewGCM(nonceSize, tagSize int) (cipher.AEAD, error) {
 	var h1, h2 uint64
-	g := &gcmAsm{cipher: c, ks: c.enc, nonceSize: nonceSize, tagSize: tagSize}
+	g := &gcmAsm{cipher: c, ks: c.enc[:c.l], nonceSize: nonceSize, tagSize: tagSize}
 
 	hle := make([]byte, gcmBlockSize)
 

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ var _ gcmAble = (*aesCipherGCM)(nil)`
`45`	`45`	`// NewGCM returns the AES cipher wrapped in Galois Counter Mode. This is only`
`46`	`46`	`// called by [crypto/cipher.NewGCM] via the gcmAble interface.`
`47`	`47`	`func (c *aesCipherGCM) NewGCM(nonceSize, tagSize int) (cipher.AEAD, error) {`
`48`		`- g := &gcmAsm{ks: c.enc, nonceSize: nonceSize, tagSize: tagSize}`
	`48`	`+ g := &gcmAsm{ks: c.enc[:c.l], nonceSize: nonceSize, tagSize: tagSize}`
`49`	`49`	`gcmAesInit(&g.productTable, g.ks)`
`50`	`50`	`return g, nil`
`51`	`51`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,8 +16,9 @@ const BlockSize = 16`
`16`	`16`
`17`	`17`	`// A cipher is an instance of AES encryption using a particular key.`
`18`	`18`	`type aesCipher struct {`
`19`		`- enc []uint32`
`20`		`- dec []uint32`
	`19`	`+ l uint8 // only this length of the enc and dec array is actually used`
	`20`	`+ enc [28 + 32]uint32`
	`21`	`+ dec [28 + 32]uint32`
`21`	`22`	`}`
`22`	`23`
`23`	`24`	`type KeySizeError int`
`@@ -47,9 +48,8 @@ func NewCipher(key []byte) (cipher.Block, error) {`
`47`	`48`	`// newCipherGeneric creates and returns a new cipher.Block`
`48`	`49`	`// implemented in pure Go.`
`49`	`50`	`func newCipherGeneric(key []byte) (cipher.Block, error) {`
`50`		`- n := len(key) + 28`
`51`		`- c := aesCipher{make([]uint32, n), make([]uint32, n)}`
`52`		`- expandKeyGo(key, c.enc, c.dec)`
	`51`	`+ c := aesCipher{l: uint8(len(key) + 28)}`
	`52`	`+ expandKeyGo(key, c.enc[:c.l], c.dec[:c.l])`
`53`	`53`	`return &c, nil`
`54`	`54`	`}`
`55`	`55`
`@@ -65,7 +65,7 @@ func (c *aesCipher) Encrypt(dst, src []byte) {`
`65`	`65`	`if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {`
`66`	`66`	`panic("crypto/aes: invalid buffer overlap")`
`67`	`67`	`}`
`68`		`- encryptBlockGo(c.enc, dst, src)`
	`68`	`+ encryptBlockGo(c.enc[:c.l], dst, src)`
`69`	`69`	`}`
`70`	`70`
`71`	`71`	`func (c *aesCipher) Decrypt(dst, src []byte) {`
`@@ -78,5 +78,5 @@ func (c *aesCipher) Decrypt(dst, src []byte) {`
`78`	`78`	`if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {`
`79`	`79`	`panic("crypto/aes: invalid buffer overlap")`
`80`	`80`	`}`
`81`		`- decryptBlockGo(c.dec, dst, src)`
	`81`	`+ decryptBlockGo(c.dec[:c.l], dst, src)`
`82`	`82`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,8 +44,9 @@ func newCipher(key []byte) (cipher.Block, error) {`
`44`	`44`	`if !supportsAES {`
`45`	`45`	`return newCipherGeneric(key)`
`46`	`46`	`}`
`47`		`- n := len(key) + 28`
`48`		`- c := aesCipherAsm{aesCipher{make([]uint32, n), make([]uint32, n)}}`
	`47`	`+ // Note that under certain circumstances, we only return the inner aesCipherAsm.`
	`48`	`+ // This avoids an unnecessary allocation of the aesCipher struct.`
	`49`	`+ c := aesCipherGCM{aesCipherAsm{aesCipher{l: uint8(len(key) + 28)}}}`
`49`	`50`	`var rounds int`
`50`	`51`	`switch len(key) {`
`51`	`52`	`case 128 / 8:`
`@@ -60,9 +61,9 @@ func newCipher(key []byte) (cipher.Block, error) {`
`60`	`61`
`61`	`62`	`expandKeyAsm(rounds, &key[0], &c.enc[0], &c.dec[0])`
`62`	`63`	`if supportsAES && supportsGFMUL {`
`63`		`- return &aesCipherGCM{c}, nil`
	`64`	`+ return &c, nil`
`64`	`65`	`}`
`65`		`- return &c, nil`
	`66`	`+ return &c.aesCipherAsm, nil`
`66`	`67`	`}`
`67`	`68`
`68`	`69`	`func (c *aesCipherAsm) BlockSize() int { return BlockSize }`
`@@ -78,7 +79,7 @@ func (c *aesCipherAsm) Encrypt(dst, src []byte) {`
`78`	`79`	`if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {`
`79`	`80`	`panic("crypto/aes: invalid buffer overlap")`
`80`	`81`	`}`
`81`		`- encryptBlockAsm(len(c.enc)/4-1, &c.enc[0], &dst[0], &src[0])`
	`82`	`+ encryptBlockAsm(int(c.l)/4-1, &c.enc[0], &dst[0], &src[0])`
`82`	`83`	`}`
`83`	`84`
`84`	`85`	`func (c *aesCipherAsm) Decrypt(dst, src []byte) {`
`@@ -92,7 +93,7 @@ func (c *aesCipherAsm) Decrypt(dst, src []byte) {`
`92`	`93`	`if alias.InexactOverlap(dst[:BlockSize], src[:BlockSize]) {`
`93`	`94`	`panic("crypto/aes: invalid buffer overlap")`
`94`	`95`	`}`
`95`		`- decryptBlockAsm(len(c.dec)/4-1, &c.dec[0], &dst[0], &src[0])`
	`96`	`+ decryptBlockAsm(int(c.l)/4-1, &c.dec[0], &dst[0], &src[0])`
`96`	`97`	`}`
`97`	`98`
`98`	`99`	`// expandKey is used by BenchmarkExpand to ensure that the asm implementation`