Add darwin AVX512 detection (#74)

* Add darwin AVX512 detection Imported from https://go-review.googlesource.com/c/sys/+/285572/ See golang/go#43089
klauspost · Mar 26, 2021 · c6a3519 · c6a3519
1 parent 3df0045
commit c6a3519
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 1 deletion.
diff --git a/cpuid.go b/cpuid.go
@@ -15,6 +15,7 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"runtime"
 	"strings"
 )
 
@@ -209,6 +210,7 @@ var cpuid func(op uint32) (eax, ebx, ecx, edx uint32)
 var cpuidex func(op, op2 uint32) (eax, ebx, ecx, edx uint32)
 var xgetbv func(index uint32) (eax, edx uint32)
 var rdtscpAsm func() (eax, ebx, ecx, edx uint32)
+var darwinHasAVX512 = func() bool { return false }
 
 // CPU contains information about the CPU as detected on startup,
 // or when Detect last was called.
@@ -922,7 +924,11 @@ func support() flagSet {
 			// Verify that XCR0[7:5] = ‘111b’ (OPMASK state, upper 256-bit of ZMM0-ZMM15 and
 			// ZMM16-ZMM31 state are enabled by OS)
 			/// and that XCR0[2:1] = ‘11b’ (XMM state and YMM state are enabled by OS).
-			if (eax>>5)&7 == 7 && (eax>>1)&3 == 3 {
+			hasAVX512 := (eax>>5)&7 == 7 && (eax>>1)&3 == 3
+			if runtime.GOOS == "darwin" {
+				hasAVX512 = fs.inSet(AVX) && darwinHasAVX512()
+			}
+			if hasAVX512 {
 				fs.setIf(ebx&(1<<16) != 0, AVX512F)
 				fs.setIf(ebx&(1<<17) != 0, AVX512DQ)
 				fs.setIf(ebx&(1<<21) != 0, AVX512IFMA)

diff --git a/cpuid_386.s b/cpuid_386.s
@@ -40,3 +40,8 @@ TEXT ·asmRdtscpAsm(SB), 7, $0
 	MOVL CX, ecx+8(FP)
 	MOVL DX, edx+12(FP)
 	RET
+
+// func asmDarwinHasAVX512() bool
+TEXT ·asmDarwinHasAVX512(SB), 7, $0
+	MOVL $0, eax+0(FP)
+	RET
diff --git a/cpuid_amd64.s b/cpuid_amd64.s
@@ -40,3 +40,33 @@ TEXT ·asmRdtscpAsm(SB), 7, $0
 	MOVL CX, ecx+8(FP)
 	MOVL DX, edx+12(FP)
 	RET
+
+// From https://go-review.googlesource.com/c/sys/+/285572/
+// func asmDarwinHasAVX512() bool
+TEXT ·asmDarwinHasAVX512(SB), 7, $0-1
+	MOVB $0, ret+0(FP) // default to false
+
+#ifdef GOOS_darwin // return if not darwin
+#ifdef GOARCH_amd64 // return if not amd64
+// These values from:
+// https://github.com/apple/darwin-xnu/blob/xnu-4570.1.46/osfmk/i386/cpu_capabilities.h
+#define commpage64_base_address         0x00007fffffe00000
+#define commpage64_cpu_capabilities64   (commpage64_base_address+0x010)
+#define commpage64_version              (commpage64_base_address+0x01E)
+#define hasAVX512F                      0x0000004000000000
+	MOVQ $commpage64_version, BX
+	MOVW (BX), AX
+	CMPW AX, $13                            // versions < 13 do not support AVX512
+	JL   no_avx512
+	MOVQ $commpage64_cpu_capabilities64, BX
+	MOVQ (BX), AX
+	MOVQ $hasAVX512F, CX
+	ANDQ CX, AX
+	JZ   no_avx512
+	MOVB $1, ret+0(FP)
+
+no_avx512:
+#endif
+#endif
+	RET
+
diff --git a/detect_x86.go b/detect_x86.go
@@ -8,12 +8,14 @@ func asmCpuid(op uint32) (eax, ebx, ecx, edx uint32)
 func asmCpuidex(op, op2 uint32) (eax, ebx, ecx, edx uint32)
 func asmXgetbv(index uint32) (eax, edx uint32)
 func asmRdtscpAsm() (eax, ebx, ecx, edx uint32)
+func asmDarwinHasAVX512() bool
 
 func initCPU() {
 	cpuid = asmCpuid
 	cpuidex = asmCpuidex
 	xgetbv = asmXgetbv
 	rdtscpAsm = asmRdtscpAsm
+	darwinHasAVX512 = asmDarwinHasAVX512
 }
 
 func addInfo(c *CPUInfo, safe bool) {