Remove unused subgraph vmap api (#52512)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52512

This API is not used at all, and is tricky to maintain. When we were using it last we ran into lifetime issues when using `Value *` as the key. In hind sight, we should have been using `value->unique()`, but regardless, this not being used and should be removed.

Test Plan: Imported from OSS

Reviewed By: navahgar

Differential Revision: D26696695

Pulled By: eellison

fbshipit-source-id: 97ed92e88ecab0085fabbac46573611666bf2420

test/cpp/jit/test_subgraph_utils.cpp

@@ -38,61 +38,6 @@ TEST(SubgraphUtilsTest, Basic) {
   ASSERT_EQ(originalNodes.size(), newNodes.size());
 }
 
-TEST(SubgraphUtilsTest, Vmap) {
-  auto graph = std::make_shared<Graph>();
-
-  std::unordered_map<std::string, Value*> parse_map;
-  parseIR(
-      R"IR(
-graph(%a : Tensor, %b : Tensor, %c : Tensor):
-  %x : Tensor = aten::tanh(%a)
-  %y : Tensor = aten::mul(%a, %b)
-
-  %p : Tensor = aten::div(%c, %b)
-  %q : Tensor = aten::mul(%p, %a)
-  return (%x, %y, %q))IR",
-      &*graph,
-      parse_map);
-
-  std::unordered_map<Value*, Value*> vmap1;
-
-  Node* subgraph1 = SubgraphUtils::createSingletonSubgraph(
-      parse_map.at("y")->node(), prim::DifferentiableGraph);
-  SubgraphUtils::mergeNodeIntoSubgraph(
-      parse_map.at("x")->node(), subgraph1, vmap1);
-  // vmap should have two entries: a mapping for the '%x' value - the output of
-  // the node we merged in, and a mapping for the '%a' value - the input of the
-  // node.
-  ASSERT_EQ(vmap1.size(), 2);
-
-  // Check that after mergeNodeIntoSubgraph we can still access the node
-  // corresponding to the original node "%x = aten::tanh(%a)".
-  //
-  // Note that parse_map["x"] points to a destroyed Value after the
-  // mergeNodeIntoSubgraph call - we cannot access its content anymore, but
-  // still can use it as a key in the value map to find the value it was moved
-  // to.
-  Node* new_tanh = vmap1.at(parse_map.at("x"))->node();
-  ASSERT_TRUE(new_tanh->kind() == aten::tanh);
-
-  Node* subgraph2 = SubgraphUtils::createSingletonSubgraph(
-      parse_map["q"]->node(), prim::DifferentiableGraph);
-  SubgraphUtils::mergeNodeIntoSubgraph(parse_map.at("p")->node(), subgraph2);
-
-  std::unordered_map<Value*, Value*> vmap2;
-  Value* new_tanh_out = new_tanh->output();
-  SubgraphUtils::mergeNodeIntoSubgraph(subgraph1, subgraph2, vmap2);
-  // vmap should have 6 entries, since we moved 4 values into the graph (the
-  // values correspond to the original values '%a', '%b', '%x', and '%y').
-  // and we map the node outputs for '%x' and '%y'
-  ASSERT_EQ(vmap2.size(), 6);
-
-  // Check that after mergeNodeIntoSubgraph we can still access the node
-  // corresponding to the original node, even if the toMerge node had a subgraph
-  // as well
-  ASSERT_TRUE(vmap2.at(new_tanh_out)->node()->kind() == aten::tanh);
-}
-
 TEST(SubgraphUtilsTest, GraphName) {
   auto graph = std::make_shared<Graph>();
 

torch/csrc/jit/passes/utils/subgraph_utils.cpp

@@ -95,73 +95,30 @@ Node* executeSubgraphMergeAndUpdateAliasing(
 
 // Combine the nodes in two subgraph together. The nodes will end up in
 // `mergeTo`, and `mergeFrom` is destroyed.
-void mergeSubgraph(
-    Node* mergeTo,
-    Node* mergeFrom,
-    std::unordered_map<Value*, Value*>& vmap) {
+void mergeSubgraph(Node* mergeTo, Node* mergeFrom) {
   Node* nodeBeforeMergeFrom = mergeFrom->prev();
   Node* nodeAfterMergeFrom = mergeFrom->next();
 
-  // will be used later to map the node outputs -> new subgraph values
-  std::unordered_map<Value*, Value*> node_outputs_to_subgraph_values;
-  for (size_t i = 0; i < mergeFrom->outputs().size(); ++i) {
-    node_outputs_to_subgraph_values[mergeFrom->output(i)] =
-        getSubgraph(mergeFrom)->outputs().at(i);
-  }
-
-  // unmerge_map will contain mapping from values from the mergeTo's subgraph
-  // (we will call them "original" values) to the corresponding values that we
-  // created in the main graph (we will call them "unmerged" values) as we
-  // unmerged the mergeTo's subgraph.
-  std::unordered_map<Value*, Value*> unmerge_vmap;
-  unmergeSubgraph(mergeFrom, unmerge_vmap);
+  unmergeSubgraph(mergeFrom);
 
   std::vector<Node*> nodes;
   const auto end_it = nodeBeforeMergeFrom->reverseIterator();
   auto it = nodeAfterMergeFrom->reverseIterator();
   ++it;
 
-  // Now we're merging the "unmerged" nodes into the mergeFrom subgraph. That
-  // will give us a new map: "unmerged" -> "merged".
-  std::unordered_map<Value*, Value*> merge_vmap;
-
-  // defer destroying nodes until after all nodes have been merged, otherwise we
-  // run into lifetime issues where the previous mapping of the merged nodes
-  // inputs/outputs can be overwritten with newly created values
+  // defer destroying nodes until after all nodes have been merged,
+  // to make iterators easier to reason about
   std::vector<Node*> merged_nodes;
   while (it != end_it) {
     Node* node = *it;
     ++it;
     merged_nodes.push_back(node);
-    mergeNodeIntoSubgraph(node, mergeTo, merge_vmap, /*destroyNode*/ false);
+    mergeNodeIntoSubgraph(node, mergeTo, /*destroyNode*/ false);
   }
 
   for (Node* n : merged_nodes) {
     n->destroy();
   }
-
-  // Vmap should contain "original" -> "merged" mapping, thus we basically need
-  // to perform the following transformation:
-  // vmap[x] = merge_vmap[unmerge_map[x]]
-  for (auto& kv : unmerge_vmap) {
-    if (merge_vmap.count(kv.second)) {
-      vmap[kv.first] = merge_vmap.at(kv.second);
-    } else {
-      vmap[kv.first] = kv.second;
-    }
-  }
-
-  // fill the value mapping with node output -> new subgraph value
-  for (const auto& mapping : node_outputs_to_subgraph_values) {
-    vmap[mapping.first] = vmap[mapping.second];
-  }
-}
-
-// Combine the nodes in two subgraph together. The nodes will end up in
-// `mergeTo`, and `mergeFrom` is destroyed.
-void mergeSubgraph(Node* mergeTo, Node* mergeFrom) {
-  std::unordered_map<Value*, Value*> vmap;
-  mergeSubgraph(mergeTo, mergeFrom, vmap);
 }
 
 } // namespace
@@ -170,14 +127,12 @@ std::shared_ptr<Graph> getSubgraph(Node* n) {
   return n->g(attr::Subgraph);
 }
 
-void unmergeSubgraph(
-    Node* subgraphNode,
-    std::unordered_map<Value*, Value*>& vmap) {
+void unmergeSubgraph(Node* subgraphNode) {
   // Inline the graph, replace uses of node outputs and destroy the node
   auto outerGraph = subgraphNode->owningGraph();
   WithInsertPoint guard(subgraphNode);
   const auto subgraphOutputs = insertGraph(
-      *outerGraph, *getSubgraph(subgraphNode), subgraphNode->inputs(), vmap);
+      *outerGraph, *getSubgraph(subgraphNode), subgraphNode->inputs());
   AT_ASSERT(subgraphOutputs.size() >= subgraphNode->outputs().size());
   for (size_t i = 0; i < subgraphNode->outputs().size(); ++i) {
     subgraphNode->outputs()[i]->replaceAllUsesWith(subgraphOutputs[i]);
@@ -185,11 +140,6 @@ void unmergeSubgraph(
   subgraphNode->destroy();
 }
 
-void unmergeSubgraph(Node* subgraphNode) {
-  std::unordered_map<Value*, Value*> vmap;
-  unmergeSubgraph(subgraphNode, vmap);
-}
-
 void collectNestedUses(
     std::unordered_set<Value*>& closed_over_values,
     std::unordered_set<Value*>& new_values,
@@ -244,11 +194,10 @@ std::unordered_set<Value*> closedOverValues(
 void mergeNodeIntoSubgraph(
     Node* toMerge,
     Node* subgraphNode,
-    std::unordered_map<Value*, Value*>& vmap,
     bool destroyNode) {
   AT_ASSERT(hasSubgraph(subgraphNode) && toMerge != subgraphNode);
   if (hasSubgraph(toMerge)) {
-    return mergeSubgraph(subgraphNode, toMerge, vmap);
+    return mergeSubgraph(subgraphNode, toMerge);
   }
 
   auto subgraph = getSubgraph(subgraphNode);
@@ -306,13 +255,6 @@ void mergeNodeIntoSubgraph(
   auto mergedNode = subgraph->insertNode(
       subgraph->createClone(toMerge, [&](Value* v) { return inputsMap[v]; }));
 
-  for (size_t idx = 0; idx < toMerge->outputs().size(); idx++) {
-    vmap[toMerge->output(idx)] = mergedNode->output(idx);
-  }
-  for (size_t idx = 0; idx < toMerge->inputs().size(); idx++) {
-    vmap[toMerge->input(idx)] = mergedNode->input(idx);
-  }
-
   // If n's outputs were inputs to `group`, remove them since we just merged
   // n in.
   //
@@ -326,7 +268,6 @@ void mergeNodeIntoSubgraph(
       size_t p = it - inputs.begin();
       subgraphNode->removeInput(p);
       subgraph->inputs()[p]->replaceAllUsesWith(mergedNode->outputs()[i]);
-      vmap[subgraph->inputs()[p]] = mergedNode->output(i);
       subgraph->eraseInput(p);
     }
   }
@@ -356,31 +297,15 @@ void mergeNodeIntoSubgraph(
   }
 }
 
-void mergeNodeIntoSubgraph(
-    Node* toMerge,
-    Node* subgraphNode,
-    bool destroyNode) {
-  std::unordered_map<Value*, Value*> vmap;
-  mergeNodeIntoSubgraph(toMerge, subgraphNode, vmap, destroyNode);
-}
-
-Node* createSingletonSubgraph(
-    Node* n,
-    Symbol subgraphKind,
-    std::unordered_map<Value*, Value*>& vmap) {
+Node* createSingletonSubgraph(Node* n, Symbol subgraphKind) {
   auto graph = n->owningGraph();
   auto subgraph = graph->create(subgraphKind, 0);
   subgraph->g_(attr::Subgraph, std::make_shared<Graph>(graph->current_scope()));
   subgraph->insertBefore(n);
-  mergeNodeIntoSubgraph(n, subgraph, vmap);
+  mergeNodeIntoSubgraph(n, subgraph);
   return subgraph;
 }
 
-Node* createSingletonSubgraph(Node* n, Symbol subgraphKind) {
-  std::unordered_map<Value*, Value*> vmap;
-  return createSingletonSubgraph(n, subgraphKind, vmap);
-}
-
 void mergeNodeIntoSubgraphAndUpdateAliasing(
     Node* to_merge,
     Node* subgraphNode,

torch/csrc/jit/passes/utils/subgraph_utils.h

@@ -19,13 +19,7 @@ namespace SubgraphUtils {
 // `n` is destroyed.
 //
 // Returns the new subgraph node.
-// An optional argument 'vmap' could be used to retrieve value mappings
-// Values will be mapped to their new subgraph values
 TORCH_API Node* createSingletonSubgraph(Node* n, Symbol subgraphKind);
-TORCH_API Node* createSingletonSubgraph(
-    Node* n,
-    Symbol subgraphKind,
-    std::unordered_map<Value*, Value*>& vmap);
 
 // Creates a new subgraph that only contains `n`, amd udpates the new outputs
 // of the subgraph to have the aliasing properties of the original `n` outputs
@@ -43,11 +37,6 @@ TORCH_API void mergeNodeIntoSubgraph(
     Node* toMerge,
     Node* subgraphNode,
     bool destroyNode = true);
-TORCH_API void mergeNodeIntoSubgraph(
-    Node* toMerge,
-    Node* subgraphNode,
-    std::unordered_map<Value*, Value*>& vmap,
-    bool destroyNode = true);
 
 // Merges a node into a subgraph node, and updates the new outputs of the
 // subgraph to have the aliasing properties of the corresponding `to_merge`
@@ -59,11 +48,7 @@ TORCH_API void mergeNodeIntoSubgraphAndUpdateAliasing(
 
 // Move nodes from a subgraph node to the outer graph.
 // `subgraphNode` is destroyed.
-// An optional argument 'vmap' could be used to retrieve value mappings.
 TORCH_API void unmergeSubgraph(Node* subgraphNode);
-TORCH_API void unmergeSubgraph(
-    Node* subgraphNode,
-    std::unordered_map<Value*, Value*>& vmap);
 
 // Convenience function
 std::shared_ptr<Graph> getSubgraph(Node* n);