Skip to content

Commit 63d817c

Browse files
inspireMeNowinspireMeNow
committed
feat: add support for NCNN_ISA environment variable detection
Changes: - Implemented `get_isa_env` function to retrieve and parse ISA flags from the environment variable. Signed-off-by: Kaiyao Duan <[email protected]>
1 parent 6211838 commit 63d817c

File tree

1 file changed

+128
-16
lines changed

1 file changed

+128
-16
lines changed

‎src/cpu.cpp

+128-16
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#endif // __wasi__
2424
#include <stdio.h>
2525
#include <string.h>
26+
#include <stdlib.h>
2627

2728
#ifdef _OPENMP
2829
#if NCNN_SIMPLEOMP
@@ -1859,6 +1860,89 @@ static int detect_cpu_is_arm_a53_a55()
18591860
#endif // __aarch64__
18601861
#endif // defined __ANDROID__ || defined __linux__
18611862

1863+
static int get_isa_env(const char* isa_flags)
1864+
{
1865+
const char* isa = getenv("NCNN_ISA");
1866+
1867+
if (!isa || strlen(isa) == 0)
1868+
{
1869+
return 0;
1870+
}
1871+
1872+
char* isa_copy = strdup(isa);
1873+
1874+
if (!isa_copy)
1875+
{
1876+
return 0;
1877+
}
1878+
1879+
char* token = strtok(isa_copy, " ,");
1880+
1881+
while (token != NULL)
1882+
{
1883+
if (strcmp(token, isa_flags) == 0)
1884+
{
1885+
if (isa_flags[0] == '+')
1886+
return false;
1887+
if (isa_flags[0] == '-')
1888+
{
1889+
memmove(token, token + 1, strlen(token));
1890+
fprintf(stderr, "warning: %s disabled via environment variable!\n", token);
1891+
return 1;
1892+
}
1893+
}
1894+
token = strtok(NULL, " ,");
1895+
}
1896+
1897+
free(isa_copy);
1898+
return 0;
1899+
}
1900+
1901+
#if (__aarch64__ || __arm__)
1902+
static int is_cpu_arm_cpuid_disabled = get_isa_env("-cpuid");
1903+
static int is_cpu_arm_asimdhp_disabled = get_isa_env("-asimdhp");
1904+
static int is_cpu_arm_asimddp_disabled = get_isa_env("-asimddp");
1905+
static int is_cpu_arm_asimdfhm_disabled = get_isa_env("-asimdfhm");
1906+
static int is_cpu_arm_bf16_disabled = get_isa_env("-bf16");
1907+
static int is_cpu_arm_i8mm_disabled = get_isa_env("-i8mm");
1908+
static int is_cpu_arm_sve_disabled = get_isa_env("-sve");
1909+
static int is_cpu_arm_sve2_disabled = get_isa_env("-sve2");
1910+
static int is_cpu_arm_svebf16_disabled = get_isa_env("-svebf16");
1911+
static int is_cpu_arm_svei8mm_disabled = get_isa_env("-svei8mm");
1912+
static int is_cpu_arm_svef32mm_disabled = get_isa_env("-svef32mm");
1913+
static int is_cpu_arm_edsp_disabled = get_isa_env("-edsp");
1914+
static int is_cpu_arm_vfpv4_disabled = get_isa_env("-vfpv4");
1915+
static int is_cpu_arm_neon_disabled = get_isa_env("-neon");
1916+
#endif
1917+
1918+
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
1919+
static int is_cpu_x86_avx_disabled = get_isa_env("-avx");
1920+
static int is_cpu_x86_fma_disabled = get_isa_env("-fma");
1921+
static int is_cpu_x86_xop_disabled = get_isa_env("-xop");
1922+
static int is_cpu_x86_f16c_disabled = get_isa_env("-f16c");
1923+
static int is_cpu_x86_avx2_disabled = get_isa_env("-avx2");
1924+
static int is_cpu_x86_avx_vnni_disabled = get_isa_env("-avx_vnni");
1925+
static int is_cpu_x86_avx512_disabled = get_isa_env("-avx512");
1926+
static int is_cpu_x86_avx512_vnni_disabled = get_isa_env("-avx512_vnni");
1927+
static int is_cpu_x86_avx512_bf16_disabled = get_isa_env("-avx512_bf16");
1928+
static int is_cpu_x86_avx512_fp16_disabled = get_isa_env("-avx512_fp16");
1929+
#endif
1930+
1931+
#if __loongarch64
1932+
static int is_cpu_loongarch_lsx_disabled = get_isa_env("-lsx");
1933+
static int is_cpu_loongarch_lasx_disabled = get_isa_env("-lasx");
1934+
#endif
1935+
1936+
#if __mips__
1937+
static int is_cpu_mips_msa_disabled = get_isa_env("-msa");
1938+
static int is_cpu_loongson_mmi_disabled = get_isa_env("-mmi");
1939+
#endif
1940+
1941+
#if __riscv
1942+
static int is_cpu_riscv_v_disabled = get_isa_env("-rvv");
1943+
static int is_cpu_riscv_zfh_disabled = get_isa_env("-zfh");
1944+
#endif
1945+
18621946
// the initialization
18631947
static void initialize_global_cpu_info()
18641948
{
@@ -2087,6 +2171,8 @@ int cpu_support_arm_edsp()
20872171
{
20882172
try_initialize_global_cpu_info();
20892173
#if __arm__ && !__aarch64__
2174+
if (is_cpu_arm_edsp_disabled)
2175+
return 0;
20902176
#if defined _WIN32
20912177
return g_cpu_support_arm_edsp;
20922178
#elif defined __ANDROID__ || defined __linux__
@@ -2107,6 +2193,8 @@ int cpu_support_arm_neon()
21072193
#if __aarch64__
21082194
return 1;
21092195
#elif __arm__
2196+
if (is_cpu_arm_neon_disabled)
2197+
return 0;
21102198
#if defined _WIN32
21112199
return g_cpu_support_arm_neon;
21122200
#elif defined __ANDROID__ || defined __linux__
@@ -2127,6 +2215,8 @@ int cpu_support_arm_vfpv4()
21272215
#if __aarch64__
21282216
return 1;
21292217
#elif __arm__
2218+
if (is_cpu_arm_vfpv4_disabled)
2219+
return 0;
21302220
#if defined _WIN32
21312221
return g_cpu_support_arm_vfpv4;
21322222
#elif defined __ANDROID__ || defined __linux__
@@ -2145,6 +2235,8 @@ int cpu_support_arm_asimdhp()
21452235
{
21462236
try_initialize_global_cpu_info();
21472237
#if __aarch64__
2238+
if (is_cpu_arm_asimdhp_disabled)
2239+
return 0;
21482240
#if defined _WIN32
21492241
return g_cpu_support_arm_asimdhp;
21502242
#elif defined __ANDROID__ || defined __linux__
@@ -2173,6 +2265,8 @@ int cpu_support_arm_cpuid()
21732265
{
21742266
try_initialize_global_cpu_info();
21752267
#if __aarch64__
2268+
if (is_cpu_arm_cpuid_disabled)
2269+
return 0;
21762270
#if defined _WIN32
21772271
return g_cpu_support_arm_cpuid;
21782272
#elif defined __ANDROID__ || defined __linux__
@@ -2191,6 +2285,8 @@ int cpu_support_arm_asimddp()
21912285
{
21922286
try_initialize_global_cpu_info();
21932287
#if __aarch64__
2288+
if (is_cpu_arm_asimddp_disabled)
2289+
return 0;
21942290
#if defined _WIN32
21952291
return g_cpu_support_arm_asimddp;
21962292
#elif defined __ANDROID__ || defined __linux__
@@ -2217,6 +2313,8 @@ int cpu_support_arm_asimdfhm()
22172313
{
22182314
try_initialize_global_cpu_info();
22192315
#if __aarch64__
2316+
if (is_cpu_arm_asimdfhm_disabled)
2317+
return 0;
22202318
#if defined _WIN32
22212319
return g_cpu_support_arm_asimdfhm;
22222320
#elif defined __ANDROID__ || defined __linux__
@@ -2243,6 +2341,8 @@ int cpu_support_arm_bf16()
22432341
{
22442342
try_initialize_global_cpu_info();
22452343
#if __aarch64__
2344+
if (is_cpu_arm_bf16_disabled)
2345+
return 0;
22462346
#if defined _WIN32
22472347
return g_cpu_support_arm_bf16;
22482348
#elif defined __ANDROID__ || defined __linux__
@@ -2267,6 +2367,8 @@ int cpu_support_arm_i8mm()
22672367
{
22682368
try_initialize_global_cpu_info();
22692369
#if __aarch64__
2370+
if (is_cpu_arm_i8mm_disabled)
2371+
return 0;
22702372
#if defined _WIN32
22712373
return g_cpu_support_arm_i8mm;
22722374
#elif defined __ANDROID__ || defined __linux__
@@ -2291,6 +2393,8 @@ int cpu_support_arm_sve()
22912393
{
22922394
try_initialize_global_cpu_info();
22932395
#if __aarch64__
2396+
if (is_cpu_arm_sve_disabled)
2397+
return 0;
22942398
#if defined _WIN32
22952399
return g_cpu_support_arm_sve;
22962400
#elif defined __ANDROID__ || defined __linux__
@@ -2309,6 +2413,8 @@ int cpu_support_arm_sve2()
23092413
{
23102414
try_initialize_global_cpu_info();
23112415
#if __aarch64__
2416+
if (is_cpu_arm_sve2_disabled)
2417+
return 0;
23122418
#if defined _WIN32
23132419
return g_cpu_support_arm_sve2;
23142420
#elif defined __ANDROID__ || defined __linux__
@@ -2327,6 +2433,8 @@ int cpu_support_arm_svebf16()
23272433
{
23282434
try_initialize_global_cpu_info();
23292435
#if __aarch64__
2436+
if (is_cpu_arm_svebf16_disabled)
2437+
return 0;
23302438
#if defined _WIN32
23312439
return g_cpu_support_arm_svebf16;
23322440
#elif defined __ANDROID__ || defined __linux__
@@ -2345,6 +2453,8 @@ int cpu_support_arm_svei8mm()
23452453
{
23462454
try_initialize_global_cpu_info();
23472455
#if __aarch64__
2456+
if (is_cpu_arm_svei8mm_disabled)
2457+
return 0;
23482458
#if defined _WIN32
23492459
return g_cpu_support_arm_svei8mm;
23502460
#elif defined __ANDROID__ || defined __linux__
@@ -2363,6 +2473,8 @@ int cpu_support_arm_svef32mm()
23632473
{
23642474
try_initialize_global_cpu_info();
23652475
#if __aarch64__
2476+
if (is_cpu_arm_svef32mm_disabled)
2477+
return 0;
23662478
#if defined _WIN32
23672479
return g_cpu_support_arm_svef32mm;
23682480
#elif defined __ANDROID__ || defined __linux__
@@ -2381,7 +2493,7 @@ int cpu_support_x86_avx()
23812493
{
23822494
try_initialize_global_cpu_info();
23832495
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2384-
return g_cpu_support_x86_avx;
2496+
return g_cpu_support_x86_avx && !is_cpu_x86_avx_disabled;
23852497
#else
23862498
return 0;
23872499
#endif
@@ -2391,7 +2503,7 @@ int cpu_support_x86_fma()
23912503
{
23922504
try_initialize_global_cpu_info();
23932505
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2394-
return g_cpu_support_x86_fma;
2506+
return g_cpu_support_x86_fma && !is_cpu_x86_fma_disabled;
23952507
#else
23962508
return 0;
23972509
#endif
@@ -2401,7 +2513,7 @@ int cpu_support_x86_xop()
24012513
{
24022514
try_initialize_global_cpu_info();
24032515
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2404-
return g_cpu_support_x86_xop;
2516+
return g_cpu_support_x86_xop && !is_cpu_x86_xop_disabled;
24052517
#else
24062518
return 0;
24072519
#endif
@@ -2411,7 +2523,7 @@ int cpu_support_x86_f16c()
24112523
{
24122524
try_initialize_global_cpu_info();
24132525
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2414-
return g_cpu_support_x86_f16c;
2526+
return g_cpu_support_x86_f16c && !is_cpu_x86_f16c_disabled;
24152527
#else
24162528
return 0;
24172529
#endif
@@ -2421,7 +2533,7 @@ int cpu_support_x86_avx2()
24212533
{
24222534
try_initialize_global_cpu_info();
24232535
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2424-
return g_cpu_support_x86_avx2;
2536+
return g_cpu_support_x86_avx2 && !is_cpu_x86_avx2_disabled;
24252537
#else
24262538
return 0;
24272539
#endif
@@ -2431,7 +2543,7 @@ int cpu_support_x86_avx_vnni()
24312543
{
24322544
try_initialize_global_cpu_info();
24332545
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2434-
return g_cpu_support_x86_avx_vnni;
2546+
return g_cpu_support_x86_avx_vnni && !is_cpu_x86_avx_vnni_disabled;
24352547
#else
24362548
return 0;
24372549
#endif
@@ -2441,7 +2553,7 @@ int cpu_support_x86_avx512()
24412553
{
24422554
try_initialize_global_cpu_info();
24432555
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2444-
return g_cpu_support_x86_avx512;
2556+
return g_cpu_support_x86_avx512 && !is_cpu_x86_avx512_disabled;
24452557
#else
24462558
return 0;
24472559
#endif
@@ -2451,7 +2563,7 @@ int cpu_support_x86_avx512_vnni()
24512563
{
24522564
try_initialize_global_cpu_info();
24532565
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2454-
return g_cpu_support_x86_avx512_vnni;
2566+
return g_cpu_support_x86_avx512_vnni && !is_cpu_x86_avx512_vnni_disabled;
24552567
#else
24562568
return 0;
24572569
#endif
@@ -2461,7 +2573,7 @@ int cpu_support_x86_avx512_bf16()
24612573
{
24622574
try_initialize_global_cpu_info();
24632575
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2464-
return g_cpu_support_x86_avx512_bf16;
2576+
return g_cpu_support_x86_avx512_bf16 && !is_cpu_x86_avx512_bf16_disabled;
24652577
#else
24662578
return 0;
24672579
#endif
@@ -2471,7 +2583,7 @@ int cpu_support_x86_avx512_fp16()
24712583
{
24722584
try_initialize_global_cpu_info();
24732585
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
2474-
return g_cpu_support_x86_avx512_fp16;
2586+
return g_cpu_support_x86_avx512_fp16 && !is_cpu_x86_avx512_fp16_disabled;
24752587
#else
24762588
return 0;
24772589
#endif
@@ -2482,7 +2594,7 @@ int cpu_support_mips_msa()
24822594
try_initialize_global_cpu_info();
24832595
#if defined __ANDROID__ || defined __linux__
24842596
#if __mips__
2485-
return g_hwcaps & HWCAP_MIPS_MSA;
2597+
return (g_hwcaps & HWCAP_MIPS_MSA) && !is_cpu_mips_msa_disabled;
24862598
#else
24872599
return 0;
24882600
#endif
@@ -2496,7 +2608,7 @@ int cpu_support_loongarch_lsx()
24962608
try_initialize_global_cpu_info();
24972609
#if defined __ANDROID__ || defined __linux__
24982610
#if __loongarch64
2499-
return g_hwcaps & HWCAP_LOONGARCH_LSX;
2611+
return (g_hwcaps & HWCAP_LOONGARCH_LSX) && !is_cpu_loongarch_lsx_disabled;
25002612
#else
25012613
return 0;
25022614
#endif
@@ -2510,7 +2622,7 @@ int cpu_support_loongarch_lasx()
25102622
try_initialize_global_cpu_info();
25112623
#if defined __ANDROID__ || defined __linux__
25122624
#if __loongarch64
2513-
return g_hwcaps & HWCAP_LOONGARCH_LASX;
2625+
return (g_hwcaps & HWCAP_LOONGARCH_LASX) && !is_cpu_loongarch_lasx_disabled;
25142626
#else
25152627
return 0;
25162628
#endif
@@ -2524,7 +2636,7 @@ int cpu_support_loongson_mmi()
25242636
try_initialize_global_cpu_info();
25252637
#if defined __ANDROID__ || defined __linux__
25262638
#if __mips__
2527-
return g_hwcaps & HWCAP_LOONGSON_MMI;
2639+
return (g_hwcaps & HWCAP_LOONGSON_MMI) && !is_cpu_loongson_mmi_disabled;
25282640
#else
25292641
return 0;
25302642
#endif
@@ -2538,7 +2650,7 @@ int cpu_support_riscv_v()
25382650
try_initialize_global_cpu_info();
25392651
#if defined __ANDROID__ || defined __linux__
25402652
#if __riscv
2541-
return g_hwcaps & COMPAT_HWCAP_ISA_V;
2653+
return (g_hwcaps & COMPAT_HWCAP_ISA_V) && !is_cpu_riscv_v_disabled;
25422654
#else
25432655
return 0;
25442656
#endif
@@ -2554,7 +2666,7 @@ int cpu_support_riscv_zfh()
25542666
#if __riscv
25552667
// v + f does not imply zfh, but how to discover zfh properly ?
25562668
// upstream issue https://github.com/riscv/riscv-isa-manual/issues/414
2557-
return g_hwcaps & COMPAT_HWCAP_ISA_V && g_hwcaps & COMPAT_HWCAP_ISA_F;
2669+
return (g_hwcaps & COMPAT_HWCAP_ISA_V && g_hwcaps & COMPAT_HWCAP_ISA_F) && !is_cpu_riscv_zfh_disabled;
25582670
#else
25592671
return 0;
25602672
#endif

0 commit comments

Comments
 (0)