When BUILD_ATEN=OFF, use ATen/core directly (#10019)

Summary:
ATenCore.h is a dummy header to just test that this is working at all.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/10019

Reviewed By: smessmer

Differential Revision: D9067262

Pulled By: ezyang

fbshipit-source-id: 58bab9c0aa83b56335e36b719b9b6505400d8dee

CMakeLists.txt

@@ -284,6 +284,8 @@ include_directories(BEFORE ${PROJECT_SOURCE_DIR})
 # in PROJECT_SOURCE_DIR.
 include_directories(BEFORE ${PROJECT_BINARY_DIR})
 
+include_directories(BEFORE ${PROJECT_SOURCE_DIR}/aten/src/)
+
 # ---[ Old caffe protobuf
 if(BUILD_CAFFE2)
   add_subdirectory(caffe/proto)

aten/src/ATen/CMakeLists.txt

@@ -44,6 +44,7 @@ CONFIGURE_FILE(cuda/CUDAConfig.h.in "${CMAKE_CURRENT_SOURCE_DIR}/cuda/CUDAConfig
 # NB: If you edit these globs, you'll have to update setup.py package_data as well
 FILE(GLOB base_h "*.h" "detail/*.h")
 FILE(GLOB base_cpp "*.cpp" "detail/*.cpp")
+add_subdirectory(core)
 FILE(GLOB cuda_h "cuda/*.h" "cuda/detail/*.h" "cuda/*.cuh" "cuda/detail/*.cuh")
 FILE(GLOB cuda_cpp "cuda/*.cpp" "cuda/detail/*.cpp")
 FILE(GLOB cuda_cu "cuda/*.cu" "cuda/detail/*.cu")
@@ -62,7 +63,7 @@ FILE(GLOB native_cuda_cpp "native/cuda/*.cpp")
 FILE(GLOB native_mkl_cpp "native/mkl/*.cpp")
 FILE(GLOB native_mkldnn_cpp "native/mkldnn/*.cpp")
 
-set(all_cpu_cpp ${base_cpp} ${native_cpp} ${native_sparse_cpp} ${native_mkl_cpp} ${native_mkldnn_cpp} ${generated_cpp} ${ATen_CPU_SRCS} ${cpu_kernel_cpp})
+set(all_cpu_cpp ${base_cpp} ${ATen_CORE_SRCS} ${native_cpp} ${native_sparse_cpp} ${native_mkl_cpp} ${native_mkldnn_cpp} ${generated_cpp} ${ATen_CPU_SRCS} ${cpu_kernel_cpp})
 if(AT_MKL_ENABLED)
   set(all_cpu_cpp ${all_cpu_cpp} ${mkl_cpp})
 endif()
@@ -393,7 +394,7 @@ INSTALL(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/ATenConfig.cmake"
   DESTINATION "${AT_INSTALL_SHARE_DIR}/cmake/ATen")
 
 # https://stackoverflow.com/questions/11096471/how-can-i-install-a-hierarchy-of-files-using-cmake
-FOREACH(HEADER ${base_h} ${cuda_h} ${cudnn_h})
+FOREACH(HEADER ${base_h} ${ATen_CORE_HEADERS} ${cuda_h} ${cudnn_h})
   string(REPLACE "${CMAKE_CURRENT_SOURCE_DIR}/" "" HEADER_SUB ${HEADER})
   GET_FILENAME_COMPONENT(DIR ${HEADER_SUB} DIRECTORY)
   INSTALL(FILES ${HEADER} DESTINATION ${AT_INSTALL_INCLUDE_DIR}/ATen/${DIR})

aten/src/ATen/Context.h

@@ -9,6 +9,9 @@
 #include "ATen/detail/CUDAHooksInterface.h"
 #include "ATen/CUDAStream.h"
 
+// This is temporary
+#include "ATen/core/ATenCoreTest.h"
+
 #include <memory>
 #include <mutex>
 #include <cstdint>

aten/src/ATen/core/ATenCoreTest.cpp

@@ -0,0 +1,10 @@
+#include <ATen/core/ATenCoreTest.h>
+
+namespace at {
+
+static int CoreTestGlobal = 0;
+int CoreTest() {
+  return CoreTestGlobal++;
+}
+
+}

aten/src/ATen/core/ATenCoreTest.h

@@ -0,0 +1,18 @@
+#pragma once
+
+// TODO: Move this to something like ATenCoreGeneral.h
+#ifdef _WIN32
+# if defined(ATen_cpu_EXPORTS) || defined(caffe2_EXPORTS)
+#  define AT_CORE_API __declspec(dllexport)
+# else
+#  define AT_CORE_API __declspec(dllimport)
+# endif
+#else
+# define AT_CORE_API
+#endif
+
+namespace at {
+
+AT_CORE_API int CoreTest();
+
+}

aten/src/ATen/core/CMakeLists.txt

@@ -0,0 +1,13 @@
+# This file solely exists to let Caffe2 Android build get at the list
+# of core files without having to trundle through all of ATen's CMakeLists.txt
+
+FILE(GLOB ATen_CORE_HEADERS "*.h")
+FILE(GLOB ATen_CORE_SRCS "*.cpp")
+
+# Pass to parent
+set(ATen_CORE_HEADERS ${ATen_CORE_HEADERS} PARENT_SCOPE)
+set(ATen_CORE_SRCS ${ATen_CORE_SRCS} PARENT_SCOPE)
+# This is a little dodgy, because it means ALL ATen headers are made
+# visible.  Fortunately, you should just get a lot of undefined symbol
+# errors if you go outside core
+set(ATen_CORE_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/../.. PARENT_SCOPE)

aten/src/ATen/core/README.md

@@ -0,0 +1,5 @@
+ATen Core
+---------
+
+ATen Core is a minimal subset of ATen which is suitable for deployment
+on mobile.  Binary size of files in this folder is an important constraint.

aten/src/ATen/test/basic.cpp

@@ -270,6 +270,10 @@ static void test(Type & type) {
     auto result = tensor.m(relu).m(mse_loss, other, Reduction::ElementwiseMean);
     REQUIRE(result.allclose(mse_loss(relu(tensor), other)));
   }
+  SECTION("core") {
+    int i = CoreTest();
+    REQUIRE(i + 1 == CoreTest());
+  }
 }
 
 TEST_CASE( "basic tests CPU", "[cpu]" ) {

caffe2/CMakeLists.txt

@@ -51,6 +51,14 @@ if(BUILD_ATEN)
     set(Caffe2_HIP_SRCS ${ATen_CUDA_SRCS})
     set(Caffe2_HIP_INCLUDES ${Caffe2_HIP_INCLUDES} ${Caffe2_GPU_INCLUDE})
   ENDIF(USE_ROCM)
+else()
+  # Only add "ATen Core", a minimal, easy-to-compile fragment of ATen.
+  # This codepath should only be exercised by the Android build.
+  add_subdirectory(../aten/src/ATen/core ATen_core)
+  list(APPEND Caffe2_CPU_SRCS ${ATen_CORE_SRCS})
+  list(APPEND Caffe2_CPU_INCLUDE ${ATen_CORE_INCLUDE})
+  # TODO: We should probably install the headers, but I don't know
+  # how to do that.
 endif()
 
 # ---[ Torch build

caffe2/core/context.h

@@ -13,6 +13,8 @@
 #include "caffe2/core/typeid.h"
 #include "caffe2/proto/caffe2.pb.h"
 
+#include "ATen/core/ATenCoreTest.h"
+
 CAFFE2_DECLARE_bool(caffe2_report_cpu_memory_usage);
 
 namespace caffe2 {

caffe2/core/context_test.cc

@@ -6,6 +6,11 @@
 
 namespace caffe2 {
 
+TEST(CPUContextTest, ATenCoreTest) {
+  int i = at::CoreTest();
+  EXPECT_EQ(i + 1, at::CoreTest());
+}
+
 TEST(CPUContextTest, TestAllocAlignment) {
   for (int i = 1; i < 10; ++i) {
     auto data = CPUContext::New(i);

cmake/Codegen.cmake

@@ -1,3 +1,9 @@
+# This ill-named file does a number of things:
+# - Installs Caffe2 header files (this has nothing to do with code generation)
+# - Configures caffe2/core/macros.h
+# - Creates an ATen target for its generated C++ files and adds it
+#   as a dependency
+
 if (DEFINED ENV{PYTORCH_PYTHON})
   message(STATUS "Using python found in $ENV{PYTORCH_PYTHON}")
   set(PYCMD "$ENV{PYTORCH_PYTHON}")
@@ -14,6 +20,11 @@ configure_file(
 install(DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/../caffe2
         DESTINATION include
         FILES_MATCHING PATTERN "*.h")
+if (NOT BUILD_ATEN)
+  install(DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/core
+          DESTINATION include/ATen/core
+          FILES_MATCHING PATTERN "*.h")
+endif()
 install(FILES ${CMAKE_BINARY_DIR}/caffe2/core/macros.h
         DESTINATION include/caffe2/core)
 

setup.py

@@ -1023,6 +1023,7 @@ if __name__ == '__main__':
                 'lib/torch_shm_manager',
                 'lib/*.h',
                 'lib/include/ATen/*.h',
+                'lib/include/ATen/core/*.h',
                 'lib/include/ATen/detail/*.h',
                 'lib/include/ATen/cuda/*.h',
                 'lib/include/ATen/cuda/*.cuh',