Windows CI integration for custom ops (pytorch#12928)

Summary:
Resubmission of pytorch#11527

ezyang orionr
Pull Request resolved: pytorch#12928

Differential Revision: D10501342

Pulled By: goldsborough

fbshipit-source-id: 7ce74795aab2f13efeb38f56ce82f53055f5eade

Author: goldsborough

.jenkins/pytorch/win-test.sh

@@ -87,14 +87,37 @@ call ci_scripts/setup_pytorch_env.bat
 cd test/ && python run_test.py --exclude nn --verbose && cd ..
 EOL
 
+cat >ci_scripts/test_custom_script_ops.bat <<EOL
+call ci_scripts/setup_pytorch_env.bat
+
+cd test/custom_operator
+
+:: Build the custom operator library.
+mkdir build
+cd build
+:: Note: Caffe2 does not support MSVC + CUDA + Debug mode (has to be Release mode)
+cmake -DCMAKE_PREFIX_PATH=%CD%\\..\\..\\torch -DCMAKE_BUILD_TYPE=Release -GNinja ..
+ninja -v
+cd ..
+
+:: Run tests Python-side and export a script module.
+python test_custom_ops.py -v
+python model.py --export-script-module="build/model.pt"
+:: Run tests C++-side and load the exported script module.
+cd build
+set PATH=C:\\Program Files\\NVIDIA Corporation\\NvToolsExt/bin/x64;%CD%\\..\\..\\torch\\lib;%PATH%
+test_custom_ops.exe model.pt
+
+EOL
+
 run_tests() {
     if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
-        ci_scripts/test_python_nn.bat && ci_scripts/test_python_all_except_nn.bat
+        ci_scripts/test_python_nn.bat && ci_scripts/test_python_all_except_nn.bat && ci_scripts/test_custom_script_ops.bat
     else
         if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
             ci_scripts/test_python_nn.bat
         elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then
-            ci_scripts/test_python_all_except_nn.bat
+            ci_scripts/test_python_all_except_nn.bat && ci_scripts/test_custom_script_ops.bat
         fi
     fi
 }

cmake/TorchConfig.cmake.in

@@ -14,6 +14,8 @@
 #
 #   torch
 
+include(FindPackageHandleStandardArgs)
+
 if ($ENV{TORCH_INSTALL_PREFIX})
   set(TORCH_INSTALL_PREFIX $ENV{TORCH_INSTALL_PREFIX})
 else()
@@ -37,7 +39,7 @@ endif()
 find_package(Caffe2 REQUIRED)
 
 find_library(TORCH_LIBRARY torch PATHS "${TORCH_INSTALL_PREFIX}/lib")
-add_library(torch SHARED IMPORTED)
+add_library(torch UNKNOWN IMPORTED)
 set(TORCH_LIBRARIES torch ${Caffe2_MAIN_LIBS})
 
 if (@USE_CUDA@)
@@ -67,11 +69,17 @@ if (@USE_CUDA@)
 endif()
 
 # 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@")
+if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
+  set(TORCH_CXX_FLAGS "-D_GLIBCXX_USE_CXX11_ABI=@GLIBCXX_USE_CXX11_ABI@")
+endif()
 
 set_target_properties(torch PROPERTIES
     IMPORTED_LOCATION "${TORCH_LIBRARY}"
     INTERFACE_INCLUDE_DIRECTORIES "${TORCH_INCLUDE_DIRS}"
-    INTERFACE_COMPILE_OPTIONS "${TORCH_CXX_FLAGS}"
     CXX_STANDARD 11
 )
+if (TORCH_CXX_FLAGS)
+  set_property(TARGET torch PROPERTY INTERFACE_COMPILE_OPTIONS "${TORCH_CXX_FLAGS}")
+endif()
+
+find_package_handle_standard_args(torch DEFAULT_MSG TORCH_LIBRARY TORCH_INCLUDE_DIRS)

test/custom_operator/CMakeLists.txt

@@ -5,7 +5,9 @@ project(custom_ops)
 find_package(Torch REQUIRED)
 
 add_library(custom_ops SHARED op.cpp)
-target_link_libraries(custom_ops ${TORCH_LIBRARIES})
+target_compile_features(custom_ops PUBLIC cxx_range_for)
+target_link_libraries(custom_ops "${TORCH_LIBRARIES}")
+target_compile_definitions(custom_ops PRIVATE custom_ops_EXPORTS)
 
 add_executable(test_custom_ops test_custom_ops.cpp)
 target_link_libraries(test_custom_ops custom_ops)

test/custom_operator/model.py

@@ -4,13 +4,16 @@
 
 import torch
 
-
-SHARED_LIBRARY_EXTENSIONS = {'linux': 'so', 'darwin': 'dylib', 'win32': 'dll'}
+SHARED_LIBRARY_NAMES = {
+    'linux': 'libcustom_ops.so',
+    'darwin': 'libcustom_ops.dylib',
+    'win32': 'custom_ops.dll'
+}
 
 
 def get_custom_op_library_path():
-    extension = SHARED_LIBRARY_EXTENSIONS[sys.platform]
-    path = os.path.abspath('build/libcustom_ops.{}'.format(extension))
+    path = os.path.abspath('build/{}'.format(
+        SHARED_LIBRARY_NAMES[sys.platform]))
     assert os.path.exists(path), path
     return path
 
@@ -27,8 +30,7 @@ def forward(self, input):
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Serialize a script module with custom ops"
-    )
+        description="Serialize a script module with custom ops")
     parser.add_argument("--export-script-module-to", required=True)
     options = parser.parse_args()
 

