[wip] update graph fuser aliasdb in-place (#37106)

suo · facebook-github-bot · commit 9e32a1f5cdce · 2020-04-30T22:21:35.000-07:00
Summary: Pull Request resolved: #37106 Recomputing the aliasdb on every fusion iteration + in every subblock is hugely expensive. Instead, update it in-place when doing fusion. The graph fuser pass operates by pushing nodes into a fusion group. So we start with ``` x, y = f(a, b, c) ``` and end with: ``` x_out, y_out = prim::fusionGroup(a, b, c) x_in, y_in = f(a_in, b_in, c_in) -> x_in, y_in ``` We destroy the `x` and `y` `Value*`s in the process. This operation is easy to express as an update to the aliasDb--`x_out` just takes on all the aliasing information `x` used to have. In particular, since we know `f` and `prim::fusionGroup` are purely functional, we don't have to mess with any write information. This PR is the bare minimum to get this working, in the interest of unscrewing the compilation times ASAP. Followups I want to do: - We don't have a way of expressing deletion of values in AliasDb. In `graph_fuser.cpp` we sometimes construct nodes that we end up throwing away, and we are littering `MemoryDAG` with references to dangling pointers. Because of the way the pass works, it's fine, but this is fragile so I want to fix it. - We should decouple alias analysis from write tracking, to simplify the job of keeping the write caches consistent as we mutate the aliasing information. - the tensorexpr fuser doesn't do this and thus is incorrect today, we need to update it to work. Test Plan: Imported from OSS Differential Revision: D21219179 Pulled By: suo fbshipit-source-id: 8ae5397b3a0ad90edec2fbc555647091f1ad5284
diff --git a/test/cpp/jit/test_fuser.cpp b/test/cpp/jit/test_fuser.cpp
@@ -176,6 +176,31 @@ void testFusion() {
   };
 }
 
