adding environment flags to disable SIMD codepaths

2 files changed, 29 insertions, 3 deletions

diff --git a/lib/TH/README.md b/lib/TH/README.md
new file mode 100644
index 0000000..c646ce9
--- /dev/null
+++ b/lib/TH/README.md
@@ -0,0 +1,7 @@
+Environment variables control the disabling of certain explicit SIMD optimizations.
+
+```
+TH_NO_AVX2=1 # disable AVX2 codepaths
+TH_NO_AVX=1  # disable AVX codepaths
+TH_NO_SSE=1  # disable SSE codepaths
+```
\ No newline at end of file
diff --git a/lib/TH/generic/simd/simd.h b/lib/TH/generic/simd/simd.h
index 2f0c1f9..19d41b1 100644
--- a/lib/TH/generic/simd/simd.h
+++ b/lib/TH/generic/simd/simd.h
@@ -2,6 +2,7 @@
 #define TH_SIMD_INC
 
 #include <stdint.h>
+#include <stdlib.h>
 #if defined(_MSC_VER)
 #include <intrin.h>
 #elif defined(HAVE_GCC_GET_CPUID) && defined(USE_GCC_GET_CPUID)
@@ -117,20 +118,38 @@ static inline uint32_t detectHostSIMDExtensions()
 {
   uint32_t eax, ebx, ecx, edx;
   uint32_t hostSimdExts = 0x0;
+  int TH_NO_AVX = 1, TH_NO_AVX2 = 1, TH_NO_SSE = 1;
+  char *evar;
+
+  evar = getenv("TH_NO_AVX2");
+  if (evar == NULL || strncmp(evar, "1", 2) != 0)
+    TH_NO_AVX2 = 0;
 
   // Check for AVX2. Requires separate CPUID
   eax = 0x7;
   ecx = 0x0;
   cpuid(&eax, &ebx, &ecx, &edx);
-  if (ebx & CPUID_AVX2_BIT)
+  if ((ebx & CPUID_AVX2_BIT) && TH_NO_AVX2 == 0) {
     hostSimdExts |= SIMDExtension_AVX2;
+  }
 
+  // Detect and enable AVX and SSE
   eax = 0x1;
   cpuid(&eax, &ebx, &ecx, &edx);
-  if (ecx & CPUID_AVX_BIT)
+
+  evar = getenv("TH_NO_AVX");
+  if (evar == NULL || strncmp(evar, "1", 2) != 0)
+    TH_NO_AVX = 0;
+  if (ecx & CPUID_AVX_BIT && TH_NO_AVX == 0) {
     hostSimdExts |= SIMDExtension_AVX;
-  if (edx & CPUID_SSE_BIT)
+  }
+
+  evar = getenv("TH_NO_SSE");
+  if (evar == NULL || strncmp(evar, "1", 2) != 0)
+    TH_NO_SSE = 0;  
+  if (edx & CPUID_SSE_BIT && TH_NO_SSE == 0) {
     hostSimdExts |= SIMDExtension_SSE;
+  }
 
   return hostSimdExts;
 }