CMakeLists and C++ files refactoring.

CMakeLists.txt

@@ -15,27 +15,22 @@ find_package(Python3 COMPONENTS Development)
 find_package(Torch REQUIRED)
 find_package(PNG REQUIRED)
 
+file(GLOB OPERATOR_SOURCES torchvision/csrc/cpu/*.cpp torchvision/csrc/*.cpp)
+file(GLOB MODELS_SOURCES torchvision/csrc/models/*.cpp)
+
+include_directories(torchvision/cinclude/${PROJECT_NAME})
 
-file(GLOB HEADERS torchvision/csrc/*.h)
-# Image extension
-file(GLOB IMAGE_HEADERS torchvision/csrc/cpu/image/*.h)
-file(GLOB IMAGE_SOURCES torchvision/csrc/cpu/image/*.cpp)
-file(GLOB OPERATOR_SOURCES torchvision/csrc/cpu/*.h torchvision/csrc/cpu/*.cpp ${IMAGE_HEADERS} ${IMAGE_SOURCES} ${HEADERS} torchvision/csrc/*.cpp)
 if(WITH_CUDA)
-  file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} torchvision/csrc/cuda/*.h torchvision/csrc/cuda/*.cu)
+  file(GLOB OPERATOR_CUDA_SOURCES torchvision/csrc/cuda/*.cu)
+  list(APPEND OPERATOR_SOURCES ${OPERATOR_CUDA_SOURCES})
+
+  include_directories(torchvision/cinclude/${PROJECT_NAME}_cuda)
 endif()
-file(GLOB MODELS_HEADERS torchvision/csrc/models/*.h)
-file(GLOB MODELS_SOURCES torchvision/csrc/models/*.h torchvision/csrc/models/*.cpp)
 
-add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES} ${IMAGE_SOURCES})
-target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} ${PNG_LIBRARY} Python3::Python)
-# target_link_libraries(${PROJECT_NAME} PRIVATE ${PNG_LIBRARY} Python3::Python)
+add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES})
+target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python)
 set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchVision)
 
-target_include_directories(${PROJECT_NAME} INTERFACE
-  $<BUILD_INTERFACE:${HEADERS}:${PNG_INCLUDE_DIR}>
-  $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
-
 include(GNUInstallDirs)
 include(CMakePackageConfigHelpers)
 
@@ -62,13 +57,13 @@ install(EXPORT TorchVisionTargets
   NAMESPACE TorchVision::
   DESTINATION ${TORCHVISION_CMAKECONFIG_INSTALL_DIR})
 
-install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME})
-install(FILES
-  torchvision/csrc/cpu/vision_cpu.h
-  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cpu)
+install(
+  DIRECTORY "torchvision/cinclude/${PROJECT_NAME}/"
+  DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}")
+
 if(WITH_CUDA)
-  install(FILES
-    torchvision/csrc/cuda/vision_cuda.h
-    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cuda)
+  install(
+    DIRECTORY "torchvision/cinclude/${PROJECT_NAME}_cuda/cuda"
+    DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}"
+  )
 endif()
-install(FILES ${MODELS_HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/models)

setup.py

@@ -158,16 +158,18 @@ def get_linux_distribution():
 
 def get_extensions():
     this_dir = os.path.dirname(os.path.abspath(__file__))
-    extensions_dir = os.path.join(this_dir, 'torchvision', 'csrc')
+    extensions_src_dir = os.path.join(this_dir, 'torchvision', 'csrc')
+    extensions_include_dir = os.path.join(this_dir, 'torchvision', 'cinclude', 'torchvision')
 
-    main_file = glob.glob(os.path.join(extensions_dir, '*.cpp'))
-    source_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', '*.cpp'))
+    main_file = glob.glob(os.path.join(extensions_src_dir, '*.cpp'))
+    source_cpu = glob.glob(os.path.join(extensions_src_dir, 'cpu', '*.cpp'))
 
     is_rocm_pytorch = False
     if torch.__version__ >= '1.5':
         from torch.utils.cpp_extension import ROCM_HOME
         is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False
 
+    # TODO fix this
     if is_rocm_pytorch:
         hipify_python.hipify(
             project_directory=this_dir,
@@ -176,13 +178,14 @@ def get_extensions():
             show_detailed=True,
             is_pytorch_extension=True,
         )
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'hip', '*.hip'))
+        source_cuda = glob.glob(os.path.join(extensions_src_dir, 'hip', '*.hip'))
         # Copy over additional files
         shutil.copy("torchvision/csrc/cuda/cuda_helpers.h", "torchvision/csrc/hip/cuda_helpers.h")
         shutil.copy("torchvision/csrc/cuda/vision_cuda.h", "torchvision/csrc/hip/vision_cuda.h")
 
     else:
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'cuda', '*.cu'))
+        source_cuda = glob.glob(os.path.join(extensions_src_dir, 'cuda', '*.cu'))
+        cuda_include_dir = os.path.join(this_dir, 'torchvision', 'cinclude', 'torchvision_cuda')
 
     sources = main_file + source_cpu
     extension = CppExtension
@@ -201,11 +204,14 @@ def get_extensions():
 
     define_macros = []
 
+    include_dirs = [extensions_include_dir]
+
     extra_compile_args = {}
     if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) \
             or os.getenv('FORCE_CUDA', '0') == '1':
         extension = CUDAExtension
         sources += source_cuda
+        include_dirs.append(cuda_include_dir)
         if not is_rocm_pytorch:
             define_macros += [('WITH_CUDA', None)]
             nvcc_flags = os.getenv('NVCC_FLAGS', '')
@@ -227,9 +233,7 @@ def get_extensions():
         extra_compile_args.setdefault('cxx', [])
         extra_compile_args['cxx'].append('/MP')
 
-    sources = [os.path.join(extensions_dir, s) for s in sources]
-
-    include_dirs = [extensions_dir]
+    sources = [os.path.join(extensions_src_dir, s) for s in sources]
 
     ext_modules = [
         extension(
@@ -263,7 +267,7 @@ def get_extensions():
 
     # Image reading extension
     image_macros = []
-    image_include = [extensions_dir]
+    image_include = [extensions_include_dir]
     image_library = []
     image_link_flags = []
 
@@ -327,7 +331,7 @@ def get_extensions():
             image_include += [png_include]
             image_link_flags.append('libpng')
 
-    image_path = os.path.join(extensions_dir, 'cpu', 'image')
+    image_path = os.path.join(extensions_src_dir, 'cpu', 'image')
     image_src = glob.glob(os.path.join(image_path, '*.cpp'))
 
     if png_found:
@@ -368,7 +372,7 @@ def get_extensions():
                     base_decoder_src_dir,
                     video_reader_src_dir,
                     ffmpeg_include_dir,
-                    extensions_dir,
+                    extensions_src_dir,
                 ],
                 libraries=[
                     'avcodec',

torchvision/csrc/cpu/ROIAlign_cpu.cpp

@@ -1,5 +1,5 @@
 #include <ATen/TensorUtils.h>
-#include "vision_cpu.h"
+#include "cpu/vision_cpu.h"
 
 // implementation taken from Caffe2
 template <typename T>

torchvision/csrc/cpu/nms_cpu.cpp

@@ -1,4 +1,4 @@
-#include "vision_cpu.h"
+#include "cpu/vision_cpu.h"
 
 template <typename scalar_t>
 at::Tensor nms_cpu_kernel(

torchvision/csrc/cuda/DeformConv_cuda.cu

@@ -72,7 +72,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <THC/THCAtomics.cuh>
 
-#include "cuda_helpers.h"
+#include "cuda/cuda_helpers.h"
 
 #include <cmath>
 #include <iostream>

torchvision/csrc/cuda/PSROIAlign_cuda.cu

@@ -5,7 +5,7 @@
 #include <THC/THCAtomics.cuh>
 #include <stdio.h>
 
-#include "cuda_helpers.h"
+#include "cuda/cuda_helpers.h"
 
 template <typename T>
 __device__ T bilinear_interpolate(

torchvision/csrc/cuda/PSROIPool_cuda.cu

@@ -4,7 +4,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <THC/THCAtomics.cuh>
 
-#include "cuda_helpers.h"
+#include "cuda/cuda_helpers.h"
 
 template <typename T>
 __global__ void PSROIPoolForward(

torchvision/csrc/cuda/ROIAlign_cuda.cu

@@ -4,7 +4,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <THC/THCAtomics.cuh>
 
-#include "cuda_helpers.h"
+#include "cuda/cuda_helpers.h"
 
 template <typename T>
 __device__ T bilinear_interpolate(

torchvision/csrc/cuda/ROIPool_cuda.cu

@@ -4,7 +4,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <THC/THCAtomics.cuh>
 
-#include "cuda_helpers.h"
+#include "cuda/cuda_helpers.h"
 
 template <typename T>
 __global__ void RoIPoolForward(

torchvision/csrc/cuda/nms_cuda.cu

@@ -2,7 +2,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 
-#include "cuda_helpers.h"
+#include "cuda/cuda_helpers.h"
 
 #include <iostream>
 #include <vector>

torchvision/csrc/models/alexnet.cpp

@@ -1,6 +1,6 @@
-#include "alexnet.h"
+#include "models/alexnet.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/densenet.cpp

@@ -1,6 +1,6 @@
-#include "densenet.h"
+#include "models/densenet.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/googlenet.cpp

@@ -1,6 +1,6 @@
-#include "googlenet.h"
+#include "models/googlenet.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/inception.cpp

@@ -1,4 +1,4 @@
-#include "inception.h"
+#include "models/inception.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/mnasnet.cpp

@@ -1,6 +1,6 @@
-#include "mnasnet.h"
+#include "models/mnasnet.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/mobilenet.cpp

@@ -1,6 +1,6 @@
-#include "mobilenet.h"
+#include "models/mobilenet.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/resnet.cpp

@@ -1,6 +1,6 @@
-#include "resnet.h"
+#include "models/resnet.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/shufflenetv2.cpp

@@ -1,6 +1,6 @@
-#include "shufflenetv2.h"
+#include "models/shufflenetv2.h"
 
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/squeezenet.cpp

@@ -1,7 +1,7 @@
-#include "squeezenet.h"
+#include "models/squeezenet.h"
 
 #include <limits>
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {

torchvision/csrc/models/vgg.cpp

@@ -1,7 +1,7 @@
-#include "vgg.h"
+#include "models/vgg.h"
 
 #include <unordered_map>
-#include "modelsimpl.h"
+#include "models/modelsimpl.h"
 
 namespace vision {
 namespace models {