From 7ee8ef40b4194eb4ac443c66bf1b970797cebf35 Mon Sep 17 00:00:00 2001
From: Anton Kuklin <anton.a.kuklin@gmail.com>
Date: Mon, 7 Aug 2023 20:35:28 +0100
Subject: [PATCH] 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 {