[hotfix] Make USE_AVX a flag in cmake to avoid compilation error for arm user (#2428)

* upd cmake

* upd

* format

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")

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)

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
 }
 
 /*!

tests/cpp/test_spmm.cc

 #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