From c7f95b7ca2fddbfdf344a29a00e7deca461c9293 Mon Sep 17 00:00:00 2001
From: Didzis Gosko <didzis@users.noreply.github.com>
Date: Mon, 15 Apr 2024 20:02:09 +0300
Subject: [PATCH] build : detect AVX512 in Makefile, add AVX512 option in CMake
 (#2043)

* make : add AVX512 detection to Makefile and CMakeLists.txt

* make : autodetect more AVX512 instruction subsets

* cmake : do not default to AVX512, must be enabled explicitly

* cmake : enable a set of AVX512 subsets, when AVX512 is turned on

* make : consolidate AVX512 subsets, add AVX512 VBMI

* cmake : revert to NO AVX512 setting, add settings for AVX512 VNNI and VBMI

* make : re-introduce AVX512VNNI back

* cmake : remove superfluous comment line
---
 CMakeLists.txt | 48 +++++++++++++++++++++++++++++++++++++-----------
 Makefile       | 18 ++++++++++++++++++
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a3c2639939e..2561dc6b65b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -59,10 +59,13 @@ option(WHISPER_BUILD_EXAMPLES         "whisper: build examples" ${WHISPER_STANDA
 
 option(WHISPER_SDL2                   "whisper: support for libSDL2" OFF)
 
-option(WHISPER_NO_AVX                 "whisper: disable AVX"  OFF)
-option(WHISPER_NO_AVX2                "whisper: disable AVX2" OFF)
-option(WHISPER_NO_FMA                 "whisper: disable FMA"  OFF)
-option(WHISPER_NO_F16C                "whisper: disable F16c" OFF)
+option(WHISPER_NO_AVX                 "whisper: disable AVX"         OFF)
+option(WHISPER_NO_AVX2                "whisper: disable AVX2"        OFF)
+option(WHISPER_NO_AVX512              "whisper: disable AVX512"      ON)
+option(WHISPER_NO_AVX512_VBMI         "whisper: disable AVX512-VBMI" ON)
+option(WHISPER_NO_AVX512_VNNI         "whisper: disable AVX512-VNNI" ON)
+option(WHISPER_NO_FMA                 "whisper: disable FMA"         OFF)
+option(WHISPER_NO_F16C                "whisper: disable F16c"        OFF)
 
 option(WHISPER_OPENVINO               "whisper: support for OpenVINO" OFF)
 
@@ -464,16 +467,30 @@ else()
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /utf-8")
         set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /utf-8")
         set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /utf-8")
-        if(NOT WHISPER_NO_AVX2)
+        if(NOT WHISPER_NO_AVX512)
+            set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX512")
+            set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /arch:AVX512")
+            set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /arch:AVX512")
+            # MSVC has no compile-time flags enabling specific
+            # AVX512 extensions, neither it defines the
+            # macros corresponding to the extensions.
+            # Do it manually.
+            if (NOT WHISPER_NO_AVX512_VBMI)
+                add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512VBMI__>)
+                add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512VBMI__>)
+            endif()
+            if (NOT WHISPER_NO_AVX512_VNNI)
+                add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512VNNI__>)
+                add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512VNNI__>)
+            endif()
+        elseif(NOT WHISPER_NO_AVX2)
             set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX2")
             set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /arch:AVX2")
             set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /arch:AVX2")
-        else()
-            if(NOT WHISPER_NO_AVX)
-                set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX")
-                set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /arch:AVX")
-                set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /arch:AVX")
-            endif()
+        elseif(NOT WHISPER_NO_AVX)
+            set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX")
+            set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /arch:AVX")
+            set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /arch:AVX")
         endif()
     else()
         if (EMSCRIPTEN)
@@ -486,6 +503,15 @@ else()
             if(NOT WHISPER_NO_AVX2)
                 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx2")
             endif()
+            if(NOT WHISPER_NO_AVX512)
+                set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw")
+            endif()
+            if(NOT WHISPER_NO_AVX512_VBMI)
+                set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx512vbmi")
+            endif()
+            if(NOT WHISPER_NO_AVX512_VNNI)
+                set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx512vnni")
+            endif()
             if(NOT WHISPER_NO_FMA)
                 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mfma")
             endif()
diff --git a/Makefile b/Makefile
index e255c52121a..b69628e0374 100644
--- a/Makefile
+++ b/Makefile
@@ -144,6 +144,24 @@ ifeq ($(UNAME_M),$(filter $(UNAME_M),x86_64 i686 amd64))
 			CXXFLAGS += -mavx2
 		endif
 
+		AVX512F_M := $(shell $(CPUINFO_CMD) | grep -iw 'AVX512F')
+		ifneq (,$(AVX512F_M))
+			CFLAGS   += -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw
+			CXXFLAGS += -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw
+		endif
+
+		AVX512VNNI_M := $(shell $(CPUINFO_CMD) | grep -iwE 'AVX512_VNNI|AVX512VNNI')
+		ifneq (,$(AVX512VNNI_M))
+			CFLAGS   += -mavx512vnni
+			CXXFLAGS += -mavx512vnni
+		endif
+
+		AVX512VBMI_M := $(shell $(CPUINFO_CMD) | grep -iw 'AVX512VBMI')
+		ifneq (,$(AVX512VBMI_M))
+			CFLAGS   += -mavx512vbmi
+			CXXFLAGS += -mavx512vbmi
+		endif
+
 		FMA_M := $(shell $(CPUINFO_CMD) | grep -iw 'FMA')
 		ifneq (,$(FMA_M))
 			CFLAGS   += -mfma