internal/chacha20: add s390x SIMD implementation

Based on the SIMD algorithm described in:
ChaCha, a variant of Salsa20 by Daniel J. Bernstein
https://cr.yp.to/chacha/chacha-20080128.pdf

Requires the vector facility (vx).

name            old speed     new speed      delta
ChaCha20/32     178MB/s ± 0%   174MB/s ± 0%    -2.10%  (p=0.000 n=9+10)
ChaCha20/63     341MB/s ± 0%   337MB/s ± 0%    -1.16%  (p=0.000 n=10+10)
ChaCha20/64     367MB/s ± 0%   335MB/s ± 0%    -8.73%  (p=0.000 n=10+10)
ChaCha20/256    404MB/s ± 0%  1448MB/s ± 0%  +258.61%  (p=0.000 n=9+10)
ChaCha20/1024   410MB/s ± 0%  1568MB/s ± 0%  +282.73%  (p=0.000 n=9+10)
ChaCha20/1350   393MB/s ± 0%  1389MB/s ± 0%  +253.58%  (p=0.000 n=10+10)
ChaCha20/65536  414MB/s ± 0%  1634MB/s ± 0%  +294.79%  (p=0.000 n=10+10)

Change-Id: I9a600fb5ae8ee3f3b81ae6b01cff139c1272d684
Reviewed-on: https://go-review.googlesource.com/35842
Run-TryBot: Michael Munday <[email protected]>
Reviewed-by: Brad Fitzpatrick <[email protected]>

Author: mundaym

internal/chacha20/asm_s390x.s

