Add in-place check to AliasDb

bwasti · facebook-github-bot · commit ff3dd724699f · 2019-08-01T12:15:52.000-07:00
Summary: Pull Request resolved: #23210 Test Plan: Imported from OSS Differential Revision: D16444529 Pulled By: bwasti fbshipit-source-id: 83af54d423989a2a726158b521093660584ee9c2
diff --git a/test/cpp/jit/test_alias_analysis.h b/test/cpp/jit/test_alias_analysis.h
@@ -539,6 +539,51 @@ void testWriteTracking() {
     ASSERT_TRUE(aliasDb.writesToAlias(
         writingNode, std::unordered_set<const Value*>{aAlias}));
   }
+  {
+    auto graph = std::make_shared<Graph>();
+    script::parseIR(
+        R"IR(
+  graph(%x: Tensor):
+    %b : (Tensor) = aten::relu_(%x)
+    return (%b)
+    )IR",
+        &*graph);
+    auto node_iter = graph->block()->nodes().begin();
+    auto relu = *node_iter;
+    AliasDb aliasDb(graph);
+    AT_ASSERT(aliasDb.isMutable(relu));
+  }
+  {
+    auto graph = std::make_shared<Graph>();
+    script::parseIR(
+        R"IR(
+  graph(%x: Tensor, %y : Tensor):
+    %b : (Tensor) = aten::mul(%x, %y)
+    return (%b)
+    )IR",
+        &*graph);
+    auto node_iter = graph->block()->nodes().begin();
+    auto mul = *node_iter;
+    AliasDb aliasDb(graph);
+    AT_ASSERT(!aliasDb.isMutable(mul));
+  }
+  {
+    auto graph = std::make_shared<Graph>();
+    std::unordered_map<std::string, Value*> vmap;
+    script::parseIR(
+        R"IR(
+  graph(%x: Tensor, %y : Tensor):
+    %c1 : int = prim::Constant[value=1]()
+    %b : (Tensor) = aten::add_(%x, %y, %c1)
+    return (%b)
+    )IR",
+        &*graph,
+        vmap);
+    auto add = vmap["b"]->node();
+    AliasDb aliasDb(graph);
+    AT_ASSERT(aliasDb.hasWriters(add));
+    AT_ASSERT(aliasDb.isMutable(add));
+  }
 }
 
 void testContainerAliasing() {
diff --git a/torch/csrc/jit/passes/alias_analysis.cpp b/torch/csrc/jit/passes/alias_analysis.cpp
@@ -60,12 +60,24 @@ AliasDb::AliasDb(std::shared_ptr<Graph> graph) : graph_(std::move(graph)) {
   GRAPH_DEBUG(toString());
 }
 
-bool AliasDb::hasWriters(const Node* n) const {
+bool AliasDb::isMutable(Node* n) const {
+  ValueSet vs;
+  for (const auto input : n->inputs()) {
+    vs.insert(input);
+  }
+  return writesToAlias(n, vs);
+}
+
+bool AliasDb::hasInputWriters(const Node* n) const {
   for (const auto input : n->inputs()) {
     if (hasWriters(input)) {
       return true;
     }
   }
+  return false;
+}
+
+bool AliasDb::hasOutputWriters(const Node* n) const {
   for (const auto output : n->outputs()) {
     if (hasWriters(output)) {
       return true;
@@ -74,6 +86,10 @@ bool AliasDb::hasWriters(const Node* n) const {
   return false;
 }
 
+bool AliasDb::hasWriters(const Node* n) const {
+  return hasInputWriters(n) || hasOutputWriters(n);
+}
+
 bool AliasDb::hasWriters(const Value* v) const {
   if (v->mustBeNone()) {
     return false;
diff --git a/torch/csrc/jit/passes/alias_analysis.h b/torch/csrc/jit/passes/alias_analysis.h
@@ -61,9 +61,19 @@ class AliasDb {
   // value in group `b`? i.e. may they overlap?
   TORCH_API bool mayAlias(const ValueSet& a, const ValueSet& b) const;
 
+  // Do any nodes write to an alias set input to `n`?
+  TORCH_API bool hasInputWriters(const Node* n) const;
+
+  // Do any nodes write to an alias set output by `n`?
+  TORCH_API bool hasOutputWriters(const Node* n) const;
+
   // Do any nodes write to an alias set inputed/outputed by `n`?
   TORCH_API bool hasWriters(const Node* n) const;
 
+  // Is the operation in-place? i.e. doesn't write anywhere but locations it
+  // reads from.
+  TORCH_API bool isMutable(Node* n) const;
+
   // Move 'n' (already in the graph) after 'movePoint' in the topological order.
   //
   // Tries to preserve value dependencies, so other nodes might be moved. We
