[dev.simd] simd, internal/cpu: support more AVX CPU Feature checks

This CL adds more checks, it also changes HasAVX512GFNI to be exactly
checking GFNI instead of being a virtual feature.

This CL copies its logic from x/sys/arch.

Change-Id: I4612b0409b8a3518928300562ae08bcf123d53a7
Reviewed-on: https://go-review.googlesource.com/c/go/+/688276
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>

3 files changed, 108 insertions, 45 deletions

diff --git a/src/internal/cpu/cpu.go b/src/internal/cpu/cpu.go
index 1eeb580711..53633c7ca8 100644
--- a/src/internal/cpu/cpu.go
+++ b/src/internal/cpu/cpu.go
@@ -26,34 +26,39 @@ var CacheLineSize uintptr = CacheLinePadSize
 // in addition to the cpuid feature bit being set.
 // The struct is padded to avoid false sharing.
 var X86 struct {
-	_             CacheLinePad
-	HasAES        bool
-	HasADX        bool
-	HasAVX        bool
-	HasAVX2       bool
-	HasAVX512GFNI bool // Virtual feature: F+CD+BW+DQ+VL+GFNI
-	HasAVX512     bool // Virtual feature: F+CD+BW+DQ+VL
-	HasAVX512F    bool
-	HasAVX512CD   bool
-	HasAVX512BW   bool
-	HasAVX512DQ   bool
-	HasAVX512VL   bool
-	HasBMI1       bool
-	HasBMI2       bool
-	HasERMS       bool
-	HasFSRM       bool
-	HasFMA        bool
-	HasGFNI       bool
-	HasOSXSAVE    bool
-	HasPCLMULQDQ  bool
-	HasPOPCNT     bool
-	HasRDTSCP     bool
-	HasSHA        bool
-	HasSSE3       bool
-	HasSSSE3      bool
-	HasSSE41      bool
-	HasSSE42      bool
-	_             CacheLinePad
+	_                  CacheLinePad
+	HasAES             bool
+	HasADX             bool
+	HasAVX             bool
+	HasAVXVNNI         bool
+	HasAVX2            bool
+	HasAVX512          bool // Virtual feature: F+CD+BW+DQ+VL
+	HasAVX512F         bool
+	HasAVX512CD        bool
+	HasAVX512BW        bool
+	HasAVX512DQ        bool
+	HasAVX512VL        bool
+	HasAVX512GFNI      bool
+	HasAVX512VNNI      bool
+	HasAVX512VBMI      bool
+	HasAVX512VBMI2     bool
+	HasAVX512BITALG    bool
+	HasAVX512VPOPCNTDQ bool
+	HasBMI1            bool
+	HasBMI2            bool
+	HasERMS            bool
+	HasFSRM            bool
+	HasFMA             bool
+	HasOSXSAVE         bool
+	HasPCLMULQDQ       bool
+	HasPOPCNT          bool
+	HasRDTSCP          bool
+	HasSHA             bool
+	HasSSE3            bool
+	HasSSSE3           bool
+	HasSSE41           bool
+	HasSSE42           bool
+	_                  CacheLinePad
 }
 
 // The booleans in ARM contain the correspondingly named cpu feature bit.
