Fix AMX detection to request OS XSAVE permission on Linux

Makefile

@@ -860,6 +860,7 @@ RUNTIME_CPP_COMPONENTS = \
   linux_arm_cpu_features \
   linux_clock \
   linux_host_cpu_count \
+  linux_x86_cpu_features \
   linux_yield \
   metal \
   metal_objc_arm \

src/LLVM_Runtime_Linker.cpp

@@ -289,6 +289,7 @@ DECLARE_NO_INITMOD(windows_vulkan)
 
 #ifdef WITH_X86
 // keep-sorted start by_regex=["INITMOD\\(.+"]
+DECLARE_CPP_INITMOD(linux_x86_cpu_features)
 DECLARE_LL_INITMOD(x86)
 DECLARE_LL_INITMOD(x86_amx)
 DECLARE_LL_INITMOD(x86_avx)
@@ -299,6 +300,7 @@ DECLARE_LL_INITMOD(x86_sse41)
 // keep-sorted end
 #else
 // keep-sorted start by_regex=["INITMOD\\(.+"]
+DECLARE_NO_INITMOD(linux_x86_cpu_features)
 DECLARE_NO_INITMOD(x86)
 DECLARE_NO_INITMOD(x86_amx)
 DECLARE_NO_INITMOD(x86_avx)
