From 3f0aab6c91302ba8f99becd13e8304fcf6041087 Mon Sep 17 00:00:00 2001
From: inspiremenow <inspiremenow@murena.io>
Date: Fri, 12 Jul 2024 18:27:57 +0800
Subject: [PATCH] add environment variable NCNN_X86_AVX512 detect

Signed-off-by: Kaiyao Duan
---
 src/cpu.cpp | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/src/cpu.cpp b/src/cpu.cpp
index b1afbba3f651..1ec82d26d0e2 100644
--- a/src/cpu.cpp
+++ b/src/cpu.cpp
@@ -23,6 +23,7 @@
 #endif // __wasi__
 #include <stdio.h>
 #include <string.h>
+#include <stdlib.h>
 
 #ifdef _OPENMP
 #if NCNN_SIMPLEOMP
@@ -620,6 +621,20 @@ static int get_cpu_support_x86_avx_vnni()
 #endif
 }
 
+static int check_avx512_flag()
+{
+    const char* ncnn_x86_avx512 = "NCNN_X86_AVX512";
+    const char* value = getenv(ncnn_x86_avx512);
+
+    if (value)
+        if (strcmp(value, "0") == 0)
+        {
+            fprintf(stderr, "Warning: AVX512 support is disabled by environment variable NCNN_X86_AVX512=0\n");
+            return 0;
+        }
+    return 1;
+}
+
 static int get_cpu_support_x86_avx512()
 {
 #if __APPLE__
@@ -628,6 +643,9 @@ static int get_cpu_support_x86_avx512()
     unsigned int cpu_info[4] = {0};
     x86_cpuid(0, cpu_info);
 
+    if (!check_avx512_flag())
+        return 0;
+
     int nIds = cpu_info[0];
     if (nIds < 7)
         return 0;
@@ -658,6 +676,9 @@ static int get_cpu_support_x86_avx512_vnni()
     unsigned int cpu_info[4] = {0};
     x86_cpuid(0, cpu_info);
 
+    if (!check_avx512_flag())
+        return 0;
+
     int nIds = cpu_info[0];
     if (nIds < 7)
         return 0;
@@ -688,6 +709,9 @@ static int get_cpu_support_x86_avx512_bf16()
     unsigned int cpu_info[4] = {0};
     x86_cpuid(0, cpu_info);
 
+    if (!check_avx512_flag())
+        return 0;
+
     int nIds = cpu_info[0];
     if (nIds < 7)
         return 0;
@@ -714,6 +738,9 @@ static int get_cpu_support_x86_avx512_fp16()
     unsigned int cpu_info[4] = {0};
     x86_cpuid(0, cpu_info);
 
+    if (!check_avx512_flag())
+        return 0;
+
     int nIds = cpu_info[0];
     if (nIds < 7)
         return 0;