fix alias annotations on to, cpu, cuda (#16460)

Summary:
Fix alias annotations for ops that may return a fresh tensor. The previous version was overly conservative.

Currently there is no actual behavior change in the alias analysis, but we may use the information in the future.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/16460

Differential Revision: D13849086

Pulled By: suo

fbshipit-source-id: cd23b314a800e5e077d866e74456d37a321439d5

diff --git a/torch/csrc/jit/passes/alias_analysis.cpp b/torch/csrc/jit/passes/alias_analysis.cpp
index 3d401c6..716762b 100644 (file)
--- a/torch/csrc/jit/passes/alias_analysis.cpp
+++ b/torch/csrc/jit/passes/alias_analysis.cpp
@@ -430,17 +430,34 @@ void AliasDb::analyze(Node* node) {
     // We don't support composite types for alias analysis yet.
     AT_ASSERT(formal->containedTypes().size() == 0);
 
-    const auto& formalAlias = formal->set();
-    auto outputAlias = formalToActual.at(formalAlias);
+    for (const auto& formalAlias : formal->sets()) {
+      // If we encounter an alias annotation that wasn't in the inputs:
+      if (!formalToActual.count(formalAlias)) {
+        // If this alias is not seen elsewhere and is the only annotation on
+        // the output, it's equivalent to being fresh:
+        //   e.g. foo(Tensor(a) self) -> Tensor(b)
+        if (formal->sets().size() == 1) {
+          giveFreshAlias(actual);
+        }
+        // Or it is the form of a|fresh, which we can ignore, taking the
+        // conservative assumption that the output must alias `a`, e.g
+        //   aten::cuda(Tensor(a) self) -> Tensor(a|fresh)
 
-    // Record writes
-    for (const auto& alias : outputAlias.sets()) {
-      if (formal->isWrite()) {
-        aliasToWrites_[alias].insert(node);
+        // Don't assign an alias set in that case.
+        continue;
       }
-    }
 
-    addAlias(actual, outputAlias);
+      auto outputAlias = formalToActual.at(formalAlias);
+
+      // Record writes
+      for (const auto& alias : outputAlias.sets()) {
+        if (formal->isWrite()) {
+          aliasToWrites_[alias].insert(node);
+        }
+      }
+
+      addAlias(actual, outputAlias);
+    }
   }
   // Keep the wildcard index up to date.
   if (hasWildcardImpl(node)) {

diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
index 573c17d..daf1b41 100644 (file)
--- a/torch/csrc/jit/register_prim_ops.cpp
+++ b/torch/csrc/jit/register_prim_ops.cpp
@@ -277,7 +277,7 @@ RegisterOperators reg({
         }),
     // reference function parse_to_conversion in python_arg_parsing.h
     Operator(
-        "aten::to(Tensor(a) self, Device? device, int? dtype=None, bool non_blocking=False, bool copy=False) -> Tensor(a)",
+        "aten::to(Tensor(a) self, Device? device, int? dtype=None, bool non_blocking=False, bool copy=False) -> Tensor(a|b)",
         [](const Node* node) -> Operation {
           return [](Stack& stack) {
             bool non_blocking;
@@ -291,7 +291,7 @@ RegisterOperators reg({
           };
         }),
     Operator(
-        "aten::to(Tensor(a) self, int? dtype=None, bool non_blocking=False, bool copy=False) -> Tensor(a)",
+        "aten::to(Tensor(a) self, int? dtype=None, bool non_blocking=False, bool copy=False) -> Tensor(a|b)",
         [](const Node* node) -> Operation {
           return [](Stack& stack) {
             bool non_blocking;
@@ -305,7 +305,7 @@ RegisterOperators reg({
           };
         }),
     Operator(
-        "aten::to(Tensor(a) self, bool non_blocking=False, bool copy=False) -> Tensor(a)",
+        "aten::to(Tensor(a) self, bool non_blocking=False, bool copy=False) -> Tensor(a|b)",
         [](const Node* node) -> Operation {
           return [](Stack& stack) {
             at::Tensor self;
@@ -377,7 +377,7 @@ RegisterOperators reg({
           };
         }),
     Operator(
-        "aten::cpu(Tensor(a) self) -> Tensor(a)",
+        "aten::cpu(Tensor(a) self) -> Tensor(a|b)",
         [](const Node* node) -> Operation {
           return [](Stack& stack) {
             at::Tensor a;
@@ -387,7 +387,7 @@ RegisterOperators reg({
           };
         }),
     Operator(
-        "aten::cuda(Tensor(a) self) -> Tensor(a)",
+        "aten::cuda(Tensor(a) self) -> Tensor(a|b)",
         [](const Node* node) -> Operation {
           return [](Stack& stack) {
             at::Tensor a;