Merge pull request #2303 from Xreki/build_drive_px2

Support native build on NVIDIA DRIVE PX2 (arm64 + GPU).
PaddlePaddle · May 31, 2017 · 5526ff8 · 5526ff8
2 parents 9bed30f + 07ac67e
commit 5526ff8
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 4 deletions.
diff --git a/cmake/cudnn.cmake b/cmake/cudnn.cmake
@@ -11,11 +11,23 @@ find_path(CUDNN_INCLUDE_DIR cudnn.h
 
 get_filename_component(__libpath_hist ${CUDA_CUDART_LIBRARY} PATH)
 
+if(NOT ${CMAKE_HOST_SYSTEM_PROCESSOR})
+    execute_process(
+        COMMAND uname -m COMMAND tr -d '\n'
+        OUTPUT_VARIABLE HOST_ARCH
+        RESULT_VARIABLE UNAME_RESULT)
+    if(${UNAME_RESULT})
+        set(HOST_ARCH "x86_64")
+    endif(${UNAME_RESULT})
+else(NOT ${CMAKE_HOST_SYSTEM_PROCESSOR})
+    set(HOST_ARCH ${CMAKE_HOST_SYSTEM_PROCESSOR})
+endif(NOT ${CMAKE_HOST_SYSTEM_PROCESSOR})
+
 list(APPEND CUDNN_CHECK_LIBRARY_DIRS
     ${CUDNN_ROOT}
     ${CUDNN_ROOT}/lib64
     ${CUDNN_ROOT}/lib
-    ${CUDNN_ROOT}/lib/x86_64-linux-gnu
+    ${CUDNN_ROOT}/lib/${HOST_ARCH}-linux-gnu
     $ENV{CUDNN_ROOT}
     $ENV{CUDNN_ROOT}/lib64
     $ENV{CUDNN_ROOT}/lib

diff --git a/paddle/utils/CpuId.cpp b/paddle/utils/CpuId.cpp
@@ -21,7 +21,7 @@ limitations under the License. */
 
 #else
 
-#if !defined(__arm__)
+#if !defined(__arm__) && !defined(__aarch64__)
 #include <cpuid.h>
 /// for GCC/Clang
 #define CPUID(info, x) __cpuid_count(x, 0, info[0], info[1], info[2], info[3])
@@ -32,7 +32,7 @@ limitations under the License. */
 namespace paddle {
 
 SIMDFlags::SIMDFlags() {
-#if defined(__arm__)
+#if defined(__arm__) || defined(__aarch64__)
   simd_flags_ = SIMD_NEON;
 #else
   unsigned int cpuInfo[4];

diff --git a/paddle/utils/tests/test_SIMDFlags.cpp b/paddle/utils/tests/test_SIMDFlags.cpp
@@ -19,7 +19,7 @@ using namespace paddle;  // NOLINT
 
 TEST(SIMDFlags, gccTest) {
 #if (defined(__GNUC__) || defined(__GNUG__)) && !(defined(__clang__)) && \
-    !defined(__arm__)
+    !defined(__arm__) && !defined(__aarch64__)
   // clang-format off
   CHECK(!__builtin_cpu_supports("sse")    != HAS_SSE);
   CHECK(!__builtin_cpu_supports("sse2")   != HAS_SSE2);