diff options
Diffstat (limited to 'intgemm/intgemm.cc')
| -rw-r--r-- | intgemm/intgemm.cc | 15 |
1 files changed, 11 insertions, 4 deletions
diff --git a/intgemm/intgemm.cc b/intgemm/intgemm.cc index 7ccd21b..d6c26b9 100644 --- a/intgemm/intgemm.cc +++ b/intgemm/intgemm.cc @@ -87,16 +87,17 @@ CPUType RealCPUID() { #endif } +#ifdef INTGEMM_CPUID_ENVIRONMENT CPUType EnvironmentCPUID() { -#if defined(_MSC_VER) +# if defined(_MSC_VER) char env_override[11]; size_t len = 0; if (getenv_s(&len, env_override, sizeof(env_override), "INTGEMM_CPUID")) return CPUType::AVX512VNNI; if (!len) return CPUType::AVX512VNNI; -#else +# else const char *env_override = getenv("INTGEMM_CPUID"); if (!env_override) return CPUType::AVX512VNNI; /* This will be capped to actual ID */ -#endif +# endif if (!strcmp(env_override, "AVX512VNNI")) return CPUType::AVX512VNNI; if (!strcmp(env_override, "AVX512BW")) return CPUType::AVX512BW; if (!strcmp(env_override, "AVX2")) return CPUType::AVX2; @@ -105,11 +106,17 @@ CPUType EnvironmentCPUID() { fprintf(stderr, "Ignoring unrecognized INTGEMM_CPUID %s\n", env_override); return CPUType::AVX512VNNI; } +#endif } // namespace CPUType GetCPUID() { - static const CPUType kLocalCPU = std::min(RealCPUID(), EnvironmentCPUID()); + static const CPUType kLocalCPU = +#ifdef INTGEMM_CPUID_ENVIRONMENT + std::min(RealCPUID(), EnvironmentCPUID()); +#else + RealCPUID(); +#endif return kLocalCPU; } |