reintroduce USE_AVX* for files which dont have -mavx* setssefix

3 files changed, 16 insertions, 14 deletions

diff --git a/lib/TH/CMakeLists.txt b/lib/TH/CMakeLists.txt
index d26d4c2..8aeb204 100644
--- a/lib/TH/CMakeLists.txt
+++ b/lib/TH/CMakeLists.txt
@@ -133,14 +133,16 @@ IF(C_SSE3_FOUND)
   MESSAGE(STATUS "SSE3 Found")
   SET(CMAKE_C_FLAGS "${C_SSE3_FLAGS} -DUSE_SSE3 ${CMAKE_C_FLAGS}")
 ENDIF(C_SSE3_FOUND)
-# we dont set AVX and AVX2 flags globally, but only for specific files
+# we dont set -mavx and -mavx2 flags globally, but only for specific files
+# however, we want to enable the AVX codepaths, so we still need to
+# add USE_AVX and USE_AVX2 macro defines
 IF(C_AVX_FOUND)
   MESSAGE(STATUS "AVX Found")
-  #   SET(CMAKE_C_FLAGS "${C_AVX_FLAGS} ${CMAKE_C_FLAGS}")
+  SET(CMAKE_C_FLAGS "-DUSE_AVX ${CMAKE_C_FLAGS}")
 ENDIF(C_AVX_FOUND)
 IF(C_AVX2_FOUND)
   MESSAGE(STATUS "AVX2 Found")
-  #   SET(CMAKE_C_FLAGS "${C_AVX2_FLAGS} ${CMAKE_C_FLAGS}")
+  SET(CMAKE_C_FLAGS "-DUSE_AVX2 ${CMAKE_C_FLAGS}")
 ENDIF(C_AVX2_FOUND)
 
 
diff --git a/lib/TH/THVector.c b/lib/TH/THVector.c
index d8493a5..4410578 100644
--- a/lib/TH/THVector.c
+++ b/lib/TH/THVector.c
@@ -15,11 +15,11 @@
 #include "vector/SSE.c"
 #endif
 
-#if defined(__AVX__) || defined(__AVX2__)
+#if defined(USE_AVX)
 #include "vector/AVX.h"
 #endif
 
-#if defined(__AVX2__)
+#if defined(USE_AVX2)
 #include "vector/AVX2.h"
 #endif
 
diff --git a/lib/TH/generic/THVectorDispatch.c b/lib/TH/generic/THVectorDispatch.c
index 6220dd6..5b88852 100644
--- a/lib/TH/generic/THVectorDispatch.c
+++ b/lib/TH/generic/THVectorDispatch.c
@@ -26,7 +26,7 @@ static FunctionDescription THVector_(fill_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(fill_AVX), SIMDExtension_AVX),
     #endif
@@ -52,13 +52,13 @@ static FunctionDescription THVector_(cadd_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX2__)
+  #if defined(USE_AVX2)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(cadd_AVX2), SIMDExtension_AVX2),
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(cadd_AVX), SIMDExtension_AVX),
     #endif
@@ -91,7 +91,7 @@ static FunctionDescription THVector_(adds_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(adds_AVX), SIMDExtension_AVX),
     #endif
@@ -119,7 +119,7 @@ static FunctionDescription THVector_(cmul_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(cmul_AVX), SIMDExtension_AVX),
     #endif
@@ -152,7 +152,7 @@ static FunctionDescription THVector_(muls_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(muls_AVX), SIMDExtension_AVX),
     #endif
@@ -179,7 +179,7 @@ static FunctionDescription THVector_(cdiv_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(cdiv_AVX), SIMDExtension_AVX),
     #endif
@@ -206,7 +206,7 @@ static FunctionDescription THVector_(divs_DISPATCHTABLE)[] = {
     #endif
   #endif
 
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(divs_AVX), SIMDExtension_AVX),
     #endif
@@ -227,7 +227,7 @@ void THVector_(divs)(real *y, const real *x, const real c, const ptrdiff_t n) {
 
 static void (*THVector_(copy_DISPATCHPTR))(real *, const real *, const ptrdiff_t) = &THVector_(copy_DEFAULT);
 static FunctionDescription THVector_(copy_DISPATCHTABLE)[] = {
-  #if defined(__AVX__)
+  #if defined(USE_AVX)
     #if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)
       FUNCTION_IMPL(THVector_(copy_AVX), SIMDExtension_AVX),
     #endif