Skip to content

Commit c24a3cf

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 c24a3cf

File tree

1 file changed

+91
-13
lines changed

1 file changed

+91
-13
lines changed

‎src/cpu.cpp

+91-13
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,55 @@ 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)
1868+
{
1869+
return false;
1870+
}
1871+
1872+
char* isa_copy = strdup(isa);
1873+
char* token = strtok(isa_copy, " ,");
1874+
1875+
while (token != NULL)
1876+
{
1877+
if (strcmp(token, isa_flags) == 0)
1878+
{
1879+
if (isa_flags[0] == '+')
1880+
return false;
1881+
if (isa_flags[0] == '-')
1882+
{
1883+
memmove(token, token + 1, strlen(token));
1884+
fprintf(stderr, "warning: %s disabled via environment variable!\n", token);
1885+
return true;
1886+
}
1887+
}
1888+
token = strtok(NULL, " ,");
1889+
}
1890+
1891+
free(isa_copy);
1892+
return false;
1893+
}
1894+
1895+
#if (__aarch64__ || __arm__)
1896+
static int is_cpu_arm_cpuid_disabled = get_isa_env("-cpuid");
1897+
static int is_cpu_arm_asimdhp_disabled = get_isa_env("-asimdhp");
1898+
static int is_cpu_arm_asimddp_disabled = get_isa_env("-asimddp");
1899+
static int is_cpu_arm_asimdfhm_disabled = get_isa_env("-asimdfhm");
1900+
static int is_cpu_arm_bf16_disabled = get_isa_env("-bf16");
1901+
static int is_cpu_arm_i8mm_disabled = get_isa_env("-i8mm");
1902+
static int is_cpu_arm_sve_disabled = get_isa_env("-sve");
1903+
static int is_cpu_arm_sve2_disabled = get_isa_env("-sve2");
1904+
static int is_cpu_arm_svebf16_disabled = get_isa_env("-svebf16");
1905+
static int is_cpu_arm_svei8mm_disabled = get_isa_env("-svei8mm");
1906+
static int is_cpu_arm_svef32mm_disabled = get_isa_env("-svef32mm");
1907+
static int is_cpu_arm_edsp_disabled = get_isa_env("-edsp");
1908+
static int is_cpu_arm_vfpv4_disabled = get_isa_env("-vfpv4");
1909+
static int is_cpu_arm_neon_disabled = get_isa_env("-neon");
1910+
#endif
1911+
18621912
// the initialization
18631913
static void initialize_global_cpu_info()
18641914
{
@@ -1909,16 +1959,16 @@ static void initialize_global_cpu_info()
19091959
#endif
19101960

19111961
#if defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
1912-
g_cpu_support_x86_avx = get_cpu_support_x86_avx();
1913-
g_cpu_support_x86_fma = get_cpu_support_x86_fma();
1914-
g_cpu_support_x86_xop = get_cpu_support_x86_xop();
1915-
g_cpu_support_x86_f16c = get_cpu_support_x86_f16c();
1916-
g_cpu_support_x86_avx2 = get_cpu_support_x86_avx2();
1917-
g_cpu_support_x86_avx_vnni = get_cpu_support_x86_avx_vnni();
1918-
g_cpu_support_x86_avx512 = get_cpu_support_x86_avx512();
1919-
g_cpu_support_x86_avx512_vnni = get_cpu_support_x86_avx512_vnni();
1920-
g_cpu_support_x86_avx512_bf16 = get_cpu_support_x86_avx512_bf16();
1921-
g_cpu_support_x86_avx512_fp16 = get_cpu_support_x86_avx512_fp16();
1962+
g_cpu_support_x86_avx = get_cpu_support_x86_avx() && !(get_isa_env("-avx"));
1963+
g_cpu_support_x86_fma = get_cpu_support_x86_fma() && !(get_isa_env("-fma"));
1964+
g_cpu_support_x86_xop = get_cpu_support_x86_xop() && !(get_isa_env("-xop"));
1965+
g_cpu_support_x86_f16c = get_cpu_support_x86_f16c() && !(get_isa_env("-f16c"));
1966+
g_cpu_support_x86_avx2 = get_cpu_support_x86_avx2() && !(get_isa_env("-avx2"));
1967+
g_cpu_support_x86_avx_vnni = get_cpu_support_x86_avx_vnni() && !(get_isa_env("-avx_vnni"));
1968+
g_cpu_support_x86_avx512 = get_cpu_support_x86_avx512() && !(get_isa_env("-avx512"));
1969+
g_cpu_support_x86_avx512_vnni = get_cpu_support_x86_avx512_vnni() && !(get_isa_env("-avx512_vnni"));
1970+
g_cpu_support_x86_avx512_bf16 = get_cpu_support_x86_avx512_bf16() && !(get_isa_env("-bf16"));
1971+
g_cpu_support_x86_avx512_fp16 = get_cpu_support_x86_avx512_fp16() && !(get_isa_env("-fp16"));
19221972
#endif // defined(__i386__) || defined(__x86_64__) || defined(_M_IX86) || defined(_M_X64)
19231973

19241974
g_cpu_level2_cachesize = get_cpu_level2_cachesize();
@@ -2087,6 +2137,8 @@ int cpu_support_arm_edsp()
20872137
{
20882138
try_initialize_global_cpu_info();
20892139
#if __arm__ && !__aarch64__
2140+
if (is_cpu_arm_edsp_disabled)
2141+
return 0;
20902142
#if defined _WIN32
20912143
return g_cpu_support_arm_edsp;
20922144
#elif defined __ANDROID__ || defined __linux__
@@ -2107,6 +2159,8 @@ int cpu_support_arm_neon()
21072159
#if __aarch64__
21082160
return 1;
21092161
#elif __arm__
2162+
if (is_cpu_arm_neon_disabled)
2163+
return 0;
21102164
#if defined _WIN32
21112165
return g_cpu_support_arm_neon;
21122166
#elif defined __ANDROID__ || defined __linux__
@@ -2127,6 +2181,8 @@ int cpu_support_arm_vfpv4()
21272181
#if __aarch64__
21282182
return 1;
21292183
#elif __arm__
2184+
if (is_cpu_arm_vfpv4_disabled)
2185+
return 0;
21302186
#if defined _WIN32
21312187
return g_cpu_support_arm_vfpv4;
21322188
#elif defined __ANDROID__ || defined __linux__
@@ -2145,6 +2201,8 @@ int cpu_support_arm_asimdhp()
21452201
{
21462202
try_initialize_global_cpu_info();
21472203
#if __aarch64__
2204+
if (is_cpu_arm_asimdhp_disabled)
2205+
return 0;
21482206
#if defined _WIN32
21492207
return g_cpu_support_arm_asimdhp;
21502208
#elif defined __ANDROID__ || defined __linux__
@@ -2173,6 +2231,8 @@ int cpu_support_arm_cpuid()
21732231
{
21742232
try_initialize_global_cpu_info();
21752233
#if __aarch64__
2234+
if (is_cpu_arm_cpuid_disabled)
2235+
return 0;
21762236
#if defined _WIN32
21772237
return g_cpu_support_arm_cpuid;
21782238
#elif defined __ANDROID__ || defined __linux__
@@ -2191,6 +2251,8 @@ int cpu_support_arm_asimddp()
21912251
{
21922252
try_initialize_global_cpu_info();
21932253
#if __aarch64__
2254+
if (is_cpu_arm_asimddp_disabled)
2255+
return 0;
21942256
#if defined _WIN32
21952257
return g_cpu_support_arm_asimddp;
21962258
#elif defined __ANDROID__ || defined __linux__
@@ -2217,6 +2279,8 @@ int cpu_support_arm_asimdfhm()
22172279
{
22182280
try_initialize_global_cpu_info();
22192281
#if __aarch64__
2282+
if (is_cpu_arm_asimdfhm_disabled)
2283+
return 0;
22202284
#if defined _WIN32
22212285
return g_cpu_support_arm_asimdfhm;
22222286
#elif defined __ANDROID__ || defined __linux__
@@ -2243,6 +2307,8 @@ int cpu_support_arm_bf16()
22432307
{
22442308
try_initialize_global_cpu_info();
22452309
#if __aarch64__
2310+
if (is_cpu_arm_bf16_disabled)
2311+
return 0;
22462312
#if defined _WIN32
22472313
return g_cpu_support_arm_bf16;
22482314
#elif defined __ANDROID__ || defined __linux__
@@ -2267,6 +2333,8 @@ int cpu_support_arm_i8mm()
22672333
{
22682334
try_initialize_global_cpu_info();
22692335
#if __aarch64__
2336+
if (is_cpu_arm_i8mm_disabled)
2337+
return 0;
22702338
#if defined _WIN32
22712339
return g_cpu_support_arm_i8mm;
22722340
#elif defined __ANDROID__ || defined __linux__
@@ -2291,6 +2359,8 @@ int cpu_support_arm_sve()
22912359
{
22922360
try_initialize_global_cpu_info();
22932361
#if __aarch64__
2362+
if (is_cpu_arm_sve_disabled)
2363+
return 0;
22942364
#if defined _WIN32
22952365
return g_cpu_support_arm_sve;
22962366
#elif defined __ANDROID__ || defined __linux__
@@ -2309,6 +2379,8 @@ int cpu_support_arm_sve2()
23092379
{
23102380
try_initialize_global_cpu_info();
23112381
#if __aarch64__
2382+
if (is_cpu_arm_sve2_disabled)
2383+
return 0;
23122384
#if defined _WIN32
23132385
return g_cpu_support_arm_sve2;
23142386
#elif defined __ANDROID__ || defined __linux__
@@ -2327,6 +2399,8 @@ int cpu_support_arm_svebf16()
23272399
{
23282400
try_initialize_global_cpu_info();
23292401
#if __aarch64__
2402+
if (is_cpu_arm_svebf16_disabled)
2403+
return 0;
23302404
#if defined _WIN32
23312405
return g_cpu_support_arm_svebf16;
23322406
#elif defined __ANDROID__ || defined __linux__
@@ -2345,6 +2419,8 @@ int cpu_support_arm_svei8mm()
23452419
{
23462420
try_initialize_global_cpu_info();
23472421
#if __aarch64__
2422+
if (is_cpu_arm_svei8mm_disabled)
2423+
return 0;
23482424
#if defined _WIN32
23492425
return g_cpu_support_arm_svei8mm;
23502426
#elif defined __ANDROID__ || defined __linux__
@@ -2363,6 +2439,8 @@ int cpu_support_arm_svef32mm()
23632439
{
23642440
try_initialize_global_cpu_info();
23652441
#if __aarch64__
2442+
if (is_cpu_arm_svef32mm_disabled)
2443+
return 0;
23662444
#if defined _WIN32
23672445
return g_cpu_support_arm_svef32mm;
23682446
#elif defined __ANDROID__ || defined __linux__
@@ -2482,7 +2560,7 @@ int cpu_support_mips_msa()
24822560
try_initialize_global_cpu_info();
24832561
#if defined __ANDROID__ || defined __linux__
24842562
#if __mips__
2485-
return g_hwcaps & HWCAP_MIPS_MSA;
2563+
return (g_hwcaps & HWCAP_MIPS_MSA) && !get_isa_env(("-msa"));
24862564
#else
24872565
return 0;
24882566
#endif
@@ -2538,7 +2616,7 @@ int cpu_support_riscv_v()
25382616
try_initialize_global_cpu_info();
25392617
#if defined __ANDROID__ || defined __linux__
25402618
#if __riscv
2541-
return g_hwcaps & COMPAT_HWCAP_ISA_V;
2619+
return (g_hwcaps & COMPAT_HWCAP_ISA_V) && !get_isa_env(("-rvv"));
25422620
#else
25432621
return 0;
25442622
#endif
@@ -2554,7 +2632,7 @@ int cpu_support_riscv_zfh()
25542632
#if __riscv
25552633
// v + f does not imply zfh, but how to discover zfh properly ?
25562634
// 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;
2635+
return (g_hwcaps & COMPAT_HWCAP_ISA_V && g_hwcaps & COMPAT_HWCAP_ISA_F) && !get_isa_env(("-rvzfh"));
25582636
#else
25592637
return 0;
25602638
#endif

0 commit comments

Comments
 (0)