Integrate custom op tests with CI (pytorch#10611)

Summary:
This PR is stacked on pytorch#10610, and only adds changes in one file `.jenkins/pytorch/test.sh`, where we now build the custom op tests and run them.

I'd also like to take this PR to discuss whether the [`TorchConfig.cmake`](https://github.com/pytorch/pytorch/blob/master/cmake/TorchConfig.cmake.in) I made is robust enough (we will also see in the CI) orionr Yangqing dzhulgakov what do you think?

Also ezyang for CI changes
Pull Request resolved: pytorch#10611

Differential Revision: D9597627

Pulled By: goldsborough

fbshipit-source-id: f5af8164c076894f448cef7e5b356a6b3159f8b3

Author: goldsborough

.gitignore

@@ -30,6 +30,7 @@ docs/cpp/html/
 docs/cpp/api/
 test/.coverage
 test/cpp/api/mnist
+test/custom_operator/model.pt
 test/data/gpu_tensors.pt
 test/data/legacy_modules.t7
 test/data/legacy_serialized.pt
@@ -66,6 +67,7 @@ torch/lib/protoc
 torch/lib/tmp_install
 torch/lib/torch_shm_manager
 torch/lib/python*
+torch/share/
 torch/version.py
 
 # IPython notebook checkpoints

.jenkins/pytorch/build.sh

@@ -138,4 +138,14 @@ if [[ "$BUILD_TEST_LIBTORCH" == "1" ]]; then
   pushd ../cpp-build/caffe2
   WERROR=1 VERBOSE=1 DEBUG=1 python $BUILD_LIBTORCH_PY
   popd
+
+  # Build custom operator tests.
+  CUSTOM_OP_BUILD="$PWD/../custom-op-build"
+  CUSTOM_OP_TEST="$PWD/test/custom_operator"
+  SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
+  mkdir "$CUSTOM_OP_BUILD"
+  pushd "$CUSTOM_OP_BUILD"
+  CMAKE_PREFIX_PATH="$SITE_PACKAGES/torch" cmake "$CUSTOM_OP_TEST"
+  make VERBOSE=1
+  popd
 fi

.jenkins/pytorch/macos-test.sh

@@ -78,13 +78,35 @@ test_cpp_api() {
   "$CPP_BUILD"/caffe2/bin/test_api
 }
 
+test_custom_script_ops() {
+  echo "Testing custom script operators"
+  pushd test/custom_operator
+  # Build the custom operator library.
+  rm -rf build && mkdir build
+  pushd build
+  SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
+  CMAKE_PREFIX_PATH="$SITE_PACKAGES/torch" cmake ..
+  make VERBOSE=1
+  popd
+
+  # Run tests Python-side and export a script module.
+  python test_custom_ops.py -v
+  python model.py --export-script-module=model.pt
+  # Run tests C++-side and load the exported script module.
+  build/test_custom_ops ./model.pt
+  popd
+}
+
+
 if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
   test_python_all
   test_cpp_api
+  test_custom_script_ops
 else
   if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
     test_python_all
   elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then
     test_cpp_api
+    test_custom_script_ops
   fi
 fi

.jenkins/pytorch/test.sh

@@ -97,7 +97,7 @@ test_aten() {
     # put the dynamic libraries somewhere were the dynamic linker can find them.
     # This is a bit of a hack.
     if [[ "$BUILD_ENVIRONMENT" == *ppc64le* ]]; then
-      SUDO=sudo 
+      SUDO=sudo
     fi
 
     ${SUDO} ln -s "$TORCH_LIB_PATH"/libcaffe2* build/bin
@@ -140,12 +140,28 @@ test_libtorch() {
   fi
 }
 
+test_custom_script_ops() {
+  if [[ "$BUILD_TEST_LIBTORCH" == "1" ]]; then
+    echo "Testing custom script operators"
+    CUSTOM_OP_BUILD="$PWD/../custom-op-build"
+    pushd test/custom_operator
+    cp -r "$CUSTOM_OP_BUILD" build
+    # Run tests Python-side and export a script module.
+    python test_custom_ops.py -v
+    python model.py --export-script-module=model.pt
+    # Run tests C++-side and load the exported script module.
+    build/test_custom_ops ./model.pt
+    popd
+  fi
+}
+
 if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
   test_python_nn
   test_python_all_except_nn
   test_aten
   test_torchvision
   test_libtorch
+  test_custom_script_ops
 else
   if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
     test_python_nn
@@ -154,5 +170,6 @@ else
     test_aten
     test_torchvision
     test_libtorch
+    test_custom_script_ops
   fi
 fi

cmake/TorchConfig.cmake.in

@@ -17,32 +17,48 @@
 #
 #   torch_add_custom_op_library(<name> <source_files>)
 
-SET(TORCH_ROOT "${CMAKE_CURRENT_LIST_DIR}/../")
+if ($ENV{TORCH_INSTALL_PREFIX})
+  set(TORCH_INSTALL_PREFIX $ENV{TORCH_INSTALL_PREFIX})
+else()
+  # Assume we are in <install-prefix>/share/cmake/Torch/TorchConfig.cmake
+  get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH)
+  get_filename_component(TORCH_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE)
+endif()
 
-set(TORCH_INCLUDE_DIRS
-  "${TORCH_ROOT}"
-  "${TORCH_ROOT}/aten/src"
-  "${CMAKE_CURRENT_LIST_DIR}/aten/src"
-  "${CMAKE_CURRENT_LIST_DIR}/caffe2/aten/src"
-  "${CMAKE_CURRENT_LIST_DIR}/caffe2/aten/src/TH"
-)
+# Include directories.
+set(TORCH_INCLUDE_DIRS "${TORCH_INSTALL_PREFIX}/lib/include")
 
-find_library(TORCH_LIBRARY torch PATHS "${CMAKE_CURRENT_LIST_DIR}/lib" NO_DEFAULT_PATH)
-find_library(CAFFE2_LIBRARY caffe2 PATHS "${CMAKE_CURRENT_LIST_DIR}/lib" NO_DEFAULT_PATH)
+# Library dependencies.
+find_package(Caffe2 REQUIRED)
+find_library(TORCH_LIBRARY torch PATHS "${TORCH_INSTALL_PREFIX}/lib")
 
+set(TORCH_LIBRARIES ${TORCH_LIBRARY} ${Caffe2_MAIN_LIBS})
 if (@USE_CUDA@)
-  find_package(CUDA REQUIRED)
-  find_library(CAFFE2_CUDA_LIBRARY caffe2_gpu PATHS "${CMAKE_CURRENT_LIST_DIR}/lib" NO_DEFAULT_PATH)
-  set(TORCH_CUDA_LIBRARIES -L${CUDA_TOOLKIT_ROOT_DIR}/lib64 cuda nvrtc cudart nvToolsExt)
-  list(APPEND TORCH_INCLUDE_DIRS ${CUDA_TOOLKIT_INCLUDE})
+  if(MSVC)
+    set(NVTOOLEXT_HOME "C:/Program Files/NVIDIA Corporation/NvToolsExt")
+    if ($ENV{NVTOOLEXT_HOME})
+      set(NVTOOLEXT_HOME $ENV{NVTOOLEXT_HOME})
+    endif()
+    set(TORCH_CUDA_LIBRARIES
+      ${NVTOOLEXT_HOME}/lib/x64/nvToolsExt64_1.lib
+      ${CUDA_LIBRARIES})
+    list(APPEND TORCH_INCLUDE_DIRS "${NVTOOLEXT_HOME}/include")
+  elseif(APPLE)
+    set(TORCH_CUDA_LIBRARIES
+      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libcudart.dylib
+      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libnvrtc.dylib
+      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libnvToolsExt.dylib
+      ${CUDA_LIBRARIES})
+  else()
+    set(TORCH_CUDA_LIBRARIES
+      ${CUDA_CUDA_LIB}
+      ${CUDA_NVRTC_LIB}
+      ${CUDA_TOOLKIT_ROOT_DIR}/lib64/libnvToolsExt.so
+      ${CUDA_LIBRARIES})
+  endif()
+  list(APPEND TORCH_LIBRARIES ${TORCH_CUDA_LIBRARIES})
 endif()
 
-set(TORCH_LIBRARIES
-  ${TORCH_LIBRARY}
-  ${CAFFE2_LIBRARY}
-  ${CAFFE2_CUDA_LIBRARY}
-  ${TORCH_CUDA_LIBRARIES})
-
 # 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
@@ -51,5 +67,5 @@ 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}")
-  target_compile_options(${name} PUBLIC -std=c++11)
+  set_property(TARGET ${name} PROPERTY CXX_STANDARD 11)
 endfunction(torch_add_custom_op_library)

