Revert "Dnt CSE across context managers"

This reverts commit 0981b01.

Reverted pytorch#76075 on behalf of https://github.com/seemethere

Author: pytorchmergebot

Diff for: test/test_jit.py

@@ -1232,19 +1232,6 @@ def fn(x, y):
 
         self.assertExportImport(g, (x, y))
 
-    def test_cse_context_managers(self):
-        def bar(x):
-            with torch.no_grad():
-                y = x * 2
-
-            z = 2 * x + y
-            return z, x.requires_grad, y.requires_grad, z.requires_grad
-
-        a = torch.rand(3, requires_grad=True)
-        b = torch.rand(3, requires_grad=True)
-
-        self.checkScript(bar, (a,))
-
     def test_cse_not_introduce_aliasing(self):
         @torch.jit.script
         def tensor_alias_outputs(x):

Diff for: torch/csrc/jit/passes/common_subexpression_elimination.cpp

@@ -5,51 +5,12 @@
 #include <torch/csrc/jit/ir/node_hashing.h>
 #include <torch/csrc/jit/jit_log.h>
 
-#include <c10/util/hash.h>
 #include <unordered_map>
 
 namespace torch {
 namespace jit {
 namespace {
 
-// There are a context managers which change global state -
-// with torch.no_grad(), with torch.cpu.amp.autocast
-// These are represented in JIT as prim::Enter and prim::Exit nodes
-// Avoid CSE across two separate with statements
-
-struct NodeAndContextNode : public std::pair<Node*, Node*> {
-  using pair::pair;
-
-  Node* node() const {
-    return this->first;
-  }
-  Node* contextNode() const {
-    return this->second;
-  }
-};
-
-struct TORCH_API HashNodeAndContext {
-  size_t operator()(const NodeAndContextNode pair) const {
-    HashNode hash;
-    // we hash on the properties of the Node to be CSE'd, and the exact node
-    // of the context its in (which may be a nullptr)
-    return c10::hash_combine(
-        hash(pair.node()), reinterpret_cast<size_t>(pair.contextNode()));
-  }
-};
-
-struct TORCH_API EqualNodeAndContext {
-  bool operator()(const NodeAndContextNode lhs, const NodeAndContextNode rhs)
-      const {
-    EqualNode eq;
-    bool nodes_equal = eq(lhs.node(), rhs.node());
-    //  similarly to equality, check equality of properties of the nodes to be
-    //  CSE'd
-    // and the exact node of its context (which may be nullptr)
-    return nodes_equal && (lhs.contextNode() == rhs.contextNode());
-  }
-};
-
 struct CommonSubexpressionEliminator {
   CommonSubexpressionEliminator(std::shared_ptr<Graph> graph)
       : graph_(std::move(graph)) {}
@@ -62,22 +23,11 @@ struct CommonSubexpressionEliminator {
   // Since the nodes are visited in topological order, one pass is enough.
   // returns true if CSE made changes to a graph
   bool run(Block* block, std::function<Node*(Node*)> parent_lookup_fn) {
-    std::unordered_set<
-        NodeAndContextNode,
-        HashNodeAndContext,
-        EqualNodeAndContext>
-        subexprs;
+    std::unordered_set<Node*, HashNode, EqualNode> subexprs;
     bool changed = false;
     for (auto it = block->nodes().begin(); it != block->nodes().end(); ++it) {
       auto node = *it;
 
-      if (node->kind() == prim::Enter) {
-        prim_enter_stack_.push_back(node);
-      }
-      if (node->kind() == prim::Exit) {
-        prim_enter_stack_.pop_back();
-      }
-
       if (node->kind() == prim::profile) {
         GRAPH_DEBUG(
             "Profiled nodes shouldn't be CSE'ed there's a separate pass that does dedup and merging:\n",
@@ -98,13 +48,9 @@ struct CommonSubexpressionEliminator {
         // Traverse sub-blocks.
         for (auto block : node->blocks()) {
           changed |= run(block, [&](Node* n) {
-            NodeAndContextNode nacn(
-                n,
-                prim_enter_stack_.size() == 0 ? nullptr
-                                              : prim_enter_stack_.back());
-            auto existing = subexprs.find(nacn);
+            auto existing = subexprs.find(n);
             if (existing != subexprs.end()) {
-              return (*existing).node();
+              return *existing;
             }
 
             return parent_lookup_fn(n);
@@ -137,12 +83,10 @@ struct CommonSubexpressionEliminator {
       }
 
       // Check whether the same subexpression already exists.
-      Node* enter_node =
-          prim_enter_stack_.size() == 0 ? nullptr : prim_enter_stack_.back();
-      auto subit = subexprs.insert(NodeAndContextNode(node, enter_node));
+      auto subit = subexprs.insert(node);
       if (!subit.second) {
         // Subexpression exists, replace the uses of node, and destroy it.
-        auto existing = (*subit.first).node();
+        auto existing = *subit.first;
 
         // don't introduce new aliasing among graph outputs
         if (getOrCreateAliasDb().mayContainAlias(
@@ -171,7 +115,6 @@ struct CommonSubexpressionEliminator {
   }
 
  private:
-  std::vector<Node*> prim_enter_stack_;
   std::unique_ptr<AliasDb> alias_db_;
   std::shared_ptr<Graph> graph_;
 };