From 7ee8ef40b4194eb4ac443c66bf1b970797cebf35 Mon Sep 17 00:00:00 2001 From: Anton Kuklin Date: Mon, 7 Aug 2023 20:35:28 +0100 Subject: [PATCH 1/3] cpu: add support for amx detection Added detection for x86 AMX detection, including AMX-Tile, AMX-INT8 and AMX-BF16 instruction sets. --- cpu/cpu.go | 4 ++++ cpu/cpu_x86.go | 11 +++++++++++ 2 files changed, 15 insertions(+) diff --git a/cpu/cpu.go b/cpu/cpu.go index 83f112c4c..2796a31ce 100644 --- a/cpu/cpu.go +++ b/cpu/cpu.go @@ -54,6 +54,10 @@ var X86 struct { 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 + HasAMX bool // Advanced Matrix Extension + HasAMXTile bool // Advanced Matrix Extension Tile instructions + HasAMXInt8 bool // Advanced Matrix Extension Int8 instructions + HasAMXBF16 bool // Advanced Matrix Extension BFloat16 instructions HasBMI1 bool // Bit manipulation instruction set 1 HasBMI2 bool // Bit manipulation instruction set 2 HasCX16 bool // Compare and exchange 16 Bytes diff --git a/cpu/cpu_x86.go b/cpu/cpu_x86.go index f5aacfc82..c6446b4c2 100644 --- a/cpu/cpu_x86.go +++ b/cpu/cpu_x86.go @@ -37,6 +37,10 @@ func initOptions() { {Name: "avx512vbmi2", Feature: &X86.HasAVX512VBMI2}, {Name: "avx512bitalg", Feature: &X86.HasAVX512BITALG}, {Name: "avx512bf16", Feature: &X86.HasAVX512BF16}, + {Name: "amx", Feature: &X86.HasAMX}, + {Name: "amxtile", Feature: &X86.HasAMXTile}, + {Name: "amxint8", Feature: &X86.HasAMXInt8}, + {Name: "amxbf16", Feature: &X86.HasAMXBF16}, {Name: "bmi1", Feature: &X86.HasBMI1}, {Name: "bmi2", Feature: &X86.HasBMI2}, {Name: "cx16", Feature: &X86.HasCX16}, @@ -138,6 +142,13 @@ func archInit() { eax71, _, _, _ := cpuid(7, 1) X86.HasAVX512BF16 = isSet(5, eax71) } + + X86.HasAMX = isSet(24, edx7) + if X86.HasAMX { + X86.HasAMXTile = true + X86.HasAMXInt8 = isSet(25, edx7) + X86.HasAMXBF16 = isSet(22, edx7) + } } func isSet(bitpos uint, value uint32) bool { From a243dce1a76dda54d275f0e6501cee0f60bcbcb1 Mon Sep 17 00:00:00 2001 From: Anton Kuklin Date: Thu, 10 Aug 2023 13:24:58 +0100 Subject: [PATCH 2/3] remove hasAMX flag --- cpu/cpu.go | 1 - cpu/cpu_x86.go | 9 +++------ 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/cpu/cpu.go b/cpu/cpu.go index 2796a31ce..3ebc32570 100644 --- a/cpu/cpu.go +++ b/cpu/cpu.go @@ -54,7 +54,6 @@ var X86 struct { 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 - HasAMX bool // Advanced Matrix Extension HasAMXTile bool // Advanced Matrix Extension Tile instructions HasAMXInt8 bool // Advanced Matrix Extension Int8 instructions HasAMXBF16 bool // Advanced Matrix Extension BFloat16 instructions diff --git a/cpu/cpu_x86.go b/cpu/cpu_x86.go index c6446b4c2..53c5309c1 100644 --- a/cpu/cpu_x86.go +++ b/cpu/cpu_x86.go @@ -143,12 +143,9 @@ func archInit() { X86.HasAVX512BF16 = isSet(5, eax71) } - X86.HasAMX = isSet(24, edx7) - if X86.HasAMX { - X86.HasAMXTile = true - X86.HasAMXInt8 = isSet(25, edx7) - X86.HasAMXBF16 = isSet(22, edx7) - } + X86.HasAMXTile = isSet(24, edx7) + X86.HasAMXInt8 = isSet(25, edx7) + X86.HasAMXBF16 = isSet(22, edx7) } func isSet(bitpos uint, value uint32) bool { From 7986ed6d4f18f81e83bb2e6a782db90443db9de0 Mon Sep 17 00:00:00 2001 From: Anton Kuklin Date: Thu, 10 Aug 2023 13:31:58 +0100 Subject: [PATCH 3/3] fix --- cpu/cpu_x86.go | 1 - 1 file changed, 1 deletion(-) diff --git a/cpu/cpu_x86.go b/cpu/cpu_x86.go index 53c5309c1..2dcde8285 100644 --- a/cpu/cpu_x86.go +++ b/cpu/cpu_x86.go @@ -37,7 +37,6 @@ func initOptions() { {Name: "avx512vbmi2", Feature: &X86.HasAVX512VBMI2}, {Name: "avx512bitalg", Feature: &X86.HasAVX512BITALG}, {Name: "avx512bf16", Feature: &X86.HasAVX512BF16}, - {Name: "amx", Feature: &X86.HasAMX}, {Name: "amxtile", Feature: &X86.HasAMXTile}, {Name: "amxint8", Feature: &X86.HasAMXInt8}, {Name: "amxbf16", Feature: &X86.HasAMXBF16},