Repository: dmlc/dgl
Pull request: [hotfix] Make USE_AVX a flag in cmake to avoid compilation error for arm user
## Description
To fix #2424 .

## Checklist
Please feel free to remove inapplicable items for your PR.
- [x] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [x] Changes are complete (i.e. I finished coding on this PR)
- [x] All changes have test coverage
- [x] Code is well-documented
- [x] To the my best knowledge, examples are either not affected by this change,
      or have been fixed to be compatible with this change
- [x] Related issue is referred in this PR


Problem statement:
Compilation on osx-arm64 apple silicon m1
## 🐛 Bug

Trying to compile from source on the new macbook air with arm64e. 

## To Reproduce

mkdir build
cd build
cmake ..
make -j 8

```
/Users/yingzhu/dgl/include/intel/cpu_support.h:70:16: error: no type named 'Reg64' in namespace 'Xbyak'
  const Xbyak::Reg64 &r_out_;
        ~~~~~~~^
/Users/yingzhu/dgl/include/intel/cpu_support.h:71:16: error: no type named 'Reg64' in namespace 'Xbyak'
  const Xbyak::Reg64 &r_left_;
        ~~~~~~~^
/Users/yingzhu/dgl/include/intel/cpu_support.h:72:16: error: no type named 'Reg64' in namespace 'Xbyak'
  const Xbyak::Reg64 &r_right;
        ~~~~~~~^
/Users/yingzhu/dgl/include/intel/cpu_support.h:73:16: error: no type named 'Reg64' in namespace 'Xbyak'
  const Xbyak::Reg64 &r_size_;
        ~~~~~~~^
/Users/yingzhu/dgl/include/intel/cpu_support.h:156:42: error: use of undeclared identifier 'r9'
    alias_load<IType>(zmm0, ptr[r_out_ + r9 * sizeof(IType)]);
                                        ^
```
## Expected behavior

Build without error

## Environment

 - DGL Version from master
 - Backend Library & Version pytorch 1.7 installed
 - OS (e.g., Linux): osx-arm64
 - How you installed DGL (`conda`, `pip`, source): source
 - Build command you used (if compiling from source): see above
 - Python version: 3.8.5
 - CUDA/cuDNN version (if applicable): na
 - GPU models and configuration (e.g. V100): na 
 - Any other relevant information:

## Additional context

It clearly tries to use intel-specifc code - is there support for arm? i would expect 80% of all mac users to switch to arm in the next 5 years

Hints from issue discussion:
Yes that's related to intel-specific optimizations (use AVX instructions, see #2309 ). Actually we should only include this header when a intel CPU with avx support is detected, I'll fix it soon.
@yzh119 thanks keep me posted. sorry for being dense but setting `DGL_CPU_INTEL_KERNEL_ENABLED=0` doesn't work for me - any other way to disable this so I can build in mean time?
`DGL_CPU_INTEL_KERNEL_ENABLED=0` only controls whether to use AVX optimized operator at runtime instead of compile time, I'll set a cmake flag to indicate whether to include the header at compile time.

Commit message: upd

Patch to review:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0abcd52fe732..9af440317ad6 100644
--- 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(WITH_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
index e4a922e4c9a8..31e14a2e1410 100644
--- 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
index 887d72178824..f67e2db4b8ea 100644
--- 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 WITH_AVX
 #include "intel/cpu_support.h"
 #endif
+#endif
 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 WITH_AVX
   typedef dgl::ElemWiseAddUpdate<Op> ElemWiseUpd;
   /* Prepare an assembler kernel */
   static std::unique_ptr<ElemWiseUpd> asm_kernel_ptr(
@@ -63,6 +66,7 @@ void SpMMSumCsr(const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat,
     }
   } else {
 #endif
+#endif
 
 #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 WITH_AVX
   }
 #endif
+#endif
 }
 
 /*!
diff --git a/tests/cpp/test_spmm.cc b/tests/cpp/test_spmm.cc
index 7b4af5a50f75..6e109bf132e7 100644
--- a/tests/cpp/test_spmm.cc
+++ b/tests/cpp/test_spmm.cc
@@ -1,4 +1,5 @@
 #if !defined(_WIN32)
+#ifdef WITH_AVX
 #include <../../src/array/cpu/spmm.h>
 #include <dgl/array.h>
 #include <gtest/gtest.h>
@@ -314,3 +315,4 @@ TEST(SpmmTest, TestSpmmDiv) {
   _TestSpmmDiv<double>();
 }
 #endif
+#endif


Code review comments:
diff_hunk: @@ -13,8 +13,10 @@
 #include <memory>
 #include "spmm_binary_ops.h"
 #if !defined(_WIN32)
+#ifdef WITH_AVX
 #include "intel/cpu_support.h"
 #endif
+#endif
path: src/array/cpu/spmm.h
text: @user1:
```c++
#endif  // WITH_AVX
#endif  // _WIN32
```

diff_hunk: @@ -108,6 +109,10 @@ if(USE_OPENMP)
   endif(OPENMP_FOUND)
 endif(USE_OPENMP)
 
+if(USE_AVX)
+  add_compile_definitions(WITH_AVX)
path: CMakeLists.txt
text: @user1:
Why not call it `USE_AVX`?