Add flag to allow avx512vbmi to be disabled

PiperOrigin-RevId: 688331132
umadevimcw · Oct 22, 2024 · d385289 · d385289
1 parent b410980
commit d385289
Show file tree

Hide file tree

Showing 8 changed files with 75 additions and 18 deletions.
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -1244,6 +1244,18 @@ config_setting(
     define_values = {"xnn_enable_avx512skx": "false"},
 )
 
+# Enables usage of Intel AVX512VBMI (evex512) kernels.
+config_setting(
+    name = "xnn_enable_avx512vbmi_explicit_true",
+    define_values = {"xnn_enable_avx512vbmi": "true"},
+)
+
+# Disables usage of Intel AVX512VBMI (evex512) kernels.
+config_setting(
+    name = "xnn_enable_avx512vbmi_explicit_false",
+    define_values = {"xnn_enable_avx512vbmi": "false"},
+)
+
 # Enables usage of Intel AVX512VNNI (evex512) kernels.
 config_setting(
     name = "xnn_enable_avx512vnni_explicit_true",
@@ -1701,6 +1713,22 @@ alias(
     }),
 )
 
+selects.config_setting_group(
+    name = "avx512vbmi_enabled_by_default",
+    match_any = [
+        "//build_config:x86",
+    ],
+)
+
+alias(
+    name = "avx512vbmi_enabled",
+    actual = select({
+        ":xnn_enable_avx512vbmi_explicit_true": ":xnn_enable_avx512vbmi_explicit_true",
+        ":xnn_enable_avx512vbmi_explicit_false": ":xnn_enable_avx512vbmi_explicit_true",
+        "//conditions:default": ":avx512vbmi_enabled_by_default",
+    }),
+)
+
 selects.config_setting_group(
     name = "avx512vnni_enabled_by_default",
     match_any = [

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -201,6 +201,16 @@ ELSEIF(CMAKE_C_COMPILER_ID STREQUAL "Clang")
     SET(XNNPACK_ENABLE_AVX512SKX OFF)
   ENDIF()
 ENDIF()
+OPTION(XNNPACK_ENABLE_AVX512VBMI "Build XNNPACK with AVX512VBMI micro-kernels" ON)
+IF(CMAKE_C_COMPILER_ID STREQUAL "GNU")
+  IF(CMAKE_C_COMPILER_VERSION VERSION_LESS "8")
+    SET(XNNPACK_ENABLE_AVX512VBMI OFF)
+  ENDIF()
+ELSEIF(CMAKE_C_COMPILER_ID STREQUAL "Clang")
+  IF(CMAKE_C_COMPILER_VERSION VERSION_LESS "6")
+    SET(XNNPACK_ENABLE_AVX512VBMI OFF)
+  ENDIF()
+ENDIF()
 OPTION(XNNPACK_ENABLE_AVX512VNNI "Build XNNPACK with AVX512VNNI micro-kernels" ON)
 IF(CMAKE_C_COMPILER_ID STREQUAL "GNU")
   IF(CMAKE_C_COMPILER_VERSION VERSION_LESS "8")
@@ -278,6 +288,7 @@ ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX256VNNI=$<BOOL:${XNNPACK_ENABLE_AVX256VNN
 ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX256VNNIGFNI=$<BOOL:${XNNPACK_ENABLE_AVX256VNNIGFNI}>")
 ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX512F=$<BOOL:${XNNPACK_ENABLE_AVX512F}>")
 ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX512SKX=$<BOOL:${XNNPACK_ENABLE_AVX512SKX}>")
+ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX512VBMI=$<BOOL:${XNNPACK_ENABLE_AVX512VBMI}>")
 ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX512VNNI=$<BOOL:${XNNPACK_ENABLE_AVX512VNNI}>")
 ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX512VNNIGFNI=$<BOOL:${XNNPACK_ENABLE_AVX512VNNIGFNI}>")
 ADD_COMPILE_DEFINITIONS("XNN_ENABLE_AVX512AMX=$<BOOL:${XNNPACK_ENABLE_AVX512AMX}>")
@@ -654,7 +665,6 @@ IF(XNNPACK_TARGET_PROCESSOR MATCHES "^x86(_64)?$")
   LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_F16C_MICROKERNEL_SRCS})
   LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_FMA3_MICROKERNEL_SRCS})
   LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_AVX2_MICROKERNEL_SRCS})
-  LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_AVX512VBMI_MICROKERNEL_SRCS})
   IF(XNNPACK_ENABLE_AVX512AMX)
     LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_AVX512AMX_MICROKERNEL_SRCS})
   ENDIF()