@@ -1245,7 +1247,11 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
             // These modules are only used for AOT compilation
             modules.push_back(get_initmod_can_use_target(c, bits_64, debug));
             if (t.arch == Target::X86) {
-                modules.push_back(get_initmod_x86_cpu_features(c, bits_64, debug));
+                if (t.os == Target::Android || t.os == Target::Linux) {
+                    modules.push_back(get_initmod_linux_x86_cpu_features(c, bits_64, debug));
+                } else {
+                    modules.push_back(get_initmod_x86_cpu_features(c, bits_64, debug));
+                }
             }
             if (t.arch == Target::ARM) {
                 if (t.bits == 64) {

src/Target.cpp

@@ -33,6 +33,10 @@
 #include <sys/types.h>
 #endif
 
+#if defined(__linux__) && (defined(__x86_64__) || defined(__i386__))
+#include <unistd.h>  // for syscall()
+#endif
+
 #if defined(__linux__) && (defined(__arm__) || defined(__aarch64__))
 #include <asm/hwcap.h>
 #include <sys/auxv.h>
@@ -105,6 +109,25 @@ struct cpuid_result {
 
 #if defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_AMD64)
 
+#ifdef __linux__
+// On Linux, the AMX XTILEDATA state component (XCR0 bit 18) is allocated
+// lazily: the kernel leaves it disabled in XCR0 for a thread until that
+// thread explicitly requests permission to use it via arch_prctl(2). Without
+// this call, xgetbv() would report AMX as OS-disabled even on hardware that
+// fully supports it. This is a no-op (and harmless) if already granted, and
+// only applies in 64-bit mode, since arch_prctl and AMX are both x86-64-only.
+constexpr long sys_arch_prctl = 158;  // x86-64 Linux syscall number
+constexpr long arch_req_xcomp_perm = 0x1023;
+constexpr long xfeature_xtiledata = 18;
+
+void request_amx_os_permission() {
+    syscall(sys_arch_prctl, arch_req_xcomp_perm, xfeature_xtiledata);
+}
+#else
+void request_amx_os_permission() {
+}
+#endif
+
 enum class VendorSignatures {
     Unknown,
     GenuineIntel,
@@ -377,11 +400,16 @@ Target calculate_host_target() {
     bool os_avx = false;
     bool os_avx512 = false;
     bool os_apx = false;
+    bool os_amx = false;
     if (have_osxsave) {
+        if (use_64_bits) {
+            request_amx_os_permission();
+        }
         uint64_t xcr0 = xgetbv(0);
         os_avx = (xcr0 & 0x6) == 0x6;                   // XMM (bit 1) + YMM (bit 2)
         os_avx512 = os_avx && ((xcr0 & 0xE0) == 0xE0);  // opmask (5) + ZMM_Hi256 (6) + Hi16_ZMM (7)
         os_apx = (xcr0 & 0x80000) == 0x80000;           // APX extended GPRs (bit 19)
+        os_amx = (xcr0 & 0x60000) == 0x60000;           // AMX XTILECFG (17) + XTILEDATA (18)
     }
 
     bool have_sse41 = (info.ecx & (1 << 19)) != 0;           // ECX[19]
@@ -490,12 +518,17 @@ Target calculate_host_target() {
             if ((info2.ebx & avx512_cannonlake) == avx512_cannonlake) {
                 initial_features.push_back(Target::AVX512_Cannonlake);
 
-                const uint32_t avxvnni = 1U << 4;     // avxvnni (note, not avx512vnni) result in eax
-                const uint32_t avx512bf16 = 1U << 5;  // bf16 result in eax, with cpuid(eax=7, ecx=1)
+                const uint32_t avxvnni = 1U << 4;    // avxvnni (note, not avx512vnni) result in eax
+                const uint32_t amx_bf16 = 1U << 22;  // amx_bf16 result in edx, with cpuid(eax=7, ecx=0)
+                const uint32_t amx_tile = 1U << 24;  // amx_tile result in edx, with cpuid(eax=7, ecx=0)
+                const uint32_t amx_int8 = 1U << 25;  // amx_int8 result in edx, with cpuid(eax=7, ecx=0)
+                const uint32_t amx = amx_bf16 | amx_tile | amx_int8;
                 // TODO: port to family/model -based detection.
                 if ((info3.eax & avxvnni) == avxvnni) {
                     initial_features.push_back(Target::AVXVNNI);
-                    if ((info3.eax & avx512bf16) == avx512bf16) {
+                    // avx512_sapphirerapids implies AMX instruction support, which
+                    // requires the OS to have enabled the AMX XCR0 state components.
+                    if ((info2.edx & amx) == amx && os_amx) {
                         initial_features.push_back(Target::AVX512_SapphireRapids);
                     }
                 }

src/runtime/CMakeLists.txt

@@ -39,6 +39,7 @@ set(RUNTIME_CPP
     linux_arm_cpu_features
     linux_clock
     linux_host_cpu_count
+    linux_x86_cpu_features
     linux_yield
     metal
     metal_objc_arm

src/runtime/linux_x86_cpu_features.cpp

@@ -0,0 +1,2 @@
+#define LINUX 1
+#include "x86_cpu_features.cpp"

src/runtime/x86_cpu_features.cpp

@@ -34,6 +34,27 @@ struct cpuid_result {
     return (uint32_t)xcr_info[0];
 }
 
+#if LINUX
+// On Linux, the AMX XTILEDATA state component (XCR0 bit 18) is allocated
+// lazily: the kernel leaves it disabled in XCR0 for a thread until that
+// thread explicitly requests permission to use it via arch_prctl(2). Without
+// this call, xgetbv() would report AMX as OS-disabled even on hardware that
+// fully supports it. This is a no-op (and harmless) if already granted, and
+// only applies in 64-bit mode, since arch_prctl and AMX are both x86-64-only.
+extern "C" long syscall(long number, ...);
+
+constexpr long sys_arch_prctl = 158;  // x86-64 Linux syscall number
+constexpr long arch_req_xcomp_perm = 0x1023;
+constexpr long xfeature_xtiledata = 18;
+
+ALWAYS_INLINE void request_amx_os_permission() {
+    syscall(sys_arch_prctl, arch_req_xcomp_perm, xfeature_xtiledata);
+}
+#else
+ALWAYS_INLINE void request_amx_os_permission() {
+}
+#endif
+
 }  // namespace
 
 extern "C" WEAK int halide_get_cpu_features(CpuFeatures *features) {
@@ -65,11 +86,16 @@ extern "C" WEAK int halide_get_cpu_features(CpuFeatures *features) {
     bool os_avx = false;
     bool os_avx512 = false;
     bool os_apx = false;
+    bool os_amx = false;
     if (have_osxsave) {
+        if (use_64_bits) {
+            request_amx_os_permission();
+        }
         uint32_t xcr0 = xgetbv(0);
         os_avx = (xcr0 & 0x6) == 0x6;                   // XMM (bit 1) + YMM (bit 2)
         os_avx512 = os_avx && ((xcr0 & 0xE0) == 0xE0);  // opmask (5) + ZMM_Hi256 (6) + Hi16_ZMM (7)
         os_apx = (xcr0 & 0x80000) == 0x80000;           // APX extended GPRs (bit 19)
+        os_amx = (xcr0 & 0x60000) == 0x60000;           // AMX XTILECFG (17) + XTILEDATA (18)
     }
 
     uint32_t family = (info.eax >> 8) & 0xF;  // Bits 8..11
@@ -166,7 +192,10 @@ extern "C" WEAK int halide_get_cpu_features(CpuFeatures *features) {
         constexpr uint32_t avx512vl = 1U << 31;
         constexpr uint32_t avx512ifma = 1U << 21;
         constexpr uint32_t avxvnni = 1U << 4;
-        constexpr uint32_t avx512bf16 = 1U << 5;  // bf16 result in eax, cpuid(eax=7, ecx=1)
+        constexpr uint32_t amx_bf16 = 1U << 22;  // amx_bf16 result in edx, cpuid(eax=7, ecx=0)
+        constexpr uint32_t amx_tile = 1U << 24;  // amx_tile result in edx, cpuid(eax=7, ecx=0)
+        constexpr uint32_t amx_int8 = 1U << 25;  // amx_int8 result in edx, cpuid(eax=7, ecx=0)
+        constexpr uint32_t amx = amx_bf16 | amx_tile | amx_int8;
         constexpr uint32_t avx512 = avx512f | avx512cd;
         constexpr uint32_t avx512_knl = avx512 | avx512pf | avx512er;
         constexpr uint32_t avx512_skylake = avx512 | avx512vl | avx512bw | avx512dq;
@@ -187,7 +216,7 @@ extern "C" WEAK int halide_get_cpu_features(CpuFeatures *features) {
 
                 if ((info3.eax & avxvnni) == avxvnni) {
                     halide_set_available_cpu_feature(features, halide_target_feature_avxvnni);
-                    if ((info3.eax & avx512bf16) == avx512bf16) {
+                    if ((info2.edx & amx) == amx && os_amx) {
                         halide_set_available_cpu_feature(features, halide_target_feature_avx512_sapphirerapids);
                     }
                 }