diff --git a/sha3/sha3_s390x.go b/sha3/sha3_s390x.go
index f1fb79c..b6cbc5c 100644
--- a/sha3/sha3_s390x.go
+++ b/sha3/sha3_s390x.go
@@ -12,6 +12,8 @@ package sha3
 
 import (
 	"hash"
+
+	"golang.org/x/sys/cpu"
 )
 
 // codes represent 7-bit KIMD/KLMD function codes as defined in
@@ -29,13 +31,6 @@ const (
 	nopad          = 0x100
 )
 
-// hasMSA6 reports whether the machine supports the SHA-3 and SHAKE function
-// codes, as defined in message-security-assist extension 6.
-func hasMSA6() bool
-
-// hasAsm caches the result of hasMSA6 (which might be expensive to call).
-var hasAsm = hasMSA6()
-
 // kimd is a wrapper for the 'compute intermediate message digest' instruction.
 // src must be a multiple of the rate for the given function code.
 //go:noescape
@@ -237,7 +232,7 @@ func (s *asmState) Clone() ShakeHash {
 // new224Asm returns an assembly implementation of SHA3-224 if available,
 // otherwise it returns nil.
 func new224Asm() hash.Hash {
-	if hasAsm {
+	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_224)
 	}
 	return nil
@@ -246,7 +241,7 @@ func new224Asm() hash.Hash {
 // new256Asm returns an assembly implementation of SHA3-256 if available,
 // otherwise it returns nil.
 func new256Asm() hash.Hash {
-	if hasAsm {
+	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_256)
 	}
 	return nil
@@ -255,7 +250,7 @@ func new256Asm() hash.Hash {
 // new384Asm returns an assembly implementation of SHA3-384 if available,
 // otherwise it returns nil.
 func new384Asm() hash.Hash {
-	if hasAsm {
+	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_384)
 	}
 	return nil
@@ -264,7 +259,7 @@ func new384Asm() hash.Hash {
 // new512Asm returns an assembly implementation of SHA3-512 if available,
 // otherwise it returns nil.
 func new512Asm() hash.Hash {
-	if hasAsm {
+	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_512)
 	}
 	return nil
@@ -273,7 +268,7 @@ func new512Asm() hash.Hash {
 // newShake128Asm returns an assembly implementation of SHAKE-128 if available,
 // otherwise it returns nil.
 func newShake128Asm() ShakeHash {
-	if hasAsm {
+	if cpu.S390X.HasSHA3 {
 		return newAsmState(shake_128)
 	}
 	return nil
@@ -282,7 +277,7 @@ func newShake128Asm() ShakeHash {
 // newShake256Asm returns an assembly implementation of SHAKE-256 if available,
 // otherwise it returns nil.
 func newShake256Asm() ShakeHash {
-	if hasAsm {
+	if cpu.S390X.HasSHA3 {
 		return newAsmState(shake_256)
 	}
 	return nil
diff --git a/sha3/sha3_s390x.s b/sha3/sha3_s390x.s
index 20978fc..b2ef69f 100644
--- a/sha3/sha3_s390x.s
+++ b/sha3/sha3_s390x.s
@@ -6,26 +6,10 @@
 
 #include "textflag.h"
 
-TEXT ·hasMSA6(SB), NOSPLIT, $16-1
-	MOVD $0, R0          // KIMD-Query function code
-	MOVD $tmp-16(SP), R1 // parameter block
-	XC   $16, (R1), (R1) // clear the parameter block
-	WORD $0xB93E0002     // KIMD --, --
-	WORD $0x91FC1004     // TM 4(R1), 0xFC (test bits [32-37])
-	BVS  yes
-
-no:
-	MOVB $0, ret+0(FP)
-	RET
-
-yes:
-	MOVB $1, ret+0(FP)
-	RET
-
-// func kimd(function code, params *[200]byte, src []byte)
+// func kimd(function code, chain *[200]byte, src []byte)
 TEXT ·kimd(SB), NOFRAME|NOSPLIT, $0-40
 	MOVD function+0(FP), R0
-	MOVD params+8(FP), R1
+	MOVD chain+8(FP), R1
 	LMG  src+16(FP), R2, R3 // R2=base, R3=len
 
 continue:
@@ -34,11 +18,11 @@ continue:
 	MOVD $0, R0      // reset R0 for pre-go1.8 compilers
 	RET
 
-// func klmd(function code, params *[200]byte, dst, src []byte)
+// func klmd(function code, chain *[200]byte, dst, src []byte)
 TEXT ·klmd(SB), NOFRAME|NOSPLIT, $0-64
 	// TODO: SHAKE support
 	MOVD function+0(FP), R0
-	MOVD params+8(FP), R1
+	MOVD chain+8(FP), R1
 	LMG  dst+16(FP), R2, R3 // R2=base, R3=len
 	LMG  src+40(FP), R4, R5 // R4=base, R5=len