Add 'optimized_graph' to Function. (pytorch#26488)

Summary:
Pull Request resolved: pytorch#26488

Currently the main use case for this graph is inlining and that's the
only optimization we perform. We probably should run more cleanups on
this graph in future.

Test Plan: Imported from OSS

Differential Revision: D17485745

Pulled By: ZolotukhinM

fbshipit-source-id: 7b30c9ba47b4e5fff3591a0063560bfeb68f2164

Author: Mikhail Zolotukhin

torch/csrc/jit/function.cpp

@@ -1,4 +1,5 @@
 #include <torch/csrc/jit/function.h>
+#include <torch/csrc/jit/passes/inliner.h>
 
 #include <torch/csrc/jit/script/error_report.h>
 
@@ -52,5 +53,12 @@ const FunctionSchema& Function::getSchema() const {
   }
   return *schema_;
 }
+
+void preoptimizeGraph(std::shared_ptr<Graph>& graph) {
+  // TODO: Invoke cleanup passes before and after inlining to reduce amount of
+  // code we're copying.
+  Inline(*graph);
+}
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/function.h

@@ -9,6 +9,8 @@ namespace jit {
 
 using Kwargs = std::unordered_map<std::string, IValue>;
 
+TORCH_API void preoptimizeGraph(std::shared_ptr<Graph>& graph);
+
 // A Function is a pure Graph with no implicit `self` object bound.
 // It contains schema information, and the executor that manages the
 // execution of the function. script::Method is a wrapper around a
@@ -40,6 +42,16 @@ struct TORCH_API Function {
     return graph_;
   }
 
+  std::shared_ptr<Graph> optimized_graph() const {
+    if (optimized_graph_) {
+      return *optimized_graph_;
+    }
+    std::lock_guard<std::mutex> lock(compile_mutex);
+    optimized_graph_ = graph_->copy();
+    preoptimizeGraph(*optimized_graph_);
+    return *optimized_graph_;
+  }
+
   const c10::QualifiedName& qualname() const {
     return name_;
   }
@@ -99,10 +111,17 @@ struct TORCH_API Function {
 
  private:
   c10::QualifiedName name_;
+  // The original, non-optimized graph
   std::shared_ptr<Graph> graph_; // for debugging and for inlining
 
+  // Optimized graph, computed lazily. Used for inlining.
+  // Note: this graph is not specialized, only generic optimizations are applied
+  // here.
+  mutable c10::optional<std::shared_ptr<Graph>> optimized_graph_;
+
   // Functions are invokable from multiple threads, so this lock needs to be
-  // held when we're initializing graph executor for the first time.
+  // held when we're initializing graph executor for the first time or computing
+  // the optimized graph.
   mutable std::mutex compile_mutex;
 
   GraphExecutor executor_; // for execution