[hotfix] Make USE_AVX a flag in cmake to avoid compilation error for …

…arm user (dmlc#2428) * upd cmake * upd * format
XINyexun · Dec 17, 2020 · e379e52 · e379e52
1 parent 50f1f4a
commit e379e52
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 4 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -24,6 +24,7 @@ endif()
 # Alernatively, use cmake -DOPTION=VALUE through command-line.
 dgl_option(USE_CUDA "Build with CUDA" OFF)
 dgl_option(USE_OPENMP "Build with OpenMP" ON)
+dgl_option(USE_AVX "Build with AVX optimization" OFF)
 dgl_option(BUILD_CPP_TEST "Build cpp unittest executables" OFF)
 dgl_option(LIBCXX_ENABLE_PARALLEL_ALGORITHMS "Enable the parallel algorithms library. This requires the PSTL to be available." OFF)
 dgl_option(USE_S3 "Build with S3 support" OFF)
@@ -108,6 +109,10 @@ if(USE_OPENMP)
   endif(OPENMP_FOUND)
 endif(USE_OPENMP)
 
+if(USE_AVX)
+  add_compile_definitions(USE_AVX)
+endif(USE_AVX)
+
 # To compile METIS correct for DGL.
 if(MSVC)
   set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /DIDXTYPEWIDTH=64 /DREALTYPEWIDTH=32")

diff --git a/cmake/config.cmake b/cmake/config.cmake
@@ -39,3 +39,6 @@ set(BUILD_CPP_TEST OFF)
 
 # Whether to enable OpenMP
 set(USE_OPENMP ON)
+
+# Whether to enable Intel's avx optimized kernel
+set(USE_AVX OFF)
diff --git a/src/array/cpu/spmm.h b/src/array/cpu/spmm.h
@@ -13,8 +13,10 @@
 #include <memory>
 #include "spmm_binary_ops.h"
 #if !defined(_WIN32)
+#ifdef USE_AVX
 #include "intel/cpu_support.h"
-#endif
+#endif  // USE_AVX
+#endif  // _WIN32
 namespace dgl {
 namespace aten {
 namespace cpu {
@@ -41,6 +43,7 @@ void SpMMSumCsr(const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat,
   int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
   DType* O = out.Ptr<DType>();
 #if !defined(_WIN32)
+#ifdef USE_AVX
   typedef dgl::ElemWiseAddUpdate<Op> ElemWiseUpd;
   /* Prepare an assembler kernel */
   static std::unique_ptr<ElemWiseUpd> asm_kernel_ptr(
@@ -62,7 +65,8 @@ void SpMMSumCsr(const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat,
       }
     }
   } else {
-#endif
+#endif  // USE_AVX
+#endif  // _WIN32
 
 #pragma omp parallel for
     for (IdType rid = 0; rid < csr.num_rows; ++rid) {
@@ -84,8 +88,10 @@ void SpMMSumCsr(const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat,
       }
     }
 #if !defined(_WIN32)
+#ifdef USE_AVX
   }
-#endif
+#endif  // USE_AVX
+#endif  // _WIN32
 }
 
 /*!

diff --git a/tests/cpp/test_spmm.cc b/tests/cpp/test_spmm.cc
@@ -1,4 +1,5 @@
 #if !defined(_WIN32)
+#ifdef USE_AVX
 #include <../../src/array/cpu/spmm.h>
 #include <dgl/array.h>
 #include <gtest/gtest.h>
@@ -313,4 +314,5 @@ TEST(SpmmTest, TestSpmmDiv) {
   _TestSpmmDiv<float>();
   _TestSpmmDiv<double>();
 }
-#endif
+#endif  // USE_AVX
+#endif  // _WIN32