Option to turn off environment variable reading (#98)

3 files changed, 41 insertions, 4 deletions

diff --git a/.github/workflows/ubuntu-no-cpuid-environment.yml b/.github/workflows/ubuntu-no-cpuid-environment.yml
new file mode 100644
index 0000000..dc1862f
--- /dev/null
+++ b/.github/workflows/ubuntu-no-cpuid-environment.yml
@@ -0,0 +1,25 @@
+name: Ubuntu No CPUID Environment Variable
+
+on:
+  push:
+    branches: [master, static]
+  pull_request:
+    branches: [master, static]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: cmake
+      run: |
+        cmake -E make_directory build
+        cd build
+        cmake -DINTGEMM_CPUID_ENVIRONMENT=OFF ..
+    - name: Compile
+      working-directory: build
+      run: cmake --build . -j2
+    - name: Test
+      working-directory: build
+      run: ctest -j2
diff --git a/CMakeLists.txt b/CMakeLists.txt
index af27542..c9f78fa 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -84,6 +84,11 @@ if (WORMHOLE)
   target_compile_definitions(intgemm PUBLIC INTGEMM_WORMHOLE)
 endif()
 
+option(INTGEMM_CPUID_ENVIRONMENT "Allow INTGEMM_CPUID environment variable to downgrade CPU model, which is mainly for testing." ON)
+if (INTGEMM_CPUID_ENVIRONMENT)
+  target_compile_definitions(intgemm PRIVATE INTGEMM_CPUID_ENVIRONMENT)
+endif()
+
 if(INTGEMM_DONT_BUILD_TESTS)
   return()
 endif()
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;
 }