Restructuring C++ project: (#3146)

Summary:
* Reduce unnecessary header inclusions in models and io.

* Move autocast to separate folder and hide autograd implementation in an anonymous namespace.

* Moving files in subfolders.

Reviewed By: fmassa

Differential Revision: D25461523

fbshipit-source-id: 756eeb6848aacaa474de4825ed4c1045d17e2cea

CMakeLists.txt

@@ -28,19 +28,6 @@ function(CUDA_CONVERT_FLAGS EXISTING_TARGET)
     endif()
 endfunction()
 
-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_HEADERS torchvision/csrc/cpu/*.h)
-file(GLOB OPERATOR_SOURCES ${OPERATOR_HEADERS} torchvision/csrc/cpu/*.cpp ${IMAGE_HEADERS} ${IMAGE_SOURCES} ${HEADERS} torchvision/csrc/*.cpp)
-if(WITH_CUDA)
-  file(GLOB OPERATOR_HEADERS ${OPERATOR_HEADERS} torchvision/csrc/cuda/*.h)
-  file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} ${OPERATOR_HEADERS} torchvision/csrc/cuda/*.cu)
-endif()
-file(GLOB MODELS_HEADERS torchvision/csrc/models/*.h)
-file(GLOB MODELS_SOURCES torchvision/csrc/models/*.h torchvision/csrc/models/*.cpp)
-
 if(MSVC)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4819")
   if(WITH_CUDA)
@@ -64,7 +51,17 @@ endif()
 include(GNUInstallDirs)
 include(CMakePackageConfigHelpers)
 
-add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES} ${IMAGE_SOURCES})
+set(TVCPP torchvision/csrc)
+list(APPEND ALLOW_LISTED ${TVCPP} ${TVCPP}/io/image ${TVCPP}/io/image/cpu ${TVCPP}/models ${TVCPP}/ops ${TVCPP}/ops/cpu)
+if(WITH_CUDA)
+    list(APPEND ALLOW_LISTED ${TVCPP}/ops/cuda ${TVCPP}/ops/autocast)
+endif()
+
+FOREACH(DIR ${ALLOW_LISTED})
+    file(GLOB ALL_SOURCES ${ALL_SOURCES} ${DIR}/*.*)
+ENDFOREACH()
+
+add_library(${PROJECT_NAME} SHARED ${ALL_SOURCES})
 target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} ${PNG_LIBRARY} ${JPEG_LIBRARIES} Python3::Python)
 set_target_properties(${PROJECT_NAME} PROPERTIES
   EXPORT_NAME TorchVision
@@ -95,13 +92,8 @@ install(EXPORT TorchVisionTargets
   NAMESPACE TorchVision::
   DESTINATION ${TORCHVISION_CMAKECONFIG_INSTALL_DIR})
 
-install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME})
-install(FILES
-  ${OPERATOR_HEADERS}
-  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cpu)
-if(WITH_CUDA)
-  install(FILES
-    ${OPERATOR_HEADERS}
-    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cuda)
-endif()
-install(FILES ${MODELS_HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/models)
+FOREACH(INPUT_DIR ${ALLOW_LISTED})
+    string(REPLACE "${TVCPP}" "${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}" OUTPUT_DIR ${INPUT_DIR})
+    file(GLOB INPUT_FILES ${INPUT_DIR}/*.*)
+    install(FILES ${INPUT_FILES} DESTINATION ${OUTPUT_DIR})
+ENDFOREACH()

examples/cpp/hello_world/main.cpp

 #include <iostream>
+#include <torch/torch.h>
+#include <torchvision/vision.h>
 #include <torchvision/models/resnet.h>
 
 int main()

setup.py

@@ -134,8 +134,9 @@ def get_extensions():
     this_dir = os.path.dirname(os.path.abspath(__file__))
     extensions_dir = os.path.join(this_dir, 'torchvision', 'csrc')
 
-    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_dir, '*.cpp')) + glob.glob(os.path.join(extensions_dir, 'ops',
+                                                                                          '*.cpp'))
+    source_cpu = glob.glob(os.path.join(extensions_dir, 'ops', 'cpu', '*.cpp'))
 
     is_rocm_pytorch = False
     if torch.__version__ >= '1.5':
@@ -146,17 +147,19 @@ def get_extensions():
         hipify_python.hipify(
             project_directory=this_dir,
             output_directory=this_dir,
-            includes="torchvision/csrc/cuda/*",
+            includes="torchvision/csrc/ops/cuda/*",
             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_dir, 'ops', 'hip', '*.hip'))
         # Copy over additional files
-        for file in glob.glob(r"torchvision/csrc/cuda/*.h"):
-            shutil.copy(file, "torchvision/csrc/hip")
+        for file in glob.glob(r"torchvision/csrc/ops/cuda/*.h"):
+            shutil.copy(file, "torchvision/csrc/ops/hip")
 
     else:
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'cuda', '*.cu'))
+        source_cuda = glob.glob(os.path.join(extensions_dir, 'ops', 'cuda', '*.cu'))
+
+    source_cuda += glob.glob(os.path.join(extensions_dir, 'ops', 'autocast', '*.cpp'))
 
     sources = main_file + source_cpu
     extension = CppExtension
@@ -309,8 +312,8 @@ def get_extensions():
             image_library += [jpeg_lib]
             image_include += [jpeg_include]
 
-    image_path = os.path.join(extensions_dir, 'cpu', 'image')
-    image_src = glob.glob(os.path.join(image_path, '*.cpp'))
+    image_path = os.path.join(extensions_dir, 'io', 'image')
+    image_src = glob.glob(os.path.join(image_path, '*.cpp')) + glob.glob(os.path.join(image_path, 'cpu', '*.cpp'))
 
     if png_found or jpeg_found:
         ext_modules.append(extension(
@@ -377,13 +380,13 @@ def get_extensions():
         print("ffmpeg library_dir: {}".format(ffmpeg_library_dir))
 
         # TorchVision base decoder + video reader
-        video_reader_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'video_reader')
+        video_reader_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'video_reader')
         video_reader_src = glob.glob(os.path.join(video_reader_src_dir, "*.cpp"))
-        base_decoder_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'decoder')
+        base_decoder_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'decoder')
         base_decoder_src = glob.glob(
             os.path.join(base_decoder_src_dir, "*.cpp"))
         # Torchvision video API
-        videoapi_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'video')
+        videoapi_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'video')
         videoapi_src = glob.glob(os.path.join(videoapi_src_dir, "*.cpp"))
         # exclude tests
         base_decoder_src = [x for x in base_decoder_src if '_test.cpp' not in x]

test/tracing/frcnn/test_frcnn_tracing.cpp

-#include <ATen/ATen.h>
 #include <torch/script.h>
 #include <torch/torch.h>
-#include <torchvision/nms.h>
-#include <torchvision/roi_align.h>
+#include <torchvision/vision.h>
+#include <torchvision/ops/nms.h>
 
 #ifdef _WIN32
 // Windows only

torchvision/csrc/cpu/image/image.h

-#pragma once
-
-// Comment
-#include <torch/script.h>
-#include <torch/torch.h>
-#include "read_image_cpu.h"
-#include "read_write_file_cpu.h"
-#include "readjpeg_cpu.h"
-#include "readpng_cpu.h"
-#include "writejpeg_cpu.h"
-#include "writepng_cpu.h"