cpu: add support for detecting AVX-512

The existing implementation hasn't AVX-512 detection support.
This patch extends cpu package with such a functionality.

Fixes golang/go#41288

Change-Id: Iaddd57c5ff6d73360b77886458b9210b6c446859
Reviewed-on: https://go-review.googlesource.com/c/sys/+/257937
Trust: Tobias Klauser <tobias.klauser@gmail.com>
Trust: Cuong Manh Le <cuong.manhle.vn@gmail.com>
Run-TryBot: Tobias Klauser <tobias.klauser@gmail.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Tobias Klauser <tobias.klauser@gmail.com>

cpu/cpu.go

@@ -34,6 +34,26 @@ var X86 struct {
 	HasADX              bool // Multi-precision add-carry instruction extensions
 	HasAVX              bool // Advanced vector extension
 	HasAVX2             bool // Advanced vector extension 2
+	HasAVX512           bool // Advanced vector extension 512
+	HasAVX512F          bool // Advanced vector extension 512 Foundation Instructions
+	HasAVX512CD         bool // Advanced vector extension 512 Conflict Detection Instructions
+	HasAVX512ER         bool // Advanced vector extension 512 Exponential and Reciprocal Instructions
+	HasAVX512PF         bool // Advanced vector extension 512 Prefetch Instructions Instructions
+	HasAVX512VL         bool // Advanced vector extension 512 Vector Length Extensions
+	HasAVX512BW         bool // Advanced vector extension 512 Byte and Word Instructions
+	HasAVX512DQ         bool // Advanced vector extension 512 Doubleword and Quadword Instructions
+	HasAVX512IFMA       bool // Advanced vector extension 512 Integer Fused Multiply Add
+	HasAVX512VBMI       bool // Advanced vector extension 512 Vector Byte Manipulation Instructions
+	HasAVX5124VNNIW     bool // Advanced vector extension 512 Vector Neural Network Instructions Word variable precision
+	HasAVX5124FMAPS     bool // Advanced vector extension 512 Fused Multiply Accumulation Packed Single precision
+	HasAVX512VPOPCNTDQ  bool // Advanced vector extension 512 Double and quad word population count instructions
+	HasAVX512VPCLMULQDQ bool // Advanced vector extension 512 Vector carry-less multiply operations
+	HasAVX512VNNI       bool // Advanced vector extension 512 Vector Neural Network Instructions
+	HasAVX512GFNI       bool // Advanced vector extension 512 Galois field New Instructions
+	HasAVX512VAES       bool // Advanced vector extension 512 Vector AES instructions
+	HasAVX512VBMI2      bool // Advanced vector extension 512 Vector Byte Manipulation Instructions 2
+	HasAVX512BITALG     bool // Advanced vector extension 512 Bit Algorithms
+	HasAVX512BF16       bool // Advanced vector extension 512 BFloat16 Instructions
 	HasBMI1             bool // Bit manipulation instruction set 1
 	HasBMI2             bool // Bit manipulation instruction set 2
 	HasERMS             bool // Enhanced REP for MOVSB and STOSB

cpu/cpu_test.go

@@ -30,6 +30,17 @@ func TestAVX2hasAVX(t *testing.T) {
 	}
 }
 