@@ -0,0 +1,283 @@
+// Copyright 2018 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build s390x,!gccgo,!appengine
+
+#include "go_asm.h"
+#include "textflag.h"
+
+// This is an implementation of the ChaCha20 encryption algorithm as
+// specified in RFC 7539. It uses vector instructions to compute
+// 4 keystream blocks in parallel (256 bytes) which are then XORed
+// with the bytes in the input slice.
+
+GLOBL ·constants<>(SB), RODATA|NOPTR, $32
+// BSWAP: swap bytes in each 4-byte element
+DATA ·constants<>+0x00(SB)/4, $0x03020100
+DATA ·constants<>+0x04(SB)/4, $0x07060504
+DATA ·constants<>+0x08(SB)/4, $0x0b0a0908
+DATA ·constants<>+0x0c(SB)/4, $0x0f0e0d0c
+// J0: [j0, j1, j2, j3]
+DATA ·constants<>+0x10(SB)/4, $0x61707865
+DATA ·constants<>+0x14(SB)/4, $0x3320646e
+DATA ·constants<>+0x18(SB)/4, $0x79622d32
+DATA ·constants<>+0x1c(SB)/4, $0x6b206574
+
+// EXRL targets:
+TEXT ·mvcSrcToBuf(SB), NOFRAME|NOSPLIT, $0
+	MVC $1, (R1), (R8)
+	RET
+
+TEXT ·mvcBufToDst(SB), NOFRAME|NOSPLIT, $0
+	MVC $1, (R8), (R9)
+	RET
+
+#define BSWAP V5
+#define J0    V6
+#define KEY0  V7
+#define KEY1  V8
+#define NONCE V9
+#define CTR   V10
+#define M0    V11
+#define M1    V12
+#define M2    V13
+#define M3    V14
+#define INC   V15
+#define X0    V16
+#define X1    V17
+#define X2    V18
+#define X3    V19
+#define X4    V20
+#define X5    V21
+#define X6    V22
+#define X7    V23
+#define X8    V24
+#define X9    V25
+#define X10   V26
+#define X11   V27
+#define X12   V28
+#define X13   V29
+#define X14   V30
+#define X15   V31
+
+#define NUM_ROUNDS 20
+
+#define ROUND4(a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3) \
+	VAF    a1, a0, a0  \
+	VAF    b1, b0, b0  \
+	VAF    c1, c0, c0  \
+	VAF    d1, d0, d0  \
+	VX     a0, a2, a2  \
+	VX     b0, b2, b2  \
+	VX     c0, c2, c2  \
+	VX     d0, d2, d2  \
+	VERLLF $16, a2, a2 \
+	VERLLF $16, b2, b2 \
+	VERLLF $16, c2, c2 \
+	VERLLF $16, d2, d2 \
+	VAF    a2, a3, a3  \
+	VAF    b2, b3, b3  \
+	VAF    c2, c3, c3  \
+	VAF    d2, d3, d3  \
+	VX     a3, a1, a1  \
+	VX     b3, b1, b1  \
+	VX     c3, c1, c1  \
+	VX     d3, d1, d1  \
+	VERLLF $12, a1, a1 \
+	VERLLF $12, b1, b1 \
+	VERLLF $12, c1, c1 \
+	VERLLF $12, d1, d1 \
+	VAF    a1, a0, a0  \
+	VAF    b1, b0, b0  \
+	VAF    c1, c0, c0  \
+	VAF    d1, d0, d0  \
+	VX     a0, a2, a2  \
+	VX     b0, b2, b2  \
+	VX     c0, c2, c2  \
+	VX     d0, d2, d2  \
+	VERLLF $8, a2, a2  \
+	VERLLF $8, b2, b2  \
+	VERLLF $8, c2, c2  \
+	VERLLF $8, d2, d2  \
+	VAF    a2, a3, a3  \
+	VAF    b2, b3, b3  \
+	VAF    c2, c3, c3  \
+	VAF    d2, d3, d3  \
+	VX     a3, a1, a1  \
+	VX     b3, b1, b1  \
+	VX     c3, c1, c1  \
+	VX     d3, d1, d1  \
+	VERLLF $7, a1, a1  \
+	VERLLF $7, b1, b1  \
+	VERLLF $7, c1, c1  \
+	VERLLF $7, d1, d1
+
+#define PERMUTE(mask, v0, v1, v2, v3) \
+	VPERM v0, v0, mask, v0 \
+	VPERM v1, v1, mask, v1 \
+	VPERM v2, v2, mask, v2 \
+	VPERM v3, v3, mask, v3
+
+#define ADDV(x, v0, v1, v2, v3) \
+	VAF x, v0, v0 \
+	VAF x, v1, v1 \
+	VAF x, v2, v2 \
+	VAF x, v3, v3
+
+#define XORV(off, dst, src, v0, v1, v2, v3) \
+	VLM  off(src), M0, M3          \
+	PERMUTE(BSWAP, v0, v1, v2, v3) \
+	VX   v0, M0, M0                \
+	VX   v1, M1, M1                \
+	VX   v2, M2, M2                \
+	VX   v3, M3, M3                \
+	VSTM M0, M3, off(dst)
+
+#define SHUFFLE(a, b, c, d, t, u, v, w) \
+	VMRHF a, c, t \ // t = {a[0], c[0], a[1], c[1]}
+	VMRHF b, d, u \ // u = {b[0], d[0], b[1], d[1]}
+	VMRLF a, c, v \ // v = {a[2], c[2], a[3], c[3]}
+	VMRLF b, d, w \ // w = {b[2], d[2], b[3], d[3]}
+	VMRHF t, u, a \ // a = {a[0], b[0], c[0], d[0]}
+	VMRLF t, u, b \ // b = {a[1], b[1], c[1], d[1]}
+	VMRHF v, w, c \ // c = {a[2], b[2], c[2], d[2]}
+	VMRLF v, w, d // d = {a[3], b[3], c[3], d[3]}
+
+// func xorKeyStreamVX(dst, src []byte, key *[8]uint32, nonce *[3]uint32, counter *uint32, buf *[256]byte, len *int)
+TEXT ·xorKeyStreamVX(SB), NOSPLIT, $0
+	MOVD $·constants<>(SB), R1
+	MOVD dst+0(FP), R2         // R2=&dst[0]
+	LMG  src+24(FP), R3, R4    // R3=&src[0] R4=len(src)
+	MOVD key+48(FP), R5        // R5=key
+	MOVD nonce+56(FP), R6      // R6=nonce
+	MOVD counter+64(FP), R7    // R7=counter
+	MOVD buf+72(FP), R8        // R8=buf
+	MOVD len+80(FP), R9        // R9=len
+
+	// load BSWAP and J0
+	VLM (R1), BSWAP, J0
+
+	// set up tail buffer
+	ADD     $-1, R4, R12
+	MOVBZ   R12, R12
+	CMPUBEQ R12, $255, aligned
+	MOVD    R4, R1
+	AND     $~255, R1
+	MOVD    $(R3)(R1*1), R1
+	EXRL    $·mvcSrcToBuf(SB), R12
+	MOVD    $255, R0
+	SUB     R12, R0
+	MOVD    R0, (R9)               // update len
+
+aligned:
+	// setup
+	MOVD  $95, R0
+	VLM   (R5), KEY0, KEY1
+	VLL   R0, (R6), NONCE
+	VZERO M0
+	VLEIB $7, $32, M0
+	VSRLB M0, NONCE, NONCE
+
+	// initialize counter values
+	VLREPF (R7), CTR
+	VZERO  INC
+	VLEIF  $1, $1, INC
+	VLEIF  $2, $2, INC
+	VLEIF  $3, $3, INC
+	VAF    INC, CTR, CTR
+	VREPIF $4, INC
+
+chacha:
+	VREPF $0, J0, X0
+	VREPF $1, J0, X1
+	VREPF $2, J0, X2
+	VREPF $3, J0, X3
+	VREPF $0, KEY0, X4
+	VREPF $1, KEY0, X5
+	VREPF $2, KEY0, X6
+	VREPF $3, KEY0, X7
+	VREPF $0, KEY1, X8
+	VREPF $1, KEY1, X9
+	VREPF $2, KEY1, X10
+	VREPF $3, KEY1, X11
+	VLR   CTR, X12
+	VREPF $1, NONCE, X13
+	VREPF $2, NONCE, X14
+	VREPF $3, NONCE, X15
+
+	MOVD $(NUM_ROUNDS/2), R1
+
+loop:
+	ROUND4(X0, X4, X12,  X8, X1, X5, X13,  X9, X2, X6, X14, X10, X3, X7, X15, X11)
+	ROUND4(X0, X5, X15, X10, X1, X6, X12, X11, X2, X7, X13, X8,  X3, X4, X14, X9)
+
+	ADD $-1, R1
+	BNE loop
+
+	// decrement length
+	ADD $-256, R4
+	BLT tail
+
+continue:
+	// rearrange vectors
+	SHUFFLE(X0, X1, X2, X3, M0, M1, M2, M3)
+	ADDV(J0, X0, X1, X2, X3)
+	SHUFFLE(X4, X5, X6, X7, M0, M1, M2, M3)
+	ADDV(KEY0, X4, X5, X6, X7)
+	SHUFFLE(X8, X9, X10, X11, M0, M1, M2, M3)
+	ADDV(KEY1, X8, X9, X10, X11)
+	VAF CTR, X12, X12
+	SHUFFLE(X12, X13, X14, X15, M0, M1, M2, M3)
+	ADDV(NONCE, X12, X13, X14, X15)
+
+	// increment counters
+	VAF INC, CTR, CTR
+
+	// xor keystream with plaintext
+	XORV(0*64, R2, R3, X0, X4,  X8, X12)
+	XORV(1*64, R2, R3, X1, X5,  X9, X13)
+	XORV(2*64, R2, R3, X2, X6, X10, X14)
+	XORV(3*64, R2, R3, X3, X7, X11, X15)
+
+	// increment pointers
+	MOVD $256(R2), R2
+	MOVD $256(R3), R3
+
+	CMPBNE  R4, $0, chacha
+	CMPUBEQ R12, $255, return
+	EXRL    $·mvcBufToDst(SB), R12 // len was updated during setup
+
+return:
+	VSTEF $0, CTR, (R7)
+	RET
+
+tail:
+	MOVD R2, R9
+	MOVD R8, R2
+	MOVD R8, R3
+	MOVD $0, R4
+	JMP  continue
+
+// func hasVectorFacility() bool
+TEXT ·hasVectorFacility(SB), NOSPLIT, $24-1
+	MOVD  $x-24(SP), R1
+	XC    $24, 0(R1), 0(R1) // clear the storage
+	MOVD  $2, R0            // R0 is the number of double words stored -1
+	WORD  $0xB2B01000       // STFLE 0(R1)
+	XOR   R0, R0            // reset the value of R0
+	MOVBZ z-8(SP), R1
+	AND   $0x40, R1
+	BEQ   novector
+
+vectorinstalled:
+	// check if the vector instruction has been enabled
+	VLEIB  $0, $0xF, V16
+	VLGVB  $0, V16, R1
+	CMPBNE R1, $0xF, novector
+	MOVB   $1, ret+0(FP)      // have vx
+	RET
+
+novector:
+	MOVB $0, ret+0(FP) // no vx
+	RET

