[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/preinline_onnx.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

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/preinline_onnx.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/api/function_impl.cpp

@@ -1,5 +1,6 @@
 #include <torch/csrc/jit/api/function_impl.h>
 #include <torch/csrc/jit/passes/inliner.h>
+#include <torch/csrc/jit/passes/onnx/preinline_onnx.h>
 
 #include <torch/csrc/jit/frontend/error_report.h>
 
@@ -64,6 +65,7 @@ const c10::FunctionSchema& GraphFunction::getSchema() const {
 void preoptimizeGraph(std::shared_ptr<Graph>& graph) {
   // TODO: Invoke cleanup passes before and after inlining to reduce amount of
   // code we're copying.
+  PreInlineONNX(*graph);
   Inline(*graph);
 }
 

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

@@ -0,0 +1,52 @@
+#include <torch/csrc/jit/passes/onnx/preinline_onnx.h>
+#include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/onnx/helper.h>
+
+namespace torch {
+namespace jit {
+
+void replaceFunctions(Node* to_replace, Function* callee) {
+  if (callee->name() == "interpolate") {
+    to_replace->removeInput(0);
+    Node* interpolate_node = to_replace->owningGraph()->create(
+        Symbol::fromQualString("aten::__interpolate"),
+        {to_replace->inputs()},
+        to_replace->outputs().size());
+    interpolate_node->output()->copyMetadata(to_replace->output());
+    interpolate_node->insertAfter(to_replace);
+    to_replace->replaceAllUsesWith(interpolate_node);
+    to_replace->removeAllInputs();
+    to_replace->destroy();
+    return;
+  }
+}
+
+void PreInlineCalls(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>();
+        replaceFunctions(cur, fun_type->function());
+      } break;
+      default: {
+        for (auto b : cur->blocks()) {
+          PreInlineCalls(b);
+        }
+      } break;
+    }
+  }
+}
+
+void PreInlineONNX(Graph& graph) {
+  GRAPH_DUMP("Before Pre-inlining: ", &graph);
+  PreInlineCalls(graph.block());
+  GRAPH_DUMP("After Pre-inlining: ", &graph);
+}
+
+} // namespace jit
+} // namespace torch

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

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch {
+namespace jit {
+
+void PreInlineONNX(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/preinline_onnx.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_preinline", PreInlineONNX)
       .def(
           "_jit_pass_onnx_peephole",
           [](std::shared_ptr<Graph>& graph,

torch/onnx/utils.py

@@ -334,7 +334,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_preinline(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)
@@ -345,8 +347,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_preinline(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)