setup.py

@@ -443,11 +443,10 @@ def check_file(f):
         # we need to find a better way to do this.
         # More information can be found in conversation thread of PR #5772
 
-        self.copy_tree('torch/csrc', 'torch/lib/include/torch/csrc/')
+        self.copy_tree('torch/lib/tmp_install/share', 'torch/share')
         self.copy_tree('third_party/pybind11/include/pybind11/',
                        'torch/lib/include/pybind11')
         self.copy_file('torch/csrc/torch.h', 'torch/lib/include/torch/torch.h')
-        self.copy_file('torch/op.h', 'torch/lib/include/torch/op.h')
 
 
 build_dep_cmds = {}
@@ -1136,27 +1135,38 @@ def make_relative_rpath(path):
                 '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',
+                'lib/include/ATen/cuda/*.h',
+                'lib/include/ATen/cuda/detail/*.cuh',
                 'lib/include/ATen/cuda/detail/*.h',
                 'lib/include/ATen/cudnn/*.h',
-                'lib/include/ATen/cuda/detail/*.cuh',
+                'lib/include/ATen/detail/*.h',
+                'lib/include/caffe2/utils/*.h',
+                'lib/include/torch/*.h',
+                'lib/include/torch/csrc/*.h',
+                'lib/include/torch/csrc/api/include/torch/detail/ordered_dict.h',
+                'lib/include/torch/csrc/autograd/*.h',
+                'lib/include/torch/csrc/autograd/generated/*.h',
+                'lib/include/torch/csrc/cuda/*.h',
+                'lib/include/torch/csrc/jit/*.h',
+                'lib/include/torch/csrc/jit/generated/*.h',
+                'lib/include/torch/csrc/jit/passes/*.h',
+                'lib/include/torch/csrc/jit/script/*.h',
+                'lib/include/torch/csrc/utils/*.h',
                 'lib/include/pybind11/*.h',
                 'lib/include/pybind11/detail/*.h',
                 'lib/include/TH/*.h*',
                 'lib/include/TH/generic/*.h*',
