From dbbe21dfd7fa04557301948aa1912e7743e9fef0 Mon Sep 17 00:00:00 2001 From: Elias Ellison Date: Mon, 1 Mar 2021 21:14:16 -0800 Subject: [PATCH] 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 | 55 ----------- .../csrc/jit/passes/utils/subgraph_utils.cpp | 95 ++----------------- torch/csrc/jit/passes/utils/subgraph_utils.h | 15 --- 3 files changed, 10 insertions(+), 155 deletions(-) diff --git a/test/cpp/jit/test_subgraph_utils.cpp b/test/cpp/jit/test_subgraph_utils.cpp index cadb143cf5d..f934db88651 100644 --- a/test/cpp/jit/test_subgraph_utils.cpp +++ b/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(); - - std::unordered_map 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 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 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(); diff --git a/torch/csrc/jit/passes/utils/subgraph_utils.cpp b/torch/csrc/jit/passes/utils/subgraph_utils.cpp index d6a0ffdd966..af7147a1743 100644 --- a/torch/csrc/jit/passes/utils/subgraph_utils.cpp +++ b/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& 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 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 unmerge_vmap; - unmergeSubgraph(mergeFrom, unmerge_vmap); + unmergeSubgraph(mergeFrom); std::vector 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 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 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 vmap; - mergeSubgraph(mergeTo, mergeFrom, vmap); } } // namespace @@ -170,14 +127,12 @@ std::shared_ptr getSubgraph(Node* n) { return n->g(attr::Subgraph); } -void unmergeSubgraph( - Node* subgraphNode, - std::unordered_map& 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 vmap; - unmergeSubgraph(subgraphNode, vmap); -} - void collectNestedUses( std::unordered_set& closed_over_values, std::unordered_set& new_values, @@ -244,11 +194,10 @@ std::unordered_set closedOverValues( void mergeNodeIntoSubgraph( Node* toMerge, Node* subgraphNode, - std::unordered_map& 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 vmap; - mergeNodeIntoSubgraph(toMerge, subgraphNode, vmap, destroyNode); -} - -Node* createSingletonSubgraph( - Node* n, - Symbol subgraphKind, - std::unordered_map& 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->current_scope())); subgraph->insertBefore(n); - mergeNodeIntoSubgraph(n, subgraph, vmap); + mergeNodeIntoSubgraph(n, subgraph); return subgraph; } -Node* createSingletonSubgraph(Node* n, Symbol subgraphKind) { - std::unordered_map vmap; - return createSingletonSubgraph(n, subgraphKind, vmap); -} - void mergeNodeIntoSubgraphAndUpdateAliasing( Node* to_merge, Node* subgraphNode, diff --git a/torch/csrc/jit/passes/utils/subgraph_utils.h b/torch/csrc/jit/passes/utils/subgraph_utils.h index f111d495c63..9c84eea4717 100644 --- a/torch/csrc/jit/passes/utils/subgraph_utils.h +++ b/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& 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& 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& vmap); // Convenience function std::shared_ptr getSubgraph(Node* n);