diff --git a/src/cpu.cpp b/src/cpu.cpp index c9307619ce9..ee646375a96 100644 --- a/src/cpu.cpp +++ b/src/cpu.cpp @@ -23,6 +23,7 @@ #endif // __wasi__ #include #include +#include #ifdef _OPENMP #if NCNN_SIMPLEOMP @@ -1980,6 +1981,91 @@ static int detect_cpu_is_arm_a53_a55() #endif // __aarch64__ #endif // defined __ANDROID__ || defined __linux__ +static int get_isa_env(const char* isa_flags) +{ + const char* isa = getenv("NCNN_ISA"); + + if (!isa || strlen(isa) == 0) + { + return 0; + } + + char* isa_copy = strdup(isa); + + if (!isa_copy) + { + return 0; + } + + char* token = strtok(isa_copy, " ,"); + + while (token != NULL) + { + if (strcmp(token + 1, isa_flags) == 0) + { + if (token[0] == '+') + { + fprintf(stderr, "warning: %s enabled via environment variable!\n", isa_flags); + return 0; + } + if (token[0] == '-') + { + fprintf(stderr, "warning: %s disabled via environment variable!\n", isa_flags); + return 1; + } + } + token = strtok(NULL, " ,"); + } + + free(isa_copy); + return 0; +} + +#if (__aarch64__ || __arm__) +static int is_cpu_arm_cpuid_disabled = get_isa_env("cpuid"); +static int is_cpu_arm_asimdhp_disabled = get_isa_env("asimdhp"); +static int is_cpu_arm_asimddp_disabled = get_isa_env("asimddp"); +static int is_cpu_arm_asimdfhm_disabled = get_isa_env("asimdfhm"); +static int is_cpu_arm_bf16_disabled = get_isa_env("bf16"); +static int is_cpu_arm_i8mm_disabled = get_isa_env("i8mm"); +static int is_cpu_arm_sve_disabled = get_isa_env("sve"); +static int is_cpu_arm_sve2_disabled = get_isa_env("sve2"); +static int is_cpu_arm_svebf16_disabled = get_isa_env("svebf16"); +static int is_cpu_arm_svei8mm_disabled = get_isa_env("svei8mm"); +static int is_cpu_arm_svef32mm_disabled = get_isa_env("svef32mm"); +static int is_cpu_arm_edsp_disabled = get_isa_env("edsp"); +static int is_cpu_arm_vfpv4_disabled = get_isa_env("vfpv4"); +static int is_cpu_arm_neon_disabled = get_isa_env("neon"); +#endif + +#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) +static int is_cpu_x86_avx_disabled = get_isa_env("avx"); +static int is_cpu_x86_fma_disabled = get_isa_env("fma"); +static int is_cpu_x86_xop_disabled = get_isa_env("xop"); +static int is_cpu_x86_f16c_disabled = get_isa_env("f16c"); +static int is_cpu_x86_avx2_disabled = get_isa_env("avx2"); +static int is_cpu_x86_avx_vnni_disabled = get_isa_env("avx_vnni"); +static int is_cpu_x86_avx512_disabled = get_isa_env("avx512"); +static int is_cpu_x86_avx512_vnni_disabled = get_isa_env("avx512_vnni"); +static int is_cpu_x86_avx512_bf16_disabled = get_isa_env("avx512_bf16"); +static int is_cpu_x86_avx512_fp16_disabled = get_isa_env("avx512_fp16"); +#endif + +#if __loongarch64 +static int is_cpu_loongarch_lsx_disabled = get_isa_env("lsx"); +static int is_cpu_loongarch_lasx_disabled = get_isa_env("lasx"); +#endif + +#if __mips__ +static int is_cpu_mips_msa_disabled = get_isa_env("msa"); +static int is_cpu_loongson_mmi_disabled = get_isa_env("mmi"); +#endif + +#if __riscv +static int is_cpu_riscv_v_disabled = get_isa_env("rvv"); +static int is_cpu_riscv_zfh_disabled = get_isa_env("zfh"); +#endif + // the initialization static void initialize_global_cpu_info() { @@ -2211,6 +2297,8 @@ int cpu_support_arm_edsp() { try_initialize_global_cpu_info(); #if __arm__ && !__aarch64__ + if (is_cpu_arm_edsp_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_edsp; #elif defined __ANDROID__ || defined __linux__ @@ -2231,6 +2319,8 @@ int cpu_support_arm_neon() #if __aarch64__ return 1; #elif __arm__ + if (is_cpu_arm_neon_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_neon; #elif defined __ANDROID__ || defined __linux__ @@ -2251,6 +2341,8 @@ int cpu_support_arm_vfpv4() #if __aarch64__ return 1; #elif __arm__ + if (is_cpu_arm_vfpv4_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_vfpv4; #elif defined __ANDROID__ || defined __linux__ @@ -2269,6 +2361,8 @@ int cpu_support_arm_asimdhp() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_asimdhp_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_asimdhp; #elif defined __ANDROID__ || defined __linux__ @@ -2297,6 +2391,8 @@ int cpu_support_arm_cpuid() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_cpuid_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_cpuid; #elif defined __ANDROID__ || defined __linux__ @@ -2315,6 +2411,8 @@ int cpu_support_arm_asimddp() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_asimddp_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_asimddp; #elif defined __ANDROID__ || defined __linux__ @@ -2341,6 +2439,8 @@ int cpu_support_arm_asimdfhm() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_asimdfhm_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_asimdfhm; #elif defined __ANDROID__ || defined __linux__ @@ -2367,6 +2467,8 @@ int cpu_support_arm_bf16() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_bf16_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_bf16; #elif defined __ANDROID__ || defined __linux__ @@ -2391,6 +2493,8 @@ int cpu_support_arm_i8mm() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_i8mm_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_i8mm; #elif defined __ANDROID__ || defined __linux__ @@ -2415,6 +2519,8 @@ int cpu_support_arm_sve() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_sve_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_sve; #elif defined __ANDROID__ || defined __linux__ @@ -2433,6 +2539,8 @@ int cpu_support_arm_sve2() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_sve2_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_sve2; #elif defined __ANDROID__ || defined __linux__ @@ -2451,6 +2559,8 @@ int cpu_support_arm_svebf16() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_svebf16_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_svebf16; #elif defined __ANDROID__ || defined __linux__ @@ -2469,6 +2579,8 @@ int cpu_support_arm_svei8mm() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_svei8mm_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_svei8mm; #elif defined __ANDROID__ || defined __linux__ @@ -2487,6 +2599,8 @@ int cpu_support_arm_svef32mm() { try_initialize_global_cpu_info(); #if __aarch64__ + if (is_cpu_arm_svef32mm_disabled) + return 0; #if defined _WIN32 return g_cpu_support_arm_svef32mm; #elif defined __ANDROID__ || defined __linux__ @@ -2505,7 +2619,7 @@ int cpu_support_x86_avx() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx; + return g_cpu_support_x86_avx && !is_cpu_x86_avx_disabled && !is_cpu_x86_avx2_disabled; #else return 0; #endif @@ -2515,7 +2629,7 @@ int cpu_support_x86_fma() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_fma; + return g_cpu_support_x86_fma && !is_cpu_x86_fma_disabled && !is_cpu_x86_avx2_disabled && !is_cpu_x86_avx_disabled; #else return 0; #endif @@ -2525,7 +2639,7 @@ int cpu_support_x86_xop() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_xop; + return g_cpu_support_x86_xop && !is_cpu_x86_xop_disabled; #else return 0; #endif @@ -2535,7 +2649,7 @@ int cpu_support_x86_f16c() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_f16c; + return g_cpu_support_x86_f16c && !is_cpu_x86_f16c_disabled; #else return 0; #endif @@ -2545,7 +2659,7 @@ int cpu_support_x86_avx2() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx2; + return g_cpu_support_x86_avx2 && !is_cpu_x86_avx2_disabled && !is_cpu_x86_avx_disabled; #else return 0; #endif @@ -2555,7 +2669,7 @@ int cpu_support_x86_avx_vnni() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx_vnni; + return g_cpu_support_x86_avx_vnni && !is_cpu_x86_avx_vnni_disabled && !is_cpu_x86_avx_disabled && !is_cpu_x86_avx2_disabled; #else return 0; #endif @@ -2595,7 +2709,7 @@ int cpu_support_x86_avx512() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx512; + return g_cpu_support_x86_avx512 && !is_cpu_x86_avx512_disabled && !is_cpu_x86_avx_disabled && !is_cpu_x86_avx2_disabled; #else return 0; #endif @@ -2605,7 +2719,7 @@ int cpu_support_x86_avx512_vnni() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx512_vnni; + return g_cpu_support_x86_avx512_vnni && !is_cpu_x86_avx512_vnni_disabled && !is_cpu_x86_avx512_disabled && !is_cpu_x86_avx_disabled && !is_cpu_x86_avx2_disabled; #else return 0; #endif @@ -2615,7 +2729,7 @@ int cpu_support_x86_avx512_bf16() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx512_bf16; + return g_cpu_support_x86_avx512_bf16 && !is_cpu_x86_avx512_bf16_disabled && !is_cpu_x86_avx512_disabled && !is_cpu_x86_avx_disabled && !is_cpu_x86_avx2_disabled; #else return 0; #endif @@ -2625,7 +2739,7 @@ int cpu_support_x86_avx512_fp16() { try_initialize_global_cpu_info(); #if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) - return g_cpu_support_x86_avx512_fp16; + return g_cpu_support_x86_avx512_fp16 && !is_cpu_x86_avx512_fp16_disabled && !is_cpu_x86_avx512_disabled && !is_cpu_x86_avx_disabled && !is_cpu_x86_avx2_disabled; #else return 0; #endif @@ -2636,7 +2750,7 @@ int cpu_support_mips_msa() try_initialize_global_cpu_info(); #if defined __ANDROID__ || defined __linux__ #if __mips__ - return g_hwcaps & HWCAP_MIPS_MSA; + return (g_hwcaps & HWCAP_MIPS_MSA) && !is_cpu_mips_msa_disabled; #else return 0; #endif @@ -2650,7 +2764,7 @@ int cpu_support_loongarch_lsx() try_initialize_global_cpu_info(); #if defined __ANDROID__ || defined __linux__ #if __loongarch64 - return g_hwcaps & HWCAP_LOONGARCH_LSX; + return (g_hwcaps & HWCAP_LOONGARCH_LSX) && !is_cpu_loongarch_lsx_disabled; #else return 0; #endif @@ -2664,7 +2778,7 @@ int cpu_support_loongarch_lasx() try_initialize_global_cpu_info(); #if defined __ANDROID__ || defined __linux__ #if __loongarch64 - return g_hwcaps & HWCAP_LOONGARCH_LASX; + return (g_hwcaps & HWCAP_LOONGARCH_LASX) && !is_cpu_loongarch_lasx_disabled && !is_cpu_loongarch_lsx_disabled; #else return 0; #endif @@ -2678,7 +2792,7 @@ int cpu_support_loongson_mmi() try_initialize_global_cpu_info(); #if defined __ANDROID__ || defined __linux__ #if __mips__ - return g_hwcaps & HWCAP_LOONGSON_MMI; + return (g_hwcaps & HWCAP_LOONGSON_MMI) && !is_cpu_loongson_mmi_disabled; #else return 0; #endif @@ -2692,7 +2806,7 @@ int cpu_support_riscv_v() try_initialize_global_cpu_info(); #if defined __ANDROID__ || defined __linux__ #if __riscv - return g_hwcaps & COMPAT_HWCAP_ISA_V; + return (g_hwcaps & COMPAT_HWCAP_ISA_V) && !is_cpu_riscv_v_disabled; #else return 0; #endif @@ -2708,7 +2822,7 @@ int cpu_support_riscv_zfh() #if __riscv // v + f does not imply zfh, but how to discover zfh properly ? // upstream issue https://github.com/riscv/riscv-isa-manual/issues/414 - return g_hwcaps & COMPAT_HWCAP_ISA_V && g_hwcaps & COMPAT_HWCAP_ISA_F; + return (g_hwcaps & COMPAT_HWCAP_ISA_V && g_hwcaps & COMPAT_HWCAP_ISA_F) && !is_cpu_riscv_zfh_disabled; #else return 0; #endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e2ddc32a00d..54e778e35e7 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -61,6 +61,58 @@ endif() ncnn_add_test(c_api) ncnn_add_test(cpu) +set(ENABLE_ENV_DETECT_TEST ON CACHE BOOL "Enable the environment detection test") +if(ENABLE_ENV_DETECT_TEST) + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + if(DEFINED CMAKE_SYSTEM_PROCESSOR) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + set(ENV{NCNN_ISA} "-avx,-fma") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + set(ENV{NCNN_ISA} "-neon") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "arm") + set(ENV{NCNN_ISA} "-neon") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "loongarch64") + set(ENV{NCNN_ISA} "-lsx,-lasx") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "mips") + set(ENV{NCNN_ISA} "-msa") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "riscv") + set(ENV{NCNN_ISA} "-rvv") + endif() + endif() + + elseif(CMAKE_SYSTEM_NAME STREQUAL "Windows") + if(DEFINED CMAKE_SYSTEM_PROCESSOR) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "AMD64" OR CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + set(ENV{NCNN_ISA} "-avx,-fma") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "ARM64") + set(ENV{NCNN_ISA} "-neon") + endif() + endif() + + elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + if(DEFINED CMAKE_SYSTEM_PROCESSOR) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + set(ENV{NCNN_ISA} "-avx,-fma") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "arm64") + set(ENV{NCNN_ISA} "-neon") + endif() + endif() + + elseif(CMAKE_SYSTEM_NAME STREQUAL "Android") + if(DEFINED CMAKE_SYSTEM_PROCESSOR) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "arm64-v8a") + set(ENV{NCNN_ISA} "-neon") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "armeabi-v7a") + set(ENV{NCNN_ISA} "-neon") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + set(ENV{NCNN_ISA} "-avx,-fma") + endif() + endif() + endif() + + ncnn_add_test(env_detect) +endif() + if(NCNN_VULKAN) ncnn_add_test(command) endif() diff --git a/tests/test_env_detect.cpp b/tests/test_env_detect.cpp new file mode 100644 index 00000000000..3e7965df67b --- /dev/null +++ b/tests/test_env_detect.cpp @@ -0,0 +1,92 @@ +#include "cpu.h" +#include +#include +#include + +// Improved version of is_option_disabled function that doesn't modify the original string +bool is_option_disabled(const char* options, const char* option) +{ + char* options_copy = strdup(options); + char* token = strtok(options_copy, ","); + bool disabled = false; + + while (token) + { + if (strcmp(token, option) == 0) + { + disabled = true; + break; + } + token = strtok(NULL, ","); + } + + free(options_copy); + return disabled; +} + +// Helper function to check and report instruction set support +bool check_instruction_disabled(const char* options, const char* option, bool cpu_support, const char* instruction_name) +{ + if (is_option_disabled(options, option) && cpu_support) + { + fprintf(stderr, "Error: %s should be disabled but it is enabled!\n", instruction_name); + return true; + } + return false; +} + +int main() +{ + const char* ncnn_isa = getenv("NCNN_ISA"); + + // Check if NCNN_ISA is set to disable certain options + if (ncnn_isa) + { +#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64) + if (check_instruction_disabled(ncnn_isa, "-avx", ncnn::cpu_support_x86_avx(), "avx")) return 1; + if (check_instruction_disabled(ncnn_isa, "-xop", ncnn::cpu_support_x86_xop(), "xop")) return 1; + if (check_instruction_disabled(ncnn_isa, "-fma", ncnn::cpu_support_x86_fma(), "fma")) return 1; + if (check_instruction_disabled(ncnn_isa, "-f16c", ncnn::cpu_support_x86_f16c(), "f16c")) return 1; + if (check_instruction_disabled(ncnn_isa, "-avx2", ncnn::cpu_support_x86_avx2(), "avx2")) return 1; + if (check_instruction_disabled(ncnn_isa, "-avx512", ncnn::cpu_support_x86_avx512(), "avx512")) return 1; + if (check_instruction_disabled(ncnn_isa, "-avx_vnni", ncnn::cpu_support_x86_avx_vnni(), "avx_vnni")) return 1; + if (check_instruction_disabled(ncnn_isa, "-avx512_vnni", ncnn::cpu_support_x86_avx512_vnni(), "avx512_vnni")) return 1; + if (check_instruction_disabled(ncnn_isa, "-avx512_bf16", ncnn::cpu_support_x86_avx512_bf16(), "avx512_bf16")) return 1; + if (check_instruction_disabled(ncnn_isa, "-avx512_fp16", ncnn::cpu_support_x86_avx512_fp16(), "avx512_fp16")) return 1; +#endif + +#if defined(__aarch64__) || defined(__arm__) + if (check_instruction_disabled(ncnn_isa, "-cpuid", ncnn::cpu_support_arm_cpuid(), "cpuid")) return 1; + if (check_instruction_disabled(ncnn_isa, "-asimdhp", ncnn::cpu_support_arm_asimdhp(), "asimdhp")) return 1; + if (check_instruction_disabled(ncnn_isa, "-asimddp", ncnn::cpu_support_arm_asimddp(), "asimddp")) return 1; + if (check_instruction_disabled(ncnn_isa, "-asimdfhm", ncnn::cpu_support_arm_asimdfhm(), "asimdfhm")) return 1; + if (check_instruction_disabled(ncnn_isa, "-bf16", ncnn::cpu_support_arm_bf16(), "bf16")) return 1; + if (check_instruction_disabled(ncnn_isa, "-i8mm", ncnn::cpu_support_arm_i8mm(), "i8mm")) return 1; + if (check_instruction_disabled(ncnn_isa, "-sve", ncnn::cpu_support_arm_sve(), "sve")) return 1; + if (check_instruction_disabled(ncnn_isa, "-sve2", ncnn::cpu_support_arm_sve2(), "sve2")) return 1; + if (check_instruction_disabled(ncnn_isa, "-svebf16", ncnn::cpu_support_arm_svebf16(), "svebf16")) return 1; + if (check_instruction_disabled(ncnn_isa, "-svei8mm", ncnn::cpu_support_arm_svei8mm(), "svei8mm")) return 1; + if (check_instruction_disabled(ncnn_isa, "-svef32mm", ncnn::cpu_support_arm_svef32mm(), "svef32mm")) return 1; + if (check_instruction_disabled(ncnn_isa, "-edsp", ncnn::cpu_support_arm_edsp(), "edsp")) return 1; + if (check_instruction_disabled(ncnn_isa, "-vfpv4", ncnn::cpu_support_arm_vfpv4(), "vfpv4")) return 1; + if (check_instruction_disabled(ncnn_isa, "-neon", ncnn::cpu_support_arm_neon(), "neon")) return 1; +#endif + +#if defined(__loongarch64) + if (check_instruction_disabled(ncnn_isa, "-lsx", ncnn::cpu_support_loongarch_lsx(), "lsx")) return 1; + if (check_instruction_disabled(ncnn_isa, "-lasx", ncnn::cpu_support_loongarch_lasx(), "lasx")) return 1; +#endif + +#if defined(__mips__) + if (check_instruction_disabled(ncnn_isa, "-msa", ncnn::cpu_support_mips_msa(), "msa")) return 1; + if (check_instruction_disabled(ncnn_isa, "-mmi", ncnn::cpu_support_loongson_mmi(), "mmi")) return 1; +#endif + +#if defined(__riscv) + if (check_instruction_disabled(ncnn_isa, "-rvv", ncnn::cpu_support_riscv_v(), "rvv")) return 1; + if (check_instruction_disabled(ncnn_isa, "-zfh", ncnn::cpu_support_riscv_zfh(), "zfh")) return 1; +#endif + } + + return 0; +}