-                'lib/include/THC/*.h*',
                 'lib/include/THC/*.cuh',
+                'lib/include/THC/*.h*',
                 'lib/include/THC/generic/*.h',
                 'lib/include/THCUNN/*.cuh',
                 'lib/include/THNN/*.h',
-                'lib/include/torch/csrc/*.h',
-                'lib/include/torch/csrc/autograd/*.h',
-                'lib/include/torch/csrc/jit/*.h',
-                'lib/include/torch/csrc/utils/*.h',
-                'lib/include/torch/csrc/cuda/*.h',
-                'lib/include/torch/torch.h',
+                'share/cmake/ATen/*.cmake',
+                'share/cmake/Caffe2/*.cmake',
+                'share/cmake/Caffe2/public/*.cmake',
+                'share/cmake/Gloo/*.cmake',
+                'share/cmake/Torch/*.cmake',
             ],
             'caffe2': [
                 rel_site_packages + '/caffe2/**/*.py'

test/custom_operator/CMakeLists.txt

@@ -4,6 +4,15 @@ 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_executable(test_custom_ops test_custom_ops.cpp)

test/custom_operator/model.py

@@ -1,9 +1,20 @@
 import argparse
 import os.path
+import sys
 
 import torch
 
 
+SHARED_LIBRARY_EXTENSIONS = {'linux': 'so', 'darwin': 'dylib', 'win32': 'dll'}
+
+
+def get_custom_op_library_path():
+    extension = SHARED_LIBRARY_EXTENSIONS[sys.platform]
+    path = os.path.abspath('build/libcustom_ops.{}'.format(extension))
+    assert os.path.exists(path), path
+    return path
+
+
 class Model(torch.jit.ScriptModule):
     def __init__(self):
         super(Model, self).__init__()
@@ -20,7 +31,7 @@ def main():
     parser.add_argument("--export-script-module-to", required=True)
     options = parser.parse_args()
 
-    torch.ops.load_library(os.path.abspath('build/libcustom_ops.so'))
+    torch.ops.load_library(get_custom_op_library_path())
 
     model = Model()
     model.save(options.export_script_module_to)

test/custom_operator/op.h

@@ -3,7 +3,7 @@
 #include <cstddef>
 #include <vector>
 
-std::vector<at::Tensor> custom_op(
+TORCH_API std::vector<at::Tensor> custom_op(
     at::Tensor tensor,
     double scalar,
     int64_t repeat);

test/custom_operator/test_custom_ops.cpp

@@ -22,9 +22,12 @@ void get_operator_from_registry_and_execute() {
   std::vector<at::Tensor> output;
   torch::jit::pop(stack, output);
 
+  const auto manual = custom_op(torch::ones(5), 2.0, 3);
+
   assert(output.size() == 3);
-  for (const auto& tensor : output) {
-    assert(tensor.allclose(torch::ones(5) * 2));
+  for (size_t i = 0; i < output.size(); ++i) {
+    assert(output[i].allclose(torch::ones(5) * 2));
+    assert(output[i].allclose(manual[i]));
   }
 }
 
@@ -71,10 +74,9 @@ void test_argument_checking_for_serialized_modules(
     module->forward({});
     assert(false);
   } catch (const at::Error& error) {
-    std::cout << error.what_without_backtrace() << std::endl;
     assert(
         std::string(error.what_without_backtrace())
-            .find("custom::op() is missing value for argument 'tensor'") == 0);
+            .find("forward() is missing value for argument 'input'") == 0);
   }
 }
 

test/custom_operator/test_custom_ops.py

@@ -5,12 +5,12 @@
 
 import torch
 
-from model import Model
+from model import Model, get_custom_op_library_path
 
 
 class TestCustomOperators(unittest.TestCase):
     def setUp(self):
-        self.library_path = os.path.abspath('build/libcustom_ops.so')
+        self.library_path = get_custom_op_library_path()
         torch.ops.load_library(self.library_path)
 
     def test_custom_library_is_loaded(self):

torch/CMakeLists.txt

@@ -26,7 +26,7 @@ if(NOT TORCH_INSTALL_BIN_DIR)
 endif()
 
 if(NOT TORCH_INSTALL_INCLUDE_DIR)
-  set(TORCH_INSTALL_INCLUDE_DIR include/libtorch)
+  set(TORCH_INSTALL_INCLUDE_DIR include)
 endif()
 
 if(NOT TORCH_INSTALL_LIB_DIR)
@@ -287,8 +287,7 @@ if (MSVC)
   target_link_libraries(torch onnx onnx_library)
 endif()
 
-target_link_libraries(torch
-  caffe2_library)
+target_link_libraries(torch caffe2_library)
 
 find_package(OpenMP)
 if(OPENMP_FOUND)
@@ -395,6 +394,8 @@ endif()
 install(DIRECTORY "${TORCH_SRC_DIR}/csrc"
         DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch
         FILES_MATCHING PATTERN "*.h")
+install(FILES "${TORCH_SRC_DIR}/op.h"
+        DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch)
 
 install(TARGETS torch
   RUNTIME DESTINATION "${TORCH_INSTALL_BIN_DIR}"
@@ -473,3 +474,7 @@ configure_file(
     ${TORCH_ROOT}/cmake/TorchConfig.cmake.in
     ${PROJECT_BINARY_DIR}/TorchConfig.cmake
     @ONLY)
+install(FILES
+    ${PROJECT_BINARY_DIR}/TorchConfigVersion.cmake
+    ${PROJECT_BINARY_DIR}/TorchConfig.cmake
+    DESTINATION share/cmake/Torch)

torch/op.h

@@ -3,6 +3,7 @@
 #include <torch/csrc/autograd/generated/variable_factories.h>
 #include <torch/csrc/jit/custom_operator.h>
 #include <torch/csrc/jit/import.h>
+#include <torch/csrc/WindowsTorchApiMacro.h>
 
 #include <ATen/ATen.h>