Open up AliasAnalysisKind for any ops (#23834)

smessmer · facebook-github-bot · commit 352032c93c76 · 2019-08-09T15:10:29.000-07:00
Summary: Pull Request resolved: #23834 Re-expot of reverted PR #23810 with the bug fixed A previous diff removed the special casing for aten:: and prim:: ops in alias analysis and implements alias analysis purely based on the AliasAnalysisKind. To be sure it doesn't break our existing code base, it added asserts that make sure that our existing aten:: and prim:: ops set the correct AliasAnalysisKind. However, we don't need that restriction for future ops. Since we are now certain all existing cases are set up correctly, we can remove these assertions. ghstack-source-id: 88050427 Differential Revision: D16657239 fbshipit-source-id: 8a7606da8e9bd961bf47e3e1587b622a9c247ec6
diff --git a/test/cpp/jit/test_alias_analysis.h b/test/cpp/jit/test_alias_analysis.h
@@ -1142,13 +1142,13 @@ void testAliasRegistration() {
   }
   {
     auto registry = torch::RegisterOperators().op(
-        "aten::rand6(Tensor arg1) -> Tensor",
+        "foo::rand6(Tensor arg1) -> Tensor",
         torch::RegisterOperators::options()
             .catchAllKernel([](at::Tensor) -> at::Tensor {
               return at::rand({2, 2});
             })
             .aliasAnalysis(AliasAnalysisKind::FROM_SCHEMA));
-    const auto rand_op = Symbol::fromQualString("aten::rand6");
+    const auto rand_op = Symbol::fromQualString("foo::rand6");
     auto graph = std::make_shared<Graph>();
     auto a = graph->addInput();
     auto b = graph->insert(rand_op, {a});
@@ -1159,11 +1159,12 @@ void testAliasRegistration() {
   }
   {
     auto registry = torch::RegisterOperators().op(
-        "aten::rand7(Tensor(a) arg1) -> Tensor(a)",
+        "foo::rand7(Tensor(a) arg1) -> Tensor(a)",
         torch::RegisterOperators::options()
             .catchAllKernel([](at::Tensor t) -> at::Tensor { return t * 2; })
             .aliasAnalysis(AliasAnalysisKind::FROM_SCHEMA));
-    const auto rand_op = Symbol::fromQualString("aten::rand7");
+    const auto rand_op = Symbol::fromQualString("foo::rand7");
+
     auto graph = std::make_shared<Graph>();
     auto a = graph->addInput();
     auto b = graph->insert(rand_op, {a});
@@ -1173,11 +1174,11 @@ void testAliasRegistration() {
   }
   {
     auto registry = torch::RegisterOperators().op(
-        "aten::rand8(Tensor(a) arg1) -> Tensor(b)",
+        "foo::rand8(Tensor(a) arg1) -> Tensor(b)",
         torch::RegisterOperators::options()
             .catchAllKernel([](at::Tensor t) -> at::Tensor { return t * 2; })
             .aliasAnalysis(AliasAnalysisKind::FROM_SCHEMA));
-    const auto rand_op = Symbol::fromQualString("aten::rand8");
+    const auto rand_op = Symbol::fromQualString("foo::rand8");
     auto graph = std::make_shared<Graph>();
     auto a = graph->addInput();
     auto b = graph->insert(rand_op, {a});
diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp
@@ -904,10 +904,12 @@ bool Node::hasSideEffects() const {
         " doesn't have one either. We don't know if this op has side effects.");
     return false;
   }
+
   if (kind_.is_prim() || kind_.is_aten()) {
-    // TODO This assert is only introduced to check that we don't break the
-    // current code base. Remove this later to allow other ops to use
-    // AliasAnalysisKind::FROM_SCHEMA
+    // TODO There is nothing in the system that relies on aten:: and prim::
+    // ops using AliasAnalysisKind::FROM_SCHEMA or AliasAnalysisKind::INTERNAL_SPECIAL_CASE,
+    // but this is the intended behavior for all current ops and a good error check.
+    // We can consider lifting this constraint later if we have a use case for it.
     TORCH_INTERNAL_ASSERT(
         op->aliasAnalysisKind() == AliasAnalysisKind::INTERNAL_SPECIAL_CASE ||
             op->aliasAnalysisKind() == AliasAnalysisKind::FROM_SCHEMA,
@@ -916,6 +918,7 @@ bool Node::hasSideEffects() const {
         " has ",
         toString(op->aliasAnalysisKind()));
   }
+
   switch (op->aliasAnalysisKind()) {
     case AliasAnalysisKind::PURE:
       return false;
diff --git a/torch/csrc/jit/passes/alias_analysis.cpp b/torch/csrc/jit/passes/alias_analysis.cpp
@@ -379,9 +379,10 @@ void AliasDb::analyzeImpl(Node* node) {
       "Special cases should be handled already if we're here.");
 
   if (node->kind().is_aten() || node->kind().is_prim()) {
-    // TODO This assert is only introduced to check that we don't break the
-    // current code base. Remove this later to allow aten:: and prim:: ops to
-    // use other alias analysis kinds.
+    // TODO There is nothing in the system that relies on aten:: and prim::
+    // ops using AliasAnalysisKind::FROM_SCHEMA or AliasAnalysisKind::INTERNAL_SPECIAL_CASE,
+    // but this is the intended behavior for all current ops and a good error check.
+    // We can consider lifting this constraint later if we have a use case for it.
     TORCH_INTERNAL_ASSERT(
         analysis == AliasAnalysisKind::FROM_SCHEMA,
         "aten:: and prim:: operators should use AliasAnalysisKind::FROM_SCHEMA but ",
@@ -418,15 +419,6 @@ void AliasDb::analyzeImpl(Node* node) {
       "AliasAnalysisKind::CONSERVATIVE/PURE/INTERNAL_SPECIAL_CASE should already have been handled above");
   const auto& schema = node->schema();
 
-  // TODO This assert is only introduced to check that we don't break the
-  // current code base. Remove this later to allow other ops to use
-  // AliasAnalysisKind::FROM_SCHEMA
-  TORCH_INTERNAL_ASSERT(
-      node->kind().is_prim() || node->kind().is_aten(),
-      "The current code base should only have AliasAnalysisKind::FROM_SCHEMA for aten:: and prim:: ops but we found it for ",
-      node->kind().toDisplayString(),
-      ". We want to open this up though.");
-
   // Bind the schema's "formal" alias annotation to the actual values those
   // schema arguments represent
   std::unordered_map<Symbol, Value*> formalToActual;
