Fixes for Torch Script C++ API (pytorch#11682)

Summary:
A couple fixes I deem necessary to the TorchScript C++ API after writing the tutorial:

1. When I was creating the custom op API, I created `torch/op.h` as the one-stop header for creating custom ops. I now notice that there is no good header for the TorchScript C++ story altogether, i.e. when you just want to load a script module in C++ without any custom ops necessarily. The `torch/op.h` header suits that purpose just as well of course, but I think we should rename it to `torch/script.h`, which seems like a great name for this feature.

2. The current API for the CMake we provided was that we defined a bunch of variables like `TORCH_LIBRARY_DIRS` and `TORCH_INCLUDES` and then expected users to add those variables to their targets. We also had a CMake function that did that for you automatically. I now realized a much smarter way of doing this is to create an `IMPORTED` target for the libtorch library in CMake, and then add all this stuff to the link interface of that target. Then all downstream users have to do is `target_link_libraries(my_target torch)` and they get all the proper includes, libraries and compiler flags added to their target. This means we can get rid of the CMake function and all that stuff. orionr  AFAIK this is a much, much better way of doing all of this, no?

3. Since we distribute libtorch with `D_GLIBCXX_USE_CXX11_ABI=0`, dependent libraries must set this flag too. I now add this to the interface compile options of this imported target.

4. Fixes to JIT docs.

These could likely be 4 different PRs but given the release I wouldn't mind landing them all asap.

zdevito dzhulgakov soumith
Pull Request resolved: pytorch#11682

Differential Revision: D9839431

Pulled By: goldsborough

fbshipit-source-id: fdc47b95f83f22d53e1995aa683e09613b4bfe65

Author: goldsborough

cmake/TorchConfig.cmake.in

@@ -7,15 +7,12 @@
 #
 #   TORCH_FOUND        -- True if the system has the Torch library
 #   TORCH_INCLUDE_DIRS -- The include directories for torch
-#   TORCH_LIBRARIES    -- Libraries to link to
+#   TORCH_LIBRARIES    -- Libraries to link against
+#   TORCH_CXX_FLAGS    -- Additional (required) compiler flags
 #
 # and the following imported targets:
 #
-#   Torch
-#
-# and the following functions:
-#
-#   torch_add_custom_op_library(<name> <source_files>)
+#   torch
 
 if ($ENV{TORCH_INSTALL_PREFIX})
   set(TORCH_INSTALL_PREFIX $ENV{TORCH_INSTALL_PREFIX})
@@ -26,13 +23,19 @@ else()
 endif()
 
 # Include directories.
-set(TORCH_INCLUDE_DIRS "${TORCH_INSTALL_PREFIX}/lib/include")
+if (EXISTS "${TORCH_INSTALL_PREFIX}/lib/include")
+  set(TORCH_INCLUDE_DIRS "${TORCH_INSTALL_PREFIX}/lib/include")
+else()
+  set(TORCH_INCLUDE_DIRS "${TORCH_INSTALL_PREFIX}/include")
+endif()
 
 # Library dependencies.
 find_package(Caffe2 REQUIRED)
+
 find_library(TORCH_LIBRARY torch PATHS "${TORCH_INSTALL_PREFIX}/lib")
+add_library(torch SHARED IMPORTED)
+set(TORCH_LIBRARIES torch ${Caffe2_MAIN_LIBS})
 
-set(TORCH_LIBRARIES ${TORCH_LIBRARY} ${Caffe2_MAIN_LIBS})
 if (@USE_CUDA@)
   if(MSVC)
     set(NVTOOLEXT_HOME "C:/Program Files/NVIDIA Corporation/NvToolsExt")
@@ -59,13 +62,12 @@ if (@USE_CUDA@)
   list(APPEND TORCH_LIBRARIES ${TORCH_CUDA_LIBRARIES})
 endif()
 
-# Creates a shared library <name> with the correct include directories
-# and linker flags set to include Torch header files and link with Torch
-# libraries. Also sets the C++ standard version to C++11. All options
-# can be override by specifying further options on the `<name>` CMake target.
-function(torch_add_custom_op_library name source_files)
-  add_library(${name} SHARED ${source_files})
-  target_include_directories(${name} PUBLIC "${TORCH_INCLUDE_DIRS}")
-  target_link_libraries(${name} "${TORCH_LIBRARIES}")
-  set_property(TARGET ${name} PROPERTY CXX_STANDARD 11)
-endfunction(torch_add_custom_op_library)
+# When we build libtorch with the old GCC ABI, dependent libraries must too.
+set(TORCH_CXX_FLAGS "-D_GLIBCXX_USE_CXX11_ABI=@GLIBCXX_USE_CXX11_ABI@")
+
+set_target_properties(torch PROPERTIES
+    IMPORTED_LOCATION ${TORCH_LIBRARY}
+    INTERFACE_INCLUDE_DIRECTORIES ${TORCH_INCLUDE_DIRS}
+    INTERFACE_COMPILE_OPTIONS ${TORCH_CXX_FLAGS}
+    CXX_STANDARD 11
+)

docs/source/jit.rst

@@ -7,13 +7,13 @@ Torch Script
 .. currentmodule:: torch.jit
 
 Torch Script is a way to create serializable and optimizable models from PyTorch code.
-Anything code written in Torch Script can be saved from your Python
-process and loaded/run a process where there is no python dependency.
+Any code written in Torch Script can be saved from your Python
+process and loaded in a process where there is no Python dependency.
 
 We provide tools to incrementally transition a model from being a pure Python program
-to a Torch Script program that can be run independently from python, for instance, in a standalone C++ process.
+to a Torch Script program that can be run independently from Python, for instance, in a standalone C++ program.
 This makes it possible to train models in PyTorch using familiar tools and then export
-the model to a production environment where it is not a good idea to run models as python programs
+the model to a production environment where it is not a good idea to run models as Python programs
 for performance and multi-threading reasons.
 
 Creating Torch Script Code
@@ -47,7 +47,7 @@ Mixing Tracing and Scripting
 ----------------------------
 
 In many cases either tracing or script is an easier approach for converting a model.
-We allow you to compose tracing and scripting to suite the particular requirements
+We allow you to compose tracing and scripting to suit the particular requirements
 of a part of a model.
 
 Scripted functions can call traced ones. This is particularly useful when you need
@@ -79,6 +79,7 @@ Example:
 ::
 
     import torch
+
     @torch.jit.script
     def foo(x, y):
         if x.max() > y.max():

test/custom_operator/CMakeLists.txt

@@ -1,19 +1,12 @@
 # Basic CMake setup
-cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
 project(custom_ops)
 
 find_package(Torch REQUIRED)
 
-# This convenience function will create a shared library target, configure
-# the right include directories and link against the right libraries. It is
-# exactly equivalent to the following lines:
-#
-# add_library(custom_ops SHARED op.cpp)
-# target_include_directories(custom_ops PUBLIC "${TORCH_INCLUDE_DIRS}")
-# target_link_libraries(custom_ops "${TORCH_LIBRARIES}")
-# set_property(TARGET custom_ops PROPERTY CXX_STANDARD 11)
-#
-torch_add_custom_op_library(custom_ops op.cpp)
+add_library(custom_ops SHARED op.cpp)
+target_compile_features(custom_ops PUBLIC cxx_range_for)
+target_link_libraries(custom_ops ${TORCH_LIBRARIES})
 
 add_executable(test_custom_ops test_custom_ops.cpp)
 target_link_libraries(test_custom_ops custom_ops)

test/custom_operator/op.cpp

@@ -1,4 +1,4 @@
-#include <torch/op.h>
+#include <torch/script.h>
 
 #include <cstddef>
 #include <vector>

test/custom_operator/op.h

@@ -1,4 +1,4 @@
-#include <torch/op.h>
+#include <torch/script.h>
 
 #include <cstddef>
 #include <vector>

test/custom_operator/test_custom_ops.cpp

@@ -1,4 +1,4 @@
-#include <torch/op.h>
+#include <torch/script.h>
 
 #include "op.h"
 

torch/CMakeLists.txt

@@ -413,7 +413,7 @@ endif()
 install(DIRECTORY "${TORCH_SRC_DIR}/csrc"
         DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch
         FILES_MATCHING PATTERN "*.h")
-install(FILES "${TORCH_SRC_DIR}/op.h"
+install(FILES "${TORCH_SRC_DIR}/script.h"
         DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch)
 
 install(TARGETS torch
@@ -488,6 +488,28 @@ if (BUILD_TEST AND NOT NO_API AND NOT USE_ROCM)
   endif()
 endif()
 
+if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
+  message(STATUS "${CMAKE_CXX_COMPILER} ${CMAKE_CURRENT_LIST_DIR}/abi-check.cpp -o ${CMAKE_BINARY_DIR}/abi-check")
+  execute_process(
+    COMMAND
+      "${CMAKE_CXX_COMPILER}"
+      "${CMAKE_CURRENT_LIST_DIR}/abi-check.cpp"
+      "-o"
+      "${CMAKE_BINARY_DIR}/abi-check"
+    RESULT_VARIABLE ABI_CHECK_COMPILE_RESULT)
+  if (ABI_CHECK_COMPILE_RESULT)
+    message(FATAL_ERROR "Could not compile ABI Check: ${ABI_CHECK_COMPILE_RESULT}")
+  endif()
+  execute_process(
+    COMMAND "${CMAKE_BINARY_DIR}/abi-check"
+    RESULT_VARIABLE ABI_CHECK_RESULT
+    OUTPUT_VARIABLE GLIBCXX_USE_CXX11_ABI)
+  if (ABI_CHECK_RESULT)
+    message(WARNING "Could not run ABI Check: ${ABI_CHECK_RESULT}")
+  endif()
+  message(STATUS "Determined _GLIBCXX_USE_CXX11_ABI=${GLIBCXX_USE_CXX11_ABI}")
+endif()
+
 # CMake config for external projects.
 configure_file(
     ${PROJECT_SOURCE_DIR}/cmake/TorchConfigVersion.cmake.in

torch/abi-check.cpp

@@ -0,0 +1,9 @@
+#include <iostream>
+
+int main() {
+#ifdef _GLIBCXX_USE_CXX11_ABI
+  std::cout << _GLIBCXX_USE_CXX11_ABI;
+#else
+  std::cout << 0;
+#endif
+}

torch/op.h torch/script.h

@@ -288,7 +288,7 @@ def _add_gnu_abi_flag_if_binary(self, extension):
         # if the extension is compiled with gcc >= 5.1,
         # then we have to define _GLIBCXX_USE_CXX11_ABI=0
         # so that the std::string in the API is resolved to
-        # non-C++11 symbols
+        # non-C++11 symbols.
         define = '-D_GLIBCXX_USE_CXX11_ABI=0'
         if is_binary_build():
             if isinstance(extension.extra_compile_args, dict):