[ONNX] Adding a pass to replace interpolate function with aten::__interpolate

caffe2/CMakeLists.txt

@@ -405,6 +405,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     ${TORCH_SRC_DIR}/csrc/jit/passes/inline_autodiff_subgraphs.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/insert_guards.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/inliner.cpp
+    ${TORCH_SRC_DIR}/csrc/jit/passes/onnx/preclude_inlining.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/lift_closures.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/inline_forked_closures.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/dead_code_elimination.cpp

test/onnx/test_pytorch_onnx_onnxruntime.py

@@ -346,7 +346,6 @@ def get_test_images(self):
         return images
 
     @skipIfUnsupportedMinOpsetVersion(11)
-    @unittest.skip("disabled due to removal of aten::__interpolate")
     def test_mask_rcnn(self):
         model = torchvision.models.detection.mask_rcnn.maskrcnn_resnet50_fpn(pretrained=True, min_size=200,
                                                                              max_size=300)
@@ -1198,8 +1197,6 @@ def forward(self, x):
         self.run_test(MyModel(), x)
 
     def _interpolate_script(self, x, mode, use_size, is_upsample, align_corners=False):
-        # test disabled
-        return 
 
         class MyModel(torch.jit.ScriptModule):
             __constants__ = ['mode', 'use_size', 'is_upsample', 'size', 'scale', 'size_array', 'scale_array', 'align_corners']

torch/CMakeLists.txt