+func TestAVX512HasAVX2AndAVX(t *testing.T) {
+	if runtime.GOARCH == "amd64" {
+		if cpu.X86.HasAVX512 && !cpu.X86.HasAVX {
+			t.Fatal("HasAVX expected true, got false")
+		}
+		if cpu.X86.HasAVX512 && !cpu.X86.HasAVX2 {
+			t.Fatal("HasAVX2 expected true, got false")
+		}
+	}
+}
+
 func TestARM64minimalFeatures(t *testing.T) {
 	if runtime.GOARCH != "arm64" || (runtime.GOOS == "darwin" || runtime.GOOS == "ios") {
 		return

cpu/cpu_x86.go

@@ -16,6 +16,26 @@ func initOptions() {
 		{Name: "aes", Feature: &X86.HasAES},
 		{Name: "avx", Feature: &X86.HasAVX},
 		{Name: "avx2", Feature: &X86.HasAVX2},
+		{Name: "avx512", Feature: &X86.HasAVX512},
+		{Name: "avx512f", Feature: &X86.HasAVX512F},
+		{Name: "avx512cd", Feature: &X86.HasAVX512CD},
+		{Name: "avx512er", Feature: &X86.HasAVX512ER},
+		{Name: "avx512pf", Feature: &X86.HasAVX512PF},
+		{Name: "avx512vl", Feature: &X86.HasAVX512VL},
+		{Name: "avx512bw", Feature: &X86.HasAVX512BW},
+		{Name: "avx512dq", Feature: &X86.HasAVX512DQ},
+		{Name: "avx512ifma", Feature: &X86.HasAVX512IFMA},
+		{Name: "avx512vbmi", Feature: &X86.HasAVX512VBMI},
+		{Name: "avx512vnniw", Feature: &X86.HasAVX5124VNNIW},
+		{Name: "avx5124fmaps", Feature: &X86.HasAVX5124FMAPS},
+		{Name: "avx512vpopcntdq", Feature: &X86.HasAVX512VPOPCNTDQ},
+		{Name: "avx512vpclmulqdq", Feature: &X86.HasAVX512VPCLMULQDQ},
+		{Name: "avx512vnni", Feature: &X86.HasAVX512VNNI},
+		{Name: "avx512gfni", Feature: &X86.HasAVX512GFNI},
+		{Name: "avx512vaes", Feature: &X86.HasAVX512VAES},
+		{Name: "avx512vbmi2", Feature: &X86.HasAVX512VBMI2},
+		{Name: "avx512bitalg", Feature: &X86.HasAVX512BITALG},
+		{Name: "avx512bf16", Feature: &X86.HasAVX512BF16},
 		{Name: "bmi1", Feature: &X86.HasBMI1},
 		{Name: "bmi2", Feature: &X86.HasBMI2},
 		{Name: "erms", Feature: &X86.HasERMS},
@@ -59,12 +79,15 @@ func archInit() {
 	X86.HasOSXSAVE = isSet(27, ecx1)
 	X86.HasRDRAND = isSet(30, ecx1)
 
-	osSupportsAVX := false
+	var osSupportsAVX, osSupportsAVX512 bool
 	// For XGETBV, OSXSAVE bit is required and sufficient.
 	if X86.HasOSXSAVE {
 		eax, _ := xgetbv()
 		// Check if XMM and YMM registers have OS support.
 		osSupportsAVX = isSet(1, eax) && isSet(2, eax)
+
+		// Check if OPMASK and ZMM registers have OS support.
+		osSupportsAVX512 = osSupportsAVX && isSet(5, eax) && isSet(6, eax) && isSet(7, eax)
 	}
 
 	X86.HasAVX = isSet(28, ecx1) && osSupportsAVX
@@ -73,7 +96,7 @@ func archInit() {
 		return
 	}
 
-	_, ebx7, _, _ := cpuid(7, 0)
+	_, ebx7, ecx7, edx7 := cpuid(7, 0)
 	X86.HasBMI1 = isSet(3, ebx7)
 	X86.HasAVX2 = isSet(5, ebx7) && osSupportsAVX
 	X86.HasBMI2 = isSet(8, ebx7)
@@ -81,6 +104,30 @@ func archInit() {
 	X86.HasRDSEED = isSet(18, ebx7)
 	X86.HasADX = isSet(19, ebx7)
 
+	X86.HasAVX512 = isSet(16, ebx7) && osSupportsAVX512 // Because avx-512 foundation is the core required extension
+	if X86.HasAVX512 {
+		X86.HasAVX512F = true
+		X86.HasAVX512CD = isSet(28, ebx7)
+		X86.HasAVX512ER = isSet(27, ebx7)
+		X86.HasAVX512PF = isSet(26, ebx7)
+		X86.HasAVX512VL = isSet(31, ebx7)
+		X86.HasAVX512BW = isSet(30, ebx7)
+		X86.HasAVX512DQ = isSet(17, ebx7)
+		X86.HasAVX512IFMA = isSet(21, ebx7)
+		X86.HasAVX512VBMI = isSet(1, ecx7)
+		X86.HasAVX5124VNNIW = isSet(2, edx7)
+		X86.HasAVX5124FMAPS = isSet(3, edx7)
+		X86.HasAVX512VPOPCNTDQ = isSet(14, ecx7)
+		X86.HasAVX512VPCLMULQDQ = isSet(10, ecx7)
+		X86.HasAVX512VNNI = isSet(11, ecx7)
+		X86.HasAVX512GFNI = isSet(8, ecx7)
+		X86.HasAVX512VAES = isSet(9, ecx7)
+		X86.HasAVX512VBMI2 = isSet(6, ecx7)
+		X86.HasAVX512BITALG = isSet(12, ecx7)
+
+		eax71, _, _, _ := cpuid(7, 1)
+		X86.HasAVX512BF16 = isSet(5, eax71)
+	}
 }
 
 func isSet(bitpos uint, value uint32) bool {