diff --git a/src/internal/cpu/cpu_x86.go b/src/internal/cpu/cpu_x86.go
index 152a08cdbf..04d89955da 100644
--- a/src/internal/cpu/cpu_x86.go
+++ b/src/internal/cpu/cpu_x86.go
@@ -18,18 +18,26 @@ func xgetbv() (eax, edx uint32)
 func getGOAMD64level() int32
 
 const (
+	// eax bits
+	cpuid_AVXVNNI = 1 << 4
+
 	// ecx bits
-	cpuid_SSE3      = 1 << 0
-	cpuid_PCLMULQDQ = 1 << 1
-	cpuid_SSSE3     = 1 << 9
-	cpuid_GFNI      = 1 << 8
-	cpuid_FMA       = 1 << 12
-	cpuid_SSE41     = 1 << 19
-	cpuid_SSE42     = 1 << 20
-	cpuid_POPCNT    = 1 << 23
-	cpuid_AES       = 1 << 25
-	cpuid_OSXSAVE   = 1 << 27
-	cpuid_AVX       = 1 << 28
+	cpuid_SSE3            = 1 << 0
+	cpuid_PCLMULQDQ       = 1 << 1
+	cpuid_AVX512VBMI      = 1 << 1
+	cpuid_AVX512VBMI2     = 1 << 6
+	cpuid_SSSE3           = 1 << 9
+	cpuid_AVX512GFNI      = 1 << 8
+	cpuid_AVX512VNNI      = 1 << 11
+	cpuid_AVX512BITALG    = 1 << 12
+	cpuid_FMA             = 1 << 12
+	cpuid_AVX512VPOPCNTDQ = 1 << 14
+	cpuid_SSE41           = 1 << 19
+	cpuid_SSE42           = 1 << 20
+	cpuid_POPCNT          = 1 << 23
+	cpuid_AES             = 1 << 25
+	cpuid_OSXSAVE         = 1 << 27
+	cpuid_AVX             = 1 << 28
 
 	// ebx bits
 	cpuid_BMI1     = 1 << 3
@@ -144,7 +152,7 @@ func doinit() {
 		return
 	}
 
-	_, ebx7, ecx7, edx7 := cpuid(7, 0)
+	eax7, ebx7, ecx7, edx7 := cpuid(7, 0)
 	X86.HasBMI1 = isSet(ebx7, cpuid_BMI1)
 	X86.HasAVX2 = isSet(ebx7, cpuid_AVX2) && osSupportsAVX
 	X86.HasBMI2 = isSet(ebx7, cpuid_BMI2)
@@ -158,10 +166,15 @@ func doinit() {
 		X86.HasAVX512BW = isSet(ebx7, cpuid_AVX512BW)
 		X86.HasAVX512DQ = isSet(ebx7, cpuid_AVX512DQ)
 		X86.HasAVX512VL = isSet(ebx7, cpuid_AVX512VL)
+		X86.HasAVX512GFNI = isSet(ecx7, cpuid_AVX512GFNI)
+		X86.HasAVX512BITALG = isSet(ecx7, cpuid_AVX512BITALG)
+		X86.HasAVX512VPOPCNTDQ = isSet(ecx7, cpuid_AVX512VPOPCNTDQ)
+		X86.HasAVX512VBMI = isSet(ecx7, cpuid_AVX512VBMI)
+		X86.HasAVX512VBMI2 = isSet(ecx7, cpuid_AVX512VBMI2)
+		X86.HasAVX512VNNI = isSet(ecx7, cpuid_AVX512VNNI)
 	}
 
 	X86.HasFSRM = isSet(edx7, cpuid_FSRM)
-	X86.HasGFNI = isSet(ecx7, cpuid_GFNI)
 
 	var maxExtendedInformation uint32
 	maxExtendedInformation, _, _, _ = cpuid(0x80000000, 0)
@@ -182,7 +195,12 @@ func doinit() {
 		// it. GOAMD64=v4 also implies exactly this set, and these are all
 		// included in AVX10.1.
 		X86.HasAVX512 = X86.HasAVX512F && X86.HasAVX512CD && X86.HasAVX512BW && X86.HasAVX512DQ && X86.HasAVX512VL
-		X86.HasAVX512GFNI = X86.HasAVX512 && X86.HasGFNI
+	}
+	if eax7 >= 1 {
+		eax71, _, _, _ := cpuid(7, 1)
+		if X86.HasAVX {
+			X86.HasAVXVNNI = isSet(4, eax71)
+		}
 	}
 }
 
diff --git a/src/simd/cpu.go b/src/simd/cpu.go
index 5ff47b8873..7bc5116525 100644
--- a/src/simd/cpu.go
+++ b/src/simd/cpu.go
@@ -11,12 +11,52 @@ package simd
 
 import "internal/cpu"
 
-// HasAVX512GFNI checks AVX512 CPU feature F+CD+BW+DQ+VL+GFNI.
-func HasAVX512GFNI() bool {
-	return cpu.X86.HasAVX512GFNI
+// HasAVX checks AVX CPU feature.
+func HasAVX() bool {
+	return cpu.X86.HasAVX
+}
+
+// HasAVXVNNI checks AVX CPU feature VNNI.
+func HasAVXVNNI() bool {
+	return cpu.X86.HasAVXVNNI
+}
+
+// HasAVX2 checks AVX2 CPU feature.
+func HasAVX2() bool {
+	return cpu.X86.HasAVX2
 }
 
 // HasAVX512 checks AVX512 CPU feature F+CD+BW+DQ+VL.
 func HasAVX512() bool {
 	return cpu.X86.HasAVX512
 }
+
+// HasAVX512GFNI checks AVX512 CPU feature GFNI.
+func HasAVX512GFNI() bool {
+	return cpu.X86.HasAVX512GFNI
+}
+
+// HasAVX512VBMI checks AVX512 CPU feature VBMI
+func HasAVX512VBMI() bool {
+	return cpu.X86.HasAVX512VBMI
+}
+
+// HasAVX512VBMI2 checks AVX512 CPU feature VBMI2
+func HasAVX512VBMI2() bool {
+	return cpu.X86.HasAVX512VBMI2
+}
+
+// HasAVX512VNNI checks AVX512 CPU feature VNNI
+func HasAVX512VNNI() bool {
+	return cpu.X86.HasAVX512VNNI
+}
+
+// HasAVX512VPOPCNTDQ checks AVX512 CPU feature VPOPCNTDQ
+func HasAVX512VPOPCNTDQ() bool {
+	return cpu.X86.HasAVX512VPOPCNTDQ
+}
+
+// HasAVX512BITALG checks AVX512 CPU feature BITALG
+func HasAVX512BITALG() bool {
+	return cpu.X86.HasAVX512BITALG
+}