From b114243b22c3fc4dcdf69c57f9f2ae09ac42726a Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Thu, 9 Jul 2026 16:03:03 -0400 Subject: [PATCH 1/3] Fix AMX detection to request OS XSAVE permission on Linux On Linux, the AMX XTILEDATA state component (XCR0 bit 18) is allocated lazily: the kernel leaves it disabled until a thread explicitly requests it via arch_prctl(ARCH_REQ_XCOMP_PERM). Without this call, xgetbv() reports AMX as OS-disabled even on hardware that fully supports it, and code generated for AVX512_SapphireRapids can fault at runtime. Request the permission and gate the feature on the resulting XCR0 bits in both the AOT runtime feature-detection module and the compiler's own host-target detection in Target.cpp. Fixes #9191 Co-Authored-By: Claude Sonnet 5 --- src/LLVM_Runtime_Linker.cpp | 8 +++++- src/Target.cpp | 39 ++++++++++++++++++++++++-- src/runtime/CMakeLists.txt | 1 + src/runtime/linux_x86_cpu_features.cpp | 2 ++ src/runtime/x86_cpu_features.cpp | 33 ++++++++++++++++++++-- 5 files changed, 77 insertions(+), 6 deletions(-) create mode 100644 src/runtime/linux_x86_cpu_features.cpp diff --git a/src/LLVM_Runtime_Linker.cpp b/src/LLVM_Runtime_Linker.cpp index 60faee4c47fe..e532124c4af6 100644 --- a/src/LLVM_Runtime_Linker.cpp +++ b/src/LLVM_Runtime_Linker.cpp @@ -295,6 +295,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) @@ -305,6 +306,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) @@ -1273,7 +1275,11 @@ std::unique_ptr 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) { diff --git a/src/Target.cpp b/src/Target.cpp index 55f2978b869c..2b37a6143e23 100644 --- a/src/Target.cpp +++ b/src/Target.cpp @@ -33,6 +33,10 @@ #include #endif +#if defined(__linux__) && (defined(__x86_64__) || defined(__i386__)) +#include // for syscall() +#endif + #if defined(__linux__) && (defined(__arm__) || defined(__aarch64__)) #include #include @@ -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); } } diff --git a/src/runtime/CMakeLists.txt b/src/runtime/CMakeLists.txt index 95e78d11e459..5f403c62d51d 100644 --- a/src/runtime/CMakeLists.txt +++ b/src/runtime/CMakeLists.txt @@ -43,6 +43,7 @@ set(RUNTIME_CPP linux_host_cpu_count linux_powerpc_thread_id linux_riscv_thread_id + linux_x86_cpu_features linux_x86_thread_id linux_yield metal diff --git a/src/runtime/linux_x86_cpu_features.cpp b/src/runtime/linux_x86_cpu_features.cpp new file mode 100644 index 000000000000..f84ee6c76fd6 --- /dev/null +++ b/src/runtime/linux_x86_cpu_features.cpp @@ -0,0 +1,2 @@ +#define LINUX 1 +#include "x86_cpu_features.cpp" diff --git a/src/runtime/x86_cpu_features.cpp b/src/runtime/x86_cpu_features.cpp index 0716d7f27078..93df6fc17419 100644 --- a/src/runtime/x86_cpu_features.cpp +++ b/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); } } From 578340c0c180214ea3b69d99f8d31597ec2b3b54 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 9 Jul 2026 13:33:31 -0700 Subject: [PATCH 2/3] Add new runtime module to makefile --- Makefile | 1 + 1 file changed, 1 insertion(+) diff --git a/Makefile b/Makefile index 67e19076c78f..a825b48365bb 100644 --- a/Makefile +++ b/Makefile @@ -866,6 +866,7 @@ RUNTIME_CPP_COMPONENTS = \ linux_riscv_thread_id \ linux_x86_thread_id \ linux_yield \ + linux_x86_cpu_features \ metal \ metal_objc_arm \ metal_objc_x86 \ From aec76c1390a8d88cdec18f193eb0a9b9c5c85a4c Mon Sep 17 00:00:00 2001 From: "halide-ci[bot]" <266445882+halide-ci[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:35:20 +0000 Subject: [PATCH 3/3] Apply pre-commit auto-fixes --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index a825b48365bb..dbb64afebccc 100644 --- a/Makefile +++ b/Makefile @@ -864,9 +864,9 @@ RUNTIME_CPP_COMPONENTS = \ linux_host_cpu_count \ linux_powerpc_thread_id \ linux_riscv_thread_id \ + linux_x86_cpu_features \ linux_x86_thread_id \ linux_yield \ - linux_x86_cpu_features \ metal \ metal_objc_arm \ metal_objc_x86 \