@@ -682,6 +692,9 @@ IF(XNNPACK_TARGET_PROCESSOR MATCHES "^x86(_64)?$")
   IF(XNNPACK_ENABLE_AVX512SKX)
     LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_AVX512SKX_MICROKERNEL_SRCS})
   ENDIF()
+  IF(XNNPACK_ENABLE_AVX512VBMI)
+    LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_AVX512VBMI_MICROKERNEL_SRCS})
+  ENDIF()
   IF(XNNPACK_ENABLE_AVX512VNNI)
     LIST(APPEND PROD_MICROKERNEL_SRCS ${PROD_AVX512VNNI_MICROKERNEL_SRCS})
   ENDIF()
@@ -699,7 +712,6 @@ IF(XNNPACK_TARGET_PROCESSOR MATCHES "^x86(_64)?$")
   LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_F16C_MICROKERNEL_SRCS})
   LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_FMA3_MICROKERNEL_SRCS})
   LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_AVX2_MICROKERNEL_SRCS})
-  LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_AVX512VBMI_MICROKERNEL_SRCS})
   IF(XNNPACK_ENABLE_AVX512AMX)
     LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_AVX512AMX_MICROKERNEL_SRCS})
   ENDIF()
@@ -727,6 +739,9 @@ IF(XNNPACK_TARGET_PROCESSOR MATCHES "^x86(_64)?$")
   IF(XNNPACK_ENABLE_AVX512SKX)
     LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_AVX512SKX_MICROKERNEL_SRCS})
   ENDIF()
+  IF(XNNPACK_ENABLE_AVX512VBMI)
+    LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_AVX512VBMI_MICROKERNEL_SRCS})
+  ENDIF()
   IF(XNNPACK_ENABLE_AVX512VNNI)
     LIST(APPEND NON_PROD_MICROKERNEL_SRCS ${NON_PROD_AVX512VNNI_MICROKERNEL_SRCS})
   ENDIF()

