From 07ac67ec7c15d0c4b93ab6d7e119570f10952a3f Mon Sep 17 00:00:00 2001
From: Liu Yiqun <liuyiqun01@baidu.com>
Date: Fri, 26 May 2017 23:38:52 -0700
Subject: [PATCH] Support native build on NVIDIA DRIVE PX2 (arm64 + GPU).

---
 cmake/cudnn.cmake                     | 14 +++++++++++++-
 paddle/utils/CpuId.cpp                |  4 ++--
 paddle/utils/tests/test_SIMDFlags.cpp |  2 +-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/cmake/cudnn.cmake b/cmake/cudnn.cmake
index af9be869618..92dce20c698 100644
--- 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
index 5abeeecae8d..7186feef041 100644
--- 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
index 185789c927b..a808d456a69 100644
--- 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);
-- 
GitLab