[JIT] Add out-of-source-tree to_backend tests (pytorch#41145)

Summary:
Pull Request resolved: pytorch#41145

**Summary**
This commit adds out-of-source-tree tests for `to_backend`. These tests check
that a Module can be lowered to a backend, exported, loaded (in both
Python and C++) and executed.

**Fixes**
This commit fixes pytorch#40067.

Test Plan: Imported from OSS

Reviewed By: jamesr66a

Differential Revision: D22510076

Pulled By: SplitInfinity

fbshipit-source-id: f65964ef3092a095740f06636ed5b1eb0884492d

Author: Meghan Lele

.jenkins/pytorch/build.sh

@@ -233,6 +233,17 @@ else
     make VERBOSE=1
     popd
     assert_git_not_dirty
+
+    # Build custom backend tests.
+    CUSTOM_BACKEND_BUILD="$PWD/../custom-backend-build"
+    CUSTOM_BACKEND_TEST="$PWD/test/custom_backend"
+    python --version
+    mkdir "$CUSTOM_BACKEND_BUILD"
+    pushd "$CUSTOM_BACKEND_BUILD"
+    cmake "$CUSTOM_BACKEND_TEST" -DCMAKE_PREFIX_PATH="$SITE_PACKAGES/torch" -DPYTHON_EXECUTABLE="$(which python)"
+    make VERBOSE=1
+    popd
+    assert_git_not_dirty
   else
     # Test standalone c10 build
     if [[ "$BUILD_ENVIRONMENT" == *xenial-cuda10.1-cudnn7-py3* ]]; then

.jenkins/pytorch/macos-test.sh

@@ -99,6 +99,26 @@ test_libtorch() {
   fi
 }
 
+test_custom_backend() {
+  echo "Testing custom backends"
+  pushd test/custom_backend
+  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 Python tests and export a lowered module.
+  python test_custom_backend.py -v
+  python backend.py --export-module-to=model.pt
+  # Run C++ tests using the exported module.
+  build/test_custom_backend ./model.pt
+  rm -f ./model.pt
+  popd
+  assert_git_not_dirty
+}
+
 test_custom_script_ops() {
   echo "Testing custom script operators"
   pushd test/custom_operator
@@ -124,11 +144,13 @@ if [ -z "${BUILD_ENVIRONMENT}" ] || [[ "${BUILD_ENVIRONMENT}" == *-test ]]; then
   test_python_all
   test_libtorch
   test_custom_script_ops
+  test_custom_backend
 else
   if [[ "${BUILD_ENVIRONMENT}" == *-test1 ]]; then
     test_python_all
   elif [[ "${BUILD_ENVIRONMENT}" == *-test2 ]]; then
     test_libtorch
     test_custom_script_ops
+    test_custom_backend
   fi
 fi

.jenkins/pytorch/test.sh

@@ -228,6 +228,23 @@ test_libtorch() {
   fi
 }
 
+test_custom_backend() {
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]] && [[ "$BUILD_ENVIRONMENT" != *asan* ]] ; then
+    echo "Testing custom backends"
+    CUSTOM_BACKEND_BUILD="$PWD/../custom-backend-build"
+    pushd test/custom_backend
+    cp -a "$CUSTOM_BACKEND_BUILD" build
+    # Run tests Python-side and export a lowered module.
+    python test_custom_backend.py -v
+    python backend.py --export-module-to=model.pt
+    # Run tests C++-side and load the exported lowered module.
+    build/test_custom_backend ./model.pt
+    rm -f ./model.pt
+    popd
+    assert_git_not_dirty
+  fi
+}
+
 test_custom_script_ops() {
   if [[ "$BUILD_ENVIRONMENT" != *rocm* ]] && [[ "$BUILD_ENVIRONMENT" != *asan* ]] ; then
     echo "Testing custom script operators"
@@ -331,6 +348,7 @@ elif [[ "${BUILD_ENVIRONMENT}" == *-test2 || "${JOB_BASE_NAME}" == *-test2 ]]; t
   test_aten
   test_libtorch
   test_custom_script_ops
+  test_custom_backend
   test_torch_function_benchmark
 elif [[ "${BUILD_ENVIRONMENT}" == *-bazel-* ]]; then
   test_bazel
@@ -346,5 +364,6 @@ else
   test_aten
   test_libtorch
   test_custom_script_ops
+  test_custom_backend
   test_torch_function_benchmark
 fi

.jenkins/pytorch/win-test-helpers/test_custom_backend.bat

@@ -0,0 +1,36 @@
+call %SCRIPT_HELPERS_DIR%\setup_pytorch_env.bat
+
+git submodule update --init --recursive third_party/pybind11
+cd test\custom_backend
+
+:: Build the custom backend library.
+mkdir build
+pushd build
+
+echo "Executing CMake for custom_backend test..."
+
+:: Note: Caffe2 does not support MSVC + CUDA + Debug mode (has to be Release mode)
+cmake -DCMAKE_PREFIX_PATH=%TMP_DIR_WIN%\build\torch -DCMAKE_BUILD_TYPE=Release -GNinja ..
+if ERRORLEVEL 1 exit /b 1
+
+echo "Executing Ninja for custom_backend test..."
+
+ninja -v
+if ERRORLEVEL 1 exit /b 1
+
+echo "Ninja succeeded for custom_backend test."
+
+popd
+
+:: Run tests Python-side and export a script module.
+python test_custom_backend.py -v
+if ERRORLEVEL 1 exit /b 1
+
+python backend.py --export-module-to="build/model.pt"
+if ERRORLEVEL 1 exit /b 1
+
+:: Run tests C++-side and load the exported script module.
+cd build
+set PATH=C:\Program Files\NVIDIA Corporation\NvToolsExt\bin\x64;%TMP_DIR_WIN%\build\torch\lib;%PATH%
+test_custom_backend.exe model.pt
+if ERRORLEVEL 1 exit /b 1

.jenkins/pytorch/win-test.sh

@@ -41,13 +41,15 @@ run_tests() {
         $SCRIPT_HELPERS_DIR/test_python_nn.bat "$DETERMINE_FROM" && \
         $SCRIPT_HELPERS_DIR/test_python_all_except_nn.bat "$DETERMINE_FROM" && \
         $SCRIPT_HELPERS_DIR/test_custom_script_ops.bat && \
+        $SCRIPT_HELPERS_DIR/test_custom_backend.bat && \
         $SCRIPT_HELPERS_DIR/test_libtorch.bat
     else
         if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
             $SCRIPT_HELPERS_DIR/test_python_nn.bat "$DETERMINE_FROM" && \
             $SCRIPT_HELPERS_DIR/test_libtorch.bat
         elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then
             $SCRIPT_HELPERS_DIR/test_python_all_except_nn.bat "$DETERMINE_FROM" && \
+            $SCRIPT_HELPERS_DIR/test_custom_backend.bat && \
             $SCRIPT_HELPERS_DIR/test_custom_script_ops.bat
         fi
     fi

test/cpp/jit/test_backend.cpp

@@ -1,4 +1,3 @@
-#include <torch/csrc/jit/api/module.h>
 #include <torch/csrc/jit/backends/backend.h>
 
 namespace torch {

test/custom_backend/CMakeLists.txt

@@ -0,0 +1,13 @@
+# Basic CMake setup
+cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
+project(custom_backend)
+
+find_package(Torch REQUIRED)
+
+add_library(custom_backend SHARED custom_backend.cpp)
+set_property(TARGET custom_backend PROPERTY CXX_STANDARD 14)
+target_link_libraries(custom_backend "${TORCH_LIBRARIES}")
+
+add_executable(test_custom_backend test_custom_backend.cpp)
+set_property(TARGET test_custom_backend PROPERTY CXX_STANDARD 14)
+target_link_libraries(test_custom_backend custom_backend)

test/custom_backend/backend.py

@@ -0,0 +1,72 @@
+import argparse
+import os.path
+import sys
+import torch
+
+
+def get_custom_backend_library_path():
+    """
+    Get the path to the library containing the custom backend.
+
+    Return:
+        The path to the custom backend object, customized by platform.
+    """
+    if sys.platform.startswith("win32"):
+        library_filename = "custom_backend.dll"
+    elif sys.platform.startswith("darwin"):
+        library_filename = "libcustom_backend.dylib"
+    else:
+        library_filename = "libcustom_backend.so"
+    path = os.path.abspath("build/{}".format(library_filename))
+    assert os.path.exists(path), path
+    return path
+
+
+def to_custom_backend(module):
+    """
+    This is a helper that wraps torch._C._jit_to_test_backend and compiles
+    only the forward method with an empty compile spec.
+
+    Args:
+        module: input ScriptModule.
+
+    Returns:
+        The module, lowered so that it can run on TestBackend.
+    """
+    lowered_module = torch._C._jit_to_backend("custom_backend", module._c, {"forward": {"": ""}})
+    return lowered_module
+
+
+class Model(torch.nn.Module):
+    """
+    Simple model used for testing that to_backend API supports saving, loading,
+    and executing in C++.
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, a, b):
+        return (a + b, a - b)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Lower a Module to a custom backend"
+    )
+    parser.add_argument("--export-module-to", required=True)
+    options = parser.parse_args()
+
+    # Load the library containing the custom backend.
+    library_path = get_custom_backend_library_path()
+    torch.ops.load_library(library_path)
+    assert library_path in torch.ops.loaded_libraries
+
+    # Lower an instance of Model to the custom backend  and export it
+    # to the specified location.
+    lowered_module = to_custom_backend(torch.jit.script(Model()))
+    torch.jit.save(lowered_module, options.export_module_to)
+
+
+if __name__ == "__main__":
+    main()

test/custom_backend/custom_backend.cpp

@@ -0,0 +1,14 @@
+#include "custom_backend.h"
+
+namespace torch {
+namespace custom_backend {
+namespace {
+constexpr auto kBackendName = "custom_backend";
+static auto cls = torch::jit::backend<CustomBackend>(kBackendName);
+}
+
+std::string getBackendName() {
+  return std::string(kBackendName);
+}
+}
+}

test/custom_backend/custom_backend.h

@@ -0,0 +1,84 @@
+#include <torch/csrc/jit/backends/backend.h>
+
+namespace torch {
+namespace custom_backend {
+// This custom JIT backend is intended to do the minimal amount of work
+// necessary to test that the JIT backend registration endpoints and
+// code generation are working correctly. It is not intended to
+// produce numerically correct results.
+class CustomBackend : public torch::jit::PyTorchBackendInterface {
+ public:
+  // Constructor.
+  explicit CustomBackend() {}
+  virtual ~CustomBackend() = default;
+
+  c10::IValue preprocess(
+      c10::IValue mod,
+      c10::impl::GenericDict method_compile_spec) override {
+    return mod;
+  }
+
+  c10::impl::GenericDict compile(
+      c10::IValue processed,
+      c10::impl::GenericDict method_compile_spec) override {
+    auto spec =
+        c10::impl::toTypedDict<std::string, at::IValue>(method_compile_spec);
+
+    // Return the same string as a value for every key in method_compile_spec.
+    auto handles = c10::Dict<std::string, std::string>();
+    for (auto it = spec.begin(), end = spec.end(); it != end; ++it) {
+      handles.insert(it->key(), it->key());
+    }
+    return c10::impl::toGenericDict(handles);
+  }
+  c10::impl::GenericList execute(
+      c10::IValue handle,
+      c10::impl::GenericList inputs) override {
+    TORCH_INTERNAL_ASSERT(handle.isString());
+    TORCH_INTERNAL_ASSERT(inputs.size() > 0);
+
+    c10::List<at::Tensor> output_list;
+
+    // Implement simple accumulator and negative accumulator (?) ops. Return one
+    // or both of them depending on the handle to make sure multiple outputs are
+    // handled.
+    c10::IValue value = inputs[0];
+    at::Tensor accum = value.toTensor();
+    accum = accum.clone();
+    at::Tensor sub_accum = value.toTensor();
+    sub_accum = sub_accum.clone();
+
+    for (size_t i = 1, e = inputs.size(); i < e; ++i) {
+      value = inputs[i];
+      accum.add_(value.toTensor(), 1.0);
+      sub_accum.sub_(value.toTensor(), 1.0);
+    }
+
+    if (handle.toStringRef() == "accum") {
+      output_list.emplace_back(accum);
+    } else if (handle.toStringRef() == "sub_accum") {
+      output_list.emplace_back(sub_accum);
+    } else if (handle.toStringRef() == "forward") {
+      output_list.emplace_back(accum);
+      output_list.emplace_back(sub_accum);
+    }
+
+    return c10::impl::toList(output_list);
+  }
+};
+
+// clang-format off
+#  if defined(_WIN32)
+#    if defined(custom_ops_EXPORTS)
+#      define CUSTOM_BACKEND_API __declspec(dllexport)
+#    else
+#      define CUSTOM_BACKEND_API __declspec(dllimport)
+#    endif
+#  else
+#    define CUSTOM_BACKEND_API
+#  endif
+// clang-format on
+
+CUSTOM_BACKEND_API std::string getBackendName();
+} // namespace custom_backend
+} // namespace torch

test/custom_backend/test_custom_backend.cpp

@@ -0,0 +1,39 @@
+#include <torch/cuda.h>
+#include <torch/script.h>
+
+#include <string>
+
+#include "custom_backend.h"
+
+// Load a module lowered for the custom backend from \p path and test that
+// it can be executed and produces correct results.
+void load_serialized_lowered_module_and_execute(const std::string& path) {
+  torch::jit::Module module = torch::jit::load(path);
+  // The custom backend is hardcoded to compute f(a, b) = (a + b, a - b).
+  auto tensor = torch::ones(5);
+  std::vector<torch::jit::IValue> inputs{tensor, tensor};
+  auto output = module.forward(inputs);
+  AT_ASSERT(output.isTuple());
+  auto output_elements = output.toTuple()->elements();
+  for (auto& e : output_elements) {
+    AT_ASSERT(e.isTensor());
+  }
+  AT_ASSERT(output_elements.size(), 2);
+  AT_ASSERT(output_elements[0].toTensor().allclose(tensor + tensor));
+  AT_ASSERT(output_elements[1].toTensor().allclose(tensor - tensor));
+}
+
+int main(int argc, const char* argv[]) {
+  if (argc != 2) {
+    std::cerr
+        << "usage: test_custom_backend <path-to-exported-script-module>\n";
+    return -1;
+  }
+  const std::string path_to_exported_script_module = argv[1];
+
+  std::cout << "Testing " << torch::custom_backend::getBackendName() << "\n";
+  load_serialized_lowered_module_and_execute(path_to_exported_script_module);
+
+  std::cout << "OK\n";
+  return 0;
+}

test/custom_backend/test_custom_backend.py

@@ -0,0 +1,54 @@
+import os
+import tempfile
+import torch
+import unittest
+
+from backend import Model, to_custom_backend, get_custom_backend_library_path
+
+
+class TestCustomBackend(unittest.TestCase):
+    def setUp(self):
+        # Load the library containing the custom backend.
+        self.library_path = get_custom_backend_library_path()
+        torch.ops.load_library(self.library_path)
+        # Create an instance of the test Module and lower it for
+        # the custom backend.
+        self.model = to_custom_backend(torch.jit.script(Model()))
+
+    def test_execute(self):
+        """
+        Test execution using the custom backend.
+        """
+        a = torch.randn(4)
+        b = torch.randn(4)
+        # The custom backend is hardcoded to compute f(a, b) = (a + b, a - b).
+        expected = (a + b, a - b)
+        out = self.model(a, b)
+        self.assertTrue(expected[0].allclose(out[0]))
+        self.assertTrue(expected[1].allclose(out[1]))
+
+    def test_save_load(self):
+        """
+        Test that a lowered module can be executed correctly
+        after saving and loading.
+        """
+        # Test execution before saving and loading to make sure
+        # the lowered module works in the first place.
+        self.test_execute()
+
+        # Save and load.
+        f = tempfile.NamedTemporaryFile(delete=False)
+        try:
+            f.close()
+            torch.jit.save(self.model, f.name)
+            loaded = torch.jit.load(f.name)
+        finally:
+            os.unlink(f.name)
+        self.model = loaded
+
+        # Test execution again.
+        self.test_execute()
+
+
+if __name__ == "__main__":
+    unittest.main()