diff --git a/src/cpu.cpp b/src/cpu.cpp index b1afbba3f651..469601f6d248 100644 --- a/src/cpu.cpp +++ b/src/cpu.cpp @@ -620,6 +620,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 = std::getenv(ncnn_x86_avx512); + + if (value) + if (std::string(value) == "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 +642,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 +675,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 +708,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 +737,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;