From 71e79b2b4343aeb6d77eb230d8096eee3a896949 Mon Sep 17 00:00:00 2001
From: Bradley Wood <bradley.wood@ibm.com>
Date: Wed, 27 Nov 2024 12:34:06 -0500
Subject: [PATCH] x86: Disable AVX-512 if zmm XCR0 flags are not set

Signed-off-by: Bradley Wood <bradley.wood@ibm.com>
---
 compiler/x/env/OMRCPU.cpp         | 37 +++++++++++++++++++++++++-----
 compiler/x/runtime/X86Runtime.hpp | 38 +++++++++++++++++++++++++------
 2 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/compiler/x/env/OMRCPU.cpp b/compiler/x/env/OMRCPU.cpp
index 036dc85d322..4ebbb181e24 100644
--- a/compiler/x/env/OMRCPU.cpp
+++ b/compiler/x/env/OMRCPU.cpp
@@ -61,14 +61,39 @@ OMR::X86::CPU::detect(OMRPortLibrary * const omrPortLib)
       processorDescription.features[i] &= featureMasks.features[i];
       }
 
+   bool disableAVX = true;
+   bool disableAVX512 = true;
+
+   // Check XCRO register for OS support of xmm/ymm/zmm
    if (TRUE == omrsysinfo_processor_has_feature(&processorDescription, OMR_FEATURE_X86_OSXSAVE))
       {
-      static const bool disableAVX = feGetEnv("TR_DisableAVX") != NULL;
-      if (((6 & _xgetbv(0)) != 6) || disableAVX) // '6' = mask for XCR0[2:1]='11b' (XMM state and YMM state are enabled)
-         {
-         // Unset OSXSAVE if not enabled via CR0
-         omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_OSXSAVE, FALSE);
-         }
+      // '6' = mask for XCR0[2:1]='11b' (XMM state and YMM state are enabled)
+      disableAVX = ((6 & _xgetbv(0)) != 6);
+      // 'e6' = (mask for XCR0[7:5]='111b' (Opmask, ZMM_Hi256, Hi16_ZMM) + XCR0[2:1]='11b' (XMM/YMM))
+      disableAVX512 = ((0xe6 & _xgetbv(0)) != 0xe6);
+      }
+
+   if(disableAVX)
+      {
+      // Unset AVX/AVX2 if not enabled via CR0 or otherwise disabled
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX2, FALSE);
+      }
+
+   if (disableAVX512)
+      {
+      // Unset AVX-512 if not enabled via CR0 or otherwise disabled
+      // If other AVX-512 extensions are supported in the port library, they need to be disabled here
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512F, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512VL, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512BW, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512CD, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512DQ, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512_BITALG, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512_VBMI, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512_VBMI2, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512_VNNI, FALSE);
+      omrsysinfo_processor_set_feature(&processorDescription, OMR_FEATURE_X86_AVX512_VPOPCNTDQ, FALSE);
       }
 
    return TR::CPU(processorDescription);
diff --git a/compiler/x/runtime/X86Runtime.hpp b/compiler/x/runtime/X86Runtime.hpp
index 7a59d52bbb4..cabde744c27 100644
--- a/compiler/x/runtime/X86Runtime.hpp
+++ b/compiler/x/runtime/X86Runtime.hpp
@@ -87,15 +87,39 @@ inline bool jitGetCPUID(TR_X86CPUIDBuffer* pBuffer)
       pBuffer->_featureFlags8      = CPUInfo[EBX];
       pBuffer->_featureFlags10     = CPUInfo[ECX];
 
-      // Check for XSAVE
+      bool disableAVX = true;
+      bool disableAVX512 = true;
+
+      // Check XCRO register for OS support of xmm/ymm/zmm
       if(pBuffer->_featureFlags2 & TR_OSXSAVE)
          {
-         static const bool disableAVX = feGetEnv("TR_DisableAVX") != NULL;
-         if(((6 & _xgetbv(0)) != 6) || disableAVX) // '6' = mask for XCR0[2:1]='11b' (XMM state and YMM state are enabled)
-            {
-            // Unset OSXSAVE if not enabled via CR0
-            pBuffer->_featureFlags2 &= ~TR_OSXSAVE;
-            }
+         // '6' = mask for XCR0[2:1]='11b' (XMM state and YMM state are enabled)
+         disableAVX = ((6 & _xgetbv(0)) != 6);
+         // 'e6' = (mask for XCR0[7:5]='111b' (Opmask, ZMM_Hi256, Hi16_ZMM) + XCR0[2:1]='11b' (XMM/YMM))
+         disableAVX512 = ((0xe6 & _xgetbv(0)) != 0xe6);
+         }
+
+      if(disableAVX)
+         {
+         // Unset AVX/AVX2 if not enabled via CR0 or otherwise disabled
+         pBuffer->_featureFlags2 &= ~TR_AVX;
+         pBuffer->_featureFlags8 &= ~TR_AVX2;
+         }
+
+      if (disableAVX512)
+         {
+         // Unset AVX-512 if not enabled via CR0 or otherwise disabled
+         // If other AVX-512 extensions are supported in the old cpuid API, they need to be disabled here
+         pBuffer->_featureFlags8 &= ~TR_AVX512F;
+         pBuffer->_featureFlags8 &= ~TR_AVX512VL;
+         pBuffer->_featureFlags8 &= ~TR_AVX512BW;
+         pBuffer->_featureFlags8 &= ~TR_AVX512CD;
+         pBuffer->_featureFlags8 &= ~TR_AVX512DQ;
+         pBuffer->_featureFlags8 &= ~TR_AVX512_BITALG;
+         pBuffer->_featureFlags8 &= ~TR_AVX512_VBMI;
+         pBuffer->_featureFlags8 &= ~TR_AVX512_VBMI2;
+         pBuffer->_featureFlags8 &= ~TR_AVX512_VNNI;
+         pBuffer->_featureFlags8 &= ~TR_AVX512_VPOPCNTDQ;
          }
 
       /* Mask out the bits the compiler does not care about.