Skip to content

Commit 7ee8ef4

Browse files
committed
cpu: add support for amx detection
Added detection for x86 AMX detection, including AMX-Tile, AMX-INT8 and AMX-BF16 instruction sets.
1 parent 5059a07 commit 7ee8ef4

File tree

2 files changed

+15
-0
lines changed

2 files changed

+15
-0
lines changed

cpu/cpu.go

+4
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@ var X86 struct {
5454
HasAVX512VBMI2 bool // Advanced vector extension 512 Vector Byte Manipulation Instructions 2
5555
HasAVX512BITALG bool // Advanced vector extension 512 Bit Algorithms
5656
HasAVX512BF16 bool // Advanced vector extension 512 BFloat16 Instructions
57+
HasAMX bool // Advanced Matrix Extension
58+
HasAMXTile bool // Advanced Matrix Extension Tile instructions
59+
HasAMXInt8 bool // Advanced Matrix Extension Int8 instructions
60+
HasAMXBF16 bool // Advanced Matrix Extension BFloat16 instructions
5761
HasBMI1 bool // Bit manipulation instruction set 1
5862
HasBMI2 bool // Bit manipulation instruction set 2
5963
HasCX16 bool // Compare and exchange 16 Bytes

cpu/cpu_x86.go

+11
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ func initOptions() {
3737
{Name: "avx512vbmi2", Feature: &X86.HasAVX512VBMI2},
3838
{Name: "avx512bitalg", Feature: &X86.HasAVX512BITALG},
3939
{Name: "avx512bf16", Feature: &X86.HasAVX512BF16},
40+
{Name: "amx", Feature: &X86.HasAMX},
41+
{Name: "amxtile", Feature: &X86.HasAMXTile},
42+
{Name: "amxint8", Feature: &X86.HasAMXInt8},
43+
{Name: "amxbf16", Feature: &X86.HasAMXBF16},
4044
{Name: "bmi1", Feature: &X86.HasBMI1},
4145
{Name: "bmi2", Feature: &X86.HasBMI2},
4246
{Name: "cx16", Feature: &X86.HasCX16},
@@ -138,6 +142,13 @@ func archInit() {
138142
eax71, _, _, _ := cpuid(7, 1)
139143
X86.HasAVX512BF16 = isSet(5, eax71)
140144
}
145+
146+
X86.HasAMX = isSet(24, edx7)
147+
if X86.HasAMX {
148+
X86.HasAMXTile = true
149+
X86.HasAMXInt8 = isSet(25, edx7)
150+
X86.HasAMXBF16 = isSet(22, edx7)
151+
}
141152
}
142153

143154
func isSet(bitpos uint, value uint32) bool {

0 commit comments

Comments
 (0)