internal/chacha20/chacha_generic.go

@@ -18,10 +18,10 @@ var _ cipher.Stream = (*Cipher)(nil)
 // and nonce. A *Cipher implements the cipher.Stream interface.
 type Cipher struct {
 	key     [8]uint32
+	counter uint32 // incremented after each block
 	nonce   [3]uint32
-	counter uint32   // incremented after each block
-	buf     [64]byte // buffer for unused keystream bytes
-	len     int      // number of unused keystream bytes at end of buf
+	buf     [bufSize]byte // buffer for unused keystream bytes
+	len     int           // number of unused keystream bytes at end of buf
 }
 
 // New creates a new ChaCha20 stream cipher with the given key and nonce.
@@ -63,6 +63,10 @@ func (s *Cipher) XORKeyStream(dst, src []byte) {
 	if len(src) == 0 {
 		return
 	}
+	if haveAsm {
+		s.xorKeyStreamAsm(dst, src)
+		return
+	}
 
 	// set up a 64-byte buffer to pad out the final block if needed
 	// (hoisted out of the main loop to avoid spills)

internal/chacha20/chacha_noasm.go

@@ -0,0 +1,16 @@
+// Copyright 2018 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build !s390x gccgo appengine
+
+package chacha20
+
+const (
+	bufSize = 64
+	haveAsm = false
+)
+
+func (*Cipher) xorKeyStreamAsm(dst, src []byte) {
+	panic("not implemented")
+}

internal/chacha20/chacha_s390x.go

@@ -0,0 +1,30 @@
+// Copyright 2018 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build s390x,!gccgo,!appengine
+
+package chacha20
+
+var haveAsm = hasVectorFacility()
+
+const bufSize = 256
+
+// hasVectorFacility reports whether the machine supports the vector
+// facility (vx).
+// Implementation in asm_s390x.s.
+func hasVectorFacility() bool
+
+// xorKeyStreamVX is an assembly implementation of XORKeyStream. It must only
+// be called when the vector facility is available.
+// Implementation in asm_s390x.s.
+//go:noescape
+func xorKeyStreamVX(dst, src []byte, key *[8]uint32, nonce *[3]uint32, counter *uint32, buf *[256]byte, len *int)
+
+func (c *Cipher) xorKeyStreamAsm(dst, src []byte) {
+	xorKeyStreamVX(dst, src, &c.key, &c.nonce, &c.counter, &c.buf, &c.len)
+}
+
+// EXRL targets, DO NOT CALL!
+func mvcSrcToBuf()
+func mvcBufToDst()