@@ -71,6 +71,7 @@ set(TORCH_PYTHON_SRCS
     ${TORCH_SRC_DIR}/csrc/jit/python/init.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/onnx.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/onnx/fixup_onnx_conditionals.cpp
+    ${TORCH_SRC_DIR}/csrc/jit/passes/onnx/preclude_inlining.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/onnx/fixup_onnx_loop.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/onnx/helper.cpp
     ${TORCH_SRC_DIR}/csrc/jit/passes/onnx/prepare_division_for_onnx.cpp

torch/csrc/jit/ir/ir.cpp

@@ -1820,16 +1820,30 @@ at::ArrayRef<Value*> createTupleUnpack(Value* v) {
   return g.insertNode(g.createTupleUnpack(v))->outputs();
 }
 
-std::vector<Value*> inlineCallTo(Node* to_replace, Function* callee) {
+// inline_optimized_graph argument is used to preclude inlining functions for
+// ONNX conversion
+std::vector<Value*> inlineCallTo(
+    Node* to_replace,
+    Function* callee,
+    bool inline_optimized_graph /*=true*/) {
   WithInsertPoint guard(to_replace);
   TORCH_INTERNAL_ASSERT(callee->isGraphFunction());
   std::unordered_map<Value*, Value*> value_map;
-  auto new_outputs = insertGraph(
-      *to_replace->owningGraph(),
-      *(callee->optimized_graph()),
-      to_replace->inputs(),
-      value_map);
-
+  std::vector<torch::jit::Value*> new_outputs;
+
+  if (inline_optimized_graph) {
+    new_outputs = insertGraph(
+        *to_replace->owningGraph(),
+        *(callee->optimized_graph()),
+        to_replace->inputs(),
+        value_map);
+  } else {
+    new_outputs = insertGraph(
+        *to_replace->owningGraph(),
+        *(callee->graph()),
+        to_replace->inputs(),
+        value_map);
+  }
   std::unordered_map<InlinedCallStack*, InlinedCallStackPtr>
       new_callstack_entries;
 

torch/csrc/jit/ir/ir.h

@@ -1367,7 +1367,10 @@ TORCH_API std::vector<Value*> insertGraph(
  * This asserts that the number of outputs of the original node and the
  * graph are the same.
  */
-TORCH_API std::vector<Value*> inlineCallTo(Node* to_replace, Function* callee);
+TORCH_API std::vector<Value*> inlineCallTo(
+    Node* to_replace,
+    Function* callee,
+    bool use_graph = true);
 
 /** If there is only one value in \p OUTPUTS and its kind is Tuple, insert a
  * tuple unpack node and return the resulting values.

torch/csrc/jit/passes/onnx/preclude_inlining.cpp

@@ -0,0 +1,87 @@
+#include <torch/csrc/jit/passes/onnx/preclude_inlining.h>
+#include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/onnx/helper.h>
+
+namespace torch {
+namespace jit {
+
+void functionCallSubstitution(Block* block) {
+  for (auto it = block->nodes().begin(), end = block->nodes().end();
+       it != end;) {
+    Node* cur = *it++;
+    switch (cur->kind()) {
+      case prim::CallFunction: {
+        AT_ASSERT(cur->input(0)->node()->kind() == prim::Constant);
+        auto function_constant = cur->input(0)->node();
+        auto fun_type =
+            function_constant->output()->type()->expect<FunctionType>();
+
+        if ((fun_type->function()->qualname().qualifiedName().find(
+                 "torch.nn.functional") != std::string::npos) &&
+            (fun_type->function()->qualname().qualifiedName().find(
+                 "interpolate") != std::string::npos)) {
+          cur->removeInput(0);
+          Node* interpolate_node = block->owningGraph()->create(
+              Symbol::fromQualString("aten::__interpolate"),
+              {cur->inputs()},
+              cur->outputs().size());
+          interpolate_node->output()->copyMetadata(cur->output());
+          interpolate_node->insertAfter(cur);
+          cur->replaceAllUsesWith(interpolate_node);
+          cur->removeAllInputs();
+          cur->destroy();
+        } else {
+          cur->removeInput(0);
+          functionCallSubstitution(fun_type->function()->graph()->block());
+          GRAPH_UPDATE(
+              "Inlining in ONNX preclude inlining function '",
+              fun_type->function()->name(),
+              "' to ",
+              *cur);
+          GRAPH_UPDATE(
+              "Function in ONNX preclude inlining body: ",
+              *fun_type->function()->optimized_graph());
+          inlineCallTo(cur, fun_type->function(), false);
+        }
+      } break;
+      case prim::CallMethod: {
+        const std::string& name = cur->s(attr::name);
+        if (auto class_type = cur->input(0)->type()->cast<ClassType>()) {
+          auto function = class_type->getMethod(name);
+          if (!function->isGraphFunction()) {
+            continue;
+          }
+          functionCallSubstitution(function->graph()->block());
+          GRAPH_UPDATE(
+              "Inlining in ONNX preclude inlining function '",
+              function->name(),
+              "' to ",
+              *cur);
+          GRAPH_UPDATE(
+              "Function in ONNX preclude inlining body: ",
+              function->optimized_graph());
+          inlineCallTo(cur, function, false);
+        }
+      } break;
+      default: {
+        for (auto b : cur->blocks()) {
+          functionCallSubstitution(b);
+        }
+      } break;
+    }
+  }
+}
+
+// This pass is to be used for ONNX conversion only. The ONNX converter depends
+// on a number of deprecated aten operators. These operators are removed from IR
+// and replaced by the compiled python function code. However, in-order to
+// maintain the behavior for ONNX conversion, we replace these function calls
+// with the aten symbolic which can still be used by the ONNX converter.
+void ONNXFunctionCallSubstitution(Graph& graph) {
+  GRAPH_DUMP("Before stop-inlining calls: ", &graph);
+  functionCallSubstitution(graph.block());
+  GRAPH_DUMP("After stop-inlining calls: ", &graph);
+}
+
+} // namespace jit
+} // namespace torch

torch/csrc/jit/passes/onnx/preclude_inlining.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch {
+namespace jit {
+
+void ONNXFunctionCallSubstitution(Graph& graph);
+
+}
+} // namespace torch

torch/csrc/jit/python/init.cpp

@@ -31,6 +31,7 @@
 #include <torch/csrc/jit/passes/onnx/fixup_onnx_conditionals.h>
 #include <torch/csrc/jit/passes/onnx/fixup_onnx_loop.h>
 #include <torch/csrc/jit/passes/onnx/peephole.h>
+#include <torch/csrc/jit/passes/onnx/preclude_inlining.h>
 #include <torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h>
 #include <torch/csrc/jit/passes/onnx/prepare_inplace_ops_for_onnx.h>
 #include <torch/csrc/jit/passes/onnx/scalar_type_analysis.h>
@@ -122,6 +123,7 @@ void initJITBindings(PyObject* module) {
       .def("_jit_pass_onnx_preprocess_caffe2", PreprocessCaffe2Ops)
       .def("_jit_pass_onnx", ToONNX)
       .def("_jit_pass_lower_all_tuples", LowerAllTuples)
+      .def("_jit_pass_onnx_function_substitution", ONNXFunctionCallSubstitution)
       .def(
           "_jit_pass_onnx_peephole",
           [](std::shared_ptr<Graph>& graph,

torch/onnx/utils.py

@@ -333,7 +333,9 @@ def _model_to_graph(model, args, verbose=False,
     if isinstance(model, torch.jit.ScriptModule):
         assert example_outputs is not None, "example_outputs must be provided when exporting a ScriptModule"
         try:
-            method_graph, params = torch._C._jit_pass_lower_graph(model.forward.graph, model._c)
+            graph = model.forward.graph
+            torch._C._jit_pass_onnx_function_substitution(graph)
+            method_graph, params = torch._C._jit_pass_lower_graph(graph, model._c)
             in_vars, in_desc = torch.jit._flatten(tuple(args) + tuple(params))
             graph = _propagate_and_assign_input_shapes(
                 method_graph, tuple(in_vars), False, propagate)
@@ -344,8 +346,10 @@ def _model_to_graph(model, args, verbose=False,
         method = model
         params = ()
         in_vars, in_desc = torch.jit._flatten(tuple(args))
+        graph = model.graph
+        torch._C._jit_pass_onnx_function_substitution(graph)
         graph = _propagate_and_assign_input_shapes(
-            model.graph, tuple(in_vars), False, propagate)
+            graph, tuple(in_vars), False, propagate)
     else:
         graph, torch_out = _trace_and_get_graph_from_model(model, args)
         state_dict = _unique_state_dict(model)
@@ -357,6 +361,7 @@ def _model_to_graph(model, args, verbose=False,
             for i, inp in enumerate(graph_inputs):
                 if i >= user_input_num:
                     inp.setDebugName(param_names[i - user_input_num])
+        torch._C._jit_pass_onnx_function_substitution(graph)
 
     input_and_param_names = [val.debugName() for val in graph.inputs()]
     param_names = input_and_param_names[len(input_and_param_names) - len(params):]