diff --git a/bench/x8-lut.cc b/bench/x8-lut.cc
@@ -74,7 +74,7 @@ static void x8_lut(
     ->UseRealTime();
 #endif  // XNN_ARCH_ARM64
 
-#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+#if XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
   BENCHMARK_CAPTURE(x8_lut, avx512vbmi_vpermx2b_u64,
                     xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u64,
                     benchmark::utils::CheckAVX512VBMI)
@@ -95,7 +95,9 @@ static void x8_lut(
                     benchmark::utils::CheckAVX512VBMI)
     ->Apply(benchmark::utils::UnaryElementwiseParameters<uint8_t, uint8_t>)
     ->UseRealTime();
+#endif  // XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
+#if XNN_ENABLE_AVX512SKX && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
   BENCHMARK_CAPTURE(x8_lut, avx512skx_vpshufb_u64,
                     xnn_x8_lut_ukernel__avx512skx_vpshufb_u64,
                     benchmark::utils::CheckAVX512SKX)
@@ -116,7 +118,9 @@ static void x8_lut(
                     benchmark::utils::CheckAVX512SKX)
     ->Apply(benchmark::utils::UnaryElementwiseParameters<uint8_t, uint8_t>)
     ->UseRealTime();
+#endif  // XNN_ENABLE_AVX512SKX && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
+#if XNN_ARCH_X86 || XNN_ARCH_X86_64
   BENCHMARK_CAPTURE(x8_lut, avx2_u32,
                     xnn_x8_lut_ukernel__avx2_u32,
                     benchmark::utils::CheckAVX2)

diff --git a/build_params.bzl b/build_params.bzl
@@ -119,6 +119,10 @@ def xnnpack_configurable_defines():
         ":avx512skx_enabled",
         ["XNN_ENABLE_AVX512SKX=1"],
         ["XNN_ENABLE_AVX512SKX=0"],
+    ) + xnnpack_select_if(
+        ":avx512vbmi_enabled",
+        ["XNN_ENABLE_AVX512VBMI=1"],
+        ["XNN_ENABLE_AVX512VBMI=0"],
     ) + xnnpack_select_if(
         ":avx512vnni_enabled",
         ["XNN_ENABLE_AVX512VNNI=1"],
@@ -645,7 +649,7 @@ XNNPACK_PARAMS_FOR_ARCH = {
         msys_copts = ["-fno-asynchronous-unwind-tables"],
     ),
     "avx512vbmi": _create_params(
-        cond = "//build_config:x86",
+        cond = "//:avx512vbmi_enabled",
         gcc_x86_copts = [
             "-mf16c",
             "-mfma",

diff --git a/src/configs/hardware-config.c b/src/configs/hardware-config.c
@@ -126,7 +126,11 @@ static void init_hardware_config(void) {
 #else
     hardware_config.use_x86_avx512skx = 0;
 #endif
+#if XNN_ENABLE_AVX512VBMI
     hardware_config.use_x86_avx512vbmi = hardware_config.use_x86_avx512skx && cpuinfo_has_x86_avx512vbmi();
+#else
+    hardware_config.use_x86_avx512vbmi = 0;
+#endif
 #if XNN_ENABLE_AVX512VNNI
     hardware_config.use_x86_avx512vnni = hardware_config.use_x86_avx512skx && cpuinfo_has_x86_avx512vnni();
 #else

diff --git a/src/configs/x8-lut-config.c b/src/configs/x8-lut-config.c
@@ -25,13 +25,14 @@ static void init_x8_lut_config(void) {
     const struct xnn_hardware_config* hardware_config = xnn_init_hardware_config();
     assert(hardware_config != NULL);
 
-    #if XNN_ENABLE_AVX512SKX
+    #if XNN_ENABLE_AVX256VBMI
+      if (!XNN_PLATFORM_MOBILE && hardware_config->use_x86_avx512vbmi) {
+        x8_lut_config.microkernel = xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u128;
+      } else
+    #endif
+    #if XNN_ENABLE_AVX256SKX
       if (!XNN_PLATFORM_MOBILE && hardware_config->use_x86_avx512skx) {
-        if (hardware_config->use_x86_avx512vbmi) {
-          x8_lut_config.microkernel = xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u128;
-        } else {
-          x8_lut_config.microkernel = xnn_x8_lut_ukernel__avx512skx_vpshufb_u64;
-        }
+        x8_lut_config.microkernel = xnn_x8_lut_ukernel__avx512skx_vpshufb_u64;
       } else
     #endif
     if (hardware_config->use_x86_avx2) {

diff --git a/test/x8-lut.cc b/test/x8-lut.cc
@@ -1040,7 +1040,7 @@ TEST(X8_LUT__SCALAR_U16, inplace) {
 #endif  // XNN_ENABLE_AVX512SKX && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
 
-#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+#if XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
   TEST(X8_LUT__AVX512VBMI_VPERMX2B_U64, batch_eq_64) {
     TEST_REQUIRES_X86_AVX512VBMI;
     LUTMicrokernelTester()
@@ -1084,10 +1084,10 @@ TEST(X8_LUT__SCALAR_U16, inplace) {
         .Test(xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u64);
     }
   }
-#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
+#endif  // XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
 
-#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+#if XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
   TEST(X8_LUT__AVX512VBMI_VPERMX2B_U128, batch_eq_128) {
     TEST_REQUIRES_X86_AVX512VBMI;
     LUTMicrokernelTester()
@@ -1131,10 +1131,10 @@ TEST(X8_LUT__SCALAR_U16, inplace) {
         .Test(xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u128);
     }
   }
-#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
+#endif  // XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
 
-#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+#if XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
   TEST(X8_LUT__AVX512VBMI_VPERMX2B_U192, batch_eq_192) {
     TEST_REQUIRES_X86_AVX512VBMI;
     LUTMicrokernelTester()
@@ -1178,10 +1178,10 @@ TEST(X8_LUT__SCALAR_U16, inplace) {
         .Test(xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u192);
     }
   }
-#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
+#endif  // XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
 
-#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+#if XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
   TEST(X8_LUT__AVX512VBMI_VPERMX2B_U256, batch_eq_256) {
     TEST_REQUIRES_X86_AVX512VBMI;
     LUTMicrokernelTester()
@@ -1225,7 +1225,7 @@ TEST(X8_LUT__SCALAR_U16, inplace) {
         .Test(xnn_x8_lut_ukernel__avx512vbmi_vpermx2b_u256);
     }
   }
-#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
+#endif  // XNN_ENABLE_AVX512VBMI && (XNN_ARCH_X86 || XNN_ARCH_X86_64)
 
 
 #if XNN_ARCH_WASMSIMD || XNN_ARCH_WASMRELAXEDSIMD

diff --git a/tools/xnncommon.py b/tools/xnncommon.py
@@ -57,6 +57,7 @@ def _remove_duplicate_newlines(text):
   "avx256vnnigfni": "XNN_ENABLE_AVX256VNNIGFNI",
   "avx512f": "XNN_ENABLE_AVX512F",
   "avx512skx": "XNN_ENABLE_AVX512SKX",
+  "avx512vbmi": "XNN_ENABLE_AVX512VBMI",
   "avx512vnni": "XNN_ENABLE_AVX512VNNI",
   "avx512vnnigfni": "XNN_ENABLE_AVX512VNNIGFNI",
   "avx512amx": "XNN_ENABLE_AVX512AMX",