Fix Caffe2 build with ATen CPU/GPU split (pytorch#7486)

Author: orionr

aten/CMakeLists.txt

@@ -544,7 +544,12 @@ ${CMAKE_CURRENT_BINARY_DIR}/src/ATen)
 if(NOT NO_CUDA)
   include_directories(${CUDA_INCLUDE_DIRS})
 endif()
-add_subdirectory(src/ATen/test)
+
+if(ATEN_NO_TEST)
+  message("disable test because ATEN_NO_TEST is set")
+else()
+  add_subdirectory(src/ATen/test)
+endif()
 
 if(ATEN_NO_CONTRIB)
   message("disable contrib because ATEN_NO_CONTRIB is set")

caffe2/CMakeLists.txt

@@ -14,6 +14,7 @@ if(BUILD_ATEN)
   if(NOT BUILD_CAFFE2)
     set(__caffe2_CMAKE_POSITION_INDEPENDENT_CODE ${CMAKE_POSITION_INDEPENDENT_CODE})
     set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+    set(ATEN_NO_TEST ON)
     set(AT_LINK_STYLE STATIC)
     add_subdirectory(../aten aten)
     set(CMAKE_POSITION_INDEPENDENT_CODE ${__caffe2_CMAKE_POSITION_INDEPENDENT_CODE})

caffe2/contrib/aten/CMakeLists.txt

@@ -3,17 +3,18 @@ if(BUILD_ATEN)
   # Unfortunately the custom command below requires this to be here
   set(__caffe2_CMAKE_POSITION_INDEPENDENT_CODE ${CMAKE_POSITION_INDEPENDENT_CODE})
   set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+  set(ATEN_NO_TEST ON)
   set(AT_LINK_STYLE STATIC)
-  add_subdirectory(${PROJECT_SOURCE_DIR}/aten aten)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../../aten aten)
   set(CMAKE_POSITION_INDEPENDENT_CODE ${__caffe2_CMAKE_POSITION_INDEPENDENT_CODE})
 
   add_custom_command(OUTPUT aten_op.h
   COMMAND
   python ${CMAKE_CURRENT_SOURCE_DIR}/gen_op.py
-  --aten_root=${PROJECT_SOURCE_DIR}/aten
-  --template_dir=${PROJECT_SOURCE_DIR}/caffe2/contrib/aten
+  --aten_root=${CMAKE_CURRENT_SOURCE_DIR}/../../../aten
+  --template_dir=${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS
-  ATen
+  ATen_cpu
   ${CMAKE_CURRENT_SOURCE_DIR}/gen_op.py
   ${CMAKE_CURRENT_SOURCE_DIR}/aten_op_template.h)
 

cmake/Dependencies.cmake

@@ -579,9 +579,9 @@ if (BUILD_ATEN)
     include_directories(${PROJECT_BINARY_DIR}/caffe2/contrib/aten)
   endif()
 
-  list(APPEND Caffe2_DEPENDENCY_LIBS ATen)
+  list(APPEND Caffe2_DEPENDENCY_LIBS ATen_cpu)
   if (USE_CUDA)
-    list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS ATen)
+    list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS ATen_cuda)
   endif()
   include_directories(${PROJECT_SOURCE_DIR}/aten/src)
 endif()