Remove 'recurse' parameter from Inline. (pytorch#26487)

Summary:
Pull Request resolved: pytorch#26487

The way it is implemented currently is bad because while we're inlining
to a graph G, we are also mutating all the graphs that are being
inlined. The problem is that the graphs we're inlining are usually the
original graphs of functions, so we're silently changing them behind the
scenes, and we don't have a way to recover 'unoptimized' graphs
afterwards.

Test Plan: Imported from OSS

Differential Revision: D17485748

Pulled By: ZolotukhinM

fbshipit-source-id: 6094ef56077240e9379d4c53680867df1b6e79ef

Author: Mikhail Zolotukhin

test/cpp/jit/test_inliner.cpp

@@ -39,26 +39,14 @@ struct InlinerGuard {
 
 void testInliner() {
   {
-    // Test that the recursive inlining works
     // disable automatic inlining so we can test it manually
     InlinerGuard guard(/*shouldInline=*/false);
 
     CompilationUnit cu(testSource);
     auto& fn = cu.get_function("foo3");
 
     auto g = fn.graph();
-    Inline(*g, /*recurse=*/true);
-    FileCheck().check_count("prim::Print", 3)->run(*g);
-  }
-  {
-    // disable automatic inlining so we can test it manually
-    InlinerGuard guard(/*shouldInline=*/false);
-
-    CompilationUnit cu(testSource);
-    auto& fn = cu.get_function("foo3");
-
-    auto g = fn.graph();
-    Inline(*g, /*recurse=*/false);
+    Inline(*g);
     FileCheck()
         .check("three")
         ->check("two")

torch/csrc/jit/passes/inliner.cpp

@@ -9,7 +9,7 @@ namespace prim {
 using namespace ::c10::prim;
 }
 
-void inlineCalls(Block* block, bool recurse) {
+void inlineCalls(Block* block) {
   for (auto it = block->nodes().begin(), end = block->nodes().end();
        it != end;) {
     Node* cur = *it++;
@@ -20,32 +20,26 @@ void inlineCalls(Block* block, bool recurse) {
         auto fun_type =
             function_constant->output()->type()->expect<FunctionType>();
         cur->removeInput(0);
-        if (recurse) {
-          Inline(*fun_type->function()->graph(), recurse);
-        }
         inlineCallTo(cur, *fun_type->function()->graph());
       } 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 (recurse) {
-            Inline(*function->graph(), recurse);
-          }
           inlineCallTo(cur, *function->graph());
         }
       } break;
       default: {
         for (auto b : cur->blocks()) {
-          inlineCalls(b, recurse);
+          inlineCalls(b);
         }
       } break;
     }
   }
 }
 
-void Inline(Graph& graph, bool recurse) {
-  inlineCalls(graph.block(), recurse);
+void Inline(Graph& graph) {
+  inlineCalls(graph.block());
 }
 
 } // namespace jit

torch/csrc/jit/passes/inliner.h

@@ -5,9 +5,8 @@
 namespace torch {
 namespace jit {
 
-// Inline function and method calls. If `recurse` is true, inline all nested
-// calls as well, resulting in a completely flattened graph.
-TORCH_API void Inline(Graph& graph, bool recurse = false);
+// Inline function and method calls.
+TORCH_API void Inline(Graph& graph);
 
 } // namespace jit
 } // namespace torch

torch/onnx/utils.py

@@ -86,8 +86,8 @@ def _split_tensor_list_constants(g, block):
 
 
 def _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop=False, fixed_batch_size=False):
-    # Inline everyting (recursively)
-    torch._C._jit_pass_inline(graph, True)
+    # Inline everyting
+    torch._C._jit_pass_inline(graph)
 
     # Remove fork/wait nodes
     torch._C._jit_pass_inline_fork_wait(graph)