+void testFusionAliasing() {
+  const auto graph_string = R"IR(
+    graph(%0 : Tensor,
+          %1 : Tensor):
+      %12 : int = prim::Constant[value=1]()
+      %2.1 : Tensor = aten::mul(%0, %1)
+      %2 : Tensor = aten::mul(%2.1, %1)
+      %3 : Tensor = aten::add_(%2, %1, %12)
+      %4 : Tensor = aten::mul(%2, %1)
+      %5 : Tensor = aten::add(%2, %4, %12)
+      return (%5))IR";
+  auto g = std::make_shared<Graph>();
+  torch::jit::parseIR(graph_string, g.get());
+
+  g->lint();
+  FuseGraph(g);
+
+  // We should not be able to fuse across the in-place operation here.
+  testing::FileCheck()
+      .check("prim::FusionGroup_0")
+      ->check("aten::add_")
+      ->check("prim::FusionGroup_1")
+      ->run(*g);
+}
+
 void testRegisterFusionCachesKernel() {
   // Constructs two functionally equivalent graphs
   const auto graph0_string = R"IR(
diff --git a/test/cpp/jit/tests.h b/test/cpp/jit/tests.h
@@ -90,7 +90,8 @@ namespace jit {
   _(LiteInterpreterParams)             \
   _(LiteInterpreterSetState)           \
   _(TorchbindIValueAPI)                \
-  _(LiteInterpreterDict)
+  _(LiteInterpreterDict)               \
+  _(FusionAliasing)
 
 #if defined(USE_CUDA)
 #define TH_FORALL_TESTS_CUDA(_)  \
diff --git a/torch/csrc/jit/ir/alias_analysis.cpp b/torch/csrc/jit/ir/alias_analysis.cpp
@@ -324,7 +324,7 @@ MemoryLocations AliasDb::getReads(Node* n) const {
 }
 
 std::string AliasDb::getElementName(const Element* e) const {
-  if (e->value == nullptr) {
+  if (e->values.empty()) {
     // not the most efficient way, but given the fact there are
     // not too many types and even fewer of them will end up in
     // wildcardIndex_, we should be fine with a linear search
@@ -336,7 +336,17 @@ std::string AliasDb::getElementName(const Element* e) const {
     }
     return "WILDCARD";
   } else {
-    return e->value->debugName();
+    std::ostringstream ss;
+    if (e->values.size() == 1) {
+      ss << "%" << (*e->values.begin())->debugName();
+      return ss.str();
+    }
+    ss << "(";
+    for (const Value* v : e->values) {
+      ss << "%" << v->debugName() << ", ";
+    }
+    ss << ")";
+    return ss.str();
   }
 }
 
@@ -1036,10 +1046,8 @@ void AliasDb::mapAliases(at::ArrayRef<Value*> from, at::ArrayRef<Value*> to) {
 // The asserts are to guard against unintentional use.
 // FIXME refactor aliasdb construction to be more robust to mutation so this
 // hack isn't necessary.
-void AliasDb::unsafeGiveFreshAlias(const Value* value) {
+void AliasDb::createValue(const Value* value) {
   TORCH_INTERNAL_ASSERT(isMutableTypeInternal(value->type()));
-  TORCH_INTERNAL_ASSERT(value->type()->containedTypes().size() == 0);
-  TORCH_INTERNAL_ASSERT(!elementMap_.count(value));
   auto new_elem = memoryDAG_->unsafeMakeFreshValue(value);
   elementMap_[value] = new_elem;
 }
@@ -1068,14 +1076,39 @@ Element* AliasDb::getOrCreateElement(const Value* value) {
   return elementMap_.at(value);
 }
 
-void AliasDb::replaceMemoryLocation(Value* existing, Value* new_value) {
+void AliasDb::replaceWithNewValue(Value* existing, Value* new_value) {
+  TORCH_INTERNAL_ASSERT(
+      *unshapedType(existing->type()) == *unshapedType(new_value->type()),
+      "Types must be strictly equal if you are replacing aliasing information. ",
+      "Got existing: '",
+      existing->type()->python_str(),
+      "', new_value: '",
+      new_value->type()->python_str(),
+      "'");
   if (!isMutableTypeInternal(existing)) {
     return;
   }
   auto existing_elem = elementMap_.at(existing);
   elementMap_[new_value] = existing_elem;
   elementMap_.erase(existing);
-  existing_elem->value = new_value;
+  existing_elem->values = {new_value};
+}
+
+void AliasDb::copyValue(Value* from, Value* to) {
+  TORCH_INTERNAL_ASSERT(
+      *unshapedType(from->type()) == *unshapedType(to->type()),
+      "Types must be strictly equal if you are copying aliasing information. ",
+      "Got from: '",
+      from->type()->python_str(),
+      "', to: '",
+      to->type()->python_str(),
+      "'");
+  if (!isMutableTypeInternal(to)) {
+    return;
+  }
+  auto origElem = elementMap_.at(from);
+  elementMap_[to] = origElem;
+  origElem->values.insert(to);
 }
 
 bool AliasDb::moveAfterTopologicallyValid(Node* n, Node* movePoint) {
@@ -1492,5 +1525,30 @@ MemoryLocations AliasDb::buildWrittenToLocationsIndex() const {
   return ret;
 }
 
+void Lint(const AliasDb* db) {
+  bool failed = false;
+
+  std::stringstream ss;
+  // Every mutable value in the system has a corresponding element.
+  for (const auto& v : db->graph_->all_values) {
+    if (!db->isMutableTypeInternal(v)) {
+      continue;
+    }
+    auto it = db->elementMap_.find(v);
+    if (it == db->elementMap_.end()) {
+      failed = true;
+      ss << "Value %" << v->debugName() << " of type "
+         << v->type()->python_str() << " wasn't found in the element map.\n"
+         << "It was defined in " << *v->node();
+    }
+  }
+  TORCH_INTERNAL_ASSERT(!failed, ss.str());
+
+  // Two checks that we want to add but can't until the mutation API is more
+  // fully developed.
+  // - Every mutable value in the aliasdb belongs to the graph
+  // - All container values have contained elements
+}
+
 } // namespace jit
 } // namespace torch
diff --git a/torch/csrc/jit/ir/alias_analysis.h b/torch/csrc/jit/ir/alias_analysis.h
@@ -117,6 +117,24 @@ class AliasDb {
   static bool isMutableType(const Value* v);
   static bool isMutableType(const TypePtr& type);
 
+  /**
+   * Mutation API
+   *
+   * These methods allow you to update AliasDb in-place if you are performing
+   * graph mutation.
+   *
+   * WARNING: These methods should be considered INTERNAL. They do not perform
+   * very many correctness checks, the user is responsible for making sure they
+   * are updating AliasDb correctly. `Lint()`ing the AliasDb can help with
+   * this.
+   */
+  // Copy `existing`s aliasing info to `new_value`, and remove `existing`.
+  void replaceWithNewValue(Value* existing, Value* new_value);
+  // Copy `from`s aliasing info to `to`.
+  void copyValue(Value* from, Value* to);
+  // Create a new `value` that does not alias anything else.
+  void createValue(const Value* value);
+
   friend struct MutationRemover;
 
  private:
@@ -187,17 +205,9 @@ class AliasDb {
       const Value* element,
       const Value* container);
   void mapAliases(at::ArrayRef<Value*> to, at::ArrayRef<Value*> from);
-  void unsafeGiveFreshAlias(const Value* value);
   void giveFreshAlias(const Value* value);
   Element* getOrCreateElement(const Value* value);
 
-  // In the Value * -> Element * map replaces the mapping
-  // of Value * existing -> Element * existing_elem with
-  // Value * new_value -> Element * existing_elem
-  // Callers are expected to maintain graph invariants & specify
-  // own correctness conditions
-  void replaceMemoryLocation(Value* existing, Value* new_value);
-
   c10::optional<TypePtr> getMutableTypePtr(const TypePtr& type) const;
 
   bool isContainerType(const TypePtr& type) const;
@@ -250,7 +260,14 @@ class AliasDb {
   std::unordered_set<const Value*> wildcards_;
 
   std::string getElementName(const Element* e) const;
+
+  friend void Lint(const AliasDb* db);
 };
 
+// Helper check that invariants over AliasDb are maintained.
+// Useful if you are using the AliasDb mutation API and want to check you did
+// the right thing.
+void Lint(const AliasDb* db);
+
 } // namespace jit
 } // namespace torch
diff --git a/torch/csrc/jit/ir/ir.h b/torch/csrc/jit/ir/ir.h
@@ -32,6 +32,7 @@ using pyobj_list = std::vector<THPObjectPtr>;
 
 namespace torch {
 namespace jit {
+class AliasDb;
 
 using ::c10::Argument;
 using ::c10::FunctionSchema;
@@ -1231,6 +1232,7 @@ struct Graph {
   TORCH_API void remapTypes(const std::function<TypePtr(TypePtr)>& type_map);
 
  private:
+  friend void Lint(const AliasDb* db);
   TORCH_API void freeNode(Node* n);
   TORCH_API void freeValue(Value* v);
   TORCH_API void freeBlock(Block* b);
diff --git a/torch/csrc/jit/passes/create_functional_graphs.cpp b/torch/csrc/jit/passes/create_functional_graphs.cpp
@@ -378,12 +378,12 @@ struct MutationRemover {
       // same aliasing relationships as the original x.
       // To avoid rebuilding the entire alias db, we can replace
       // the memory dag element of x with x0.
-      aliasDb_->replaceMemoryLocation(mutated_value, new_node->output());
+      aliasDb_->replaceWithNewValue(mutated_value, new_node->output());
 
       // it is an invariant that all mutable types have an element in the memory
       // dag so we must regive x an alias db element. We have already verified
       // that the mutated value is a fresh alias with a single use.
-      aliasDb_->unsafeGiveFreshAlias(mutated_value);
+      aliasDb_->createValue(mutated_value);
 
       // We must erase the destroyed node from the AliasDb lists of writes
       aliasDb_->writeIndex_->erase(node);
diff --git a/torch/csrc/jit/passes/graph_fuser.cpp b/torch/csrc/jit/passes/graph_fuser.cpp
diff --git a/torch/csrc/jit/passes/utils/memory_dag.cpp b/torch/csrc/jit/passes/utils/memory_dag.cpp
diff --git a/torch/csrc/jit/passes/utils/memory_dag.h b/torch/csrc/jit/passes/utils/memory_dag.h