test/custom_operator/op.cpp

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

test/custom_operator/op.h

@@ -3,7 +3,19 @@
 #include <cstddef>
 #include <vector>
 
-TORCH_API std::vector<at::Tensor> custom_op(
+// clang-format off
+#  if defined(_WIN32)
+#    if defined(custom_ops_EXPORTS)
+#      define CUSTOM_OP_API __declspec(dllexport)
+#    else
+#      define CUSTOM_OP_API __declspec(dllimport)
+#    endif
+#  else
+#    define CUSTOM_OP_API
+#  endif
+// clang-format on
+
+CUSTOM_OP_API std::vector<at::Tensor> custom_op(
     at::Tensor tensor,
     double scalar,
     int64_t repeat);

test/custom_operator/test_custom_ops.cpp

@@ -2,7 +2,6 @@
 
 #include "op.h"
 
-#include <cassert>
 #include <memory>
 #include <string>
 #include <vector>
@@ -26,10 +25,10 @@ void check_all_parameters(
 void get_operator_from_registry_and_execute() {
   auto& ops = torch::jit::getAllOperatorsFor(
       torch::jit::Symbol::fromQualString("custom::op"));
-  assert(ops.size() == 1);
+  AT_ASSERT(ops.size() == 1);
 
   auto& op = ops.front();
-  assert(op->schema().name == "custom::op");
+  AT_ASSERT(op->schema().name == "custom::op");
 
   torch::jit::Stack stack;
   torch::jit::push(stack, torch::ones(5), 2.0, 3);
@@ -39,57 +38,57 @@ void get_operator_from_registry_and_execute() {
 
   const auto manual = custom_op(torch::ones(5), 2.0, 3);
 
-  assert(output.size() == 3);
+  AT_ASSERT(output.size() == 3);
   for (size_t i = 0; i < output.size(); ++i) {
-    assert(output[i].allclose(torch::ones(5) * 2));
-    assert(output[i].allclose(manual[i]));
+    AT_ASSERT(output[i].allclose(torch::ones(5) * 2));
+    AT_ASSERT(output[i].allclose(manual[i]));
   }
 }
 
 void load_serialized_module_with_custom_op_and_execute(
     const std::string& path_to_exported_script_module) {
   std::shared_ptr<torch::jit::script::Module> module =
       torch::jit::load(path_to_exported_script_module);
-  assert(module != nullptr);
+  AT_ASSERT(module != nullptr);
 
   std::vector<torch::jit::IValue> inputs;
   inputs.push_back(torch::ones(5));
   auto output = module->forward(inputs).toTensor();
 
-  assert(output.allclose(torch::ones(5) + 1));
+  AT_ASSERT(output.allclose(torch::ones(5) + 1));
 }
 
 void test_argument_checking_for_serialized_modules(
     const std::string& path_to_exported_script_module) {
   std::shared_ptr<torch::jit::script::Module> module =
       torch::jit::load(path_to_exported_script_module);
-  assert(module != nullptr);
+  AT_ASSERT(module != nullptr);
 
   try {
     module->forward({torch::jit::IValue(1), torch::jit::IValue(2)});
-    assert(false);
+    AT_ASSERT(false);
   } catch (const c10::Error& error) {
-    assert(
+    AT_ASSERT(
         std::string(error.what_without_backtrace())
             .find("Expected at most 1 argument(s) for operator 'forward', "
                   "but received 2 argument(s)") == 0);
   }
 
   try {
     module->forward({torch::jit::IValue(5)});
-    assert(false);
+    AT_ASSERT(false);
   } catch (const c10::Error& error) {
-    assert(
+    AT_ASSERT(
         std::string(error.what_without_backtrace())
             .find("Expected value of type Dynamic for argument 'input' in "
                   "position 0, but instead got value of type int") == 0);
   }
 
   try {
     module->forward({});
-    assert(false);
+    AT_ASSERT(false);
   } catch (const c10::Error& error) {
-    assert(
+    AT_ASSERT(
         std::string(error.what_without_backtrace())
             .find("forward() is missing value for argument 'input'") == 0);
   }

torch/csrc/jit/operator.h

@@ -20,7 +20,7 @@
 
 namespace torch { namespace jit {
 
-FunctionSchema parseSchema(const std::string& schema);
+TORCH_API FunctionSchema parseSchema(const std::string& schema);
 
 using OperationCreator = std::function<Operation(Node*)>;
 
@@ -90,7 +90,7 @@ inline Operation getOperation(Node* node) {
   return getOperatorFor(node).getOperation(node);
 }
 
-void registerOperator(Operator&& op);
+TORCH_API void registerOperator(Operator&& op);
 
 // XXX: this function is meant to be used with string literals only!
 Operator& sig(const char *signature_literal);

torch/csrc/jit/script/module.cpp

@@ -13,30 +13,6 @@ void placeholderCreator(Method&) {
   throw RecursiveMethodCallError();
 }
 
-static FunctionSchema defaultSchemaFor(const Method& method) {
-  std::vector<Argument> args;
-  std::vector<Argument> returns;
-  Graph& g = *method.graph();
-  size_t num_inputs = method.num_inputs();
-  for(size_t i = 0; i < num_inputs; ++i) {
-    const Value* v = g.inputs().at(i);
-    std::string name = v->hasUniqueName() ? v->uniqueName() : ("argument_"  + std::to_string(i));
-    args.push_back({std::move(name), unshapedType(g.inputs()[i]->type())});
-  }
-  for(size_t i = 0; i < g.outputs().size(); ++i) {
-    returns.push_back({"", unshapedType(g.outputs()[i]->type())});
-  }
-  return { method.name(), std::move(args), std::move(returns) };
-}
-
-
-const FunctionSchema& Method::getSchema() const {
-  if(schema == nullptr) {
-    schema.reset(new FunctionSchema(defaultSchemaFor(*this)));
-  }
-  return *schema;
-}
-
 c10::optional<std::vector<Value*>> try_emit_call_to(
     Graph& graph,
     SourceRange loc,

torch/csrc/jit/script/module.h

@@ -11,6 +11,7 @@
 
 #include <torch/csrc/api/include/torch/detail/ordered_dict.h>
 #include <torch/csrc/utils/memory.h>
+#include <torch/csrc/WindowsTorchApiMacro.h>
 
 #include <ATen/core/ArrayRef.h>
 #include "c10/util/Optional.h"
@@ -156,7 +157,12 @@ struct Method {
     return *this;
   }
 
-  const FunctionSchema& getSchema() const;
+  const FunctionSchema& getSchema() const {
+    if(schema == nullptr) {
+      schema.reset(new FunctionSchema(defaultSchemaFor(*this)));
+    }
+    return *schema;
+  }
 
   std::string pretty_print_schema() const {
     JIT_ASSERT(schema);
@@ -178,6 +184,23 @@ struct Method {
   }
 
 private:
+
+  static FunctionSchema defaultSchemaFor(const Method& method) {
+    std::vector<Argument> args;
+    std::vector<Argument> returns;
+    Graph& g = *method.graph();
+    size_t num_inputs = method.num_inputs();
+    for(size_t i = 0; i < num_inputs; ++i) {
+      const Value* v = g.inputs().at(i);
+      std::string name = v->hasUniqueName() ? v->uniqueName() : ("argument_"  + std::to_string(i));
+      args.push_back({std::move(name), unshapedType(g.inputs()[i]->type())});
+    }
+    for(size_t i = 0; i < g.outputs().size(); ++i) {
+      returns.push_back({"", unshapedType(g.outputs()[i]->type())});
+    }
+    return { method.name(), std::move(args), std::move(returns) };
+  }
+
   std::string name_;
   std::shared_ptr<Graph> graph_; // for debugging and for inlining
   bool optimize;
@@ -368,7 +391,7 @@ struct Module {
   /// destination is on the GPU or vice versa, the copy is performed
   /// asynchronously with respect to the host. Otherwise, the argument has no
   /// effect.
-  void to(
+  TORCH_API void to(
       at::Device device,
       at::ScalarType dtype,
       bool non_blocking = false);
@@ -379,15 +402,15 @@ struct Module {
   /// destination is on the GPU or vice versa, the copy is performed
   /// asynchronously with respect to the host. Otherwise, the argument has no
   /// effect.
-  void to(at::ScalarType dtype, bool non_blocking = false);
+  TORCH_API void to(at::ScalarType dtype, bool non_blocking = false);
 
   /// Recursively moves all parameters to the given device.
   ///
   /// If `non_blocking` is true and the source is in pinned memory and
   /// destination is on the GPU or vice versa, the copy is performed
   /// asynchronously with respect to the host. Otherwise, the argument has no
   /// effect.
-  void to(at::Device device, bool non_blocking = false);
+  TORCH_API void to(at::Device device, bool non_blocking = false);
 
   /// Run a method from this module.
   ///

torch/script.h

@@ -3,6 +3,5 @@
 #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>