From 92e236dd24a5e52b44fbfccd75f4833b68e14b75 Mon Sep 17 00:00:00 2001 From: Tong Shen Date: Fri, 30 Nov 2018 11:12:00 -0800 Subject: [PATCH] Fix for outside compilation. When generating Placeholder node for host computation/outside compilation nodes, we have to generate different Placeholder node for different src_output edges. PiperOrigin-RevId: 223544322 --- tensorflow/compiler/jit/encapsulate_util.cc | 21 +++++---- .../compiler/jit/encapsulate_util_test.cc | 46 +++++++++++++++++-- 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/tensorflow/compiler/jit/encapsulate_util.cc b/tensorflow/compiler/jit/encapsulate_util.cc index bcc3213285b..1f4b9c90a4f 100644 --- a/tensorflow/compiler/jit/encapsulate_util.cc +++ b/tensorflow/compiler/jit/encapsulate_util.cc @@ -241,7 +241,7 @@ Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( // Remove the edge from host to outside compilation. Add a placeholder as // outside compilation node input. - std::map placeholders; + std::map, Node*> placeholders; for (int i = 0; i < edges.size(); i++) { Node* dst = g->FindNodeId(edges[i].dst_node_id); const Edge* e; @@ -253,9 +253,10 @@ Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( // Find or create placeholder node. string new_name = edges[i].is_host_to_outside_compilation - ? absl::StrCat(src->name(), "_host_to_oc_placeholder") - : absl::StrCat(src->name(), "_oc_to_host_placeholder"); - auto iter = placeholders.find(new_name); + ? absl::StrCat(src->name(), "_host_to_oc_placeholder_", src_output) + : absl::StrCat(src->name(), "_oc_to_host_placeholder_", src_output); + auto placeholder_index = std::make_pair(src->name(), src_output); + auto iter = placeholders.find(placeholder_index); Node* placeholder_node; if (iter == placeholders.end()) { NodeDefBuilder placeholder_builder(new_name, "Placeholder"); @@ -288,7 +289,7 @@ Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( Status s; placeholder_node = g->AddNode(placeholder_def, &s); TF_RETURN_IF_ERROR(s); - placeholders[new_name] = placeholder_node; + placeholders[placeholder_index] = placeholder_node; } else { placeholder_node = iter->second; } @@ -642,7 +643,7 @@ Status PreprocessDataEdgesBetweenOutsideCompilations( // Remove the edge from host to outside compilation. Add a placeholder as // outside compilation node input. - std::map placeholders; + std::map, Node*> placeholders; for (int i = 0; i < edges.size(); i++) { Node* dst = g->FindNodeId(edges[i].dst_node_id); const Edge* e; @@ -652,8 +653,10 @@ Status PreprocessDataEdgesBetweenOutsideCompilations( g->RemoveEdge(e); // Find or create placeholder node. - string new_name = absl::StrCat(src->name(), "_oc_to_oc_placeholder"); - auto iter = placeholders.find(new_name); + string new_name = + absl::StrCat(src->name(), "_oc_to_oc_placeholder_", src_output); + auto placeholder_index = std::make_pair(src->name(), src_output); + auto iter = placeholders.find(placeholder_index); Node* placeholder_node; if (iter == placeholders.end()) { NodeDefBuilder placeholder_builder(new_name, "Placeholder"); @@ -673,7 +676,7 @@ Status PreprocessDataEdgesBetweenOutsideCompilations( Status s; placeholder_node = g->AddNode(placeholder_def, &s); TF_RETURN_IF_ERROR(s); - placeholders[new_name] = placeholder_node; + placeholders[placeholder_index] = placeholder_node; } else { placeholder_node = iter->second; } diff --git a/tensorflow/compiler/jit/encapsulate_util_test.cc b/tensorflow/compiler/jit/encapsulate_util_test.cc index 25c32cef01d..3b8b49cb92f 100644 --- a/tensorflow/compiler/jit/encapsulate_util_test.cc +++ b/tensorflow/compiler/jit/encapsulate_util_test.cc @@ -153,23 +153,33 @@ TEST(PreprocessForEncapsulationTest, ControlEdges) { TEST(PreprocessForEncapsulationTest, DataEdges) { // Build the graph: // "const_0" and "const_1" in host computation + // "identityn0" = ("const_0", "const_1") in host computation 0 // "add0" = "const_0" + "const_1" in XLA computation 0 // "add1" = "add0" + "const_0" in XLA computation 0 & outside compilation 0 // "identity0" = "add1" in XLA computation 0 // "add2" = "add1" + "identity0" in host computation // "add3" = "add1" + "add2" in XLA computation 1 - // "add4" = "identity0" + "add2" in XLA computation 1 & outside compilation 1 + // "add4" = "identity0" + "add2" in XLA computation 1 & outside compilation 0 + // "add5" = "identityn0"[0] + "identityn0"[1] in XLA computation 1 & + // outside compilation 0 + // "identityn1" = ("identityn0"[0], "identityn0"[1]) in XLA computation 1 & + // outside compilation 0 // "identity1" = "add4" in XLA computation 1 // "identity2" = "identity1" in host computation tensorflow::Scope s = tensorflow::Scope::NewRootScope(); Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {}); Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {}); + auto identityn0 = + ops::IdentityN(s.WithOpName("identityn_0"), {const_0, const_1}); Output add0 = ops::Add(s.WithOpName("add0"), const_0, const_1); Output add1 = ops::Add(s.WithOpName("add1"), add0, const_0); Output identity0 = ops::Identity(s.WithOpName("identity0"), add1); Output add2 = ops::Add(s.WithOpName("add2"), add1, identity0); Output add3 = ops::Add(s.WithOpName("add3"), add1, add2); Output add4 = ops::Add(s.WithOpName("add4"), identity0, add2); + Output add5 = ops::Add(s.WithOpName("add5"), identityn0[0], identityn0[1]); + auto identityn1 = ops::IdentityN(s.WithOpName("identityn_1"), + {identityn0[0], identityn0[1]}); Output identity1 = ops::Identity(s.WithOpName("identity1"), add4); Output identity2 = ops::Identity(s.WithOpName("identity2"), add4); Graph g(OpRegistry::Global()); @@ -180,6 +190,8 @@ TEST(PreprocessForEncapsulationTest, DataEdges) { Node *add0_node = node_index["add0"], *add1_node = node_index["add1"], *identity0_node = node_index["identity0"], *add3_node = node_index["add3"], *add4_node = node_index["add4"], + *add5_node = node_index["add5"], + *identityn1_node = node_index["identityn_1"], *identity1_node = node_index["identity1"]; add0_node->AddAttr("_xla", "0"); add1_node->AddAttr("_xla", "0"); @@ -188,6 +200,10 @@ TEST(PreprocessForEncapsulationTest, DataEdges) { add3_node->AddAttr("_xla", "1"); add4_node->AddAttr("_xla", "1"); add4_node->AddAttr("_oc", "0"); + add5_node->AddAttr("_xla", "1"); + add5_node->AddAttr("_oc", "0"); + identityn1_node->AddAttr("_xla", "1"); + identityn1_node->AddAttr("_oc", "0"); identity1_node->AddAttr("_xla", "1"); TF_CHECK_OK(PreprocessForEncapsulation(&g, "_xla", "_oc")); @@ -205,8 +221,9 @@ TEST(PreprocessForEncapsulationTest, DataEdges) { EXPECT_NE(bridge_identity0_add4, nullptr); // Step 3: add placeholder for edges between host computation and outside // compilation. - EXPECT_EQ(bridge_add1_add3->def().input(0), "add1_oc_to_host_placeholder"); - Node *add1_oc_to_host_placeholder = node_index["add1_oc_to_host_placeholder"]; + EXPECT_EQ(bridge_add1_add3->def().input(0), "add1_oc_to_host_placeholder_0"); + Node *add1_oc_to_host_placeholder = + node_index["add1_oc_to_host_placeholder_0"]; TF_CHECK_OK(GetNodeAttr(add1_oc_to_host_placeholder->attrs(), kOutsideCompilationToHostOriginalNodeAttrName, &str)); EXPECT_EQ(str, "add1"); @@ -217,15 +234,34 @@ TEST(PreprocessForEncapsulationTest, DataEdges) { add4_node = node_index["add4"]; ASSERT_NE(add4_node, nullptr); EXPECT_EQ(add4_node->def().input(0), - "bridge_identity0_add4_host_to_oc_placeholder"); + "bridge_identity0_add4_host_to_oc_placeholder_0"); Node *identity0_host_to_oc_placeholder = - node_index["bridge_identity0_add4_host_to_oc_placeholder"]; + node_index["bridge_identity0_add4_host_to_oc_placeholder_0"]; TF_CHECK_OK(GetNodeAttr(identity0_host_to_oc_placeholder->attrs(), kHostToOutsideCompilationOriginalNodeAttrName, &str)); EXPECT_EQ(str, "bridge_identity0_add4"); TF_CHECK_OK(GetNodeAttr(identity0_host_to_oc_placeholder->attrs(), kHostToOutsideCompilationSrcOutputAttrName, &i)); EXPECT_EQ(i, 0); + + // Check different placeholder nodes are created for different src_output. + Node *placeholder0 = node_index["identityn_0_host_to_oc_placeholder_0"], + *placeholder1 = node_index["identityn_0_host_to_oc_placeholder_1"]; + EXPECT_NE(placeholder0, nullptr); + EXPECT_NE(placeholder1, nullptr); + // Check we only have 2 placeholder nodes created for "identityn_0". + int placeholder_count = 0; + for (Node *n : g.nodes()) { + if (HasNodeAttr(n->def(), kHostToOutsideCompilationOriginalNodeAttrName)) { + string attr; + TF_CHECK_OK(GetNodeAttr( + n->attrs(), kHostToOutsideCompilationOriginalNodeAttrName, &attr)); + if (attr == "identityn_0") { + ++placeholder_count; + } + } + } + EXPECT_EQ(placeholder_count, 2); } TEST(PostprocessForEncapsulationTest, ControlEdges) {